diff --git a/.circleci/config.yml b/.circleci/config.yml index 6e368a3debe..32d2cf0390c 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1440,6 +1440,7 @@ jobs: TEST_FILES=$(printf "%s\n" \ tests/local_testing/test_dual_cache.py \ tests/local_testing/test_redis_batch_optimizations.py \ + tests/local_testing/test_redis_increment_with_floor.py \ tests/local_testing/test_router_utils.py) echo "$TEST_FILES" | circleci tests run \ --verbose \ @@ -2648,6 +2649,19 @@ jobs: name: Start mock LLM server command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true + - run: + name: Start mock Presidio server + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py + background: true + - run: + name: Wait for mock Presidio server + command: | + for i in $(seq 1 30); do + if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi + sleep 1 + done + echo "Mock Presidio server never answered /health on port 8091" >&2 + exit 1 - run: name: Start LiteLLM proxy environment: @@ -2778,6 +2792,19 @@ jobs: name: Start mock LLM server command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true + - run: + name: Start mock Presidio server + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py + background: true + - run: + name: Wait for mock Presidio server + command: | + for i in $(seq 1 30); do + if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi + sleep 1 + done + echo "Mock Presidio server never answered /health on port 8091" >&2 + exit 1 - run: name: Start LiteLLM proxy under a server root path environment: diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml index cbf380bac01..fc86c229c81 100644 --- a/.github/ISSUE_TEMPLATE/config.yml +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -1,4 +1,4 @@ -blank_issues_enabled: true +blank_issues_enabled: false contact_links: - name: Schedule Demo url: https://enterprise.litellm.ai/demo diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 0da07038152..672f102eeb1 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -106,11 +106,6 @@ dockerfiles: and lint workflows already exercise that output, so building the image adds no signal about it paths: - ui/Dockerfile - - reason: >- - The Rust gateway ships as its own chart and package with a separate release pipeline, so its - image is not part of this repo's Python image set - paths: - - litellm-rust/crates/ai-gateway/Dockerfile - reason: >- An example image under cookbook/ that is documentation rather than a shipped artifact paths: diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py new file mode 100644 index 00000000000..c4348c20873 --- /dev/null +++ b/.github/e2e-stack/assert_tests_ran.py @@ -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()) diff --git a/.github/e2e-stack/down.sh b/.github/e2e-stack/down.sh new file mode 100755 index 00000000000..9f72f2d6e64 --- /dev/null +++ b/.github/e2e-stack/down.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +set -uo pipefail + +STACK_DIR="${E2E_STACK_DIR:-${RUNNER_TEMP:-/tmp}/litellm-e2e-stack}" + +for pid_file in "${STACK_DIR}"/pids/*.pid; do + [[ -f "${pid_file}" ]] || continue + pkill -TERM -P "$(cat "${pid_file}")" 2>/dev/null + kill -TERM "$(cat "${pid_file}")" 2>/dev/null + rm -f "${pid_file}" +done + +for container in e2e-nginx e2e-valkey e2e-jaeger e2e-postgres; do + docker rm -f "${container}" >/dev/null 2>&1 +done + +exit 0 diff --git a/.github/e2e-stack/secrets_to_env.py b/.github/e2e-stack/secrets_to_env.py new file mode 100644 index 00000000000..7913b25918b --- /dev/null +++ b/.github/e2e-stack/secrets_to_env.py @@ -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()) diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py new file mode 100644 index 00000000000..a62358f81ff --- /dev/null +++ b/.github/e2e-stack/select_tests.py @@ -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()) diff --git a/.github/e2e-stack/up.sh b/.github/e2e-stack/up.sh new file mode 100755 index 00000000000..2f2e6c6f9a8 --- /dev/null +++ b/.github/e2e-stack/up.sh @@ -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" </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" <&2 + docker logs ai-gateway >&2 + exit 1 + - name: Assert the gateway loaded the baked config + run: | + docker logs ai-gateway 2>&1 | tee gateway.log + grep 'via python config reader' gateway.log + - name: Stop the gateway + if: always() + run: docker rm -f ai-gateway || true diff --git a/.github/workflows/auto-close-duplicates.yml b/.github/workflows/auto-close-duplicates.yml index d8256917805..9a362b4c85c 100644 --- a/.github/workflows/auto-close-duplicates.yml +++ b/.github/workflows/auto-close-duplicates.yml @@ -22,7 +22,7 @@ on: permissions: {} jobs: - test: + sweep-tests: if: github.event_name == 'pull_request' runs-on: ubuntu-latest timeout-minutes: 5 diff --git a/.github/workflows/cost-map-guard.yml b/.github/workflows/cost-map-guard.yml new file mode 100644 index 00000000000..61a56f1f47e --- /dev/null +++ b/.github/workflows/cost-map-guard.yml @@ -0,0 +1,45 @@ +name: Cost map guard + +on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed + pull_request_target: + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + cost-map-guard: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + - name: Fetch the pull request head and its merge base + id: revisions + env: + GH_TOKEN: ${{ github.token }} + BASE_SHA: ${{ github.event.pull_request.base.sha }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${BASE_SHA}...${HEAD_SHA}" --jq '.merge_base_commit.sha')" + git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA" + echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT" + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + - name: Run the guard + env: + MERGE_BASE: ${{ steps.revisions.outputs.merge_base }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + HEAD_REF: ${{ github.event.pull_request.head.ref }} + run: | + uv run --frozen python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF" diff --git a/.github/workflows/publish-basedpyright-base-counts.yml b/.github/workflows/publish-basedpyright-base-counts.yml index cd443a8e9db..27d4682dbd9 100644 --- a/.github/workflows/publish-basedpyright-base-counts.yml +++ b/.github/workflows/publish-basedpyright-base-counts.yml @@ -1,6 +1,6 @@ name: Publish basedpyright base counts -# Every commit on litellm_internal_staging is some branch's future merge-base. +# Every commit on main or litellm_internal_staging can become a future merge-base. # Publishing its per-rule basedpyright counts as an artifact lets # scripts/type_check_gate.py download them in seconds instead of paying a # 60-110s second basedpyright pass on every fresh worktree or moved merge-base. @@ -10,13 +10,13 @@ name: Publish basedpyright base counts on: push: branches: + - main - litellm_internal_staging workflow_dispatch: inputs: ref: - description: "Ref to compute and publish base counts for" + description: "Ref to compute and publish base counts for (defaults to the workflow run's commit)" required: false - default: litellm_internal_staging permissions: contents: read diff --git a/.github/workflows/report-rust-release-wheel.yml b/.github/workflows/report-rust-release-wheel.yml deleted file mode 100644 index 74e6be69604..00000000000 --- a/.github/workflows/report-rust-release-wheel.yml +++ /dev/null @@ -1,129 +0,0 @@ -name: Report LiteLLM Rust release wheel - -on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs - workflow_run: - workflows: - - LiteLLM Rust - types: - - completed - -permissions: {} - -concurrency: - group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }} - cancel-in-progress: false - -jobs: - report-release-wheel: - name: report release wheel - if: >- - github.event.workflow_run.event == 'pull_request' && - github.event.workflow_run.path == '.github/workflows/test-rust.yml' && - github.event.workflow_run.head_repository.full_name == github.repository && - github.event.workflow_run.pull_requests[0].number != null - runs-on: ubuntu-latest - timeout-minutes: 5 - permissions: - pull-requests: write - - steps: - - name: Link release wheel report on PR - uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 - env: - COMMENT_MARKER: "" - with: - script: | - const marker = process.env.COMMENT_MARKER; - const workflowRun = context.payload.workflow_run; - const allowedConclusions = new Set([ - "action_required", - "cancelled", - "failure", - "neutral", - "skipped", - "stale", - "startup_failure", - "success", - "timed_out", - ]); - if ( - !allowedConclusions.has(workflowRun.conclusion) || - workflowRun.event !== "pull_request" || - workflowRun.path !== ".github/workflows/test-rust.yml" || - workflowRun.head_repository?.full_name !== - `${context.repo.owner}/${context.repo.repo}` || - workflowRun.pull_requests?.length !== 1 - ) { - throw new Error("unexpected source workflow"); - } - const pullRequest = workflowRun.pull_requests[0]; - const pullRequestNumber = pullRequest.number; - const headSha = workflowRun.head_sha; - const runId = workflowRun.id; - if ( - !Number.isSafeInteger(pullRequestNumber) || - pullRequestNumber <= 0 || - !Number.isSafeInteger(runId) || - runId <= 0 || - !/^[0-9a-f]{40}$/.test(headSha) || - pullRequest.head?.sha !== headSha - ) { - throw new Error("invalid source workflow metadata"); - } - const runUrl = - `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` + - `/actions/runs/${runId}`; - const result = - workflowRun.conclusion === "success" - ? "successfully" - : `with \`${workflowRun.conclusion}\``; - const body = [ - marker, - "## LiteLLM Rust workflow", - "", - `Workflow completed ${result} for \`${headSha}\``, - "", - `[View workflow run](${runUrl})`, - ].join("\n"); - const comments = await github.paginate(github.rest.issues.listComments, { - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: pullRequestNumber, - per_page: 100, - }); - const existing = comments.find( - (comment) => - comment.user?.login === "github-actions[bot]" && - comment.body?.startsWith(marker), - ); - const currentPullRequest = ( - await github.rest.pulls.get({ - owner: context.repo.owner, - repo: context.repo.repo, - pull_number: pullRequestNumber, - }) - ).data; - if ( - currentPullRequest.state !== "open" || - currentPullRequest.head.repo?.full_name !== - `${context.repo.owner}/${context.repo.repo}` || - currentPullRequest.head.sha !== headSha - ) { - core.info("source workflow no longer matches the current pull request head"); - return; - } - if (existing) { - await github.rest.issues.updateComment({ - owner: context.repo.owner, - repo: context.repo.repo, - comment_id: existing.id, - body, - }); - } else { - await github.rest.issues.createComment({ - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: pullRequestNumber, - body, - }); - } diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml index 1daaadeabe2..18d5e3de1eb 100644 --- a/.github/workflows/sync-together-ai-models.yml +++ b/.github/workflows/sync-together-ai-models.yml @@ -13,10 +13,12 @@ jobs: sync_together_ai_models: if: github.repository == 'BerriAI/litellm' runs-on: ubuntu-latest + env: + BASE_BRANCH: ${{ github.event.repository.default_branch }} steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: - ref: litellm_internal_staging + ref: ${{ env.BASE_BRANCH }} persist-credentials: false - name: Set up uv uses: ./.github/actions/setup-uv-with-retries @@ -63,6 +65,6 @@ jobs: gh pr create --title "feat(models): sync together_ai model registry" \ --body-file "$RUNNER_TEMP/pr_body.md" \ --head "$branch" \ - --base litellm_internal_staging + --base "$BASE_BRANCH" env: GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index d02f5878396..e6d2264fbf0 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -74,6 +74,15 @@ jobs: - name: check_workflow_startup_safety run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py + - name: check_workflow_job_name_collisions + run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_job_name_collisions.py + + - name: test_workflow_job_name_collisions + run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.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 diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml new file mode 100644 index 00000000000..23ab6dfcfe4 --- /dev/null +++ b/.github/workflows/test-e2e-changed.yml @@ -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 diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 9f1283da19e..06d369eabcd 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -113,7 +113,7 @@ jobs: - name: Check ruff format if: steps.changes.outputs.decision != 'skip' run: | - git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true + git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- ':(glob)litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then echo "No changed litellm Python files to check with ruff format." exit 0 @@ -172,7 +172,7 @@ jobs: - name: Check tests/e2e basedpyright (zero errors) if: steps.changes.outputs.decision != 'skip' run: | - if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- 'tests/e2e/**/*.py' | grep -q .; then + if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- ':(glob)tests/e2e/**/*.py' | grep -q .; then uv run --no-sync basedpyright tests/e2e else echo "No changed tests/e2e Python files; skipping." diff --git a/.github/workflows/test-litellm-ui-unit.yml b/.github/workflows/test-litellm-ui-unit.yml index 314efcc49d5..b93bf84320d 100644 --- a/.github/workflows/test-litellm-ui-unit.yml +++ b/.github/workflows/test-litellm-ui-unit.yml @@ -12,6 +12,7 @@ on: - "litellm_**" push: branches: + - main - litellm_internal_staging concurrency: @@ -65,7 +66,7 @@ jobs: BASE_SHA: ${{ github.event.pull_request.base.sha }} HEAD_SHA: ${{ github.event.pull_request.head.sha }} run: | - full_suite() { npm run test -- --run --pool forks --poolOptions.forks.maxForks=14; } + full_suite() { npm run test -- --run --pool forks --maxWorkers=14; } if [ -z "$BASE_SHA" ]; then echo "Push to $GITHUB_REF_NAME: running the full suite" @@ -94,4 +95,4 @@ jobs: echo "Pull request: running tests related to ${#changed_files[@]} changed UI files" npm run test -- related "${changed_files[@]}" --run --passWithNoTests \ - --pool forks --poolOptions.forks.maxForks=14 + --pool forks --maxWorkers=14 diff --git a/.github/workflows/test-model-map.yml b/.github/workflows/test-model-map.yml deleted file mode 100644 index c2770e5da4c..00000000000 --- a/.github/workflows/test-model-map.yml +++ /dev/null @@ -1,37 +0,0 @@ -name: Validate model_prices_and_context_window.json - -on: - pull_request: - branches: - - main - - litellm_internal_staging - - litellm_oss_staging - - "litellm_**" - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} - cancel-in-progress: ${{ github.event_name == 'pull_request' }} - -jobs: - validate-model-prices-json: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Validate model_prices_and_context_window.json - run: | - jq empty model_prices_and_context_window.json - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Check model_prices_and_context_window.schema.json is in sync - run: | - uv run --frozen python ci_cd/generate_model_prices_schema.py --check diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 9b8b132df62..17b6481a2bf 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,9 +4,24 @@ on: push: paths: - "litellm-rust/**" + - "litellm/rust_bridge/**" + - "tests/test_litellm_rust/**" + - "litellm/integrations/custom_logger.py" + - "litellm/litellm_core_utils/litellm_logging.py" + - "litellm/litellm_core_utils/logging_worker.py" + - "litellm/proxy/guardrails/**" + - "litellm/utils.py" + - "litellm/ocr/**" + - "litellm/llms/base_llm/ocr/**" + - "litellm/llms/custom_httpx/llm_http_handler.py" + - "tests/test_litellm/ocr/**" + - "tests/test_litellm/conftest.py" + - "Makefile" - ".cargo/**" - "pyproject.toml" + - "uv.lock" - "rust-toolchain.toml" + - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" @@ -19,9 +34,24 @@ on: - "litellm_**" paths: - "litellm-rust/**" + - "litellm/rust_bridge/**" + - "tests/test_litellm_rust/**" + - "litellm/integrations/custom_logger.py" + - "litellm/litellm_core_utils/litellm_logging.py" + - "litellm/litellm_core_utils/logging_worker.py" + - "litellm/proxy/guardrails/**" + - "litellm/utils.py" + - "litellm/ocr/**" + - "litellm/llms/base_llm/ocr/**" + - "litellm/llms/custom_httpx/llm_http_handler.py" + - "tests/test_litellm/ocr/**" + - "tests/test_litellm/conftest.py" + - "Makefile" - ".cargo/**" - "pyproject.toml" + - "uv.lock" - "rust-toolchain.toml" + - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" @@ -34,102 +64,92 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true +env: + CARGO_TERM_COLOR: always + jobs: - rust-checks: - name: rustfmt, clippy, test + rust-lint: runs-on: ubuntu-latest timeout-minutes: 10 defaults: run: working-directory: litellm-rust - env: - CARGO_TERM_COLOR: always - steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: persist-credentials: false - - name: Set up Rust - run: rustup toolchain install + - run: rustup toolchain install --no-self-update - - name: Cache Cargo registry and target - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + - run: cargo fmt --check + + - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 with: path: | ~/.cargo/registry ~/.cargo/git litellm-rust/target - key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }} + key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} restore-keys: | - ${{ runner.os }}-cargo- + ${{ runner.os }}-cargo-${{ github.job }}- - - name: Check Rust formatting - run: cargo fmt --check + - run: cargo clippy --workspace --all-targets --locked -- -D warnings - - name: Run Clippy - run: cargo clippy --workspace --all-targets --locked -- -D warnings + - run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings - - name: Run Clippy with Bedrock auth - run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings + - run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings - - name: Run Clippy with all gateway features - run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings - - - name: Run Rust tests - run: cargo test --workspace --locked - - - name: Run core tests with Bedrock auth - run: cargo test -p litellm-core --features bedrock-auth --locked - - # Not --all-features: python-config links libpython, which this job does not install. - - name: Run gateway tests with the server feature - run: cargo test -p litellm-ai-gateway --features server --locked - - release-wheel: - name: release wheel + rust-test: runs-on: ubuntu-latest - timeout-minutes: 20 - permissions: - contents: read - env: - CARGO_TERM_COLOR: always + timeout-minutes: 30 steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: persist-credentials: false - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: "3.12" - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries + - uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" - - name: Set up Rust - run: rustup toolchain install + - run: rustup toolchain install --no-self-update - - name: Build release wheel - run: uv build --wheel --out-dir dist + - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + litellm-rust/target + key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo-${{ github.job }}- - - name: Build panic contract wheel - run: >- + - run: cargo test --workspace --locked + working-directory: litellm-rust + + - run: cargo test -p litellm-core --features bedrock-auth --locked + working-directory: litellm-rust + + - run: cargo test -p litellm-ai-gateway --features server --locked + working-directory: litellm-rust + + - run: uv build --wheel --out-dir dist + + - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl + env: + RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + + - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + + - name: Run pytest tests/test_litellm_rust with the compiled extension + run: make test-rust-extension + + - run: >- uv build --wheel --out-dir panic-dist --config-setting "maturin.build-args=--features panic-test,extension-module" - - name: Smoke-test native panic unwinding - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - - - name: Verify stripped native extension - env: - RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl - - - name: Test native route wheel - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl diff --git a/.github/workflows/test-terraform-modules.yml b/.github/workflows/test-terraform-modules.yml index 52006d9b578..e6896604b7f 100644 --- a/.github/workflows/test-terraform-modules.yml +++ b/.github/workflows/test-terraform-modules.yml @@ -25,13 +25,17 @@ concurrency: cancel-in-progress: true jobs: - aws-module: - name: fmt, validate, test (aws) + module: + name: fmt, validate, test (${{ matrix.module }}) runs-on: ubuntu-latest timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + module: [aws, gcp] defaults: run: - working-directory: terraform/litellm/aws + working-directory: terraform/litellm/${{ matrix.module }} steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: @@ -51,35 +55,7 @@ jobs: - name: validate run: terraform validate - # Plan-only, mock_provider-backed: no AWS credentials, no API calls. + # Plan-only, mock_provider-backed: no cloud credentials, no API calls. - name: test run: terraform test - gcp-module: - name: fmt, validate, test (gcp) - runs-on: ubuntu-latest - timeout-minutes: 15 - defaults: - run: - working-directory: terraform/litellm/gcp - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2 - with: - terraform_version: 1.13.3 - terraform_wrapper: false - - - name: fmt - run: terraform fmt -recursive -check -diff - - - name: init - run: terraform init -backend=false -input=false - - - name: validate - run: terraform validate - - - name: test - run: terraform test diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 33245ec5b5f..f55c87c2ae5 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -116,6 +116,7 @@ jobs: tests/test_litellm/rerank_api tests/test_litellm/rust_bridge tests/test_litellm/sandbox + tests/test_litellm/skills tests/test_litellm/test_router tests/test_litellm/vector_stores tests/test_litellm/videos @@ -199,6 +200,7 @@ jobs: tests/test_litellm/proxy/types_utils tests/test_litellm/proxy/logging_endpoints tests/test_litellm/proxy/test_*.py + tests/test_gateway workers: 4 reruns: 2 timeout-minutes: 20 diff --git a/AGENTS.md b/AGENTS.md index 41921fdff4d..a1e8f6f618d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1 +1,3 @@ Read @CLAUDE.md for coding guidelines + +Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 3d2fa3e51c8..b04e004aa1a 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -149,7 +149,7 @@ graph TD | `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user | | `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation | | `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation | -| `litellm_skills` | `proxy/hooks/skills_injection.py` | Skills injection | +| `litellm_skills` | `proxy/hooks/litellm_skills/main.py` | Skills injection | To add a new proxy hook, implement `CustomLogger` and register in `PROXY_HOOKS`. @@ -220,20 +220,20 @@ graph LR | Job | Interval | Purpose | Key Files | |-----|----------|---------|-----------| | `update_spend` | 60s | Batch write spend logs to PostgreSQL | `proxy/db/db_spend_update_writer.py` | -| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/management_helpers/budget_reset_job.py` | +| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/common_utils/reset_budget_job.py` | | `add_deployment` | 10s | Sync new model deployments from DB | `proxy/proxy_server.py` (`ProxyConfig`) | -| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/management_helpers/spend_log_cleanup.py` | -| `check_batch_cost` | 30min | Calculate costs for batch jobs | `proxy/management_helpers/check_batch_cost_job.py` | -| `check_responses_cost` | 30min | Calculate costs for responses API | `proxy/management_helpers/check_responses_cost_job.py` | -| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/management_helpers/key_rotation_manager.py` | +| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/db/db_transaction_queue/spend_log_cleanup.py` | +| `check_batch_cost` | 30min | Calculate costs for batch jobs | `enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py` | +| `check_responses_cost` | 30min | Calculate costs for responses API | `enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py` | +| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/common_utils/key_rotation_manager.py` | | `_run_background_health_check` | continuous | Health check model deployments | `proxy/proxy_server.py` | | `send_weekly_spend_report` | weekly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) | | `send_monthly_spend_report` | monthly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) | **Cost Attribution Flow:** 1. LLM response returns to `utils.py` wrapper after `litellm.acompletion()` completes -2. `update_response_metadata()` (`llm_response_utils/response_metadata.py`) is called -3. `logging_obj._response_cost_calculator()` (`litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`) +2. `update_response_metadata()` (`litellm_core_utils/llm_response_utils/response_metadata.py`) is called +3. `logging_obj._response_cost_calculator()` (`litellm_core_utils/litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`) 4. Cost is stored in `response._hidden_params["response_cost"]` 5. `proxy/common_request_processing.py` extracts cost from `hidden_params` and adds to response headers (`x-litellm-response-cost`) 6. `logging_obj.async_success_handler()` triggers callbacks including `_ProxyDBLogger.async_log_success_event()` diff --git a/CLAUDE.md b/CLAUDE.md index d9e9e8f1586..41678432989 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -29,7 +29,7 @@ Never test structure of code only function of it End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md` -When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions +When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD` When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule @@ -37,7 +37,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank -Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR +Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis @@ -52,7 +52,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a Python max line length is 120, not 88 -When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing +Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR `make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice @@ -70,7 +70,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests, Always pull before starting any work. The checkout or worktree may be sitting on a stale branch -If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names +If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9ef1d5ae2b8..0443f1bed75 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -315,10 +315,12 @@ Ensure the UI builds successfully before submitting your PR: npm run build ``` +Local lint and budget checks follow origin's current default branch. They refresh it from the remote instead of trusting cached `origin/HEAD`. For an intentional comparison against another branch or commit, use `make check BASE_REF=` or the standalone gate's `--base ` option. An explicit ref can also be used offline once it has been fetched locally. Without an override, unavailable remote metadata stops the check + ## Submitting Your PR 1. **Push your branch**: `git push origin your-feature-branch` -2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`. +2. **Create a PR**: Go to GitHub and open a pull request against the repository's current default branch. Run `python3 scripts/default_branch.py --branch` to check its name 3. **Fill out the PR template**: Provide clear description of changes 4. **Wait for review**: Maintainers will review and provide feedback 5. **Address feedback**: Make requested changes and push updates diff --git a/Dockerfile b/Dockerfile index 0a92aa9a68c..759dac76795 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 +# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) +ARG PGBOUNCER_VERSION=1.25.2 +ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332 FROM $UV_IMAGE AS uvbin +FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder +ARG PGBOUNCER_VERSION +ARG PGBOUNCER_SHA256 +USER root +RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl +WORKDIR /build +RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ + echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ + tar xzf pgbouncer.tar.gz --strip-components=1 && \ + ./configure --prefix=/usr/local --with-openssl=/usr && \ + make -j"$(nproc)" pgbouncer && \ + install -m 0755 pgbouncer /usr/local/bin/pgbouncer + # Admin UI builder. Pinned to the build platform so the architecture-independent # Next.js static export compiles once natively even in a multi-arch build, # instead of once per target arch under QEMU. @@ -110,7 +126,8 @@ USER root RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile +RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent +COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer WORKDIR /app ENV PATH="/app/.venv/bin:${PATH}" \ diff --git a/Makefile b/Makefile index e17fdba3c85..d360074ea4e 100644 --- a/Makefile +++ b/Makefile @@ -4,6 +4,7 @@ .PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \ test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \ test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \ + test-rust-extension \ info lint lint-inner lint-dev lint-checks format \ lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \ lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \ @@ -34,7 +35,7 @@ help: @echo " make lint-basedpyright-budget-update - Ratchet basedpyright limits down by what this branch fixed" @echo " make lint-format - Check ruff format formatting (matches CI)" @echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit" - @echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)" + @echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches the default branch, simulates the merge)" @echo " make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed" @echo " make lint-test-quality - Gate the test suite against test-quality-budget.json" @echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + test quality + basedpyright)" @@ -54,12 +55,16 @@ help: @echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)" @echo " make test-integration - Run integration tests" @echo " make test-unit-helm - Run helm unit tests" + @echo " make test-rust-extension - Build the Rust extension and run its public Python tests" @echo "" @echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide" @echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine." UV := uv UV_RUN := $(UV) run --no-sync +BASE_REF ?= +export BASE_REF +RESOLVE_BASE = python3 scripts/default_branch.py --base "$(BASE_REF)" # Machine-wide slot queue for the heavy targets below; python3 + stdlib only, so # it runs before any venv exists. See scripts/gate_slot_lock.py. @@ -67,7 +72,7 @@ GATE_SLOT_LOCK := python3 scripts/gate_slot_lock.py LINT_DEP_INSTALL ?= install-dev LINT_E2E_DEP_INSTALL ?= lint-install -LINT_DEP_BASE ?= lint-fetch-base +LINT_DEP_BASE ?= LINT_JOBS := $(shell sysctl -n hw.ncpu 2>/dev/null || nproc 2>/dev/null || echo 4) LINT_OUTPUT_SYNC := $(if $(filter output-sync,$(.FEATURES)),--output-sync=target,) @@ -130,10 +135,8 @@ format: install-dev format-check: install-dev cd litellm && $(UV_RUN) ruff format --check --exclude '/enterprise/' . && cd .. -# Single fetch of the PR base so the delta-based gates below share one network round -# trip instead of each re-fetching when chained from `lint`. lint-fetch-base: - git fetch origin litellm_internal_staging + @$(RESOLVE_BASE) # Mirror test-linting.yml's lint job environment: the proxy-dev group plus a generated # Prisma client, so `basedpyright tests/e2e` resolves the same modules CI does. The @@ -147,10 +150,12 @@ lint-install: # Diff-scoped format check, mirroring test-linting.yml's "Check ruff format" step: # only the litellm Python files changed vs the base are checked, so a pre-existing # format issue elsewhere doesn't block an unrelated commit. Git pathspecs match -# recursively, so 'litellm/*.py' covers nested modules and the top-level files that -# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step. +# recursively, so 'litellm/*.py' covers top-level files and nested modules alike, +# the same set CI's ':(glob)litellm/**/*.py' selects. lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - @files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \ + @base_ref=$$($(RESOLVE_BASE)) && \ + changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \ + files=$$(printf '%s\n' "$$changed" | grep -v '^litellm/enterprise/' || true) || exit $$?; \ if [ -z "$$files" ]; then \ echo "No changed litellm Python files to format-check."; \ else \ @@ -167,7 +172,9 @@ lint-ruff: $(LINT_DEP_INSTALL) # https://github.com/astral-sh/ruff/discussions/10977 # https://github.com/astral-sh/ruff/discussions/4049 lint-format-changed: install-dev - @git diff origin/main --unified=0 --no-color -- '*.py' | \ + @base_ref=$$($(RESOLVE_BASE)) && \ + diff=$$(git diff "$$base_ref" --unified=0 --no-color -- '*.py') && \ + printf '%s\n' "$$diff" | \ perl -ne '\ if (/^diff --git a\/(.*) b\//) { $$file = $$1; } \ if (/^@@ .* \+(\d+)(?:,(\d+))? @@/) { \ @@ -182,20 +189,22 @@ lint-format-changed: install-dev done lint-ruff-dev: install-dev - @tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \ + @base_ref=$$($(RESOLVE_BASE)) || exit $$?; \ + tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \ cd litellm && \ ($(UV_RUN) ruff check . --output-format=pylint || true) > "$$tmpfile" && \ - $(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch=origin/main && \ + $(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch="$$base_ref" && \ cd .. ; \ rm -f "$$tmpfile" lint-ruff-FULL-dev: install-dev - @files=$$(git diff --name-only origin/main -- '*.py'); \ + @base_ref=$$($(RESOLVE_BASE)) && \ + files=$$(git diff --name-only "$$base_ref" -- '*.py') || exit $$?; \ if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \ else echo "No changed .py files to check."; fi lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging + $(UV_RUN) python scripts/type_check_gate.py --base "$(BASE_REF)" lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL) $(UV_RUN) basedpyright tests/e2e @@ -203,37 +212,37 @@ lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL) # Type-discipline budget (mutable collections / casts / type guards / kwargs / # unexplained suppressions), the test-linting.yml step `make lint` used to omit. lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - $(UV_RUN) python scripts/type_discipline_gate.py --base origin/litellm_internal_staging + $(UV_RUN) python scripts/type_discipline_gate.py --base "$(BASE_REF)" # Test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes, # litellm module-global mutation, credential-gated skips, conftest snapshot # inventory), counted across tests/ the same delta-vs-base way. lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - $(UV_RUN) python scripts/test_quality_gate.py --base origin/litellm_internal_staging + $(UV_RUN) python scripts/test_quality_gate.py --base "$(BASE_REF)" # --update lowers each limit by what this branch fixed since its branch point, so # it needs the base ref fetched to resolve the merge-base. -lint-basedpyright-budget-update: install-dev lint-fetch-base - $(UV_RUN) python scripts/type_check_gate.py --update +lint-basedpyright-budget-update: install-dev + $(UV_RUN) python scripts/type_check_gate.py --update --base "$(BASE_REF)" lint-format: format-check lint-ruff-budget: install-dev - $(UV_RUN) python scripts/ruff_strict_gate.py + $(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)" # Strict gate, invoked the same way CI does in test-linting.yml so a local pass # means the CI check will pass too. lint-gate: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - $(UV_RUN) python scripts/ruff_strict_gate.py --base origin/litellm_internal_staging + $(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)" -lint-ruff-budget-update: install-dev lint-fetch-base - $(UV_RUN) python scripts/ruff_strict_gate.py --update +lint-ruff-budget-update: install-dev + $(UV_RUN) python scripts/ruff_strict_gate.py --update --base "$(BASE_REF)" -lint-type-discipline-budget-update: install-dev lint-fetch-base - $(UV_RUN) python scripts/type_discipline_gate.py --update +lint-type-discipline-budget-update: install-dev + $(UV_RUN) python scripts/type_discipline_gate.py --update --base "$(BASE_REF)" -lint-test-quality-budget-update: install-dev lint-fetch-base - $(UV_RUN) python scripts/test_quality_gate.py --update +lint-test-quality-budget-update: install-dev + $(UV_RUN) python scripts/test_quality_gate.py --update --base "$(BASE_REF)" # Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright) lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-test-quality-budget-update lint-basedpyright-budget-update @@ -249,14 +258,15 @@ check-import-safety: $(LINT_DEP_INSTALL) # runs the diff-scoped ruff format check, whole-tree ruff check, the strict-rule / # type-discipline / basedpyright budgets as a delta vs the base, then the circular-import # and import-safety checks. Steps that compare against the base resolve it the same way CI -# does (merge-base with origin/litellm_internal_staging). Setup (env sync, Prisma client, +# does (merge-base with origin's current default branch). Setup (env sync, Prisma client, # base fetch) runs once up front; the checks themselves are independent, so a sub-make # fans them out with -j and the fast ones finish under basedpyright's shadow. lint: @$(GATE_SLOT_LOCK) $(MAKE) lint-inner -lint-inner: lint-install lint-fetch-base - $(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks +lint-inner: lint-install + @base_ref=$$($(RESOLVE_BASE)) && \ + $(MAKE) BASE_REF="$$base_ref" -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety @@ -281,6 +291,17 @@ pre-commit: @$(MAKE) check # Testing targets +test-rust-extension: + @temporary=$$(mktemp -d) && \ + trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \ + $(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \ + set -- "$$temporary"/wheels/*.whl && \ + [ "$$#" -eq 1 ] && \ + UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \ + $(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \ + LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \ + "$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust + test: install-test-deps $(UV_RUN) pytest tests/ diff --git a/README.md b/README.md index 92757fcbbc1..901cc5b0cea 100644 --- a/README.md +++ b/README.md @@ -262,6 +262,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ } ``` +For MCP OAuth, an upstream may advertise dynamic client registration but refuse requests with HTTP 401 or 403. If the provider requires a pre-registered OAuth app, configure its `credentials.client_id` and, when required, `credentials.client_secret` on the MCP server. This skips dynamic registration in the gateway sign-in flow. The provider must approve the app for MCP access; reaching its authorization page does not establish that login or tool calls will succeed + [**Docs: MCP Gateway**](https://docs.litellm.ai/docs/mcp) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0b0a61192e6..26e4e06a796 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1808 + "limit": 1804 }, "reportRedeclaration": { "limit": 8 @@ -105,7 +105,7 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38271 + "limit": 38269 }, "reportUnknownParameterType": { "limit": 19584 @@ -135,7 +135,7 @@ "limit": 21 }, "reportUnusedFunction": { - "limit": 138 + "limit": 136 }, "reportUnusedImport": { "limit": 542 diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py new file mode 100644 index 00000000000..50aa40ba220 --- /dev/null +++ b/ci_cd/cost_map_guard.py @@ -0,0 +1,146 @@ +"""Guard the cost map on pull requests. + +Every pull request gets the file checks: the three cost map files parse, the backup copy matches the root file, +and the JSON schema is in sync and validates the map. Pull requests from the cost map sync bot (branches named +litellm_cost_map_sync_*) additionally may only touch those three files and may only add or update models. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Final + +from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors + +COST_MAP_PATH: Final = "model_prices_and_context_window.json" +BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" +SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" +GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH) +BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_" + +CostMap = dict[str, object] + + +@dataclass(frozen=True, slots=True) +class Snapshot: + cost_map: str + backup: str + schema: str + + +def _parse_object(text: str, path: str) -> CostMap | str: + try: + parsed: Final = json.loads(text) + except json.JSONDecodeError as error: + return f"{path} is not valid JSON: {error}" + return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root" + + +def _rendered_schema(cost_map: CostMap) -> str: + try: + return render(build_schema(cost_map)) + except SystemExit as error: + return str(error) + + +def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]: + schema_text: Final = _rendered_schema(head_map) + if not schema_text.startswith("{"): + return (schema_text,) + backup_failure: Final = ( + () + if head.backup == head.cost_map + else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",) + ) + schema_failure: Final = ( + () + if head.schema == schema_text + else ( + f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; " + "run `python ci_cd/generate_model_prices_schema.py` and commit the result", + ) + ) + return ( + *backup_failure, + *schema_failure, + *( + f"{COST_MAP_PATH} does not validate against its schema: {error}" + for error in validation_errors(head_map, json.loads(schema_text))[:20] + ), + ) + + +def _entries(cost_map: CostMap) -> dict[str, dict[str, object]]: + return {key: entry for key, entry in cost_map.items() if isinstance(entry, dict)} + + +def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str]) -> tuple[str, ...]: + base_map: Final = _parse_object(base.cost_map, COST_MAP_PATH) + if isinstance(base_map, str): + return (f"merge base: {base_map}",) + base_entries: Final = _entries(base_map) + head_entries: Final = _entries(head_map) + removed_fields: Final = tuple( + f"{key}.{field}" + for key, entry in base_entries.items() + if key in head_entries + for field in entry + if field not in head_entries[key] + ) + return ( + *( + f"bot PRs may only change the cost map files, not {path}" + for path in changed_files + if path not in GUARDED_PATHS + ), + *(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map), + *(f"bot PRs may not remove fields: {ref}" for ref in removed_fields), + *( + f"bot PRs may not change {key}" + for key in sorted(SPECIAL_ROOT_KEYS) + if base_map.get(key) != head_map.get(key) + ), + ) + + +def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]: + head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH) + if isinstance(head_map, str): + return (head_map,) + return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ())) + + +def _git(*args: str) -> str: + result: Final = subprocess.run(("git", *args), check=False, capture_output=True, text=True) + return result.stdout if result.returncode == 0 else "" + + +def snapshot(revision: str) -> Snapshot: + return Snapshot(*(_git("show", f"{revision}:{path}") for path in GUARDED_PATHS)) + + +def main(argv: Sequence[str]) -> int: + parser: Final = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", required=True, help="merge base of the pull request") + parser.add_argument("--head", required=True, help="head commit of the pull request") + parser.add_argument("--head-ref", required=True, help="head branch name of the pull request") + args: Final = parser.parse_args(argv) + bot: Final = args.head_ref.startswith(BOT_BRANCH_PREFIX) + changed_files: Final = tuple(_git("diff", "--name-only", args.base, args.head).splitlines()) + failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed_files, bot) + contract: Final = "bot contract enforced" if bot else "human PR, file checks only" + if failures: + print(f"cost map guard failed ({contract}):") + print("\n".join(f"- {failure}" for failure in failures)) + return 1 + print(f"cost map guard passed ({contract})") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 57cc742d5c4..ab29b70bdd4 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -58,6 +58,11 @@ OBJECT_KEYS: dict[str, JsonSchema] = { } ARRAY_KEYS: dict[str, JsonSchema] = { + "supported_audio_formats": { + "type": "array", + "description": "Audio container formats the model can return.", + "items": {"type": "string", "enum": ["mp3", "wav"]}, + }, "supported_endpoints": { "type": "array", "description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.", @@ -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"], + }, } diff --git a/ci_cd/run_migration.py b/ci_cd/run_migration.py index feec4046ee1..c737050f3a8 100644 --- a/ci_cd/run_migration.py +++ b/ci_cd/run_migration.py @@ -6,12 +6,9 @@ import subprocess import sys from datetime import datetime from pathlib import Path - -import testing.postgresql - +from typing import Final DESTRUCTIVE_PATTERN = re.compile(r"\bDROP\s+(COLUMN|TABLE|INDEX)\b", re.IGNORECASE) -DEFAULT_BASE_BRANCH = "litellm_internal_staging" def _find_destructive_statements(sql: str) -> list: @@ -94,31 +91,57 @@ def _print_stale_branch_refusal(base_branch: str, behind: int) -> None: print(banner, file=out) -def _check_branch_freshness(root_dir: Path, base_branch: str) -> None: +def _default_base_branch(root_dir: Path) -> str: + try: + result: Final = subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve().parents[1] / "scripts" / "default_branch.py"), + "--repo-root", + str(root_dir), + "--branch", + ], + check=True, + capture_output=True, + text=True, + timeout=90, + ) + except (OSError, subprocess.SubprocessError) as exc: + _print_freshness_failure( + "default branch", + "Could not discover origin's default branch. Pass --base-branch to choose one.", + exc.stderr if isinstance(exc, subprocess.CalledProcessError) else str(exc), + ) + sys.exit(3) + return result.stdout.strip() + + +def _check_branch_freshness(root_dir: Path, base_branch: str | None = None) -> None: """Fetch origin/ and exit 3 if HEAD is behind it.""" + resolved_branch: Final = base_branch or _default_base_branch(root_dir) cwd = str(root_dir) try: subprocess.run( - ["git", "fetch", "origin", base_branch], + ["git", "fetch", "origin", f"+refs/heads/{resolved_branch}:refs/remotes/origin/{resolved_branch}"], check=True, capture_output=True, text=True, cwd=cwd, ) except FileNotFoundError: - _print_freshness_failure(base_branch, "git executable not found on PATH") + _print_freshness_failure(resolved_branch, "git executable not found on PATH") sys.exit(3) except subprocess.CalledProcessError as e: _print_freshness_failure( - base_branch, - f"`git fetch origin {base_branch}` failed", + resolved_branch, + f"`git fetch origin {resolved_branch}` failed", e.stderr or "", ) sys.exit(3) try: result = subprocess.run( - ["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"], + ["git", "rev-list", "--count", f"HEAD..origin/{resolved_branch}"], check=True, capture_output=True, text=True, @@ -127,23 +150,23 @@ def _check_branch_freshness(root_dir: Path, base_branch: str) -> None: behind = int(result.stdout.strip()) except subprocess.CalledProcessError as e: _print_freshness_failure( - base_branch, - f"`git rev-list HEAD..origin/{base_branch}` failed", + resolved_branch, + f"`git rev-list HEAD..origin/{resolved_branch}` failed", e.stderr or "", ) sys.exit(3) except ValueError: _print_freshness_failure( - base_branch, + resolved_branch, "could not parse commit count from `git rev-list`", ) sys.exit(3) if behind > 0: - _print_stale_branch_refusal(base_branch, behind) + _print_stale_branch_refusal(resolved_branch, behind) sys.exit(3) - print(f"Branch freshness OK: up to date with origin/{base_branch}.") + print(f"Branch freshness OK: up to date with origin/{resolved_branch}.") def _print_destructive_refusal(destructive_lines: list) -> None: @@ -198,7 +221,7 @@ def _print_destructive_refusal(destructive_lines: list) -> None: def create_migration( migration_name: str = None, allow_destructive: bool = False, - base_branch: str = DEFAULT_BASE_BRANCH, + base_branch: str | None = None, skip_freshness_check: bool = False, ): """ @@ -211,7 +234,7 @@ def create_migration( DROP COLUMN, DROP TABLE, or DROP INDEX statements. Without this flag, the script exits non-zero and prints guidance. base_branch (str): Branch to check freshness against - (default: "litellm_internal_staging"). + (default: origin's current default branch). skip_freshness_check (bool): Skip the "branch is up to date" check. Only for intentional migrations against an older base. """ @@ -225,6 +248,8 @@ def create_migration( else: _check_branch_freshness(root_dir, base_branch) + import testing.postgresql + try: migrations_dir = ( root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" @@ -342,9 +367,8 @@ if __name__ == "__main__": ) parser.add_argument( "--base-branch", - default=DEFAULT_BASE_BRANCH, help=( - f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). " + "Branch to check freshness against (default: origin's current default branch). " "The script fetches origin/ and refuses to run if HEAD " "is behind it." ), diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index e9ad2849bb2..b0bf935c616 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 +# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) +ARG PGBOUNCER_VERSION=1.25.2 +ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332 FROM $UV_IMAGE AS uvbin +FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder +ARG PGBOUNCER_VERSION +ARG PGBOUNCER_SHA256 +USER root +RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl +WORKDIR /build +RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ + echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ + tar xzf pgbouncer.tar.gz --strip-components=1 && \ + ./configure --prefix=/usr/local --with-openssl=/usr && \ + make -j"$(nproc)" pgbouncer && \ + install -m 0755 pgbouncer /usr/local/bin/pgbouncer + # Admin UI builder. Pinned to the build platform so the architecture-independent # Next.js static export compiles once natively even in a multi-arch build, # instead of once per target arch under QEMU. @@ -101,7 +117,8 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile +RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent +COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer WORKDIR /app ENV PATH="/app/.venv/bin:${PATH}" \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index edf20e8bbff..5d729046678 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -7,9 +7,25 @@ ARG PROXY_EXTRAS_SOURCE=published ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 +# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) +ARG PGBOUNCER_VERSION=1.25.2 +ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332 FROM $UV_IMAGE AS uvbin +FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder +ARG PGBOUNCER_VERSION +ARG PGBOUNCER_SHA256 +USER root +RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl +WORKDIR /build +RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ + echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ + tar xzf pgbouncer.tar.gz --strip-components=1 && \ + ./configure --prefix=/usr/local --with-openssl=/usr && \ + make -j"$(nproc)" pgbouncer && \ + install -m 0755 pgbouncer /usr/local/bin/pgbouncer + # Admin UI builder. Pinned to the build platform so the architecture-independent # Next.js static export compiles once natively even in a multi-arch build, # instead of once per target arch under QEMU. @@ -128,8 +144,9 @@ RUN for i in 1 2 3; do \ apk upgrade --no-cache && break || sleep 5; \ done && \ for i in 1 2 3; do \ - apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs && break || sleep 5; \ + apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \ done +COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer # Copy only what runtime needs. The application is installed inside the venv; # the rest of the builder's /app is source and build metadata that must not diff --git a/docker/component_entrypoint.sh b/docker/component_entrypoint.sh index 413957b9929..173afafe1ad 100755 --- a/docker/component_entrypoint.sh +++ b/docker/component_entrypoint.sh @@ -1,5 +1,11 @@ #!/bin/sh +# stale samples from a previous container incarnation would be summed into the aggregate +if [ -n "$PROMETHEUS_MULTIPROC_DIR" ]; then + mkdir -p "$PROMETHEUS_MULTIPROC_DIR" + rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db +fi + case "$USE_DDTRACE" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED="False" diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index e6f00877a26..13e9e5093a8 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -142,6 +142,40 @@ class CheckBatchCost: verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}") return None + async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None: + org_id = getattr(job, "org_id", None) + if org_id: + return org_id + api_key = getattr(job, "api_key", None) + team_id = getattr(job, "team_id", None) + if api_key: + try: + key_row: prisma_models.LiteLLM_VerificationToken | None = ( + await self.prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": api_key} + ) + ) + key_org_id = getattr(key_row, "organization_id", None) if key_row is not None else None + if key_org_id: + return key_org_id + except Exception as e: + verbose_proxy_logger.error( + f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, " + f"still trying the team's: {e}" + ) + if not team_id: + return None + try: + team_row: prisma_models.LiteLLM_TeamTable | None = ( + await self.prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + ) + return getattr(team_row, "organization_id", None) if team_row is not None else None + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}") + return None + async def _build_creator_attribution_metadata( self, job: "LiteLLM_ManagedObjectTable", batch_id: str ) -> dict[str, object]: @@ -153,6 +187,10 @@ class CheckBatchCost: user_api_key_alias; when it has no alias, or the key has since been rotated or deleted, the field keeps the creating user's alias that _get_user_info filled in, because a resolvable name is more useful on the spend row than a null. + + user_api_key_org_id must be resolved here too: the spend update writer reads it + off this metadata to increment organization spend, so leaving it out silently + drops batch cost from org accounting for keys and teams that belong to one. """ api_key = getattr(job, "api_key", None) team_id = getattr(job, "team_id", None) @@ -172,6 +210,9 @@ class CheckBatchCost: team_alias = await self._get_team_alias(team_id) if team_alias is not None: metadata["user_api_key_team_alias"] = team_alias + org_id: Final = await self._get_org_id(job, batch_id) + if org_id is not None: + metadata["user_api_key_org_id"] = org_id if isinstance(request_tags, list) and request_tags: metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)] @@ -641,7 +682,7 @@ class CheckBatchCost: from litellm.files.main import afile_content from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info + from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, ) @@ -805,6 +846,7 @@ class CheckBatchCost: function_id=str(uuid.uuid4()), ) + deployment_api_base: Final = deployment_info.litellm_params.api_base logging_obj.update_environment_variables( litellm_params={ # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks @@ -813,9 +855,17 @@ class CheckBatchCost: "user-agent": CHECK_BATCH_COST_USER_AGENT, } }, - "metadata": await self._build_creator_attribution_metadata(job, batch_id), + **({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}), + "metadata": { + **(await self._build_creator_attribution_metadata(job, batch_id)), + # spend logs read the deployment identity off these metadata keys, so + # without them the batch cost row carries no model_id or model_group + "model_info": {"id": model_id}, + "model_group": deployment_info.model_name, + }, }, optional_params={}, + custom_llm_provider=str(llm_provider) if llm_provider else None, ) if not await self._claim_job_for_costing(job): @@ -833,6 +883,8 @@ class CheckBatchCost: batch_models=batch_result.models, batch_successful_requests=batch_result.successful_requests, batch_failed_requests=batch_result.failed_requests, + batch_prompt_cost=batch_result.prompt_cost, + batch_completion_cost=batch_result.completion_cost, ) except Exception: await self._release_job_claim(job) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index bc1eb6cebc2..c7e1b94a2ef 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -280,6 +280,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}") + async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]: + if user_api_key_dict.org_id: + return user_api_key_dict.org_id + if not user_api_key_dict.team_id: + return None + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + try: + team: Final = await get_team_object( + team_id=user_api_key_dict.team_id, + prisma_client=self.prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return team.organization_id + except Exception as e: + verbose_logger.warning(f"could not resolve org for managed object attribution: {e}") + return None + async def store_unified_object_id( self, unified_object_id: str, @@ -352,6 +373,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "file_purpose": file_purpose, "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, + "org_id": await self._resolve_creator_org_id(user_api_key_dict), "updated_by": user_api_key_dict.user_id, "status": file_object.status, **attribution_columns, @@ -1779,7 +1801,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Remove conflicting keys from data to avoid duplicate keyword arguments filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")} for model_id, model_file_id in specific_model_file_id_mapping.items(): - delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) + delete_data = { + **{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"}, + **( + {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))} + if credentials is not None + else {} + ), + } + delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data) stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span) @@ -1790,7 +1821,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): prom_logger.record_managed_file_deleted(result="success") if stored_file_object: - return stored_file_object + return OpenAIFileObject.model_validate(stored_file_object).model_copy(update={"id": file_id}) elif delete_response: delete_response.id = file_id return delete_response diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 3699087dbfa..903c5155a12 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.65" +version = "0.1.66" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.65" +version = "0.1.66" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 308d70a6b26..33d3791dbba 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,9 +1,25 @@ ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a +# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) +ARG PGBOUNCER_VERSION=1.25.2 +ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332 FROM $UV_IMAGE AS uvbin +FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder +ARG PGBOUNCER_VERSION +ARG PGBOUNCER_SHA256 +USER root +RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl +WORKDIR /build +RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ + echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ + tar xzf pgbouncer.tar.gz --strip-components=1 && \ + ./configure --prefix=/usr/local --with-openssl=/usr && \ + make -j"$(nproc)" pgbouncer && \ + install -m 0755 pgbouncer /usr/local/bin/pgbouncer + # ---------- Builder ---------- FROM $LITELLM_BUILD_IMAGE AS builder @@ -61,6 +77,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra bedrock-realtime \ --python python3.13 +# PYTHONPATH=/app makes the source tree shadow the installed package, so the +# compiled Rust extension must live next to the source or it is never imported. +RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/ + RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ prisma generate --schema=./schema.prisma @@ -73,7 +93,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \ + apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -90,15 +110,17 @@ ENV HOME=/home/nonroot \ COPY --from=builder --chown=nonroot:nonroot /app /app COPY --from=builder /opt/prisma /opt/prisma +COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ find /app/.venv -type d -path "*/tornado/test" -delete && \ chmod -R a+rX /opt/prisma && \ - python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" + python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \ + python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()" USER nonroot EXPOSE 4000/tcp -ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"] +ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh python -m gateway.launch --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"] CMD ["--host", "0.0.0.0", "--port", "4000"] diff --git a/gateway/launch.py b/gateway/launch.py new file mode 100644 index 00000000000..d67432caee1 --- /dev/null +++ b/gateway/launch.py @@ -0,0 +1,77 @@ +"""Gateway supervisor: assemble DATABASE_URL, start the in-container PgBouncer, then run uvicorn. + +``gateway/main.py`` assembles ``DATABASE_URL`` inside every uvicorn worker, which +is fine for a plain Postgres URL but not for the pooler: PgBouncer must be +started exactly once per pod, before the workers fork, and the workers must be +handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in +``DatabaseURLSettings.apply_to_env`` under password auth, and one marked pooled +wins under token auth too, so exporting it here is enough for every worker to +pick the pooled URL up unchanged. + +Run with: + python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000 +""" + +import os +import sys +from collections.abc import Callable, Mapping, Sequence +from typing import Final + +from uvicorn.main import main as uvicorn_main + +from litellm.proxy.db.db_url_settings import DatabaseURLSettings +from litellm.proxy.db.pgbouncer import ( + PgBouncerError, + PgBouncerSettings, + export_pooled_database_url, + start_in_container_pgbouncer, +) + +GATEWAY_APP: Final = "gateway.main:app" +KEEPALIVE_FLAG: Final = "--timeout-keep-alive" + + +def uvicorn_argv(argv: Sequence[str], environ: Mapping[str, str]) -> tuple[str, ...]: + """Honor ``KEEPALIVE_TIMEOUT`` like ``proxy_cli.py`` does, unless the flag was passed explicitly.""" + keepalive: Final = environ.get("KEEPALIVE_TIMEOUT") + if keepalive is None or any(arg == KEEPALIVE_FLAG or arg.startswith(f"{KEEPALIVE_FLAG}=") for arg in argv): + return (GATEWAY_APP, *argv) + return (GATEWAY_APP, *argv, KEEPALIVE_FLAG, keepalive) + + +def pool_database_url( + settings: DatabaseURLSettings, + pgbouncer: PgBouncerSettings, + environ: Mapping[str, str], +) -> str | PgBouncerError | None: + """Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off. + + The upstream URL is whatever ``apply_to_env`` assembled from the discrete + ``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Under token + auth the pooler mints and renews the upstream token itself. + """ + if not pgbouncer.enabled: + return None + upstream_url: Final = environ.get("DATABASE_URL") + if upstream_url is None: + return PgBouncerError("LITELLM_PGBOUNCER_ENABLED is set but no DATABASE_URL could be assembled") + return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth()) + + +def _serve(argv: Sequence[str]) -> None: + uvicorn_main(tuple(argv), prog_name="uvicorn") + + +def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: + settings: Final = DatabaseURLSettings.from_env() + settings.apply_to_env() + pooled_url: Final = pool_database_url(settings, PgBouncerSettings(), os.environ) + if isinstance(pooled_url, PgBouncerError): + sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}") + if pooled_url is not None: + export_pooled_database_url(pooled_url) + serve(uvicorn_argv(argv, os.environ)) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/helm/litellm-helm/templates/_helpers.tpl b/helm/litellm-helm/templates/_helpers.tpl index 8f2acb20fce..9630633912e 100644 --- a/helm/litellm-helm/templates/_helpers.tpl +++ b/helm/litellm-helm/templates/_helpers.tpl @@ -161,3 +161,163 @@ taken before the change, which by that point no longer exists. {{- fail (printf "postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got %q). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore." $tag) -}} {{- end -}} {{- end -}} + +{{/* +Environment shared by the proxy container and the opt-in collector sidecar: +database, pgbouncer, master key, redis, user envVars. Both containers must see +the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer +and the same spend transaction buffer. +*/}} +{{- define "litellm.proxyEnv" -}} +- name: HOST + value: "{{ .Values.listen | default "0.0.0.0" }}" +- name: PORT + value: {{ .Values.service.port | quote}} +{{- if .Values.db.deployStandalone }} +- name: DATABASE_USERNAME + valueFrom: + secretKeyRef: + name: {{ include "litellm.fullname" . }}-dbcredentials + key: username +- name: DATABASE_PASSWORD + valueFrom: + secretKeyRef: + name: {{ include "litellm.fullname" . }}-dbcredentials + key: password +- name: DATABASE_HOST + value: {{ .Release.Name }}-postgresql +- name: DATABASE_NAME + value: litellm +{{- else if .Values.db.useExisting }} +- name: DATABASE_USERNAME + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.usernameKey }} +- name: DATABASE_PASSWORD + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.passwordKey }} +- name: DATABASE_HOST + {{- if .Values.db.secret.endpointKey }} + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.endpointKey }} + {{- else }} + value: {{ .Values.db.endpoint }} + {{- end }} +- name: DATABASE_NAME + value: {{ .Values.db.database }} +- name: DATABASE_URL + value: {{ .Values.db.url | quote }} +{{- end }} +{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }} +- name: DATABASE_READER_HOST + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.readReplicaEndpointKey }} +{{- end }} +{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }} +- name: DATABASE_URL_READ_REPLICA + valueFrom: + secretKeyRef: + name: {{ .Values.db.secret.name }} + key: {{ .Values.db.secret.readReplicaUrlKey }} +{{- else if .Values.db.readReplicaUrl }} +- name: DATABASE_URL_READ_REPLICA + value: {{ .Values.db.readReplicaUrl | quote }} +{{- end }} +{{- if .Values.db.connectionPool.enabled }} +- name: LITELLM_PGBOUNCER_ENABLED + value: "true" +- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: {{ .Values.db.connectionPool.maxDbConnections | quote }} +- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: {{ .Values.db.connectionPool.maxClientConn | quote }} +{{- end }} +- name: PROXY_MASTER_KEY + valueFrom: + secretKeyRef: + name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }} + key: {{ .Values.masterkeySecretKey | default "masterkey" }} +{{- if .Values.redis.enabled }} +- name: REDIS_HOST + value: {{ include "litellm.redis.serviceName" . }} +- name: REDIS_PORT + value: {{ include "litellm.redis.port" . | quote }} +- name: REDIS_PASSWORD + valueFrom: + secretKeyRef: + name: {{ include "redis.secretName" .Subcharts.redis }} + key: {{include "redis.secretPasswordKey" .Subcharts.redis }} +{{- end }} +{{- /* + Inject LITELLM_LOG only when envVars does not already define it. +*/}} +{{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }} +- name: LITELLM_LOG + value: {{ .Values.logLevel | quote }} +{{- end }} +{{- if .Values.envVars }} +{{- range $key, $val := .Values.envVars }} +- name: {{ $key }} + value: {{ $val | quote }} +{{- end }} +{{- end }} +{{- with .Values.extraEnvVars }} +{{ toYaml . }} +{{- end }} +{{- if .Values.migrationJob.enabled }} +# Schema updates are owned by the dedicated migrations Job; skip +# the proxy's startup `prisma db push` so N replicas don't race +# one DB on every rollout. Placed last (after envVars and +# extraEnvVars) so this override can't be silently shadowed by a +# user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env +# semantics — same pattern the migrations Job uses. +- name: DISABLE_SCHEMA_UPDATE + value: "true" +{{- end }} +{{- end -}} + +{{/* +Proxy-only metering and metrics env. The collector sidecar serves no HTTP +traffic, so it gets neither. +*/}} +{{- define "litellm.proxyMetricsEnv" -}} +{{- if .Values.billingMetrics.enabled }} +{{ include "litellm.billingMetricsEnv" . }} +{{- end }} +{{- if .Values.metricsServer.enabled }} +{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }} +{{- fail "metricsServer.port must differ from service.port" }} +{{- end }} +- name: PROMETHEUS_METRICS_PORT + value: {{ .Values.metricsServer.port | quote }} +{{- end }} +{{- end -}} + +{{/* +Directory of the collector's unix socket, shared between the two containers +through an emptyDir. Empty when the sidecar is off or uses 127.0.0.1 TCP. +*/}} +{{- define "litellm.collector.socketDir" -}} +{{- if and .Values.collector.enabled (hasPrefix "unix://" .Values.collector.address) -}} +{{- dir (trimPrefix "unix://" .Values.collector.address) -}} +{{- end -}} +{{- end -}} + +{{- define "litellm.collectorEnv" -}} +- name: LITELLM_COLLECTOR_ENABLED + value: "true" +- name: LITELLM_COLLECTOR_ADDRESS + value: {{ .Values.collector.address | quote }} +- name: LITELLM_COLLECTOR_BUFFER_SIZE + value: {{ .Values.collector.bufferSize | quote }} +- name: LITELLM_COLLECTOR_ON_UNAVAILABLE + value: {{ .Values.collector.onUnavailable | quote }} +- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS + value: {{ .Values.collector.drainTimeoutSeconds | quote }} +{{- end -}} diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index 52ffd117535..cf7b3f8a38d 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -56,111 +56,10 @@ spec: image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" imagePullPolicy: {{ .Values.image.pullPolicy }} env: - - name: HOST - value: "{{ .Values.listen | default "0.0.0.0" }}" - - name: PORT - value: {{ .Values.service.port | quote}} - {{- if .Values.db.deployStandalone }} - - name: DATABASE_USERNAME - valueFrom: - secretKeyRef: - name: {{ include "litellm.fullname" . }}-dbcredentials - key: username - - name: DATABASE_PASSWORD - valueFrom: - secretKeyRef: - name: {{ include "litellm.fullname" . }}-dbcredentials - key: password - - name: DATABASE_HOST - value: {{ .Release.Name }}-postgresql - - name: DATABASE_NAME - value: litellm - {{- else if .Values.db.useExisting }} - - name: DATABASE_USERNAME - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.usernameKey }} - - name: DATABASE_PASSWORD - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.passwordKey }} - - name: DATABASE_HOST - {{- if .Values.db.secret.endpointKey }} - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.endpointKey }} - {{- else }} - value: {{ .Values.db.endpoint }} - {{- end }} - - name: DATABASE_NAME - value: {{ .Values.db.database }} - - name: DATABASE_URL - value: {{ .Values.db.url | quote }} - {{- end }} - {{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }} - - name: DATABASE_READER_HOST - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.readReplicaEndpointKey }} - {{- end }} - {{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }} - - name: DATABASE_URL_READ_REPLICA - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.readReplicaUrlKey }} - {{- else if .Values.db.readReplicaUrl }} - - name: DATABASE_URL_READ_REPLICA - value: {{ .Values.db.readReplicaUrl | quote }} - {{- end }} - - name: PROXY_MASTER_KEY - valueFrom: - secretKeyRef: - name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }} - key: {{ .Values.masterkeySecretKey | default "masterkey" }} - {{- if .Values.redis.enabled }} - - name: REDIS_HOST - value: {{ include "litellm.redis.serviceName" . }} - - name: REDIS_PORT - value: {{ include "litellm.redis.port" . | quote }} - - name: REDIS_PASSWORD - valueFrom: - secretKeyRef: - name: {{ include "redis.secretName" .Subcharts.redis }} - key: {{include "redis.secretPasswordKey" .Subcharts.redis }} - {{- end }} - {{- /* - Inject LITELLM_LOG only when envVars does not already define it. - */}} - {{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }} - - name: LITELLM_LOG - value: {{ .Values.logLevel | quote }} - {{- end }} - {{- if .Values.envVars }} - {{- range $key, $val := .Values.envVars }} - - name: {{ $key }} - value: {{ $val | quote }} - {{- end }} - {{- end }} - {{- with .Values.extraEnvVars }} - {{- toYaml . | nindent 12 }} - {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsEnv" . | nindent 12 }} - {{- end }} - {{- if .Values.migrationJob.enabled }} - # Schema updates are owned by the dedicated migrations Job; skip - # the proxy's startup `prisma db push` so N replicas don't race - # one DB on every rollout. Placed last (after envVars and - # extraEnvVars) so this override can't be silently shadowed by a - # user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env - # semantics — same pattern the migrations Job uses. - - name: DISABLE_SCHEMA_UPDATE - value: "true" + {{- include "litellm.proxyEnv" . | nindent 12 }} + {{- include "litellm.proxyMetricsEnv" . | nindent 12 }} + {{- if .Values.collector.enabled }} + {{- include "litellm.collectorEnv" . | nindent 12 }} {{- end }} envFrom: {{- range .Values.environmentSecrets }} @@ -189,6 +88,11 @@ spec: - name: http containerPort: {{ .Values.service.port }} protocol: TCP + {{- if .Values.metricsServer.enabled }} + - name: metrics + containerPort: {{ .Values.metricsServer.port }} + protocol: TCP + {{- end }} livenessProbe: httpGet: path: {{ .Values.livenessProbe.path | quote }} @@ -233,6 +137,10 @@ spec: {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} {{- end }} + {{- if include "litellm.collector.socketDir" . }} + - name: collector-socket + mountPath: {{ include "litellm.collector.socketDir" . }} + {{- end }} {{- with .Values.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -240,6 +148,53 @@ spec: lifecycle: {{- toYaml . | nindent 12 }} {{- end }} + {{- if .Values.collector.enabled }} + - name: {{ include "litellm.name" . }}-collector + securityContext: + {{- toYaml .Values.securityContext | nindent 12 }} + image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + command: {{ toYaml .Values.collector.command | nindent 12 }} + env: + {{- include "litellm.proxyEnv" . | nindent 12 }} + {{- include "litellm.collectorEnv" . | nindent 12 }} + - name: LITELLM_JOB_ROLE + value: collector + {{- if not (hasKey (default dict .Values.envVars) "CONFIG_FILE_PATH") }} + - name: CONFIG_FILE_PATH + value: /etc/litellm/config.yaml + {{- end }} + envFrom: + {{- range .Values.environmentSecrets }} + - secretRef: + name: {{ . }} + {{- end }} + {{- range .Values.environmentConfigMaps }} + - configMapRef: + name: {{ . }} + {{- end }} + resources: + {{- toYaml .Values.collector.resources | nindent 12 }} + volumeMounts: + - name: litellm-config + mountPath: /etc/litellm/config.yaml + subPath: config.yaml + {{- if include "litellm.collector.socketDir" . }} + - name: collector-socket + mountPath: {{ include "litellm.collector.socketDir" . }} + {{- end }} + {{ if .Values.securityContext.readOnlyRootFilesystem }} + - name: tmp + mountPath: /tmp + - name: cache + mountPath: /.cache + - name: npm + mountPath: /.npm + {{- end }} + {{- with .Values.volumeMounts }} + {{- toYaml . | nindent 12 }} + {{- end }} + {{- end }} {{- with .Values.extraContainers }} {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} @@ -268,6 +223,11 @@ spec: {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} {{- end }} + {{- if include "litellm.collector.socketDir" . }} + - name: collector-socket + emptyDir: + sizeLimit: 1Mi + {{- end }} {{- with .Values.volumes }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/helm/litellm-helm/templates/hpa.yaml b/helm/litellm-helm/templates/hpa.yaml index fec4d1f5c5e..a651f916d21 100644 --- a/helm/litellm-helm/templates/hpa.yaml +++ b/helm/litellm-helm/templates/hpa.yaml @@ -18,6 +18,15 @@ spec: {{- end }} metrics: {{- if .Values.autoscaling.targetCPUUtilizationPercentage }} + {{- if and .Values.collector.enabled .Values.collector.scaleOnProxyContainerCpu }} + - type: ContainerResource + containerResource: + name: cpu + container: {{ include "litellm.name" . }} + target: + type: Utilization + averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }} + {{- else }} - type: Resource resource: name: cpu @@ -25,6 +34,7 @@ spec: type: Utilization averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }} {{- end }} + {{- end }} {{- if .Values.autoscaling.targetMemoryUtilizationPercentage }} - type: Resource resource: @@ -33,4 +43,22 @@ spec: type: Utilization averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }} {{- end }} + {{- with .Values.autoscaling.targetRequestsPerSecond }} + - type: Pods + pods: + metric: + name: litellm_requests_per_second + target: + type: AverageValue + averageValue: {{ toJson . | trimAll "\"" | quote }} + {{- end }} + {{- with .Values.autoscaling.targetTokensPerSecond }} + - type: Pods + pods: + metric: + name: litellm_tokens_per_second + target: + type: AverageValue + averageValue: {{ toJson . | trimAll "\"" | quote }} + {{- end }} {{- end }} diff --git a/helm/litellm-helm/templates/keda.yaml b/helm/litellm-helm/templates/keda.yaml index fe5190fffc6..bf585d0d4be 100644 --- a/helm/litellm-helm/templates/keda.yaml +++ b/helm/litellm-helm/templates/keda.yaml @@ -23,6 +23,27 @@ spec: triggers: {{- with .Values.keda.triggers }} {{- toYaml . | nindent 2 }} +{{- end }} +{{- $prom := .Values.keda.prometheus }} +{{- if or $prom.requestsPerSecond $prom.tokensPerSecond }} +{{- if not $prom.serverAddress }} +{{- fail "keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set" }} +{{- end }} +{{- $selector := printf "namespace=%q,job=%q" .Release.Namespace (printf "%s%s" (include "litellm.fullname" .) (ternary "-metrics" "" .Values.metricsServer.enabled)) }} +{{- with $prom.requestsPerSecond }} + - type: prometheus + metadata: + serverAddress: {{ $prom.serverAddress | quote }} + threshold: {{ toJson . | trimAll "\"" | quote }} + query: {{ printf "sum(rate(litellm_proxy_total_requests_metric_total{%s}[1m]))" $selector | quote }} +{{- end }} +{{- with $prom.tokensPerSecond }} + - type: prometheus + metadata: + serverAddress: {{ $prom.serverAddress | quote }} + threshold: {{ toJson . | trimAll "\"" | quote }} + query: {{ printf "sum(rate(litellm_total_tokens_metric_total{%s}[1m]))" $selector | quote }} +{{- end }} {{- end }} advanced: restoreToOriginalReplicaCount: {{ .Values.keda.restoreToOriginalReplicaCount }} diff --git a/helm/litellm-helm/templates/service-metrics.yaml b/helm/litellm-helm/templates/service-metrics.yaml new file mode 100644 index 00000000000..1d23fe39606 --- /dev/null +++ b/helm/litellm-helm/templates/service-metrics.yaml @@ -0,0 +1,17 @@ +{{- if .Values.metricsServer.enabled }} +apiVersion: v1 +kind: Service +metadata: + name: {{ include "litellm.fullname" . }}-metrics + labels: + {{- include "litellm.labels" . | nindent 4 }} +spec: + type: ClusterIP + ports: + - port: {{ .Values.metricsServer.port }} + targetPort: metrics + protocol: TCP + name: metrics + selector: + {{- include "litellm.selectorLabels" . | nindent 4 }} +{{- end }} diff --git a/helm/litellm-helm/templates/servicemonitor.yaml b/helm/litellm-helm/templates/servicemonitor.yaml index 743098deb3f..68083d0da61 100644 --- a/helm/litellm-helm/templates/servicemonitor.yaml +++ b/helm/litellm-helm/templates/servicemonitor.yaml @@ -26,7 +26,7 @@ spec: {{- toYaml .namespaceSelector.matchNames | nindent 4 }} {{- end }} endpoints: - - port: http + - port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }} path: /metrics/ interval: {{ .interval }} scrapeTimeout: {{ .scrapeTimeout }} diff --git a/helm/litellm-helm/tests/collector_tests.yaml b/helm/litellm-helm/tests/collector_tests.yaml new file mode 100644 index 00000000000..0340b1161b7 --- /dev/null +++ b/helm/litellm-helm/tests/collector_tests.yaml @@ -0,0 +1,272 @@ +suite: test collector sidecar +templates: + - deployment.yaml + - hpa.yaml + - configmap-litellm.yaml +tests: + - it: should run the proxy alone with no collector env by default + template: deployment.yaml + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 1 + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ENABLED + value: "true" + - notContains: + path: spec.template.spec.volumes + content: + name: collector-socket + any: true + + - it: should add the sidecar on the same image and point both containers at the unix socket + template: deployment.yaml + set: + image.tag: test + db.connectionPool.enabled: true + collector.enabled: true + collector.resources: + requests: + cpu: 500m + memory: 1Gi + limits: + cpu: "1" + memory: 2Gi + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 2 + - equal: + path: spec.template.spec.containers[1].name + value: litellm-collector + - equal: + path: spec.template.spec.containers[1].image + value: ghcr.io/berriai/litellm:test + - equal: + path: spec.template.spec.containers[1].command + value: [python, -m, litellm.proxy.collector] + - equal: + path: spec.template.spec.containers[1].resources.requests.cpu + value: 500m + - equal: + path: spec.template.spec.containers[1].resources.limits.memory + value: 2Gi + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ADDRESS + value: unix:///var/run/litellm/collector.sock + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_BUFFER_SIZE + value: "1000" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ON_UNAVAILABLE + value: fallback + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_JOB_ROLE + value: collector + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_JOB_ROLE + value: collector + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_COLLECTOR_ADDRESS + value: unix:///var/run/litellm/collector.sock + - contains: + path: spec.template.spec.containers[1].env + content: + name: CONFIG_FILE_PATH + value: /etc/litellm/config.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: DATABASE_HOST + value: RELEASE-NAME-postgresql + - contains: + path: spec.template.spec.containers[1].env + content: + name: DATABASE_PASSWORD + valueFrom: + secretKeyRef: + name: RELEASE-NAME-litellm-dbcredentials + key: password + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "20" + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: collector-socket + mountPath: /var/run/litellm + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: collector-socket + mountPath: /var/run/litellm + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: litellm-config + mountPath: /etc/litellm/config.yaml + subPath: config.yaml + - contains: + path: spec.template.spec.volumes + content: + name: collector-socket + emptyDir: + sizeLimit: 1Mi + + - it: should skip the socket volume and pass the policy through on tcp transport + template: deployment.yaml + set: + collector.enabled: true + collector.address: tcp://127.0.0.1:4100 + collector.onUnavailable: drop + collector.bufferSize: 50 + envVars: + CONFIG_FILE_PATH: /custom/config.yaml + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 2 + - notContains: + path: spec.template.spec.containers[1].env + content: + name: CONFIG_FILE_PATH + value: /etc/litellm/config.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: CONFIG_FILE_PATH + value: /custom/config.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ADDRESS + value: tcp://127.0.0.1:4100 + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ON_UNAVAILABLE + value: drop + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_BUFFER_SIZE + value: "50" + - notContains: + path: spec.template.spec.volumes + content: + name: collector-socket + any: true + + - it: should keep metrics and billing env on the proxy container only + template: deployment.yaml + set: + collector.enabled: true + metricsServer.enabled: true + metricsServer.port: 9090 + billingMetrics.enabled: true + billingMetrics.endpoint: https://metering.example.com + billingMetrics.secretName: billing-mtls + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROMETHEUS_METRICS_PORT + value: "9090" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://metering.example.com + - notContains: + path: spec.template.spec.containers[1].env + content: + name: PROMETHEUS_METRICS_PORT + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + any: true + - notContains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: billing-metrics-mtls + any: true + + - it: should give the sidecar the same scratch mounts as the proxy on a read-only root + template: deployment.yaml + set: + collector.enabled: true + securityContext.readOnlyRootFilesystem: true + asserts: + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: npm + mountPath: /.npm + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: cache + mountPath: /.cache + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: tmp + mountPath: /tmp + + - it: should keep the pod-wide cpu metric unless asked to scale on the proxy container + template: hpa.yaml + set: + autoscaling.enabled: true + collector.enabled: true + asserts: + - equal: { path: "spec.metrics[0].type", value: Resource } + - equal: { path: "spec.metrics[0].resource.name", value: cpu } + + - it: should scale on the proxy container's cpu only when opted in + template: hpa.yaml + set: + autoscaling.enabled: true + collector.enabled: true + collector.scaleOnProxyContainerCpu: true + asserts: + - equal: { path: "spec.metrics[0].type", value: ContainerResource } + - equal: { path: "spec.metrics[0].containerResource.name", value: cpu } + - equal: { path: "spec.metrics[0].containerResource.container", value: litellm } + - equal: { path: "spec.metrics[0].containerResource.target.averageUtilization", value: 60 } + - isNull: { path: "spec.metrics[0].resource" } + + - it: should not switch to the container metric while the sidecar is off + template: hpa.yaml + set: + autoscaling.enabled: true + collector.scaleOnProxyContainerCpu: true + asserts: + - equal: { path: "spec.metrics[0].type", value: Resource } diff --git a/helm/litellm-helm/tests/connection_pool_tests.yaml b/helm/litellm-helm/tests/connection_pool_tests.yaml new file mode 100644 index 00000000000..203082f27ba --- /dev/null +++ b/helm/litellm-helm/tests/connection_pool_tests.yaml @@ -0,0 +1,112 @@ +suite: test in-container connection pool +templates: + - deployment.yaml + - configmap-litellm.yaml +tests: + - it: should not emit pgbouncer env vars by default + template: deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "20" + + - it: should enable the pool with the default sizing when connectionPool.enabled is set + template: deployment.yaml + set: + db.connectionPool.enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "20" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "1000" + + - it: should pass custom sizing through as strings next to the worker count + template: deployment.yaml + set: + numWorkers: 4 + db.connectionPool.enabled: true + db.connectionPool.maxDbConnections: 8 + db.connectionPool.maxClientConn: 400 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "400" + - contains: + path: spec.template.spec.containers[0].args + content: "4" + + - it: should give the collector sidecar the same pool env as the proxy container + template: deployment.yaml + set: + collector.enabled: true + db.connectionPool.enabled: true + db.connectionPool.maxDbConnections: 8 + db.connectionPool.maxClientConn: 400 + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: litellm-collector + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "400" + + - it: should give the collector sidecar no pool env when the pool is off + template: deployment.yaml + set: + collector.enabled: true + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: litellm-collector + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + any: true diff --git a/helm/litellm-helm/tests/hpa_tests.yaml b/helm/litellm-helm/tests/hpa_tests.yaml index cd062dd5971..e446f58c8fe 100644 --- a/helm/litellm-helm/tests/hpa_tests.yaml +++ b/helm/litellm-helm/tests/hpa_tests.yaml @@ -61,6 +61,84 @@ tests: - equal: { path: "spec.metrics[1].resource.name", value: memory } - equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 } + - it: "renders no workload metrics by default" + set: + autoscaling.enabled: true + autoscaling.targetMemoryUtilizationPercentage: 80 + asserts: + - lengthEqual: { path: spec.metrics, count: 2 } + - notContains: { path: spec.metrics, content: { type: Pods }, any: true } + + - it: "adds a requests-per-second Pods metric after the cpu metric" + set: + autoscaling.enabled: true + autoscaling.targetRequestsPerSecond: 90 + asserts: + - lengthEqual: { path: spec.metrics, count: 2 } + - equal: { path: "spec.metrics[0].resource.name", value: cpu } + - equal: + path: "spec.metrics[1]" + value: + type: Pods + pods: + metric: { name: litellm_requests_per_second } + target: { type: AverageValue, averageValue: "90" } + + - it: "adds a tokens-per-second Pods metric on its own" + set: + autoscaling.enabled: true + autoscaling.targetTokensPerSecond: 6M + asserts: + - lengthEqual: { path: spec.metrics, count: 2 } + - equal: + path: "spec.metrics[1]" + value: + type: Pods + pods: + metric: { name: litellm_tokens_per_second } + target: { type: AverageValue, averageValue: "6M" } + - notContains: + path: spec.metrics + content: { type: Pods, pods: { metric: { name: litellm_requests_per_second } } } + any: true + + - it: "renders requests, tokens, cpu and memory metrics together" + set: + autoscaling.enabled: true + autoscaling.targetMemoryUtilizationPercentage: 80 + autoscaling.targetRequestsPerSecond: 90 + autoscaling.targetTokensPerSecond: 6000000 + asserts: + - lengthEqual: { path: spec.metrics, count: 4 } + - equal: { path: "spec.metrics[0].resource.name", value: cpu } + - equal: { path: "spec.metrics[1].resource.name", value: memory } + - equal: { path: "spec.metrics[2].pods.metric.name", value: litellm_requests_per_second } + - equal: { path: "spec.metrics[2].pods.target.averageValue", value: "90" } + - equal: { path: "spec.metrics[3].pods.metric.name", value: litellm_tokens_per_second } + - equal: { path: "spec.metrics[3].pods.target.averageValue", value: "6000000" } + + - it: "scales on workload metrics alone when the cpu target is cleared" + set: + autoscaling.enabled: true + autoscaling.targetCPUUtilizationPercentage: null + autoscaling.targetRequestsPerSecond: 90 + autoscaling.targetTokensPerSecond: 6000000 + asserts: + - lengthEqual: { path: spec.metrics, count: 2 } + - notContains: { path: spec.metrics, content: { type: Resource }, any: true } + - equal: { path: "spec.metrics[0].pods.metric.name", value: litellm_requests_per_second } + - equal: { path: "spec.metrics[1].pods.metric.name", value: litellm_tokens_per_second } + - notMatchRegexRaw: { pattern: per_minute } + + - it: "ignores the per-minute keys, which the chart never shipped" + set: + autoscaling.enabled: true + autoscaling.targetRequestsPerMinute: 5400 + autoscaling.targetTokensPerMinute: 360000000 + asserts: + - lengthEqual: { path: spec.metrics, count: 1 } + - notContains: { path: spec.metrics, content: { type: Pods }, any: true } + - it: "renders no hpa when autoscaling is disabled" asserts: - hasDocuments: { count: 0 } diff --git a/helm/litellm-helm/tests/keda_tests.yaml b/helm/litellm-helm/tests/keda_tests.yaml new file mode 100644 index 00000000000..c9598646223 --- /dev/null +++ b/helm/litellm-helm/tests/keda_tests.yaml @@ -0,0 +1,106 @@ +suite: "keda" +templates: + - keda.yaml +release: + name: rel + namespace: llm +tests: + - it: "renders no scaled object by default" + asserts: + - hasDocuments: { count: 0 } + + - it: "passes user triggers through and adds no prometheus triggers by default" + set: + keda.enabled: true + keda.triggers: + - type: cpu + metricType: Utilization + metadata: { value: "60" } + asserts: + - isKind: { of: ScaledObject } + - equal: + path: spec.triggers + value: + - type: cpu + metricType: Utilization + metadata: { value: "60" } + + - it: "scales on release-wide requests per second divided by the per-replica target" + set: + keda.enabled: true + keda.prometheus.serverAddress: http://prometheus-operated.monitoring.svc:9090 + keda.prometheus.requestsPerSecond: 90 + asserts: + - lengthEqual: { path: spec.triggers, count: 1 } + - equal: + path: "spec.triggers[0]" + value: + type: prometheus + metadata: + serverAddress: http://prometheus-operated.monitoring.svc:9090 + threshold: "90" + query: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm"}[1m])) + + - it: "scales on tokens per second on its own" + set: + keda.enabled: true + keda.prometheus.serverAddress: http://prom:9090 + keda.prometheus.tokensPerSecond: 6000000 + asserts: + - lengthEqual: { path: spec.triggers, count: 1 } + - equal: { path: "spec.triggers[0].type", value: prometheus } + - equal: { path: "spec.triggers[0].metadata.threshold", value: "6000000" } + - equal: + path: "spec.triggers[0].metadata.query" + value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm"}[1m])) + + - it: "appends requests and tokens triggers after user triggers and selects the metrics service job" + set: + keda.enabled: true + metricsServer.enabled: true + keda.triggers: + - type: cpu + metricType: Utilization + metadata: { value: "60" } + keda.prometheus.serverAddress: http://prom:9090 + keda.prometheus.requestsPerSecond: 90 + keda.prometheus.tokensPerSecond: 6000000 + asserts: + - lengthEqual: { path: spec.triggers, count: 3 } + - equal: { path: "spec.triggers[0].type", value: cpu } + - equal: { path: "spec.triggers[1].metadata.threshold", value: "90" } + - equal: + path: "spec.triggers[1].metadata.query" + value: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m])) + - equal: { path: "spec.triggers[2].metadata.threshold", value: "6000000" } + - equal: + path: "spec.triggers[2].metadata.query" + value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m])) + - notMatchRegexRaw: { pattern: "\\* *60|per_minute|PerMinute" } + + - it: "ignores the per-minute keys, which the chart never shipped" + set: + keda.enabled: true + keda.prometheus.serverAddress: http://prom:9090 + keda.prometheus.requestsPerMinute: 5400 + keda.prometheus.tokensPerMinute: 360000000 + asserts: + - isKind: { of: ScaledObject } + - isNullOrEmpty: { path: spec.triggers } + + - it: "refuses a workload target without a prometheus server address" + set: + keda.enabled: true + keda.prometheus.requestsPerSecond: 90 + asserts: + - failedTemplate: + errorMessage: keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set + + - it: "yields to the hpa when both autoscalers are enabled" + set: + autoscaling.enabled: true + keda.enabled: true + keda.prometheus.serverAddress: http://prom:9090 + keda.prometheus.requestsPerSecond: 90 + asserts: + - hasDocuments: { count: 0 } diff --git a/helm/litellm-helm/tests/metrics_server_tests.yaml b/helm/litellm-helm/tests/metrics_server_tests.yaml new file mode 100644 index 00000000000..085d69ac640 --- /dev/null +++ b/helm/litellm-helm/tests/metrics_server_tests.yaml @@ -0,0 +1,106 @@ +suite: separate metrics server +templates: + - configmap-litellm.yaml + - deployment.yaml + - service.yaml + - service-metrics.yaml + - servicemonitor.yaml +tests: + - it: should not expose a metrics port or PROMETHEUS_METRICS_PORT by default + asserts: + - notContains: + path: spec.template.spec.containers[0].ports + content: + name: metrics + any: true + template: deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: PROMETHEUS_METRICS_PORT + any: true + template: deployment.yaml + - lengthEqual: + path: spec.ports + count: 1 + template: service.yaml + - hasDocuments: + count: 0 + template: service-metrics.yaml + + - it: should scrape the proxy port when the metrics server is disabled + template: servicemonitor.yaml + set: + serviceMonitor.enabled: true + asserts: + - equal: + path: spec.endpoints[0].port + value: http + + - it: should wire the separate metrics server through container, a ClusterIP metrics service and servicemonitor + set: + metricsServer.enabled: true + metricsServer.port: 4101 + serviceMonitor.enabled: true + service.type: LoadBalancer + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROMETHEUS_METRICS_PORT + value: "4101" + template: deployment.yaml + - contains: + path: spec.template.spec.containers[0].ports + content: + name: metrics + containerPort: 4101 + protocol: TCP + template: deployment.yaml + - lengthEqual: + path: spec.ports + count: 1 + template: service.yaml + - equal: + path: spec.type + value: LoadBalancer + template: service.yaml + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-metrics + template: service-metrics.yaml + - equal: + path: spec.type + value: ClusterIP + template: service-metrics.yaml + - equal: + path: spec.ports + value: + - port: 4101 + targetPort: metrics + protocol: TCP + name: metrics + template: service-metrics.yaml + - equal: + path: spec.selector + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + template: service-metrics.yaml + - equal: + path: spec.endpoints[0].port + value: metrics + template: servicemonitor.yaml + - equal: + path: spec.endpoints[0].path + value: /metrics/ + template: servicemonitor.yaml + + - it: should reject a metrics port equal to the proxy port + template: deployment.yaml + set: + metricsServer.enabled: true + metricsServer.port: 4000 + asserts: + - failedTemplate: + errorMessage: metricsServer.port must differ from service.port diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 637be2322e3..fcee331a5aa 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -180,6 +180,58 @@ proxy_config: general_settings: master_key: os.environ/PROXY_MASTER_KEY +# Serve Prometheus /metrics from a separate process (PROMETHEUS_METRICS_PORT) +# so a scrape never runs on an inference worker. Adds a `metrics` port to the +# container and a dedicated ClusterIP `-metrics` Service, and the +# ServiceMonitor scrapes it instead of the proxy port. The separate port has +# no virtual-key auth: keep it off public ingress. Needs the proxy image +# v1.101.0 or newer. +metricsServer: + enabled: false + port: 4001 + +# Opt-in sidecar that runs the post-response spend pipeline (cost calculation, +# spend logs, spend counters, budget reservation reconciliation) so the proxy's +# uvicorn workers only serialise a compact typed event and go back to serving +# inference. Same image and tag as the proxy, second container in the same pod, +# fed over loopback (a unix socket on a shared emptyDir, or 127.0.0.1 TCP). It +# reuses the pod's in-container pgbouncer (db.connectionPool) and the same Redis +# spend transaction buffer, so the per-pod DB connection budget is unchanged. +# Delivery is at-most-once inside the pod: events already handed to the sidecar +# are lost if it crashes before writing them; events the workers could not hand +# over follow onUnavailable. Both containers drain on SIGTERM within +# terminationGracePeriodSeconds +collector: + enabled: false + # unix:////.sock (the becomes a shared emptyDir) or tcp://127.0.0.1: + address: unix:///var/run/litellm/collector.sock + # Events each uvicorn worker holds in memory while the sidecar is slow or restarting + bufferSize: 1000 + # fallback: run the pipeline in the worker when the sidecar is unreachable or the + # buffer is full (spend stays exact, that request costs proxy CPU again) + # drop: count and discard the event instead (spend under-reports) + onUnavailable: fallback + # How long the workers keep pushing buffered events on shutdown, and how long the + # sidecar keeps serving its open connections after SIGTERM + drainTimeoutSeconds: 10 + command: + - python + - -m + - litellm.proxy.collector + # Sized independently of the proxy container; the pipeline is CPU bound + resources: {} + # requests: + # cpu: 500m + # memory: 1Gi + # limits: + # cpu: "1" + # memory: 2Gi + # When autoscaling.enabled, swap the pod-wide cpu Resource metric for an + # autoscaling/v2 ContainerResource metric on the proxy container only, so the + # sidecar's CPU never scales inference replicas. Needs Kubernetes 1.30+ (or the + # HPAContainerMetrics feature gate on 1.27 to 1.29) + scaleOnProxyContainerCpu: false + resources: {} # Unset by default so the chart installs on small clusters such as Minikube, and so an @@ -212,6 +264,25 @@ autoscaling: # Memory is a floor to provision under 'resources', not a signal to scale on. # targetMemoryUtilizationPercentage: 80 # behavior: {} + # Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics + # named `litellm_requests_per_second` and `litellm_tokens_per_second` with an + # AverageValue target, alongside whichever resource targets are set (the HPA + # follows the metric asking for the most replicas). A Prometheus Adapter must + # serve those two names on custom.metrics.k8s.io from the proxy's counters, + # grouped by the scrape target's `pod` label (enable serviceMonitor below so + # every pod is scraped on its own): + # litellm_requests_per_second: + # sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>) + # litellm_tokens_per_second: + # sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>) + # rate() over [1m] is already per second, so no `* 60`. How fast the HPA + # reacts is set by that window, the scrape interval and the HPA sync period + # (15s by default), not by the unit: keep serviceMonitor.interval at 15s or + # faster so a 1m window holds at least 4 samples. averageValue takes SI + # suffixes, so "6M" is six million tokens per second per pod. Tokens are + # counted when a response completes, so TPS trails long streams. + targetRequestsPerSecond: "" + targetTokensPerSecond: "" # Autoscaling with keda is mutually exclusive with hpa keda: @@ -233,6 +304,23 @@ keda: # metricName: http_requests_total # threshold: '100' # query: sum(rate(http_requests_total{deployment="my-deployment"}[2m])) + # First-class Prometheus triggers on the proxy's own request and token + # counters, appended to `triggers`. Each target is the per-second load one + # replica should carry: KEDA divides the release-wide + # `sum(rate([1m]))` by it to pick the replica count. Thresholds + # are plain numbers (KEDA parses them as floats, no SI suffixes). The + # queries select samples by the release namespace and the `job` label the + # chart's ServiceMonitor produces (the metrics Service name), so enable + # serviceMonitor below together with metricsServer: the http port serves + # /metrics/ behind virtual-key auth and answers an unauthenticated scrape + # with 401. Reaction time comes from the [1m] window, the scrape interval + # and pollingInterval above, so keep both at 15s or faster. Tokens are + # counted at completion, so TPS trails long streams. serverAddress is + # required once either target is set. + prometheus: + serverAddress: "" + requestsPerSecond: "" + tokensPerSecond: "" behavior: {} # scaleDown: # stabilizationWindowSeconds: 300 @@ -309,6 +397,20 @@ db: # only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime). readReplicaUrl: "" + # In-container connection pool (PgBouncer, transaction mode) shared by every + # worker in the pod. Without it each --num_workers worker opens its own + # connection_limit connections to Postgres, so a pod's footprint against the + # database's connection ceiling is workers x connection_limit and grows with + # every replica. With it, the pod holds at most maxDbConnections upstream + # connections no matter how many workers run; the workers connect to the pool + # over loopback, with no extra network hop. Migrations still go straight to + # Postgres. Starting profile for numWorkers: 4 is maxDbConnections: 20, so + # a database with a 5000-connection ceiling fits roughly 200 replicas. + connectionPool: + enabled: false + maxDbConnections: 20 + maxClientConn: 1000 + # Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster. # The Stackgres Operator must already be installed within the target # Kubernetes cluster. diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index be6b9093f53..692a799e783 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -360,6 +360,20 @@ harmless no-op for the Job and authoritative for the app pods. {{- end }} {{- end -}} +{{/* +In-container PgBouncer env for the gateway container. Under IAM or Entra auth the pooler mints and renews the database token itself. +*/}} +{{- define "litellm.connectionPoolEnv" -}} +{{- with .Values.database.connectionPool -}} +- name: LITELLM_PGBOUNCER_ENABLED + value: "true" +- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: {{ required "database.connectionPool.maxDbConnections is required when the pool is enabled" .maxDbConnections | quote }} +- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: {{ required "database.connectionPool.maxClientConn is required when the pool is enabled" .maxClientConn | quote }} +{{- end }} +{{- end -}} + {{/* PodDisruptionBudget shared by gateway, backend, and ui. @@ -441,3 +455,36 @@ ImplementationSpecific {{- .pathType -}} {{- end -}} {{- end -}} + +{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}} + +{{/* +Directory of the collector's unix socket, shared by the gateway and +collector containers through an emptyDir. Empty when the sidecar is off +or gateway.collector.address is a tcp://127.0.0.1: address. +*/}} +{{- define "litellm.gateway.collectorSocketDir" -}} +{{- if and .Values.gateway.collector.enabled (hasPrefix "unix://" .Values.gateway.collector.address) -}} +{{- dir (trimPrefix "unix://" .Values.gateway.collector.address) -}} +{{- end -}} +{{- end -}} + +{{/* +LITELLM_COLLECTOR_* env shared by the producer (gateway container) and the +consumer (collector container), so both agree on the transport and the +shutdown drain window. +*/}} +{{- define "litellm.gateway.collectorEnv" -}} +{{- with .Values.gateway.collector }} +- name: LITELLM_COLLECTOR_ENABLED + value: "true" +- name: LITELLM_COLLECTOR_ADDRESS + value: {{ .address | quote }} +- name: LITELLM_COLLECTOR_BUFFER_SIZE + value: {{ .bufferSize | quote }} +- name: LITELLM_COLLECTOR_ON_UNAVAILABLE + value: {{ .onUnavailable | quote }} +- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS + value: {{ .drainTimeoutSeconds | quote }} +{{- end }} +{{- end -}} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 5030ba2c9dc..c06cc9583a0 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -61,17 +61,38 @@ spec: - name: NUM_WORKERS value: {{ .Values.gateway.numWorkers | quote }} {{- end }} + {{- if .Values.database.connectionPool.enabled }} + {{- include "litellm.connectionPoolEnv" $ | nindent 12 }} + {{- end }} {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsEnv" . | nindent 12 }} {{- end }} + {{- if .Values.gateway.metricsServer.enabled }} + {{- if eq (int .Values.gateway.metricsServer.port) 4000 }} + {{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }} + {{- end }} + - name: PROMETHEUS_MULTIPROC_DIR + value: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + {{- end }} + {{- if .Values.gateway.collector.enabled }} + {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} + {{- end }} {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} volumeMounts: {{- if .Values.gateway.config.create }} - name: gateway-config mountPath: /app/config/config.yaml subPath: config.yaml {{- end }} + {{- if .Values.gateway.metricsServer.enabled }} + - name: prometheus-multiproc + mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + {{- end }} + {{- if include "litellm.gateway.collectorSocketDir" . }} + - name: collector-socket + mountPath: {{ include "litellm.gateway.collectorSocketDir" . }} + {{- end }} {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} {{- end }} @@ -97,16 +118,102 @@ spec: {{- end }} resources: {{- toYaml .Values.gateway.resources | nindent 12 }} + {{- if .Values.gateway.metricsServer.enabled }} + - name: metrics + image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.gateway.image.pullPolicy }} + {{- with .Values.gateway.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} + command: + - python + - -m + - litellm.proxy.prometheus_metrics_server + - --port + - {{ .Values.gateway.metricsServer.port | quote }} + env: + - name: PROMETHEUS_MULTIPROC_DIR + value: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + ports: + - name: metrics + containerPort: {{ .Values.gateway.metricsServer.port }} + protocol: TCP + volumeMounts: + - name: prometheus-multiproc + mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + readinessProbe: + tcpSocket: { port: metrics } + periodSeconds: 10 + livenessProbe: + tcpSocket: { port: metrics } + periodSeconds: 15 + failureThreshold: 6 + resources: + {{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }} + {{- end }} + {{- if .Values.gateway.collector.enabled }} + - name: collector + image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.gateway.image.pullPolicy }} + {{- with .Values.gateway.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} + command: + - python + - -m + - litellm.proxy.collector + env: + {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }} + {{- if .Values.gateway.config.create }} + - name: CONFIG_FILE_PATH + value: /app/config/config.yaml + {{- end }} + {{- if .Values.database.connectionPool.enabled }} + {{- include "litellm.connectionPoolEnv" $ | nindent 12 }} + {{- end }} + {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} + - name: LITELLM_JOB_ROLE + value: collector + {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} + {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts (include "litellm.gateway.collectorSocketDir" .) }} + volumeMounts: + {{- if .Values.gateway.config.create }} + - name: gateway-config + mountPath: /app/config/config.yaml + subPath: config.yaml + {{- end }} + {{- if include "litellm.gateway.collectorSocketDir" . }} + - name: collector-socket + mountPath: {{ include "litellm.gateway.collectorSocketDir" . }} + {{- end }} + {{- with .Values.gateway.volumeMounts }} + {{- toYaml . | nindent 12 }} + {{- end }} + {{- end }} + resources: + {{- toYaml .Values.gateway.collector.resources | nindent 12 }} + {{- end }} {{- with .Values.gateway.extraContainers }} {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} - {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} volumes: {{- if .Values.gateway.config.create }} - name: gateway-config configMap: name: {{ include "litellm.gateway.fullname" . }}-config {{- end }} + {{- if .Values.gateway.metricsServer.enabled }} + - name: prometheus-multiproc + emptyDir: {} + {{- end }} + {{- if include "litellm.gateway.collectorSocketDir" . }} + - name: collector-socket + emptyDir: + sizeLimit: 1Mi + {{- end }} {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} {{- end }} diff --git a/helm/litellm/templates/gateway/hpa.yaml b/helm/litellm/templates/gateway/hpa.yaml index e97cef95ffb..e7094e96106 100644 --- a/helm/litellm/templates/gateway/hpa.yaml +++ b/helm/litellm/templates/gateway/hpa.yaml @@ -15,6 +15,15 @@ spec: maxReplicas: {{ .Values.gateway.hpa.maxReplicas }} metrics: {{- if .Values.gateway.hpa.targetCPUUtilizationPercentage }} + {{- if and .Values.gateway.collector.enabled .Values.gateway.collector.scaleOnGatewayContainerCpu }} + - type: ContainerResource + containerResource: + name: cpu + container: gateway + target: + type: Utilization + averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }} + {{- else }} - type: Resource resource: name: cpu @@ -22,6 +31,7 @@ spec: type: Utilization averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }} {{- end }} + {{- end }} {{- if .Values.gateway.hpa.targetMemoryUtilizationPercentage }} - type: Resource resource: @@ -30,6 +40,24 @@ spec: type: Utilization averageUtilization: {{ .Values.gateway.hpa.targetMemoryUtilizationPercentage }} {{- end }} + {{- with .Values.gateway.hpa.targetRequestsPerSecond }} + - type: Pods + pods: + metric: + name: litellm_requests_per_second + target: + type: AverageValue + averageValue: {{ toJson . | trimAll "\"" | quote }} + {{- end }} + {{- with .Values.gateway.hpa.targetTokensPerSecond }} + - type: Pods + pods: + metric: + name: litellm_tokens_per_second + target: + type: AverageValue + averageValue: {{ toJson . | trimAll "\"" | quote }} + {{- end }} {{- with .Values.gateway.hpa.behavior }} behavior: {{- toYaml . | nindent 4 }} diff --git a/helm/litellm/templates/gateway/service-metrics.yaml b/helm/litellm/templates/gateway/service-metrics.yaml new file mode 100644 index 00000000000..ad9bc05a9fd --- /dev/null +++ b/helm/litellm/templates/gateway/service-metrics.yaml @@ -0,0 +1,18 @@ +{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }} +apiVersion: v1 +kind: Service +metadata: + name: {{ include "litellm.gateway.fullname" . }}-metrics + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: gateway +spec: + type: ClusterIP + ports: + - port: {{ .Values.gateway.metricsServer.port }} + targetPort: metrics + protocol: TCP + name: metrics + selector: + {{- include "litellm.gateway.selectorLabels" . | nindent 4 }} +{{- end }} diff --git a/helm/litellm/templates/gateway/servicemonitor.yaml b/helm/litellm/templates/gateway/servicemonitor.yaml new file mode 100644 index 00000000000..e1bafa6e388 --- /dev/null +++ b/helm/litellm/templates/gateway/servicemonitor.yaml @@ -0,0 +1,28 @@ +{{- if and .Values.gateway.enabled .Values.gateway.serviceMonitor.enabled }} +{{- if not .Values.gateway.metricsServer.enabled }} +{{- fail "gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled: the http port serves /metrics/ behind virtual-key auth, so an unauthenticated scrape gets 401" }} +{{- end }} +apiVersion: monitoring.coreos.com/v1 +kind: ServiceMonitor +metadata: + name: {{ include "litellm.gateway.fullname" . }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: gateway + {{- with .Values.gateway.serviceMonitor.labels }} + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + selector: + matchLabels: + {{- include "litellm.gateway.selectorLabels" . | nindent 6 }} + namespaceSelector: + matchNames: + - {{ .Release.Namespace | quote }} + endpoints: + - port: metrics + path: /metrics/ + interval: {{ .Values.gateway.serviceMonitor.interval }} + scrapeTimeout: {{ .Values.gateway.serviceMonitor.scrapeTimeout }} + scheme: http +{{- end }} diff --git a/helm/litellm/tests/collector_tests.yaml b/helm/litellm/tests/collector_tests.yaml new file mode 100644 index 00000000000..4ef7e3c8ca4 --- /dev/null +++ b/helm/litellm/tests/collector_tests.yaml @@ -0,0 +1,216 @@ +suite: test gateway collector sidecar +templates: + - gateway/configmap.yaml + - gateway/deployment.yaml + - gateway/hpa.yaml +values: + - ./values/required.yaml +tests: + - it: adds no sidecar, env, volume or container metric when the collector is off + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 1 + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ENABLED + value: "true" + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.volumes + content: + name: collector-socket + any: true + template: gateway/deployment.yaml + - equal: + path: spec.metrics[0].type + value: Resource + template: gateway/hpa.yaml + + - it: runs the collector as a sidecar sharing env, config, the pod pool and a unix socket emptyDir, and scales on the gateway container only + set: + gateway.collector.enabled: true + gateway.collector.bufferSize: 250 + gateway.collector.onUnavailable: drop + gateway.image.tag: v1.102.0 + gateway.numWorkers: 4 + database.connectionPool.enabled: true + database.connectionPool.maxDbConnections: 8 + database.connectionPool.maxClientConn: 250 + gateway.envSecrets: + - litellm-license + gateway.volumes: + - name: redis-ca + secret: + secretName: redis-ca + gateway.volumeMounts: + - name: redis-ca + mountPath: /etc/litellm/redis-ca + readOnly: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ADDRESS + value: unix:///var/run/litellm/collector.sock + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_BUFFER_SIZE + value: "250" + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ON_UNAVAILABLE + value: drop + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: collector-socket + mountPath: /var/run/litellm + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].name + value: collector + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].image + value: ghcr.io/berriai/litellm-gateway:v1.102.0 + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].command + value: + - python + - -m + - litellm.proxy.collector + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_JOB_ROLE + value: collector + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: CONFIG_FILE_PATH + value: /app/config/config.yaml + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "250" + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: DATABASE_HOST + value: postgres.example.com + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_COLLECTOR_ADDRESS + value: unix:///var/run/litellm/collector.sock + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[1].env + content: + name: NUM_WORKERS + any: true + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].envFrom + value: + - secretRef: + name: litellm-license + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: gateway-config + mountPath: /app/config/config.yaml + subPath: config.yaml + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: collector-socket + mountPath: /var/run/litellm + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: redis-ca + mountPath: /etc/litellm/redis-ca + readOnly: true + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].resources.limits.cpu + value: "1" + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.volumes + content: + name: collector-socket + emptyDir: + sizeLimit: 1Mi + template: gateway/deployment.yaml + - equal: + path: spec.metrics[0] + value: + type: ContainerResource + containerResource: + name: cpu + container: gateway + target: + type: Utilization + averageUtilization: 70 + template: gateway/hpa.yaml + + - it: uses loopback tcp without a socket volume and keeps the pod-wide cpu metric when asked + set: + gateway.collector.enabled: true + gateway.collector.address: tcp://127.0.0.1:4010 + gateway.collector.scaleOnGatewayContainerCpu: false + asserts: + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_COLLECTOR_ADDRESS + value: tcp://127.0.0.1:4010 + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.volumes + content: + name: collector-socket + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[1].volumeMounts + content: + name: collector-socket + any: true + template: gateway/deployment.yaml + - equal: + path: spec.metrics[0].type + value: Resource + template: gateway/hpa.yaml diff --git a/helm/litellm/tests/connection_pool_tests.yaml b/helm/litellm/tests/connection_pool_tests.yaml new file mode 100644 index 00000000000..c39651a52c9 --- /dev/null +++ b/helm/litellm/tests/connection_pool_tests.yaml @@ -0,0 +1,204 @@ +suite: test in-container connection pool env vars +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - backend/configmap.yaml +values: + - ./values/required.yaml +tests: + - it: renders no pool env by default + template: gateway/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + any: true + + - it: enabled pool renders the three pgbouncer vars with the configured sizes + template: gateway/deployment.yaml + set: + gateway.numWorkers: 4 + database.connectionPool.enabled: true + database.connectionPool.maxDbConnections: 8 + database.connectionPool.maxClientConn: 250 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "250" + - contains: + path: spec.template.spec.containers[0].env + content: + name: NUM_WORKERS + value: "4" + + - it: enabled pool uses the chart default sizes + template: gateway/deployment.yaml + set: + database.connectionPool.enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "20" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "1000" + + - it: backend never gets the pool env + template: backend/deployment.yaml + set: + database.connectionPool.enabled: true + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + + - it: collector sidecar gets the same pool env as the gateway container, the metrics sidecar none + template: gateway/deployment.yaml + set: + gateway.collector.enabled: true + gateway.metricsServer.enabled: true + database.connectionPool.enabled: true + database.connectionPool.maxDbConnections: 8 + database.connectionPool.maxClientConn: 250 + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: metrics + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - equal: + path: spec.template.spec.containers[2].name + value: collector + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "250" + + - it: collector sidecar gets no pool env when the pool is off + template: gateway/deployment.yaml + set: + gateway.collector.enabled: true + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: collector + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + any: true + + - it: pool with IAM auth renders both the pool and the token auth flag + template: gateway/deployment.yaml + set: + gateway.collector.enabled: true + database.connectionPool.enabled: true + database.writer.useIAMAuth: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[0].env + content: + name: IAM_TOKEN_DB_AUTH + value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: IAM_TOKEN_DB_AUTH + value: "true" + + - it: pool with Entra auth renders both the pool and the token auth flag + template: gateway/deployment.yaml + set: + database.connectionPool.enabled: true + database.writer.useAzureEntraAuth: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[0].env + content: + name: AZURE_POSTGRESQL_AUTH + value: "true" + + - it: IAM auth without the pool still renders + template: gateway/deployment.yaml + set: + database.writer.useIAMAuth: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: IAM_TOKEN_DB_AUTH + value: "true" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true diff --git a/helm/litellm/tests/hpa_workload_metrics_tests.yaml b/helm/litellm/tests/hpa_workload_metrics_tests.yaml new file mode 100644 index 00000000000..a29c53c5ed9 --- /dev/null +++ b/helm/litellm/tests/hpa_workload_metrics_tests.yaml @@ -0,0 +1,213 @@ +suite: test gateway HPA per-pod requests-per-second and tokens-per-second targets +templates: + - gateway/hpa.yaml + - gateway/servicemonitor.yaml +values: + - ./values/required.yaml +tests: + - it: scales on CPU and memory only by default + template: gateway/hpa.yaml + asserts: + - equal: + path: spec.metrics + value: + - type: Resource + resource: + name: cpu + target: + type: Utilization + averageUtilization: 70 + - type: Resource + resource: + name: memory + target: + type: Utilization + averageUtilization: 80 + + - it: adds a requests-per-second Pods metric next to the resource metrics + template: gateway/hpa.yaml + set: + gateway.hpa.targetRequestsPerSecond: 90 + asserts: + - lengthEqual: + path: spec.metrics + count: 3 + - contains: + path: spec.metrics + content: + type: Resource + resource: + name: cpu + target: + type: Utilization + averageUtilization: 70 + - equal: + path: spec.metrics[2] + value: + type: Pods + pods: + metric: + name: litellm_requests_per_second + target: + type: AverageValue + averageValue: "90" + - notContains: + path: spec.metrics + content: + type: Pods + pods: + metric: + name: litellm_tokens_per_second + any: true + + - it: adds a tokens-per-second Pods metric on its own + template: gateway/hpa.yaml + set: + gateway.hpa.targetTokensPerSecond: 6M + asserts: + - lengthEqual: + path: spec.metrics + count: 3 + - equal: + path: spec.metrics[2] + value: + type: Pods + pods: + metric: + name: litellm_tokens_per_second + target: + type: AverageValue + averageValue: "6M" + - notContains: + path: spec.metrics + content: + type: Pods + pods: + metric: + name: litellm_requests_per_second + any: true + + - it: renders requests and tokens targets together and keeps CPU and memory + template: gateway/hpa.yaml + set: + gateway.hpa.targetRequestsPerSecond: 90 + gateway.hpa.targetTokensPerSecond: 6000000 + asserts: + - lengthEqual: + path: spec.metrics + count: 4 + - equal: + path: spec.metrics[0].resource.name + value: cpu + - equal: + path: spec.metrics[1].resource.name + value: memory + - equal: + path: spec.metrics[2].pods.metric.name + value: litellm_requests_per_second + - equal: + path: spec.metrics[3].pods.metric.name + value: litellm_tokens_per_second + - equal: + path: spec.metrics[3].pods.target.averageValue + value: "6000000" + + - it: scales on workload metrics alone when the resource targets are cleared + template: gateway/hpa.yaml + set: + gateway.hpa.targetCPUUtilizationPercentage: null + gateway.hpa.targetMemoryUtilizationPercentage: null + gateway.hpa.targetRequestsPerSecond: 90 + gateway.hpa.targetTokensPerSecond: 6000000 + asserts: + - lengthEqual: + path: spec.metrics + count: 2 + - notContains: + path: spec.metrics + content: + type: Resource + any: true + - equal: + path: spec.metrics[0].pods.metric.name + value: litellm_requests_per_second + - equal: + path: spec.metrics[1].pods.metric.name + value: litellm_tokens_per_second + - notMatchRegexRaw: + pattern: per_minute + + - it: ignores the per-minute keys, which the chart never shipped + template: gateway/hpa.yaml + set: + gateway.hpa.targetRequestsPerMinute: 5400 + gateway.hpa.targetTokensPerMinute: 360000000 + asserts: + - lengthEqual: + path: spec.metrics + count: 2 + - notContains: + path: spec.metrics + content: + type: Pods + any: true + + - it: renders no ServiceMonitor by default + template: gateway/servicemonitor.yaml + asserts: + - hasDocuments: + count: 0 + + - it: refuses a ServiceMonitor without the metrics server, whose http port needs a bearer token + template: gateway/servicemonitor.yaml + set: + gateway.serviceMonitor.enabled: true + asserts: + - failedTemplate: + errorPattern: gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled + + - it: scrapes each gateway pod through the metrics port + template: gateway/servicemonitor.yaml + release: + name: rel + namespace: llm + set: + gateway.serviceMonitor.enabled: true + gateway.metricsServer.enabled: true + gateway.serviceMonitor.labels: + release: kube-prometheus-stack + asserts: + - isKind: + of: ServiceMonitor + - equal: + path: metadata.labels.release + value: kube-prometheus-stack + - equal: + path: spec.selector.matchLabels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: rel + app.kubernetes.io/component: gateway + - equal: + path: spec.namespaceSelector.matchNames + value: + - llm + - equal: + path: spec.endpoints + value: + - port: metrics + path: /metrics/ + interval: 15s + scrapeTimeout: 10s + scheme: http + + - it: honours a custom scrape interval + template: gateway/servicemonitor.yaml + set: + gateway.serviceMonitor.enabled: true + gateway.serviceMonitor.interval: 30s + gateway.metricsServer.enabled: true + asserts: + - equal: + path: spec.endpoints[0].interval + value: 30s diff --git a/helm/litellm/tests/metrics_server_tests.yaml b/helm/litellm/tests/metrics_server_tests.yaml new file mode 100644 index 00000000000..0e7d9d9e9ee --- /dev/null +++ b/helm/litellm/tests/metrics_server_tests.yaml @@ -0,0 +1,148 @@ +suite: test gateway metrics sidecar +templates: + - gateway/configmap.yaml + - gateway/deployment.yaml + - gateway/service.yaml + - gateway/service-metrics.yaml +values: + - ./values/required.yaml +tests: + - it: adds no sidecar, volume, env or service port when the metrics server is off + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 1 + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: PROMETHEUS_MULTIPROC_DIR + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.volumes + content: + name: prometheus-multiproc + any: true + template: gateway/deployment.yaml + - lengthEqual: + path: spec.ports + count: 1 + template: gateway/service.yaml + - hasDocuments: + count: 0 + template: gateway/service-metrics.yaml + + - it: runs the metrics server as a sidecar over a shared multiproc dir and exposes it on a ClusterIP metrics service + set: + gateway.metricsServer.enabled: true + gateway.metricsServer.port: 4101 + gateway.service.type: LoadBalancer + gateway.image.tag: v1.101.0 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROMETHEUS_MULTIPROC_DIR + value: /tmp/litellm_prometheus_multiproc + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: prometheus-multiproc + mountPath: /tmp/litellm_prometheus_multiproc + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].name + value: metrics + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].image + value: ghcr.io/berriai/litellm-gateway:v1.101.0 + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].command + value: + - python + - -m + - litellm.proxy.prometheus_metrics_server + - --port + - "4101" + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].env + value: + - name: PROMETHEUS_MULTIPROC_DIR + value: /tmp/litellm_prometheus_multiproc + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].ports + value: + - name: metrics + containerPort: 4101 + protocol: TCP + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].volumeMounts + value: + - name: prometheus-multiproc + mountPath: /tmp/litellm_prometheus_multiproc + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].readinessProbe.tcpSocket.port + value: metrics + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].livenessProbe.tcpSocket.port + value: metrics + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].resources.requests.cpu + value: 50m + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.volumes + content: + name: prometheus-multiproc + emptyDir: {} + template: gateway/deployment.yaml + - lengthEqual: + path: spec.ports + count: 1 + template: gateway/service.yaml + - equal: + path: spec.type + value: LoadBalancer + template: gateway/service.yaml + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-gateway-metrics + template: gateway/service-metrics.yaml + - equal: + path: spec.type + value: ClusterIP + template: gateway/service-metrics.yaml + - equal: + path: spec.ports + value: + - port: 4101 + targetPort: metrics + protocol: TCP + name: metrics + template: gateway/service-metrics.yaml + - equal: + path: spec.selector + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: gateway + template: gateway/service-metrics.yaml + + - it: rejects a metrics port equal to the gateway port + template: gateway/deployment.yaml + set: + gateway.metricsServer.enabled: true + gateway.metricsServer.port: 4000 + asserts: + - failedTemplate: + errorMessage: gateway.metricsServer.port must differ from the gateway port 4000 diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 6c9fb9440c7..1873219d1ea 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -225,6 +225,26 @@ database: usernameKey: username passwordKey: password + # In-container connection pool (PgBouncer, transaction mode) shared by every + # gateway worker in the pod. Without it each of the `gateway.numWorkers` + # workers opens its own Prisma pool straight to Postgres, so a pod's + # footprint against the database's connection ceiling is + # numWorkers x connection_limit and grows with every replica. With it, the + # pod holds at most maxDbConnections upstream connections no matter how many + # workers run; the workers connect to the pool over loopback, with no extra + # network hop. The chart emits LITELLM_PGBOUNCER_ENABLED / + # LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / LITELLM_PGBOUNCER_MAX_CLIENT_CONN on + # the gateway container and its collector sidecar only: the backend runs a + # single worker and the migrations Job must keep a direct connection. With + # `database.writer.useIAMAuth` or `useAzureEntraAuth` the pool mints and + # renews the database token itself, so the workers never see it. Starting profile for + # `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a + # 5000-connection ceiling fits roughly 200 gateway replicas. + connectionPool: + enabled: false + maxDbConnections: 20 + maxClientConn: 1000 + # Optional Redis. Leave host empty to disable. # # This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend @@ -268,6 +288,68 @@ gateway: config: create: true proxy_config: {} + # Serve Prometheus /metrics from a `metrics` sidecar container (same image, + # `python -m litellm.proxy.prometheus_metrics_server`) that aggregates the + # workers' PROMETHEUS_MULTIPROC_DIR samples over a shared emptyDir, so a + # scrape never runs on an inference worker. Adds a `metrics` port to the pod + # and a dedicated ClusterIP `-metrics` Service; point your scrape + # config at it. The port has no virtual-key auth: keep it off public ingress. + # Needs the gateway image v1.101.0 or newer. + metricsServer: + enabled: false + port: 4001 + resources: + requests: + cpu: 50m + memory: 128Mi + limits: + memory: 512Mi + # Prometheus Operator ServiceMonitor for the gateway pods. Scrapes the + # `-metrics` Service, so it requires metricsServer above (the http + # port serves /metrics/ behind virtual-key auth). Every pod is its own scrape + # target, so the samples carry the `pod` label the per-pod autoscaling + # queries below group by. + serviceMonitor: + enabled: false + labels: {} + interval: 15s + scrapeTimeout: 10s + # Opt-in `collector` sidecar (same image, `python -m litellm.proxy.collector`) + # that runs the post-response spend pipeline (cost calculation, spend logs, + # spend counters, budget reservation reconciliation) so the uvicorn workers + # only serialise a compact event over loopback and go back to serving + # requests. It shares the pod's env, proxy config, in-container pgbouncer and + # Redis spend buffer, so the per-pod DB connection budget is unchanged. + # Delivery is at-most-once inside the pod: events already handed over are + # lost if the sidecar dies before writing them; events the workers cannot + # hand over follow `onUnavailable`. + collector: + enabled: false + # unix:////.sock (the becomes a shared emptyDir) or + # tcp://127.0.0.1: + address: unix:///var/run/litellm/collector.sock + # Events each uvicorn worker holds in memory while the sidecar is slow or + # restarting. + bufferSize: 1000 + # fallback: run the pipeline in the worker when the sidecar is unreachable + # or the buffer is full (spend stays exact, that request costs gateway CPU + # again). drop: count and discard the event instead (spend under-reports). + onUnavailable: fallback + # How long the workers keep pushing buffered events on shutdown, and how + # long the sidecar keeps serving open connections after SIGTERM. + drainTimeoutSeconds: 10 + # Sized independently of the gateway container; the pipeline is CPU bound. + resources: + requests: + cpu: 500m + memory: 1Gi + limits: + cpu: "1" + memory: 2Gi + # With hpa.targetCPUUtilizationPercentage set, scale on an autoscaling/v2 + # ContainerResource metric of the `gateway` container only, so the + # sidecar's CPU never drives inference replicas. Needs Kubernetes 1.30+. + scaleOnGatewayContainerCpu: true image: repository: ghcr.io/berriai/litellm-gateway tag: "" # defaults to .Chart.AppVersion @@ -324,6 +406,25 @@ gateway: # policies: # - { type: Percent, value: 100, periodSeconds: 30 } behavior: {} + # Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics + # named `litellm_requests_per_second` and `litellm_tokens_per_second` with an + # AverageValue target. They coexist with the CPU/memory targets above: the + # HPA scales on whichever metric asks for the most replicas. Kubernetes has + # no idea what a token is, so a Prometheus Adapter must serve those two + # names on custom.metrics.k8s.io from the proxy's counters, grouped by the + # scrape target's `pod` label (enable serviceMonitor above): + # litellm_requests_per_second: + # sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>) + # litellm_tokens_per_second: + # sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>) + # rate() over [1m] is already per second, so no `* 60`. How fast the HPA + # reacts is set by that window, the scrape interval and the HPA sync period + # (15s by default), not by the unit: keep serviceMonitor.interval at 15s or + # faster so a 1m window holds at least 4 samples. averageValue takes SI + # suffixes, so "6M" is six million tokens per second per pod. Tokens are + # counted when a response completes, so TPS trails long streams. + targetRequestsPerSecond: "" + targetTokensPerSecond: "" # PodDisruptionBudget for the gateway pods. Set exactly one of # `minAvailable` / `maxUnavailable` (minAvailable wins if both are set; # enabling without either falls back to `maxUnavailable: 1`). Disabled by diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql new file mode 100644 index 00000000000..e5d48abcb52 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql @@ -0,0 +1,15 @@ +-- DropForeignKey +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN + ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey"; + END IF; +END $$; + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN + ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql new file mode 100644 index 00000000000..bbe980bb66f --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql @@ -0,0 +1,4 @@ +-- Add org_id column to LiteLLM_ManagedObjectTable +-- Snapshots the creating key's organization at submission time, like team_id, +-- so CheckBatchCost can bill organization spend hours later without re-resolving +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "org_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql new file mode 100644 index 00000000000..2fa1234bbfc --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "per_server_oauth_discovery" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql new file mode 100644 index 00000000000..5503167ce09 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql @@ -0,0 +1,3 @@ +ALTER TABLE "LiteLLM_AutoRouterSession" +ADD COLUMN IF NOT EXISTS "classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0, +ADD COLUMN IF NOT EXISTS "classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_add_autorouter_session_baseline_models/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_add_autorouter_session_baseline_models/migration.sql new file mode 100644 index 00000000000..e7ce1a3180b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_add_autorouter_session_baseline_models/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "baseline_models" JSONB NOT NULL DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql new file mode 100644 index 00000000000..c982bc38a69 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "skills" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 1c43668f227..7d521d54791 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + skills String[] @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -343,6 +344,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? @@ -491,7 +493,7 @@ model LiteLLM_JWTKeyMapping { updated_at DateTime @default(now()) @updatedAt updated_by String? - litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) @@unique([jwt_claim_name, jwt_claim_value]) @@index([jwt_claim_name, jwt_claim_value, is_active]) @@ -1035,6 +1037,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") updated_at DateTime @updatedAt @@ -1508,7 +1511,10 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") + baseline_models Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 71e7e9c683b..2145f891318 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -8,14 +8,10 @@ import tempfile import time from dataclasses import dataclass, replace from pathlib import Path -from typing import Optional +from typing import TYPE_CHECKING, Final, Optional from litellm_proxy_extras import prisma_toolchain from litellm_proxy_extras._logging import logger -from litellm_proxy_extras.replica_identity import ( - REPLICA_IDENTITY_FULL_ENV_VAR, - apply_replica_identity_full, -) from litellm_proxy_extras.prisma_toolchain import ( PRISMA_COMMAND_TIMEOUT_ENV_VAR, PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, @@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import ( prisma_command_timeout, prisma_migrate_deploy_timeout, ) +from litellm_proxy_extras.replica_identity import ( + REPLICA_IDENTITY_FULL_ENV_VAR, + apply_replica_identity_full, +) + +if TYPE_CHECKING: + import psycopg + import psycopg.sql def str_to_bool(value: Optional[str]) -> bool: @@ -46,6 +50,28 @@ def _get_prisma_env() -> dict: _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") _MIGRATION_DEADLOCK_MARKER = "deadlock detected" +INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big") +_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$") +_INVALID_LITELLM_INDEXES_SQL: Final = ( + "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " + "FROM pg_index i " + "JOIN pg_class c ON c.oid = i.indexrelid " + "JOIN pg_class t ON t.oid = i.indrelid " + "JOIN pg_namespace n ON n.oid = t.relnamespace " + "WHERE NOT i.indisvalid " + " AND c.relkind = 'i' " + " AND n.nspname = %s " + " AND t.relname LIKE %s " + " AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) " + "ORDER BY c.relname" +) + + +@dataclass(frozen=True, slots=True) +class _InvalidIndex: + schema: str + name: str + table_size: str MAX_MIGRATE_DEPLOY_ATTEMPTS = 4 @@ -624,7 +650,7 @@ class ProxyExtrasDBManager: def _strip_prisma_query_params(url: str) -> str: """Remove Prisma-specific query params (connection_limit, pool_timeout, schema, etc.) from DATABASE_URL so psycopg can parse it.""" - from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode + from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse parsed = urlparse(url) if not parsed.query: @@ -645,7 +671,7 @@ class ProxyExtrasDBManager: "target_session_attrs", } kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params] - return urlunparse(parsed._replace(query=urlencode(kept))) + return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote))) @staticmethod def _warn_if_db_ahead_of_head(migrations_dir: str) -> None: @@ -719,6 +745,95 @@ class ProxyExtrasDBManager: ", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""), ) + @staticmethod + def _invalid_litellm_indexes( + conn: "psycopg.Connection[tuple[str, str, str]]", schema: str + ) -> tuple[_InvalidIndex, ...]: + rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall() + return tuple(_InvalidIndex(*row) for row in rows) + + @staticmethod + def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]: + from psycopg import sql + + target: Final = sql.Identifier(index.schema, index.name) + if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name): + return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover" + return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt" + + @staticmethod + def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None: + import psycopg + + statement, action = ProxyExtrasDBManager._index_repair(index) + try: + conn.execute(statement) + except psycopg.Error as e: + logger.warning( + "Could not repair invalid index %s.%s, will retry on the next startup. " + "If this keeps happening, run `%s` by hand as the index owner. Error: %s", + index.schema, + index.name, + statement.as_string(conn), + e, + ) + return + logger.info("%s invalid index %s.%s", action, index.schema, index.name) + + @staticmethod + def repair_invalid_indexes(lock_timeout: str = "30s") -> bool: + """Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left + INVALID (a migration deadlock between replicas is the usual cause; the + retried migration skips them because of IF NOT EXISTS). Never raises: + returns True when no invalid index remains, False when the repair was + skipped or failed and will be retried on the next startup. Looks in the + schema DATABASE_URL names, the only URL Prisma migrates through, but + connects over DIRECT_URL when set: the session settings, the advisory + lock and REINDEX CONCURRENTLY all need one server session, which a + transaction pooler does not give.""" + prisma_url: Final = os.getenv("DATABASE_URL") + if not prisma_url: + return False + + try: + import psycopg + from psycopg import sql + except ImportError: + logger.warning( + "psycopg is not installed; skipping the invalid index check. " + "Install the litellm[extra_proxy] extra, which includes psycopg." + ) + return False + + schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public" + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url) + try: + with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: + conn.execute("SET statement_timeout = 0") + conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout))) + found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + if not found: + return True + logger.warning( + "Found %d invalid index(es) left by an interrupted CREATE INDEX " + "CONCURRENTLY, rebuilding: %s", + len(found), + ", ".join(f"{index.name} (table size {index.table_size})" for index in found), + ) + lock_row: Final = conn.execute( + "SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,) + ).fetchone() + if lock_row is None or not lock_row[0]: + logger.info("Another replica is already rebuilding the invalid indexes, skipping") + return False + for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema): + ProxyExtrasDBManager._repair_index(conn, index) + remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + except psycopg.Error as e: + logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e) + return False + return not remaining + @staticmethod def _setup_database_v2(use_migrate: bool) -> bool: """ @@ -994,6 +1109,7 @@ class ProxyExtrasDBManager: use_migrate=use_migrate, use_v2_resolver=use_v2_resolver ) if migrated: + ProxyExtrasDBManager.repair_invalid_indexes() ProxyExtrasDBManager.apply_replica_identity_full_if_requested() return migrated diff --git a/litellm-proxy-extras/migration_runbook.md b/litellm-proxy-extras/migration_runbook.md index a277441b164..b1e9236e520 100644 --- a/litellm-proxy-extras/migration_runbook.md +++ b/litellm-proxy-extras/migration_runbook.md @@ -48,7 +48,7 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n ## What It Does -1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check)) +1. **Verifies the current branch is up to date with origin's current default branch** (see [Branch freshness](#branch-freshness-check)) 2. Creates temp PostgreSQL DB 3. Applies existing migrations 4. Compares with `schema.prisma` @@ -57,11 +57,11 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n ## Branch Freshness Check -Before generating anything, `run_migration.py` runs `git fetch origin ` and refuses to proceed if `HEAD` is behind `origin/`. Default base is `litellm_internal_staging` (the branch PRs target). A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense. +Before generating anything, `run_migration.py` runs `git fetch origin ` and refuses to proceed if `HEAD` is behind `origin/`. The default base is discovered from origin's advertised HEAD on each run, so an existing clone follows a default-branch change without trusting cached `origin/HEAD`. If discovery or fetching fails, migration generation stops. A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense. Flags: -- `--base-branch ` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`. +- `--base-branch ` — check against a different base (e.g. a release branch). Defaults to origin's current default branch - `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base. When the guard fires: @@ -69,8 +69,9 @@ When the guard fires: 1. Update your branch: ```bash - git fetch origin && git rebase origin/litellm_internal_staging - # or git merge origin/litellm_internal_staging — whichever matches your workflow + base_branch=$(python3 scripts/default_branch.py --branch) && + git fetch origin "+refs/heads/$base_branch:refs/remotes/origin/$base_branch" && + git rebase "origin/$base_branch" ``` 2. Re-run `run_migration.py`. diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 82d31fec373..7d4c78088f1 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.94" +version = "0.4.96" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.94" +version = "0.4.96" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index b8d1f2db4d7..17856218e60 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,18 +1,19 @@ # AGENTS.md -litellm-rust has five crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. +litellm-rust has six crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. ## Crates | Crate | Role | |-------|------| | litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. | +| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. | | litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. | | litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. | | litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`. The interop foundation depends on no LiteLLM domain crate. +Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers, `litellm-token-counter`, and `litellm-python-interop`. The token counter and interop foundations depend on no LiteLLM domain crate. ## Where a route lives diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 62e943d0f42..cf000e85d68 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,26 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.5" @@ -20,6 +40,15 @@ dependencies = [ "cc", ] +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + [[package]] name = "anes" version = "0.1.6" @@ -41,6 +70,29 @@ dependencies = [ "rustversion", ] +[[package]] +name = "async-compression" +version = "0.4.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f10dafd0c8d2e51ae9a748805777613ed0bbe17bf586b76c8311f45c020a32f" +dependencies = [ + "compression-codecs", + "compression-core", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-trait" version = "0.1.91" @@ -418,7 +470,7 @@ checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" dependencies = [ "async-trait", "axum-core", - "base64", + "base64 0.22.1", "bytes", "futures-util", "http 1.4.2", @@ -468,12 +520,76 @@ dependencies = [ "tracing", ] +[[package]] +name = "azure_core" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e41cbd819986ba41904c207d8ffc4106f8f8352a548d773e9554906379bb2fb" +dependencies = [ + "async-lock", + "async-trait", + "azure_core_macros", + "bytes", + "futures", + "pin-project", + "rustc_version", + "serde", + "serde_json", + "tokio", + "tracing", + "typespec", + "typespec_client_core", +] + +[[package]] +name = "azure_core_macros" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9b52dba6a345f3ad2d42ff8d0d63df9d0994cfa29657bf18ffdbf149f78a4f5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "tracing", +] + +[[package]] +name = "azure_identity" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32edf96b356ca7c51d7590c4925cc36efc3947a5da4468e8e0b25c56ecbb3de5" +dependencies = [ + "async-lock", + "async-trait", + "azure_core", + "futures", + "pin-project", + "serde", + "serde_json", + "time", + "tokio", + "tracing", + "url", +] + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + [[package]] name = "base64" version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64-simd" version = "0.8.0" @@ -542,6 +658,15 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + [[package]] name = "cc" version = "1.3.0" @@ -577,6 +702,20 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", +] + [[package]] name = "ciborium" version = "0.2.2" @@ -644,6 +783,48 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" +[[package]] +name = "combine" +version = "4.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" +dependencies = [ + "bytes", + "memchr", +] + +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "compression-codecs" +version = "0.4.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58a6d0db8759036a783bc7c3f7a07f8cef3bf9470eb1db3bc86e8bcd1c5d0fe8" +dependencies = [ + "compression-core", + "flate2", + "memchr", +] + +[[package]] +name = "compression-core" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e8ccc4ea9f6acc32d102c0f6d471d11d913ad15f20c04de743374861fa1d414" + [[package]] name = "const-oid" version = "0.10.2" @@ -684,6 +865,15 @@ dependencies = [ "libc", ] +[[package]] +name = "crc32fast" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" +dependencies = [ + "cfg-if", +] + [[package]] name = "criterion" version = "0.8.2" @@ -696,7 +886,7 @@ dependencies = [ "ciborium", "clap", "criterion-plot", - "itertools", + "itertools 0.13.0", "num-traits", "oorandom", "page_size", @@ -716,7 +906,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" dependencies = [ "cast", - "itertools", + "itertools 0.13.0", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98b0cc327b5bc766e7fda9c9260cc0fa81b43a8e240440422dff70788e3f9ef1" +dependencies = [ + "crossbeam-utils", ] [[package]] @@ -778,17 +977,107 @@ dependencies = [ "cmov", ] +[[package]] +name = "daachorse" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5614204febbc33cc07a2806aa6440b904ac012b68eecc37f4493ea4a76455a3d" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.119", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + [[package]] name = "data-encoding" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +[[package]] +name = "data-url" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be1e0bca6c3637f992fc1cc7cbc52a78c1ef6db076dbf1059c4323d6a2048376" + [[package]] name = "deranged" version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "serde_core", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn 2.0.119", +] [[package]] name = "digest" @@ -829,6 +1118,12 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "either" version = "1.16.0" @@ -841,6 +1136,42 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + +[[package]] +name = "event-listener" +version = "5.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" +dependencies = [ + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + [[package]] name = "fastrand" version = "2.5.0" @@ -853,6 +1184,17 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "flate2" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" +dependencies = [ + "crc32fast", + "miniz_oxide", + "zlib-rs", +] + [[package]] name = "fnv" version = "1.0.7" @@ -874,6 +1216,21 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" +[[package]] +name = "futures" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.33" @@ -890,6 +1247,17 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" +[[package]] +name = "futures-executor" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6754879cc9f2c66f88c6e5c35344bb0bdb0708b0352b1201815667c7eabc7458" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.33" @@ -931,6 +1299,7 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -941,6 +1310,33 @@ dependencies = [ "slab", ] +[[package]] +name = "gcp_auth" +version = "0.12.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d27dbcc645b60b8e7f6e2868a9d7102ece97d1bb49c1288b5321fcc67f7260" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "chrono", + "http 1.4.2", + "http-body-util", + "hyper 1.10.1", + "hyper-rustls 0.27.9", + "hyper-util", + "ring", + "rustls 0.23.42", + "rustls-pki-types", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tracing", + "tracing-futures", + "url", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -964,6 +1360,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + [[package]] name = "getrandom" version = "0.4.3" @@ -973,7 +1381,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", + "r-efi 6.0.0", "rand_core 0.10.1", "wasm-bindgen", ] @@ -1220,7 +1628,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -1237,6 +1645,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "2.2.0" @@ -1319,6 +1751,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "idna" version = "1.1.0" @@ -1348,6 +1786,8 @@ checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", "hashbrown", + "serde", + "serde_core", ] [[package]] @@ -1365,12 +1805,70 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys", + "log", + "simd_cesu8", + "thiserror 2.0.19", + "walkdir", + "windows-link", +] + +[[package]] +name = "jni-macros" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "simd_cesu8", + "syn 2.0.119", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn 2.0.119", +] + [[package]] name = "jobserver" version = "0.1.35" @@ -1409,12 +1907,14 @@ name = "litellm-ai-gateway" version = "0.1.0" dependencies = [ "axum", - "base64", + "base64 0.22.1", "futures-channel", "futures-util", "litellm-config", "litellm-core", - "reqwest", + "reqwest 0.12.28", + "rustls 0.23.42", + "rustls-native-certs", "serde", "serde_json", "sha2 0.10.9", @@ -1445,17 +1945,27 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", - "base64", + "azure_core", + "azure_identity", + "base64 0.22.1", + "data-url", + "gcp_auth", + "moka", "rand 0.8.7", - "reqwest", + "reqwest 0.12.28", "rstest", "serde", "serde_json", + "serde_path_to_error", "sha2 0.10.9", + "strum", + "subtle", "thiserror 2.0.19", "tokio", "tracing", "tracing-subscriber", + "url", + "veil", ] [[package]] @@ -1467,6 +1977,7 @@ dependencies = [ "litellm-ai-gateway", "litellm-core", "litellm-python-interop", + "litellm-token-counter", "pyo3", "pyo3-async-runtimes", "serde", @@ -1487,12 +1998,37 @@ dependencies = [ "serde_json", ] +[[package]] +name = "litellm-token-counter" +version = "0.1.0" +dependencies = [ + "criterion", + "indexmap", + "itoa", + "rand 0.8.7", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokenizers", + "unicode-normalization-alignments", +] + [[package]] name = "litemap" version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.33" @@ -1505,6 +2041,22 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "macro_rules_attribute" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3ae8f6d608c795738406608304d30a2dfbdc8e58e44f7ba43236da5208ded3c" +dependencies = [ + "macro_rules_attribute-proc_macro", + "pastey", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c" + [[package]] name = "matchit" version = "0.7.3" @@ -1533,6 +2085,22 @@ dependencies = [ "unicase", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.2.2" @@ -1544,6 +2112,58 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "moka" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9" +dependencies = [ + "async-lock", + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "event-listener", + "futures-util", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -1574,6 +2194,28 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "oorandom" version = "11.1.5" @@ -1602,12 +2244,73 @@ dependencies = [ "winapi", ] +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + [[package]] name = "percent-encoding" version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -1809,6 +2512,7 @@ version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ + "aws-lc-rs", "bytes", "getrandom 0.4.3", "lru-slab", @@ -1848,6 +2552,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r-efi" version = "6.0.0" @@ -1861,10 +2571,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", - "rand_chacha", + "rand_chacha 0.3.1", "rand_core 0.6.4", ] +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.2" @@ -1886,6 +2606,16 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + [[package]] name = "rand_core" version = "0.6.4" @@ -1895,6 +2625,15 @@ dependencies = [ "getrandom 0.2.17", ] +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + [[package]] name = "rand_core" version = "0.10.1" @@ -1920,6 +2659,17 @@ dependencies = [ "rayon-core", ] +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools 0.14.0", + "rayon", +] + [[package]] name = "rayon-core" version = "1.13.0" @@ -1930,6 +2680,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + [[package]] name = "regex" version = "1.13.1" @@ -1977,7 +2736,7 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-core", @@ -2010,11 +2769,49 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams", + "wasm-streams 0.4.2", "web-sys", "webpki-roots", ] +[[package]] +name = "reqwest" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029" +dependencies = [ + "base64 0.23.1", + "bytes", + "futures-core", + "futures-util", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "hyper 1.10.1", + "hyper-rustls 0.27.9", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls 0.23.42", + "rustls-pki-types", + "rustls-platform-verifier", + "sync_wrapper", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams 0.5.0", + "web-sys", +] + [[package]] name = "ring" version = "0.17.14" @@ -2122,6 +2919,33 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-platform-verifier" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" +dependencies = [ + "core-foundation", + "core-foundation-sys", + "jni", + "log", + "once_cell", + "rustls 0.23.42", + "rustls-native-certs", + "rustls-platform-verifier-android", + "rustls-webpki 0.103.13", + "security-framework", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls-platform-verifier-android" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" + [[package]] name = "rustls-webpki" version = "0.101.7" @@ -2174,6 +2998,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "sct" version = "0.7.1" @@ -2327,6 +3157,38 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + +[[package]] +name = "simd_cesu8" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11031e251abf8611c80f460e19dbdeb54a66db918e49c65a7065b46ac7aec520" +dependencies = [ + "rustc_version", + "simdutf8", +] + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + [[package]] name = "slab" version = "0.4.12" @@ -2359,12 +3221,57 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom", + "serde", + "unicode-segmentation", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "strum" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "subtle" version = "2.6.1" @@ -2413,6 +3320,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "target-lexicon" version = "0.13.5" @@ -2533,6 +3446,39 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tokenizers" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7afbf6e88718afcc138bad01d6ccc3051dbbc3b2ce9793d8b8a3aeb610969cfc" +dependencies = [ + "ahash", + "compact_str", + "daachorse", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "itertools 0.14.0", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.9.5", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.19", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokio" version = "1.53.0" @@ -2543,6 +3489,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2 0.6.5", "tokio-macros", "windows-sys 0.61.2", @@ -2660,12 +3607,17 @@ version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ + "async-compression", "bitflags", "bytes", + "futures-core", "futures-util", "http 1.4.2", "http-body 1.1.0", + "http-body-util", "pin-project-lite", + "tokio", + "tokio-util", "tower", "tower-layer", "tower-service", @@ -2716,6 +3668,16 @@ dependencies = [ "once_cell", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "pin-project", + "tracing", +] + [[package]] name = "tracing-subscriber" version = "0.3.23" @@ -2759,6 +3721,57 @@ version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" +[[package]] +name = "typespec" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "753a2fe021e407d4fc9ee6f4f0a33403cc306d5c54c4e4ebe1b8cbde0ca052b9" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures", + "serde", + "serde_json", + "url", +] + +[[package]] +name = "typespec_client_core" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0373af0f9d4f580b3a1a9d9639cedaabe015ed262b35bfbe13941bfb14fe1ea6" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "dyn-clone", + "futures", + "pin-project", + "rand 0.10.2", + "reqwest 0.13.5", + "serde", + "serde_json", + "time", + "tokio", + "tracing", + "typespec", + "typespec_macros", + "url", + "uuid", +] + +[[package]] +name = "typespec_macros" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c608f4427943f8adb211abc95c87672b1b98847152783507d54e3246e502f60" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "syn 2.0.119", +] + [[package]] name = "unicase" version = "2.9.0" @@ -2771,6 +3784,27 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "untrusted" version = "0.9.0" @@ -2813,10 +3847,32 @@ version = "1.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" dependencies = [ + "getrandom 0.4.3", "js-sys", "wasm-bindgen", ] +[[package]] +name = "veil" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7352f0bbf3ab98911b0c0277065094c1b1ec79bbc85fa3b7d16bf1859c3d96f" +dependencies = [ + "once_cell", + "veil-macros", +] + +[[package]] +name = "veil-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47a3f4f06d904eb789b935253752ba6bcc1dfa61349f8d5341c66abe070b44e5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "version_check" version = "0.9.5" @@ -2854,6 +3910,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "wasm-bindgen" version = "0.2.126" @@ -2922,6 +3987,19 @@ dependencies = [ "web-sys", ] +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "web-sys" version = "0.3.103" @@ -2942,6 +4020,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-root-certs" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "webpki-roots" version = "1.0.9" @@ -2982,12 +4069,65 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.52.0" @@ -3079,6 +4219,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + [[package]] name = "writeable" version = "0.6.3" @@ -3194,6 +4340,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "zlib-rs" +version = "0.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34b31d188d9d685a4f9c7b46d6e36631b07058d2cfe190267adce54dc230bf12" + [[package]] name = "zmij" version = "1.0.23" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 720c4545181..5f25e69a1f8 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "crates/core", + "crates/token-counter", "crates/config", "crates/ai-gateway", "crates/python-interop", @@ -18,6 +19,7 @@ repository = "https://github.com/BerriAI/litellm" tracing = "0.1" tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } litellm-core = { path = "crates/core" } +litellm-token-counter = { path = "crates/token-counter" } litellm-config = { path = "crates/config" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } litellm-python-interop = { path = "crates/python-interop" } @@ -28,6 +30,8 @@ pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" +rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } +rustls-native-certs = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } sha2 = "0.10" @@ -37,6 +41,14 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" +gcp_auth = "0.12.7" +azure_core = "1.0.0" +azure_identity = { version = "1.0.0", features = ["tokio"] } +moka = { version = "0.12.16", features = ["future"] } +strum = { version = "0.28.0", features = ["derive"] } +url = "2.5.8" +criterion = "0.8.2" +veil = "0.3.0" [profile.release] opt-level = 3 diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 10369fa3bfd..74cf66e88a2 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -20,6 +20,10 @@ litellm-config.workspace = true # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. reqwest.workspace = true +# rustls and its root store are direct dependencies so `io::tls` can build the +# one TLS config the outbound dials use; see that module for why it has to. +rustls.workspace = true +rustls-native-certs.workspace = true # `sync` powers the bounded mpsc channel the realtime logger drains. tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] } tokio-tungstenite.workspace = true diff --git a/litellm-rust/crates/ai-gateway/Dockerfile b/litellm-rust/crates/ai-gateway/Dockerfile index 2bc3c05ad7e..72ac25ce1d6 100644 --- a/litellm-rust/crates/ai-gateway/Dockerfile +++ b/litellm-rust/crates/ai-gateway/Dockerfile @@ -14,15 +14,20 @@ # ---- Chef ------------------------------------------------------------------- # cargo-chef caches the dependency build so only the gateway crate recompiles on # a source-only change. python3-dev is present in every rust stage because the -# `python-config` feature links libpython via pyo3 (even in the cook step). -FROM rust:1.90-slim-bookworm AS chef +# `python-config` feature links libpython via pyo3 (even in the cook step), and +# python3-pip builds the litellm wheel in the builder stage. +FROM rust:1.98-slim-bookworm AS chef ENV PYO3_PYTHON=python3.11 +# rustup reads rust-toolchain.toml from any parent of the working directory, so +# copying it in is what keeps every cargo call below on the repo's pinned +# channel rather than on whatever the base image happens to ship. +COPY rust-toolchain.toml /build/rust-toolchain.toml +WORKDIR /build/litellm-rust RUN apt-get update \ && apt-get install -y --no-install-recommends \ - python3 python3-dev pkg-config libssl-dev clang \ + python3 python3-dev python3-pip pkg-config libssl-dev clang \ && rm -rf /var/lib/apt/lists/* \ && cargo install cargo-chef --locked --version 0.1.77 -WORKDIR /build/litellm-rust # ---- Planner ---------------------------------------------------------------- # Produce the dependency recipe from the rust workspace manifests + Cargo.lock. @@ -43,6 +48,19 @@ RUN cargo chef cook --locked --release \ COPY litellm-rust/ . RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config +# The root pyproject builds with maturin against litellm-rust/crates/python-bridge, +# so the wheel is built here, next to the crate sources and the cargo toolchain, +# and the runtime stage installs the artifact instead of compiling anything. +# litellm[proxy] pins litellm-enterprise and litellm-proxy-extras to the versions +# in this repo, and those hit PyPI hours after every version bump merges, so both +# wheels are built from the repo too instead of being resolved from PyPI. +COPY pyproject.toml README.md LICENSE /build/ +COPY litellm/ /build/litellm/ +COPY enterprise/ /build/enterprise/ +COPY litellm-proxy-extras/ /build/litellm-proxy-extras/ +RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist \ + /build /build/enterprise /build/litellm-proxy-extras + # ---- Runtime ---------------------------------------------------------------- # python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3 # 3.11 ABI so the embedded interpreter links and imports cleanly. @@ -56,11 +74,16 @@ RUN apt-get update \ WORKDIR /app # Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so -# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the -# package + packaging metadata, then pip install the proxy extra. -COPY pyproject.toml README.md LICENSE ./ -COPY litellm/ ./litellm/ -RUN pip install --no-cache-dir ".[proxy]" +# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. The two +# sibling wheels come from the builder as well, so the pins in litellm[proxy] +# resolve against them and never wait on a PyPI publish. +COPY --from=builder /build/dist/*.whl /tmp/wheels/ +RUN wheel="$(ls /tmp/wheels/litellm-*.whl)" \ + && pip install --no-cache-dir \ + /tmp/wheels/litellm_enterprise-*.whl \ + /tmp/wheels/litellm_proxy_extras-*.whl \ + "${wheel}[proxy]" \ + && rm -rf /tmp/wheels # The compiled gateway binary (pure-Rust realtime hot path; Python is load-time # only). diff --git a/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore index 030ee6a37c5..d1386ff684d 100644 --- a/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore +++ b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore @@ -9,19 +9,28 @@ # Strategy: ignore everything, then re-include only what the build needs: # - litellm/ (pip install . needs the full package + proxy reader) # - litellm-rust/ (the rust workspace; Cargo.lock + crate sources) -# - pyproject.toml / README.md / LICENSE (packaging metadata for pip install) +# - enterprise/ (litellm/proxy/enterprise symlinks into it; maturin walks it) +# - litellm-proxy-extras/ (built into a wheel alongside enterprise/ for litellm[proxy]) +# - pyproject.toml / README.md / LICENSE (packaging metadata for the wheel build) +# - rust-toolchain.toml (the pinned channel every cargo call in the build uses) * # --- re-include the build inputs --- !litellm/ !litellm-rust/ +!enterprise/ +!litellm-proxy-extras/ !pyproject.toml +!rust-toolchain.toml !README.md !LICENSE # --- prune heavy / irrelevant subpaths back out of the re-included trees --- # Rust build artifacts (huge; regenerated in the builder). **/target/ +# Committed python distribution artifacts; the wheel build does not read them. +enterprise/dist/ +litellm-proxy-extras/dist/ # Python caches and compiled bytecode. **/__pycache__/ **/*.pyc diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 9fef59a277d..cbcd8119546 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -6,17 +6,18 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. ## Crates -`litellm-rust` has five crates. A crate is a layer or shared foundation, not a route: +`litellm-rust` has six crates. A crate is a layer or shared foundation, not a route: | Crate | Role | |-------|------| | litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. | +| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. | | litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. | | litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | | litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. | -Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop. +Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers, token counter, and Python interop. - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index dbe2d3a325b..6f48f38c9f6 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -269,7 +269,11 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error { fn core_error_kind(error: &Error) -> &'static str { match error { - Error::Auth(_) => "AuthError", + Error::Auth(_) + | Error::MissingApiKey { .. } + | Error::MissingAzureAiCredentials + | Error::MissingAzureDocumentIntelligenceCredentials + | Error::MissingReductoApiKey => "AuthError", Error::InvalidProvider(_) => "InvalidProvider", Error::InvalidRequest(_) => "InvalidRequest", Error::InvalidType { .. } => "InvalidType", diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index cce56dd2121..7098d67993f 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -3,3 +3,4 @@ pub mod ocr; pub mod realtime; pub mod realtime_pool; pub mod responses_ws; +pub(crate) mod tls; diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 662f7328982..1aa31adcc38 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -15,6 +15,8 @@ use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{Sink, SinkExt, Stream, StreamExt}; +use litellm_core::AuthError; +use litellm_core::auth::error::MissingCredential; use litellm_core::error::Error; use litellm_core::realtime::transformation::RealtimeProviderConfig; use litellm_core::realtime::types::RealtimeEvent; @@ -23,15 +25,15 @@ use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; +use crate::io::tls::connect_upstream; + /// Environment variable holding the OpenAI API key (last-resort fallback). const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; - /// Default **idle** timeout: if neither side sends a frame for this long, the /// session is reaped. It resets on any activity, so it does not cap a healthy /// (continuously streaming) session — it only frees a stalled one (e.g. a @@ -57,7 +59,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { .ok() .filter(|key| !key.trim().is_empty()) }) - .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiRealtimeApiKey))) } /// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`. @@ -84,7 +86,7 @@ pub(crate) async fn dial_upstream( .map_err(|err| Error::Auth(err.to_string()))?, ); - let (upstream, _response) = connect_async(request) + let (upstream, _response) = connect_upstream(request) .await .map_err(|err| Error::Network(err.to_string()))?; Ok(upstream) @@ -284,6 +286,33 @@ mod tests { serde_json::from_str(raw).expect("valid event json") } + /// The realtime dial has to reach a `wss://` upstream without a process-wide + /// crypto provider installed, which is what dialing through `io::tls` buys. + #[tokio::test] + async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + let result = dial_upstream( + "gpt-realtime", + "sk-test", + Some(&format!("wss://127.0.0.1:{port}")), + ) + .await; + + assert!(matches!(result, Err(Error::Network(_)))); + } + #[test] fn resolve_api_key_prefers_param_then_blank_falls_through() { assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test"); diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 0b01747b1a5..7f3b6b0650f 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -4,7 +4,9 @@ use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{Sink, SinkExt, Stream, StreamExt}; +use litellm_core::AuthError; use litellm_core::Error; +use litellm_core::auth::error::MissingCredential; use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG; use litellm_core::responses::types::ResponsesWsEvent; use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig; @@ -14,15 +16,15 @@ use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName}; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +use crate::io::tls::connect_upstream; use crate::constants::{ DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS, }; const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; - pub type ResponsesUpstreamWs = WebSocketStream>; type UpstreamTx = SplitSink; type UpstreamRx = SplitStream; @@ -49,14 +51,14 @@ impl ResponsesWebSocketConnection { .map_err(|error| Error::InvalidRequest(error.to_string()))?; request.headers_mut().insert(header_name, header_value); } - let connect = connect_async(request); + let connect = connect_upstream(request); let result = match timeout { Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { Error::Network("Responses WebSocket connection timed out".to_string()) })?, None => connect.await, }; - let (socket, _) = result.map_err(|error| match error { + let (socket, _) = result.map_err(|error| match *error { tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), @@ -118,7 +120,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { .ok() .filter(|value| !value.trim().is_empty()) }) - .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiResponsesApiKey))) } async fn dial_upstream( @@ -138,13 +140,13 @@ async fn dial_upstream( ); let result = tokio::time::timeout( Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS), - connect_async(request), + connect_upstream(request), ) .await .map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?; result .map(|(socket, _)| socket) - .map_err(|error| match error { + .map_err(|error| match *error { tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), @@ -324,6 +326,29 @@ mod tests { use tokio::net::TcpListener; use tokio_tungstenite::accept_async; + /// The Responses dial has to reach a `wss://` upstream without a process-wide + /// crypto provider installed, which is what dialing through `io::tls` buys. + #[tokio::test] + async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + let result = + dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await; + + assert!(matches!(result, Err(Error::Network(_)))); + } + async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); let address = listener.local_addr().expect("local address"); diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs new file mode 100644 index 00000000000..a2562f60345 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs @@ -0,0 +1,80 @@ +//! Outbound WebSocket dials over a TLS config this crate builds once and owns. +//! +//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth` +//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that +//! `tokio-tungstenite` uses when handed no connector panics rather than guess +//! between them. Naming ring on a connector of our own settles that for these +//! dials without touching the process-wide default, and building the config +//! once keeps the platform trust store, which `tokio-tungstenite` would +//! otherwise re-read on every dial, off the dial path. + +use std::io; +use std::sync::{Arc, OnceLock}; + +use rustls::{ClientConfig, RootCertStore}; +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::Error; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::error::TlsError; +use tokio_tungstenite::tungstenite::handshake::client::Response; +use tokio_tungstenite::{ + Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, +}; + +static TLS_CONFIG: OnceLock> = OnceLock::new(); + +fn build_config() -> Result> { + let native = rustls_native_certs::load_native_certs(); + let roots = { + let mut store = RootCertStore::empty(); + let (added, _ignored) = store.add_parsable_certificates(native.certs); + if added == 0 { + return Err(Box::new(Error::Io(io::Error::other(format!( + "no usable native root certificates: {:?}", + native.errors + ))))); + } + store + }; + + ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) + .with_safe_default_protocol_versions() + .map(|builder| builder.with_root_certificates(roots).with_no_client_auth()) + .map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error)))) +} + +fn tls_config() -> Result, Box> { + if let Some(config) = TLS_CONFIG.get() { + return Ok(Arc::clone(config)); + } + let built = Arc::new(build_config()?); + Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) +} + +pub(crate) async fn connect_upstream( + request: R, +) -> Result<(WebSocketStream>, Response), Box> +where + R: IntoClientRequest + Unpin, +{ + let request = request.into_client_request().map_err(Box::new)?; + let connector = match request.uri().scheme_str() { + Some("wss") => Some(Connector::Rustls(tls_config()?)), + _ => None, + }; + connect_async_tls_with_config(request, None, false, connector) + .await + .map_err(Box::new) +} + +#[cfg(test)] +mod tests { + use super::build_config; + + #[test] + fn builds_a_usable_config_with_both_provider_features_enabled() { + let config = build_config().expect("a client config"); + + assert!(!config.crypto_provider().cipher_suites.is_empty()); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs deleted file mode 100644 index d2be17260a3..00000000000 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ /dev/null @@ -1,529 +0,0 @@ -use std::net::IpAddr; -use std::time::{Duration, Instant}; - -use base64::Engine; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use litellm_core::error::Error; -use litellm_core::ocr::transformation::OcrProviderConfig; -use reqwest::Url; -use serde_json::{Map, Value}; - -use litellm_core::providers::azure_ai::ocr::transformation::{ - AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG, -}; -use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; -use litellm_core::providers::reducto::ocr::transformation as reducto; -use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai; -use litellm_core::providers::vertex_ai::ocr::transformation::{ - VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, -}; - -use crate::client::http_client; - -const ERROR_BODY_MAX_CHARS: usize = 256; -const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; -const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0; -const MAX_SAFE_FETCH_REDIRECTS: usize = 10; - -pub(super) fn truncate_error_body(body: &str) -> String { - if body.chars().count() <= ERROR_BODY_MAX_CHARS { - return body.to_string(); - } - let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect(); - format!("{truncated}... (truncated)") -} - -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub(super) fn ocr_provider_config( - provider: &str, - model: &str, -) -> Option<&'static dyn OcrProviderConfig> { - match provider { - "mistral" => Some(&MISTRAL_OCR_CONFIG), - "reducto" => reducto::config_for_model(model), - "azure_ai" if is_azure_document_intelligence_model(model) => { - Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG) - } - "azure_ai" => Some(&AZURE_AI_OCR_CONFIG), - "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG), - "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG), - _ => None, - } -} - -fn is_azure_document_intelligence_model(model: &str) -> bool { - let model = model.to_ascii_lowercase(); - model.contains("doc-intelligence") || model.contains("documentintelligence") -} - -pub(super) fn string_headers( - extra_headers: Option>, -) -> Result, Error> { - extra_headers - .unwrap_or_default() - .into_iter() - .map(|(key, value)| { - value - .as_str() - .map(|value| (key.clone(), value.to_string())) - .ok_or_else(|| { - Error::InvalidRequest(format!( - "OCR extra_headers.{key} must be a string, got {}", - litellm_core::error::json_type_name(&value) - )) - }) - }) - .collect() -} - -fn document_url_field(document: &Value) -> Result, Error> { - let Some(object) = document.as_object() else { - return Ok(None); - }; - let Some(doc_type) = object.get("type").and_then(Value::as_str) else { - return Ok(None); - }; - let field = match doc_type { - "document_url" => "document_url", - "image_url" => "image_url", - _ => return Ok(None), - }; - let Some(url) = object.get(field).and_then(Value::as_str) else { - return Ok(None); - }; - Ok(Some((field, url))) -} - -fn is_url_requiring_fetch(url: &str) -> bool { - !url.starts_with("data:") && (url.starts_with("http://") || url.starts_with("https://")) -} - -fn max_document_download_bytes() -> u64 { - let max_size_mb = std::env::var("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB") - .ok() - .and_then(|value| value.parse::().ok()) - .unwrap_or(DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB); - (max_size_mb.max(0.0) * 1024.0 * 1024.0) as u64 -} - -fn is_blocked_ip(ip: IpAddr) -> bool { - match ip { - IpAddr::V4(ip) => { - ip.is_private() - || ip.is_loopback() - || ip.is_link_local() - || ip.is_broadcast() - || ip.is_multicast() - || ip.is_unspecified() - } - IpAddr::V6(ip) => { - let first_segment = ip.segments()[0]; - let is_unique_local = (first_segment & 0xfe00) == 0xfc00; - let is_link_local = (first_segment & 0xffc0) == 0xfe80; - ip.is_loopback() - || ip.is_unspecified() - || ip.is_multicast() - || is_unique_local - || is_link_local - || ip - .to_ipv4_mapped() - .or_else(|| ip.to_ipv4()) - .map(|v4| is_blocked_ip(IpAddr::V4(v4))) - .unwrap_or(false) - } - } -} - -fn blocked_url_error(url: &Url) -> Error { - Error::InvalidRequest(format!( - "OCR document URL rejected by SSRF protection: {url}" - )) -} - -async fn validate_safe_fetch_url(url: &Url) -> Result<(), Error> { - if !matches!(url.scheme(), "http" | "https") { - return Err(blocked_url_error(url)); - } - - let host = url.host_str().ok_or_else(|| blocked_url_error(url))?; - if let Ok(ip) = host.parse::() { - if is_blocked_ip(ip) { - return Err(blocked_url_error(url)); - } - return Ok(()); - } - - let port = url - .port_or_known_default() - .ok_or_else(|| blocked_url_error(url))?; - let addresses = tokio::net::lookup_host((host, port)) - .await - .map_err(|err| Error::Network(err.to_string()))?; - let mut saw_address = false; - for address in addresses { - saw_address = true; - if is_blocked_ip(address.ip()) { - return Err(blocked_url_error(url)); - } - } - if !saw_address { - return Err(blocked_url_error(url)); - } - Ok(()) -} - -fn redirect_location(response: &reqwest::Response, url: &Url) -> Result { - let location = response - .headers() - .get(reqwest::header::LOCATION) - .and_then(|value| value.to_str().ok()) - .ok_or_else(|| { - Error::InvalidResponse("OCR document redirect missing Location header".to_string()) - })?; - url.join(location) - .map_err(|err| Error::InvalidResponse(format!("invalid OCR document redirect: {err}"))) -} - -async fn safe_get_document_url(url: &str) -> Result<(Url, reqwest::Response), Error> { - let client = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .map_err(|err| Error::Network(err.to_string()))?; - let mut current_url = Url::parse(url) - .map_err(|err| Error::InvalidRequest(format!("invalid OCR document URL: {err}")))?; - - for _ in 0..MAX_SAFE_FETCH_REDIRECTS { - validate_safe_fetch_url(¤t_url).await?; - let response = client - .get(current_url.clone()) - .send() - .await - .map_err(|err| Error::Network(err.to_string()))?; - if !response.status().is_redirection() { - return Ok((current_url, response)); - } - current_url = redirect_location(&response, ¤t_url)?; - } - - Err(Error::InvalidRequest( - "Too many redirects while fetching OCR document URL".to_string(), - )) -} - -fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Result<(), Error> { - if max_bytes == 0 { - return Err(Error::InvalidRequest(format!( - "OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}" - ))); - } - if content_length > max_bytes { - let size_mb = content_length as f64 / (1024.0 * 1024.0); - let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0); - return Err(Error::InvalidRequest(format!( - "OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}" - ))); - } - Ok(()) -} - -async fn read_response_with_limit( - mut response: reqwest::Response, - url: &Url, -) -> Result, Error> { - let max_bytes = max_document_download_bytes(); - if let Some(content_length) = response.content_length() { - enforce_download_size(content_length, max_bytes, url)?; - } else { - enforce_download_size(0, max_bytes, url)?; - } - - let mut bytes = Vec::new(); - let mut bytes_downloaded: u64 = 0; - while let Some(chunk) = response - .chunk() - .await - .map_err(|err| Error::Network(err.to_string()))? - { - bytes_downloaded += chunk.len() as u64; - enforce_download_size(bytes_downloaded, max_bytes, url)?; - bytes.extend_from_slice(&chunk); - } - Ok(bytes) -} - -pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result { - let Some((field, url)) = document_url_field(&document)? else { - return Ok(document); - }; - if !is_url_requiring_fetch(url) { - return Ok(document); - } - - let (final_url, response) = safe_get_document_url(url).await?; - let status = response.status(); - if !status.is_success() { - let body = response.text().await.unwrap_or_default(); - return Err(Error::Http { - status: status.as_u16(), - body: truncate_error_body(&body), - }); - } - let content_type = response - .headers() - .get(reqwest::header::CONTENT_TYPE) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.split(';').next()) - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or("application/octet-stream") - .to_string(); - let bytes = read_response_with_limit(response, &final_url).await?; - let data_uri = format!( - "data:{content_type};base64,{}", - BASE64_STANDARD.encode(bytes) - ); - - let mut transformed = document - .as_object() - .cloned() - .ok_or_else(|| Error::InvalidRequest("OCR document must be an object".to_string()))?; - transformed.insert(field.to_string(), Value::String(data_uri)); - Ok(Value::Object(transformed)) -} - -fn same_origin(left: &str, right: &str) -> bool { - let Ok(left) = reqwest::Url::parse(left) else { - return false; - }; - let Ok(right) = reqwest::Url::parse(right) else { - return false; - }; - left.scheme() == right.scheme() - && left.host_str() == right.host_str() - && left.port_or_known_default() == right.port_or_known_default() -} - -fn retry_after_secs(response: &reqwest::Response) -> u64 { - response - .headers() - .get(reqwest::header::RETRY_AFTER) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.parse::().ok()) - .unwrap_or(2) -} - -fn operation_status(response_json: &Value) -> Result<&str, Error> { - let status = response_json - .get("status") - .and_then(Value::as_str) - .ok_or(Error::MissingField("status"))?; - match status { - "succeeded" => Ok("succeeded"), - "running" | "notStarted" => Ok("running"), - "failed" => { - let message = response_json - .get("error") - .and_then(|error| error.get("message")) - .and_then(Value::as_str) - .unwrap_or("Unknown error"); - Err(Error::InvalidResponse(format!( - "Azure Document Intelligence analysis failed: {message}" - ))) - } - other => Err(Error::InvalidResponse(format!( - "Unknown operation status: {other}" - ))), - } -} - -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub(super) async fn poll_document_intelligence( - operation_url: &str, - original_url: &str, - headers: &[(String, String)], - timeout: Option, -) -> Result { - if !same_origin(operation_url, original_url) { - return Err(Error::InvalidResponse( - "Azure Document Intelligence: rejected cross-origin polling URL".to_string(), - )); - } - - let start = Instant::now(); - let timeout = timeout.unwrap_or(Duration::from_secs( - AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS, - )); - loop { - if start.elapsed() > timeout { - return Err(Error::Network(format!( - "Azure Document Intelligence operation polling timed out after {} seconds", - timeout.as_secs() - ))); - } - - let mut request_builder = http_client().get(operation_url); - for (key, value) in headers { - if key.eq_ignore_ascii_case("ocp-apim-subscription-key") { - request_builder = request_builder.header(key, value); - } - } - let response = request_builder - .send() - .await - .map_err(|err| Error::Network(err.to_string()))?; - let retry_after = retry_after_secs(&response); - let status = response.status(); - let text = response - .text() - .await - .map_err(|err| Error::Network(err.to_string()))?; - if !status.is_success() { - return Err(Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }); - } - let response_json: Value = serde_json::from_str(&text).map_err(|err| { - Error::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}")) - })?; - if operation_status(&response_json)? == "succeeded" { - return Ok(response_json); - } - tokio::time::sleep(Duration::from_secs(retry_after)).await; - } -} - -#[cfg(test)] -mod tests { - use litellm_core::ocr::transformation::OcrResponseHandling; - use serde_json::json; - - use super::*; - - #[test] - fn blocks_private_and_metadata_ips() { - assert!(is_blocked_ip("127.0.0.1".parse().unwrap())); - assert!(is_blocked_ip("10.0.0.1".parse().unwrap())); - assert!(is_blocked_ip("169.254.169.254".parse().unwrap())); - assert!(is_blocked_ip("::1".parse().unwrap())); - assert!(is_blocked_ip("fd00::1".parse().unwrap())); - assert!(is_blocked_ip("fe80::1".parse().unwrap())); - assert!(is_blocked_ip("::ffff:169.254.169.254".parse().unwrap())); - assert!(is_blocked_ip("::ffff:10.0.0.1".parse().unwrap())); - assert!(!is_blocked_ip("8.8.8.8".parse().unwrap())); - assert!(!is_blocked_ip("::ffff:8.8.8.8".parse().unwrap())); - } - - #[tokio::test] - async fn convert_document_url_rejects_loopback_fetch() { - let error = convert_document_url_to_data_uri(json!({ - "type": "image_url", - "image_url": "http://127.0.0.1/image.png" - })) - .await - .unwrap_err(); - - assert!(matches!( - error, - Error::InvalidRequest(message) - if message.contains("SSRF protection") - )); - } - - #[tokio::test] - async fn convert_document_url_leaves_data_uri_untouched() { - let document = json!({ - "type": "image_url", - "image_url": "data:image/png;base64,abcd" - }); - - let transformed = convert_document_url_to_data_uri(document.clone()) - .await - .unwrap(); - - assert_eq!(transformed, document); - } - - #[test] - fn truncate_error_body_passes_short_strings_through() { - let body = "Unauthorized"; - assert_eq!(truncate_error_body(body), "Unauthorized"); - } - - #[test] - fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(306); - let truncated = truncate_error_body(&body); - - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); - } - - #[test] - fn truncate_error_body_does_not_split_multibyte_chars() { - let body = "é".repeat(266); - let truncated = truncate_error_body(&body); - assert!(truncated.is_char_boundary(truncated.len())); - } - - #[test] - fn ocr_dispatch_supports_migrated_providers() { - assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); - assert!( - ocr_provider_config("azure_ai", "pixtral-12b-2409") - .expect("azure ai config resolves") - .requires_data_uri_document() - ); - assert_eq!( - ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") - .expect("document intelligence config resolves") - .response_handling(), - OcrResponseHandling::AzureDocumentIntelligencePoll - ); - assert!( - ocr_provider_config("vertex_ai", "deepseek-ocr-maas") - .expect("vertex deepseek config resolves") - .supported_ocr_params() - .contains(&"temperature") - ); - assert!(ocr_provider_config("openai", "gpt-4o").is_none()); - } - - #[test] - fn string_headers_accepts_string_values() { - let headers = json!({ - "x-trace-id": "trace-1" - }) - .as_object() - .unwrap() - .clone(); - - assert_eq!( - string_headers(Some(headers)).expect("string headers accepted"), - vec![("x-trace-id".to_string(), "trace-1".to_string())] - ); - } - - #[test] - fn string_headers_rejects_non_string_values() { - let headers = json!({ - "x-retry-count": 3 - }) - .as_object() - .unwrap() - .clone(); - - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert_eq!( - err, - Error::InvalidRequest( - "OCR extra_headers.x-retry-count must be a string, got number".to_string() - ) - ); - } -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs deleted file mode 100644 index 6c6e12724cd..00000000000 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ /dev/null @@ -1,84 +0,0 @@ -use litellm_core::error::Error; -use litellm_core::http_utils::http_request; -use litellm_core::ocr::transformation::OcrResponseHandling; -use serde_json::Value; - -use super::common_utils::{poll_document_intelligence, truncate_error_body}; -use super::hooks::OcrLifecycleHooks; -use super::types::PreparedOcrRequest; -use crate::client::http_client; - -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub(crate) async fn execute_ocr_provider_call( - request: PreparedOcrRequest, - hooks: &OcrLifecycleHooks, -) -> Result { - let request = hooks.prepare_provider_request(request).await?; - let mut request_builder = http_client().post(&request.url).json(&request.body); - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = http_request(request_builder) - .await - .map_err(|err| Error::Network(err.to_string()))?; - - let status = response.status(); - if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll - && status.as_u16() == 202 - { - let operation_url = response - .headers() - .get("operation-location") - .and_then(|value| value.to_str().ok()) - .map(str::to_string) - .ok_or_else(|| { - Error::InvalidResponse( - "Azure Document Intelligence returned 202 but no Operation-Location header found" - .to_string(), - ) - })?; - let response_json = poll_document_intelligence( - &operation_url, - &request.url, - &request.upstream_headers, - request.timeout, - ) - .await?; - return Ok(request - .config - .transform_ocr_response_with_params( - &request.model, - response_json, - &request.optional_params, - )? - .into_json()); - } - - let text = response - .text() - .await - .map_err(|err| Error::Network(err.to_string()))?; - - if !status.is_success() { - return Err(Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }); - } - - let response_json: Value = serde_json::from_str(&text) - .map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; - - Ok(request - .config - .transform_ocr_response_with_params( - &request.model, - response_json, - &request.optional_params, - )? - .into_json()) -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs deleted file mode 100644 index 446b323db3a..00000000000 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ /dev/null @@ -1,389 +0,0 @@ -use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; -use litellm_core::error::Error; -use litellm_core::providers::reducto::ocr::transformation::{ - build_upload_request, extract_document_source, extract_upload_file_id, -}; -use serde_json::{Map, Value, json}; -use std::future::Future; -use std::pin::Pin; - -use super::common_utils::{convert_document_url_to_data_uri, string_headers, truncate_error_body}; -use super::types::{PreparedOcrRequest, ProviderOcrRequest}; -use crate::client::http_client; -use crate::integrations::custom_guardrail::{ - CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, -}; -use crate::integrations::custom_logger::{ - CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails, -}; -use crate::integrations::types::{ - RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, -}; - -pub(crate) struct OcrLifecycleHooks { - logger_runner: CustomLoggerRunner, - guardrail_runner: CustomGuardrailRunner, - request_metadata: RequestMetadata, -} - -type OcrFuture<'a, T> = Pin> + Send + 'a>>; -type OcrLogFuture<'a> = Pin + Send + 'a>>; - -impl OcrLifecycleHooks { - pub(crate) fn new( - logger_runner: CustomLoggerRunner, - guardrail_runner: CustomGuardrailRunner, - request_metadata: RequestMetadata, - ) -> Self { - Self { - logger_runner, - guardrail_runner, - request_metadata, - } - } - - async fn run_pre_call_guardrails( - &self, - request: PreparedOcrRequest, - ) -> Result { - if self.guardrail_runner.is_empty() { - return Ok(request); - } - - let context = guardrail_context(&self.request_metadata); - let guardrail_request = GuardrailRequest::new(json!({ - "model": request.model, - "custom_llm_provider": request.custom_llm_provider, - "document": request.document, - "optional_params": request.optional_params, - })); - let (guardrail_request, _) = self - .guardrail_runner - .run_pre_call(&context, guardrail_request) - .await - .map_err(guardrail_error_to_core_error)?; - let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?; - let optional_params = match &request.config { - Ok(config) => config.map_ocr_params(&optional_params), - Err(_) => optional_params, - }; - Ok(PreparedOcrRequest { - document, - optional_params, - ..request - }) - } - - pub(crate) async fn prepare_provider_request( - &self, - request: PreparedOcrRequest, - ) -> Result { - let config = request.config?; - let env_lookup = |key: &str| std::env::var(key).ok(); - let upstream_headers = config.validate_environment( - string_headers(request.extra_headers)?, - request.api_key.as_deref(), - &env_lookup, - )?; - let url = config.complete_url( - request.api_base.as_deref(), - &request.model, - &request.optional_params, - &env_lookup, - )?; - let model = request.model.clone(); - let custom_llm_provider = request.custom_llm_provider.clone(); - let is_reducto = custom_llm_provider == "reducto"; - let document = if is_reducto { - let guarded_document = self - .run_during_call_guardrails(&model, &custom_llm_provider, &url, request.document) - .await?; - upload_reducto_document( - &guarded_document, - request.api_base.as_deref(), - request.timeout, - &upstream_headers, - ) - .await? - } else if config.requires_data_uri_document() { - convert_document_url_to_data_uri(request.document).await? - } else { - request.document - }; - let optional_params = request.optional_params; - let body = config - .transform_ocr_request(&request.model, document, optional_params.clone())? - .data; - let body = if is_reducto { - body - } else { - self.run_during_call_guardrails(&model, &custom_llm_provider, &url, body) - .await? - }; - Ok(ProviderOcrRequest { - model, - config, - url, - body, - optional_params, - upstream_headers, - timeout: request.timeout, - }) - } - - async fn run_during_call_guardrails( - &self, - model: &str, - custom_llm_provider: &str, - url: &str, - body: Value, - ) -> Result { - if self.guardrail_runner.is_empty() { - return Ok(body); - } - - let context = guardrail_context(&self.request_metadata); - let guardrail_request = GuardrailRequest::new(json!({ - "model": model, - "custom_llm_provider": custom_llm_provider, - "url": url, - "body": body, - })); - let (guardrail_request, _) = self - .guardrail_runner - .run_during_call(&context, guardrail_request) - .await - .map_err(guardrail_error_to_core_error)?; - parse_ocr_during_call_guardrail_request(guardrail_request) - } - - fn standard_logging_payload( - &self, - context: &CallLifecycleContext, - timing: &CallLifecycleTiming, - ) -> StandardLoggingPayload { - StandardLoggingPayload { - id: context.litellm_call_id.clone(), - litellm_call_id: context.litellm_call_id.clone(), - call_type: context.call_type.clone(), - model: context.model.clone(), - custom_llm_provider: context.custom_llm_provider.clone(), - response_cost: 0.0, - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - start_time: timing.start_time, - end_time: timing.end_time, - stream: false, - metadata: StandardLoggingMetadata { - user_api_key_hash: self.request_metadata.user_api_key_hash.clone(), - user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(), - user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(), - ..Default::default() - }, - messages: None, - } - } -} - -async fn upload_reducto_document( - document: &Value, - api_base: Option<&str>, - timeout: Option, - upstream_headers: &[(String, String)], -) -> Result { - let source = extract_document_source(document)?; - let Some(authorization) = upstream_headers - .iter() - .find(|(name, _)| name.eq_ignore_ascii_case("authorization")) - .map(|(_, value)| value.as_str()) - else { - return Err(Error::Auth( - "Reducto upload requires an Authorization header".to_string(), - )); - }; - let Some(upload) = build_upload_request(source, authorization, api_base) else { - return Ok(document.clone()); - }; - let part = reqwest::multipart::Part::bytes(upload.bytes) - .file_name(upload.file_name) - .mime_str(&upload.mime_type) - .map_err(|error| Error::InvalidRequest(error.to_string()))?; - let form = reqwest::multipart::Form::new().part("file", part); - let mut request_builder = http_client().post(upload.url).multipart(form); - for (name, value) in upstream_headers { - if !name.eq_ignore_ascii_case("content-type") - && !name.eq_ignore_ascii_case("content-length") - { - request_builder = request_builder.header(name, value); - } - } - if let Some(timeout) = timeout { - request_builder = request_builder.timeout(timeout); - } - let response = request_builder - .send() - .await - .map_err(|error| Error::Network(error.to_string()))?; - let status = response.status(); - let body = response - .text() - .await - .map_err(|error| Error::Network(error.to_string()))?; - if !status.is_success() { - return Err(Error::Http { - status: status.as_u16(), - body: truncate_error_body(&body), - }); - } - let response_json: Value = serde_json::from_str(&body).map_err(|error| { - Error::InvalidResponse(format!("invalid Reducto upload response JSON: {error}")) - })?; - let file_id = extract_upload_file_id(&response_json)?; - Ok(json!({"type": "document_url", "document_url": file_id})) -} - -impl CallLifecycleHooks for OcrLifecycleHooks { - type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; - type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; - type SuccessFuture<'a> = OcrLogFuture<'a>; - type FailureFuture<'a> = OcrLogFuture<'a>; - - fn async_pre_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: PreparedOcrRequest, - ) -> Self::PreCallFuture<'a> { - Box::pin(async move { self.run_pre_call_guardrails(request).await }) - } - - fn async_during_call_hook<'a>( - &'a self, - _context: &'a CallLifecycleContext, - request: PreparedOcrRequest, - ) -> Self::DuringCallFuture<'a> { - Box::pin(async move { Ok(request) }) - } - - fn async_log_success_event<'a>( - &'a self, - context: &'a CallLifecycleContext, - response: &'a Value, - timing: &'a CallLifecycleTiming, - ) -> Self::SuccessFuture<'a> { - Box::pin(async move { - if self.logger_runner.is_empty() { - return; - } - let response_obj = CallbackValue::new("ocr", response.clone()); - self.logger_runner - .async_log_success_event( - &ModelCallDetails::from_standard_logging_payload( - self.standard_logging_payload(context, timing), - ), - &response_obj, - CallbackTiming::new(timing.start_time, timing.end_time), - ) - .await; - }) - } - - fn async_log_failure_event<'a>( - &'a self, - context: &'a CallLifecycleContext, - error: &'a Error, - timing: &'a CallLifecycleTiming, - ) -> Self::FailureFuture<'a> { - Box::pin(async move { - if self.logger_runner.is_empty() { - return; - } - let logging_error = LoggingError { - message: error.to_string(), - kind: core_error_kind(error).to_string(), - }; - let response_obj = CallbackValue::new( - "error", - json!({ - "message": logging_error.message, - "kind": logging_error.kind, - }), - ); - self.logger_runner - .async_log_failure_event( - &ModelCallDetails::from_standard_logging_payload( - self.standard_logging_payload(context, timing), - ) - .with_failure_error(logging_error), - Some(&response_obj), - CallbackTiming::new(timing.start_time, timing.end_time), - ) - .await; - }) - } -} - -fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { - GuardrailContext { - call_type: CallType::Ocr, - selected_guardrails: Vec::new(), - metadata: std::collections::HashMap::new(), - user_api_key_hash: metadata.user_api_key_hash.clone(), - user_api_key_user_id: metadata.user_api_key_user_id.clone(), - user_api_key_team_id: metadata.user_api_key_team_id.clone(), - trace_parent: None, - } -} - -fn parse_ocr_pre_call_guardrail_request( - request: GuardrailRequest, -) -> Result<(Value, Map), Error> { - let Value::Object(mut data) = request.data else { - return Err(Error::InvalidRequest( - "OCR pre_call guardrail must return an object".to_string(), - )); - }; - let document = data.remove("document").ok_or_else(|| { - Error::InvalidRequest("OCR pre_call guardrail removed document".to_string()) - })?; - let optional_params = match data.remove("optional_params") { - Some(Value::Object(params)) => params, - Some(_) => { - return Err(Error::InvalidRequest( - "OCR pre_call guardrail optional_params must be an object".to_string(), - )); - } - None => Map::new(), - }; - Ok((document, optional_params)) -} - -fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result { - let Value::Object(mut data) = request.data else { - return Err(Error::InvalidRequest( - "OCR during_call guardrail must return an object".to_string(), - )); - }; - data.remove("body") - .ok_or_else(|| Error::InvalidRequest("OCR during_call guardrail removed body".to_string())) -} - -fn guardrail_error_to_core_error(error: GuardrailError) -> Error { - Error::InvalidRequest(format!("{}: {}", error.kind, error.message)) -} - -fn core_error_kind(error: &Error) -> &'static str { - match error { - Error::Auth(_) => "AuthError", - Error::InvalidProvider(_) => "InvalidProvider", - Error::InvalidRequest(_) => "InvalidRequest", - Error::InvalidType { .. } => "InvalidType", - Error::MissingField(_) => "MissingField", - Error::Http { .. } => "HttpError", - Error::InvalidResponse(_) => "InvalidResponse", - Error::Network(_) => "NetworkError", - Error::Connect(_) => "ConnectError", - Error::Routing(_) => "RoutingError", - Error::Unsupported(_) => "UnsupportedRequest", - } -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index 2acdd232c80..fb63a02f7ad 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -1,174 +1,127 @@ use litellm_core::Error; -use litellm_core::call_lifecycle::CallLifecycle; +use litellm_core::ocr::{ + OcrClient, + wire::{OcrWireRequest, decode_request}, +}; use serde_json::Value; -mod common_utils; -mod handler; -mod hooks; -mod prepare; mod types; pub use types::OcrRequest; -use handler::execute_ocr_provider_call; -use prepare::{PreparedOcrCall, prepare_ocr_call}; - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn ocr(request: OcrRequest<'_>) -> Result { - let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); - CallLifecycle::default() - .run_request(request, &hooks, |request| { - execute_ocr_provider_call(request, &hooks) - }) + core_ocr(request).await +} + +async fn core_ocr(request: OcrRequest<'_>) -> Result { + validate_host_hooks(&request)?; + let client = OcrClient::new(crate::client::http_client().clone())?; + let core_request = decode_request(OcrWireRequest { + model: request.model.to_string(), + document: request.document, + api_key: request.api_key.map(str::to_string), + api_base: request.api_base.map(str::to_string), + custom_llm_provider: request.custom_llm_provider.map(str::to_string), + extra_headers: request.extra_headers, + optional_params: request.optional_params, + input_sources: Default::default(), + timeout_seconds: request.timeout.map(|timeout| timeout.as_secs_f64()), + })?; + client + .perform(core_request) .await + .map(|response| response.into_json()) +} + +fn validate_host_hooks(request: &OcrRequest<'_>) -> Result<(), Error> { + if !request.guardrails.is_empty() { + return Err(Error::Unsupported( + "OCR host guardrails are not wired to the core path", + )); + } + if !request.callbacks.is_empty() { + return Err(Error::Unsupported( + "OCR host callbacks are not wired to the core path", + )); + } + Ok(()) } #[cfg(test)] mod tests { + use std::sync::Arc; + + use litellm_core::ocr::wire::is_supported_request; use serde_json::{Map, json}; - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::{TcpListener, TcpStream}; - use super::{OcrRequest, ocr}; - use crate::integrations::types::RequestMetadata; + use super::{OcrRequest, validate_host_hooks}; + use crate::integrations::custom_guardrail::{CustomGuardrail, GuardrailEventHook}; + use crate::integrations::custom_logger::CustomLogger; - async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); + struct TestGuardrail; + + impl CustomGuardrail for TestGuardrail { + fn guardrail_name(&self) -> &str { + "test" + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &[] } - String::from_utf8(request).expect("request is utf8") } - fn base_ocr_request(model: &str) -> OcrRequest<'_> { + struct TestLogger; + + impl CustomLogger for TestLogger {} + + fn request() -> OcrRequest<'static> { OcrRequest { - model, - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("sk-test"), + model: "model", + document: json!({"type":"image_url","image_url":"data:image/png;base64,YQ=="}), + api_key: None, api_base: None, - custom_llm_provider: None, + custom_llm_provider: Some("mistral"), extra_headers: None, optional_params: Map::new(), timeout: None, callbacks: Vec::new(), guardrails: Vec::new(), - request_metadata: RequestMetadata::default(), + request_metadata: Default::default(), litellm_call_id: None, } } - #[tokio::test] - async fn reducto_file_upload_then_parse_maps_response() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let address = listener.local_addr().expect("listener has local address"); - let server = tokio::spawn(async move { - let (mut upload_socket, _) = listener.accept().await.expect("accepts upload request"); - let upload_request = read_http_request(&mut upload_socket).await; - let upload_body = r#"{"file_id":"reducto://uploaded.pdf"}"#; - let upload_response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - upload_body.len(), - upload_body - ); - upload_socket - .write_all(upload_response.as_bytes()) - .await - .expect("writes upload response"); + #[test] + fn core_activation_includes_migrated_providers() { + assert!(is_supported_request("model", Some("mistral"))); + assert!(is_supported_request("pixtral-12b", Some("azure_ai"))); + assert!(is_supported_request( + "doc-intelligence/prebuilt-layout", + Some("azure_ai") + )); + assert!(is_supported_request("parse-v3", Some("reducto"))); + assert!(is_supported_request("mistral-ocr", Some("vertex_ai"))); + assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); + } - let (mut parse_socket, _) = listener.accept().await.expect("accepts parse request"); - let parse_request = read_http_request(&mut parse_socket).await; - let parse_body = r#"{"job_id":"job_123","usage":{"num_pages":3,"credits":3},"result":{"chunks":[{"content":"Page 1 block A","blocks":[{"content":"Page 1 block A","bbox":{"page":1},"kind":"text"}]},{"content":"Page 2 block A","blocks":[{"content":"Page 2 block A","bbox":{"page":2},"kind":"table"}]},{"content":"Page 1 block B","blocks":[{"content":"Page 1 block B","bbox":{"page":1},"kind":"text"}]},{"content":"Page 3 block A","blocks":[{"content":"Page 3 block A","bbox":{"page":3},"kind":"figure"}]}]}}"#; - let parse_response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - parse_body.len(), - parse_body - ); - parse_socket - .write_all(parse_response.as_bytes()) - .await - .expect("writes parse response"); - (upload_request, parse_request) - }); - let api_base = format!("http://{address}"); - let mut request = base_ocr_request("reducto/parse-v3"); - request.api_base = Some(&api_base); - request.api_key = None; - request.extra_headers = Some(Map::from_iter([ - ("Authorization".to_string(), json!("Bearer test-key")), - ("x-trace-id".to_string(), json!("trace-1")), - ])); - request.document = json!({ - "type": "document_url", - "document_url": "data:application/pdf;base64,JVBERi0xLjQ=" - }); - request.optional_params = Map::from_iter([ - ( - "formatting".to_string(), - json!({"table_output_format": "html"}), - ), - ("retrieval".to_string(), json!({"chunk_mode": "section"})), - ("settings".to_string(), json!({"ocr_system": "standard"})), - ]); + #[test] + fn core_path_rejects_unwired_guardrails() { + let request = OcrRequest { + guardrails: vec![Arc::new(TestGuardrail)], + ..request() + }; + let error = validate_host_hooks(&request).unwrap_err(); + assert!(error.to_string().contains("guardrails are not wired")); + } - let response = ocr(request).await.expect("Reducto OCR succeeds"); - - assert_eq!(response["pages"].as_array().map(Vec::len), Some(3)); - assert_eq!( - response["pages"][0]["markdown"], - "Page 1 block A\n\nPage 1 block B" - ); - assert_eq!(response["pages"][1]["markdown"], "Page 2 block A"); - assert_eq!(response["pages"][2]["markdown"], "Page 3 block A"); - assert_eq!(response["usage_info"]["pages_processed"], 3); - assert_eq!(response["usage_info"]["credits"], 3); - assert_eq!(response["provider_native_response"]["job_id"], "job_123"); - let (upload_request, parse_request) = server.await.expect("server task completes"); - assert!( - upload_request - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert!(upload_request.contains("application/pdf")); - assert!(upload_request.contains("%PDF-1.4")); - assert!(upload_request.contains("x-trace-id: trace-1")); - assert!( - parse_request - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert!(parse_request.contains(r#""input":"reducto://uploaded.pdf""#)); - assert!(parse_request.contains(r#""table_output_format":"html""#)); - assert!(parse_request.contains(r#""chunk_mode":"section""#)); - assert!(parse_request.contains(r#""ocr_system":"standard""#)); + #[test] + fn core_path_rejects_unwired_callbacks() { + let request = OcrRequest { + callbacks: vec![Arc::new(TestLogger)], + ..request() + }; + let error = validate_host_hooks(&request).unwrap_err(); + assert!(error.to_string().contains("callbacks are not wired")); } } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs deleted file mode 100644 index fa9ca1a193e..00000000000 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ /dev/null @@ -1,163 +0,0 @@ -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::{SystemTime, UNIX_EPOCH}; - -use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; -use serde_json::{Map, Value}; - -use super::common_utils::ocr_provider_config; -use super::hooks::OcrLifecycleHooks; -use super::types::{OcrRequest, PreparedOcrRequest}; -use crate::integrations::custom_guardrail::CustomGuardrailRunner; -use crate::integrations::custom_logger::CustomLoggerRunner; - -pub(crate) struct PreparedOcrCall { - pub(crate) request: PreparedOcrRequest, - pub(crate) hooks: OcrLifecycleHooks, -} - -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { - let call_id = request - .litellm_call_id - .map(str::to_string) - .unwrap_or_else(new_ocr_call_id); - let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) - .unwrap_or(CustomLlmProvider { - model: request.model, - custom_llm_provider: "mistral", - }); - let model = provider_info.model.to_string(); - let custom_llm_provider = provider_info.custom_llm_provider.to_string(); - let config = ocr_provider_config(&custom_llm_provider, &model) - .ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone())) - .and_then(|config| { - validate_request_format(config, &request.optional_params, &custom_llm_provider)?; - Ok(config) - }); - let optional_params = match &config { - Ok(config) => { - let supported = config.supported_ocr_params(); - let mut mapped = config.map_ocr_params( - &request - .optional_params - .iter() - .filter(|(name, _)| supported.contains(&name.as_str())) - .map(|(name, value)| (name.clone(), value.clone())) - .collect(), - ); - for name in [ - "vertex_project", - "vertex_ai_project", - "vertex_location", - "vertex_ai_location", - ] { - if let Some(value) = request.optional_params.get(name) { - mapped.insert(name.to_string(), value.clone()); - } - } - mapped - } - Err(_) => request.optional_params, - }; - - PreparedOcrCall { - request: PreparedOcrRequest { - config, - model, - custom_llm_provider, - litellm_call_id: call_id, - document: request.document, - api_key: request.api_key.map(str::to_string), - api_base: request.api_base.map(str::to_string), - extra_headers: request.extra_headers, - optional_params, - timeout: request.timeout, - }, - hooks: OcrLifecycleHooks::new( - CustomLoggerRunner::new(request.callbacks), - CustomGuardrailRunner::new(request.guardrails), - request.request_metadata, - ), - } -} - -fn validate_request_format( - config: &'static dyn litellm_core::ocr::transformation::OcrProviderConfig, - optional_params: &Map, - provider: &str, -) -> Result<(), litellm_core::Error> { - let Some(format) = optional_params.get("req_format") else { - return Ok(()); - }; - match format.as_str() { - Some("litellm") => Ok(()), - Some("native") if config.supported_ocr_params().contains(&"req_format") => Ok(()), - Some("native") => Err(litellm_core::Error::InvalidRequest(format!( - "`req_format=native` is not supported for provider {provider}" - ))), - _ => Err(litellm_core::Error::InvalidRequest(format!( - "Invalid `req_format`: {format}. Expected `litellm` or `native`" - ))), - } -} - -fn new_ocr_call_id() -> String { - static COUNTER: AtomicU64 = AtomicU64::new(1); - let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_nanos()) - .unwrap_or(0); - format!("ocr-{timestamp}-{sequence}") -} - -#[cfg(test)] -mod tests { - use litellm_core::error::Error; - use serde_json::{Map, json}; - - use super::{OcrRequest, prepare_ocr_call}; - use crate::integrations::types::RequestMetadata; - - fn base_ocr_request(model: &str) -> OcrRequest<'_> { - OcrRequest { - model, - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("sk-test"), - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Map::new(), - timeout: None, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: RequestMetadata::default(), - litellm_call_id: None, - } - } - - fn request_with_format(format: &str) -> OcrRequest<'_> { - let mut request = base_ocr_request("mistral/mistral-ocr-latest"); - request.optional_params = Map::from_iter([("req_format".to_string(), json!(format))]); - request - } - - #[test] - fn native_format_rejected_for_provider_without_support_as_bad_request() { - let prepared = prepare_ocr_call(request_with_format("native")); - assert!( - matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("not supported for provider")) - ); - } - - #[test] - fn unknown_format_rejected_for_provider_without_support_as_bad_request() { - let prepared = prepare_ocr_call(request_with_format("raw")); - assert!( - matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("Invalid `req_format`")) - ); - } -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index 75a8e61ddbf..e96d2df1adb 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -1,8 +1,6 @@ use std::sync::Arc; use std::time::Duration; -use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; -use litellm_core::ocr::transformation::OcrProviderConfig; use serde_json::{Map, Value}; use crate::integrations::custom_guardrail::CustomGuardrail; @@ -23,37 +21,3 @@ pub struct OcrRequest<'a> { pub request_metadata: RequestMetadata, pub litellm_call_id: Option<&'a str>, } - -pub(crate) struct PreparedOcrRequest { - pub(crate) config: Result<&'static dyn OcrProviderConfig, litellm_core::Error>, - pub(crate) model: String, - pub(crate) custom_llm_provider: String, - pub(crate) litellm_call_id: String, - pub(crate) document: Value, - pub(crate) api_key: Option, - pub(crate) api_base: Option, - pub(crate) extra_headers: Option>, - pub(crate) optional_params: Map, - pub(crate) timeout: Option, -} - -impl CallLifecycleRequest for PreparedOcrRequest { - fn lifecycle_context(&self) -> CallLifecycleContext { - CallLifecycleContext::new( - "ocr", - self.model.clone(), - self.custom_llm_provider.clone(), - self.litellm_call_id.clone(), - ) - } -} - -pub(crate) struct ProviderOcrRequest { - pub(crate) model: String, - pub(crate) config: &'static dyn OcrProviderConfig, - pub(crate) url: String, - pub(crate) body: Value, - pub(crate) optional_params: Map, - pub(crate) upstream_headers: Vec<(String, String)>, - pub(crate) timeout: Option, -} diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index bb9f3851a77..39465e28e84 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -105,7 +105,11 @@ impl IntoResponse for MessagesRouteError { StatusCode::NOT_FOUND, "no messages deployment is configured for this model".to_string(), ), - Error::Auth(_) => ( + Error::Auth(_) + | Error::MissingApiKey { .. } + | Error::MissingAzureAiCredentials + | Error::MissingAzureDocumentIntelligenceCredentials + | Error::MissingReductoApiKey => ( StatusCode::BAD_GATEWAY, "messages provider authentication failed".to_string(), ), diff --git a/litellm-rust/crates/ai-gateway/src/trace_parity.rs b/litellm-rust/crates/ai-gateway/src/trace_parity.rs index 614852c541d..21123df3f1c 100644 --- a/litellm-rust/crates/ai-gateway/src/trace_parity.rs +++ b/litellm-rust/crates/ai-gateway/src/trace_parity.rs @@ -47,10 +47,10 @@ pub async fn messages_request( .header(CONTENT_TYPE, "application/json") .body(Body::from(body.to_string())) .map_err(|error| Error::InvalidRequest(error.to_string()))?; - let response = routes::app(state) - .oneshot(request) - .await - .map_err(|error| match error {})?; + let response = match routes::app(state).oneshot(request).await { + Ok(response) => response, + Err(error) => match error {}, + }; let status: StatusCode = response.status(); let bytes = to_bytes(response.into_body(), usize::MAX) .await diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs new file mode 100644 index 00000000000..05f7d9610d5 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs @@ -0,0 +1,48 @@ +//! Guards the wiring, not just the helper: a `wss://` dial through the public +//! API has to resolve its own crypto provider, in a test binary where nothing +//! has installed a process-wide one, and has to leave it uninstalled. + +use std::collections::HashMap; +use std::time::Duration; + +use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection; +use tokio::net::TcpListener; + +async fn dead_tls_server() -> u16 { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + port +} + +#[tokio::test] +async fn dialing_wss_returns_an_error_instead_of_panicking() { + let port = dead_tls_server().await; + + let result = ResponsesWebSocketConnection::connect_url( + &format!("wss://127.0.0.1:{port}/"), + &HashMap::new(), + Some(Duration::from_secs(10)), + ) + .await; + + assert!( + result.is_err(), + "a plain TCP server cannot finish a TLS handshake" + ); + assert!( + rustls::crypto::CryptoProvider::get_default().is_none(), + "the dial settles its provider on its own connector, not process-wide" + ); +} diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs deleted file mode 100644 index c3a89f4394d..00000000000 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ /dev/null @@ -1,609 +0,0 @@ -use std::sync::{Arc, Mutex}; -use std::time::Duration; - -use litellm_ai_gateway::integrations::custom_guardrail::{ - CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, - GuardrailFuture, GuardrailRequest, -}; -use litellm_ai_gateway::integrations::custom_logger::{ - CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, -}; -use litellm_ai_gateway::integrations::types::RequestMetadata; -use litellm_ai_gateway::ocr::{OcrRequest, ocr}; -use litellm_core::error::Error; -use serde_json::{Map, Value, json}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::{TcpListener, TcpStream}; - -async fn read_http_headers(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - if request.windows(4).any(|window| window == b"\r\n\r\n") { - break; - } - } - String::from_utf8(request).expect("request is utf8") -} - -async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") -} - -#[derive(Clone, Debug, PartialEq)] -struct RecordedLogEvent { - hook: &'static str, - model: String, - call_type: String, - user_id: Option, - response_object: Option, - error_kind: Option, -} - -#[derive(Default)] -struct RecordingOcrLogger { - events: Mutex>, -} - -impl RecordingOcrLogger { - fn events(&self) -> Vec { - self.events.lock().unwrap().clone() - } -} - -impl CustomLogger for RecordingOcrLogger { - fn async_log_success_event<'a>( - &'a self, - model_call_details: &'a ModelCallDetails, - response_obj: &'a CallbackValue, - _timing: CallbackTiming, - ) -> LogFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push(RecordedLogEvent { - hook: "async_log_success_event", - model: model_call_details.model.clone(), - call_type: model_call_details.call_type.to_string(), - user_id: model_call_details.metadata.user_api_key_user_id.clone(), - response_object: Some(response_obj.object.clone()), - error_kind: None, - }); - Ok(()) - }) - } - - fn async_log_failure_event<'a>( - &'a self, - model_call_details: &'a ModelCallDetails, - response_obj: Option<&'a CallbackValue>, - _timing: CallbackTiming, - ) -> LogFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push(RecordedLogEvent { - hook: "async_log_failure_event", - model: model_call_details.model.clone(), - call_type: model_call_details.call_type.to_string(), - user_id: model_call_details.metadata.user_api_key_user_id.clone(), - response_object: response_obj.map(|value| value.object.clone()), - error_kind: model_call_details - .failure_error - .as_ref() - .map(|error| error.kind.clone()), - }); - Ok(()) - }) - } -} - -struct RecordingOcrGuardrail { - hooks: Vec, - events: Mutex>, - block_pre_call: bool, - block_during_call: bool, -} - -impl RecordingOcrGuardrail { - fn new(hooks: Vec) -> Self { - Self { - hooks, - events: Mutex::new(Vec::new()), - block_pre_call: false, - block_during_call: false, - } - } - - fn blocking_pre_call() -> Self { - Self { - hooks: vec![GuardrailEventHook::PreCall], - events: Mutex::new(Vec::new()), - block_pre_call: true, - block_during_call: false, - } - } - - fn blocking_during_call() -> Self { - Self { - hooks: vec![GuardrailEventHook::DuringCall], - events: Mutex::new(Vec::new()), - block_pre_call: false, - block_during_call: true, - } - } - - fn events(&self) -> Vec<&'static str> { - self.events.lock().unwrap().clone() - } -} - -impl CustomGuardrail for RecordingOcrGuardrail { - fn guardrail_name(&self) -> &str { - "recording-ocr-guardrail" - } - - fn supported_event_hooks(&self) -> &[GuardrailEventHook] { - &self.hooks - } - - fn async_pre_call_hook<'a>( - &'a self, - _context: &'a GuardrailContext, - mut request: GuardrailRequest, - ) -> GuardrailFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("async_pre_call_hook"); - if self.block_pre_call { - return Ok(GuardrailDecision::Block(GuardrailError::blocked( - "blocked before provider", - ))); - } - request.data["document"]["guarded_pre"] = json!(true); - Ok(GuardrailDecision::Mask(request)) - }) - } - - fn async_moderation_hook<'a>( - &'a self, - _context: &'a GuardrailContext, - mut request: GuardrailRequest, - ) -> GuardrailFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("async_moderation_hook"); - if self.block_during_call { - return Ok(GuardrailDecision::Block(GuardrailError::blocked( - "blocked before provider", - ))); - } - request.data["body"]["guarded_during"] = json!(true); - Ok(GuardrailDecision::Mask(request)) - }) - } -} - -fn base_ocr_request(model: &str) -> OcrRequest<'_> { - OcrRequest { - model, - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("sk-test"), - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Map::new(), - timeout: None, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: RequestMetadata::default(), - litellm_call_id: None, - } -} - -#[tokio::test] -async fn reducto_during_call_guardrail_blocks_before_upload() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let address = listener.local_addr().expect("listener has local address"); - let api_base = format!("http://{address}"); - let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call()); - let mut request = base_ocr_request("reducto/parse-v3"); - request.api_base = Some(&api_base); - request.document = json!({ - "type": "document_url", - "document_url": "data:application/pdf;base64,JVBERi0xLjQ=" - }); - request.guardrails = vec![guardrail.clone()]; - - let error = ocr(request).await.expect_err("guardrail blocks upload"); - - assert!(matches!(error, Error::InvalidRequest(_))); - assert_eq!(guardrail.events(), vec!["async_moderation_hook"]); - let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await; - assert!(accepted.is_err(), "upload socket should not be touched"); -} - -#[tokio::test] -async fn reducto_upload_error_body_is_truncated() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let address = listener.local_addr().expect("listener has local address"); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts upload request"); - let _request = read_http_request(&mut socket).await; - let body = "x".repeat(300); - let response = format!( - "HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ); - socket - .write_all(response.as_bytes()) - .await - .expect("writes upload response"); - }); - let api_base = format!("http://{address}"); - let mut request = base_ocr_request("reducto/parse-v3"); - request.api_base = Some(&api_base); - request.document = json!({ - "type": "document_url", - "document_url": "data:application/pdf;base64,JVBERi0xLjQ=" - }); - - let error = ocr(request).await.expect_err("upload should fail"); - - assert!( - matches!(error, Error::Http { status: 500, body } if body.chars().count() < 300 && body.ends_with("... (truncated)")) - ); - server.await.expect("server task completes"); -} - -#[tokio::test] -async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let addr = listener.local_addr().expect("listener has local addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts one request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; - let response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - response_body.len(), - response_body - ); - socket - .write_all(response.as_bytes()) - .await - .expect("writes response"); - request - }); - - let logger = Arc::new(RecordingOcrLogger::default()); - let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![ - GuardrailEventHook::PreCall, - GuardrailEventHook::DuringCall, - ])); - let response = ocr(OcrRequest { - model: "mistral-ocr-latest", - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("sk-test"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("mistral"), - extra_headers: None, - optional_params: Map::new(), - timeout: Some(Duration::from_secs(5)), - callbacks: vec![logger.clone()], - guardrails: vec![guardrail.clone()], - request_metadata: RequestMetadata { - user_api_key_user_id: Some("user-1".to_string()), - ..Default::default() - }, - litellm_call_id: Some("ocr-call-1"), - }) - .await - .expect("ocr request succeeds"); - - assert_eq!(response["pages"][0]["markdown"], "ok"); - assert_eq!( - guardrail.events(), - vec!["async_pre_call_hook", "async_moderation_hook"] - ); - assert_eq!( - logger.events(), - vec![RecordedLogEvent { - hook: "async_log_success_event", - model: "mistral-ocr-latest".to_string(), - call_type: "ocr".to_string(), - user_id: Some("user-1".to_string()), - response_object: Some("ocr".to_string()), - error_kind: None, - }] - ); - - let request = server.await.expect("server task completes"); - assert!(request.contains(r#""guarded_pre":true"#), "{request}"); - assert!(request.contains(r#""guarded_during":true"#), "{request}"); -} - -#[tokio::test] -async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let addr = listener.local_addr().expect("listener has local addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts one request"); - let _request = read_http_request(&mut socket).await; - let response_body = "provider failed"; - let response = format!( - "HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - response_body.len(), - response_body - ); - socket - .write_all(response.as_bytes()) - .await - .expect("writes response"); - }); - - let logger = Arc::new(RecordingOcrLogger::default()); - let err = ocr(OcrRequest { - model: "mistral-ocr-latest", - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("sk-test"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("mistral"), - extra_headers: None, - optional_params: Map::new(), - timeout: Some(Duration::from_secs(5)), - callbacks: vec![logger.clone()], - guardrails: Vec::new(), - request_metadata: RequestMetadata::default(), - litellm_call_id: Some("ocr-call-2"), - }) - .await - .expect_err("provider error propagates"); - - assert!(matches!(err, Error::Http { status: 500, .. })); - server.await.expect("server task completes"); - assert_eq!( - logger.events(), - vec![RecordedLogEvent { - hook: "async_log_failure_event", - model: "mistral-ocr-latest".to_string(), - call_type: "ocr".to_string(), - user_id: None, - response_object: Some("error".to_string()), - error_kind: Some("HttpError".to_string()), - }] - ); -} - -#[tokio::test] -async fn ocr_lifecycle_pre_call_block_skips_provider_socket() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let addr = listener.local_addr().expect("listener has local addr"); - let logger = Arc::new(RecordingOcrLogger::default()); - let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call()); - - let err = ocr(OcrRequest { - model: "mistral-ocr-latest", - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("sk-test"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("mistral"), - extra_headers: None, - optional_params: Map::new(), - timeout: Some(Duration::from_millis(100)), - callbacks: vec![logger.clone()], - guardrails: vec![guardrail.clone()], - request_metadata: RequestMetadata::default(), - litellm_call_id: Some("ocr-call-3"), - }) - .await - .expect_err("guardrail blocks request"); - - assert!(matches!(err, Error::InvalidRequest(_))); - assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]); - assert_eq!( - logger.events(), - vec![RecordedLogEvent { - hook: "async_log_failure_event", - model: "mistral-ocr-latest".to_string(), - call_type: "ocr".to_string(), - user_id: None, - response_object: Some("error".to_string()), - error_kind: Some("InvalidRequest".to_string()), - }] - ); - let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await; - assert!(accepted.is_err(), "provider socket should not be touched"); -} - -#[tokio::test] -async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let addr = listener.local_addr().expect("listener has local addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts one request"); - let request = read_http_headers(&mut socket).await; - let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; - let response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - response_body.len(), - response_body - ); - socket - .write_all(response.as_bytes()) - .await - .expect("writes response"); - request - }); - - let mut headers = Map::new(); - headers.insert( - "Authorization".to_string(), - Value::String("Bearer sk-from-python".to_string()), - ); - headers.insert( - "x-trace-id".to_string(), - Value::String("trace-1".to_string()), - ); - - let response = ocr(OcrRequest { - model: "mistral-ocr-latest", - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("sk-for-rust-fallback"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("mistral"), - extra_headers: Some(headers), - optional_params: Map::new(), - timeout: Some(Duration::from_secs(5)), - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: RequestMetadata::default(), - litellm_call_id: None, - }) - .await - .expect("ocr request succeeds"); - - assert_eq!(response["pages"][0]["markdown"], "ok"); - - let request = server.await.expect("server task completes"); - let authorization_count = request - .lines() - .filter(|line| line.to_ascii_lowercase().starts_with("authorization:")) - .count(); - assert_eq!(authorization_count, 1, "{request}"); - assert!( - request.contains("authorization: Bearer sk-from-python") - || request.contains("Authorization: Bearer sk-from-python"), - "{request}" - ); -} - -#[tokio::test] -async fn document_intelligence_poll_uses_resolved_subscription_key() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let addr = listener.local_addr().expect("listener has local addr"); - let operation_url = format!("http://{addr}/operations/1"); - - let server = tokio::spawn(async move { - let (mut post_socket, _) = listener.accept().await.expect("accepts post request"); - let post_request = read_http_headers(&mut post_socket).await; - let post_response = format!( - "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" - ); - post_socket - .write_all(post_response.as_bytes()) - .await - .expect("writes post response"); - - let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request"); - let poll_request = read_http_headers(&mut poll_socket).await; - let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#; - let poll_response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - response_body.len(), - response_body - ); - poll_socket - .write_all(poll_response.as_bytes()) - .await - .expect("writes poll response"); - (post_request, poll_request) - }); - - let response = ocr(OcrRequest { - model: "doc-intelligence/prebuilt-read", - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("di-key"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: None, - optional_params: Map::new(), - timeout: Some(Duration::from_secs(5)), - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: RequestMetadata::default(), - litellm_call_id: None, - }) - .await - .expect("document intelligence request succeeds"); - - assert_eq!(response["pages"][0]["markdown"], "ok"); - - let (post_request, poll_request) = server.await.expect("server task completes"); - assert!( - post_request - .to_ascii_lowercase() - .contains("ocp-apim-subscription-key: di-key"), - "{post_request}" - ); - assert!( - poll_request - .to_ascii_lowercase() - .contains("ocp-apim-subscription-key: di-key"), - "{poll_request}" - ); -} diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index c0de7ff3977..a2433435e34 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -4,17 +4,33 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true +autotests = false + +[[test]] +name = "workspace_crate_allowlist" +path = "tests/workspace_crate_allowlist.rs" [dependencies] base64.workspace = true +azure_core.workspace = true +azure_identity.workspace = true +data-url = "0.3.2" +gcp_auth.workspace = true +moka.workspace = true rand.workspace = true reqwest.workspace = true serde.workspace = true serde_json.workspace = true +serde_path_to_error = "0.1" +strum.workspace = true +subtle.workspace = true +tokio.workspace = true thiserror.workspace = true tracing.workspace = true tracing-subscriber = { workspace = true, optional = true } sha2.workspace = true +url.workspace = true +veil.workspace = true aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true } aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } @@ -36,5 +52,4 @@ observability = ["dep:tracing-subscriber"] [dev-dependencies] rstest.workspace = true -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } tracing-subscriber.workspace = true diff --git a/litellm-rust/crates/core/src/auth/credential.rs b/litellm-rust/crates/core/src/auth/credential.rs new file mode 100644 index 00000000000..b5235b6780c --- /dev/null +++ b/litellm-rust/crates/core/src/auth/credential.rs @@ -0,0 +1,168 @@ +use std::future::Future; +use std::path::PathBuf; +use std::pin::Pin; +use std::sync::Arc; + +use veil::Redact; + +use crate::AuthError; + +use super::{ResolvedCredential, SecretValue, TokenProviderHandle}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialFileRef { + Path(PathBuf), + EnvironmentVariable(String), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialRef { + Explicit(SecretValue), + Env(String), + File(CredentialFileRef), + Request(String), + Host(String), + None, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialLookup { + Found(SecretValue), + Missing, + Declined, +} + +pub type CredentialLookupFuture<'a> = + Pin> + Send + 'a>>; + +pub trait CredentialResolver: std::fmt::Debug + Send + Sync { + fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a>; +} + +#[derive(Clone, Redact)] +pub struct CredentialResolverHandle(#[redact(with = "[REDACTED]")] Arc); + +impl CredentialResolverHandle { + pub fn new(resolver: Arc) -> Self { + Self(resolver) + } + + pub async fn resolve(&self, reference: &CredentialRef) -> Result { + self.0.resolve(reference).await + } +} + +#[derive(Clone, Debug)] +pub enum CredentialPlan { + Static(CredentialRef), + Caller(TokenProviderHandle), + None, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialPlanResolution { + Resolved(ResolvedCredential), + Unavailable, +} + +impl CredentialPlan { + pub async fn resolve( + &self, + resolver: &CredentialResolverHandle, + ) -> Result { + match self { + Self::Static(CredentialRef::Explicit(secret)) => Ok( + CredentialPlanResolution::Resolved(ResolvedCredential::Static(secret.clone())), + ), + Self::Static(CredentialRef::None) | Self::None => { + Ok(CredentialPlanResolution::Unavailable) + } + Self::Static(reference) => match resolver.resolve(reference).await? { + CredentialLookup::Found(secret) => Ok(CredentialPlanResolution::Resolved( + ResolvedCredential::Static(secret), + )), + CredentialLookup::Missing | CredentialLookup::Declined => { + Ok(CredentialPlanResolution::Unavailable) + } + }, + Self::Caller(caller) => { + let credential = caller.acquire().await?; + if credential.secret().expose().is_empty() { + return Err(AuthError::EmptyCallerCredential); + } + Ok(CredentialPlanResolution::Resolved(credential)) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::{ + CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, + CredentialRef, CredentialResolver, CredentialResolverHandle, + }; + use crate::AuthError; + use crate::auth::SecretValue; + + #[derive(Debug)] + struct HostResolver; + + impl CredentialResolver for HostResolver { + fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { + Box::pin(async move { + Ok(match reference { + CredentialRef::Host(name) if name == "rotating-token" => { + CredentialLookup::Found(SecretValue::new("resolved")) + } + _ => CredentialLookup::Declined, + }) + }) + } + } + + #[tokio::test] + async fn static_host_reference_resolves_at_acquisition_time() { + let resolver = CredentialResolverHandle::new(Arc::new(HostResolver)); + let plan = CredentialPlan::Static(CredentialRef::Host("rotating-token".to_string())); + + let resolved = plan.resolve(&resolver).await.unwrap(); + + assert!(matches!(resolved, CredentialPlanResolution::Resolved(_))); + } + + #[tokio::test] + async fn declined_reference_is_available_for_pre_acquisition_fallback() { + let resolver = CredentialResolverHandle::new(Arc::new(HostResolver)); + let plan = CredentialPlan::Static(CredentialRef::Request("api-key".to_string())); + + assert_eq!( + plan.resolve(&resolver).await.unwrap(), + CredentialPlanResolution::Unavailable + ); + } + + #[derive(Debug)] + struct FailingResolver; + + impl CredentialResolver for FailingResolver { + fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { + Box::pin(async { Err(AuthError::UnresolvedOidcReference) }) + } + } + + #[tokio::test] + async fn acquisition_failure_is_terminal() { + let resolver = CredentialResolverHandle::new(Arc::new(FailingResolver)); + let plan = CredentialPlan::Static(CredentialRef::Host("token".to_string())); + + let error = plan + .resolve(&resolver) + .await + .expect_err("acquisition errors cannot become fallback"); + + assert_eq!(error, AuthError::UnresolvedOidcReference); + } +} diff --git a/litellm-rust/crates/core/src/auth/error.rs b/litellm-rust/crates/core/src/auth/error.rs new file mode 100644 index 00000000000..e7027c0df10 --- /dev/null +++ b/litellm-rust/crates/core/src/auth/error.rs @@ -0,0 +1,128 @@ +use thiserror::Error; + +#[derive(Clone, Debug, Error, PartialEq, Eq)] +pub enum AuthError { + #[error("invalid authentication configuration: {0}")] + Configuration(#[from] AuthConfigurationError), + #[error("credential acquisition failed: {0}")] + AzureTokenAcquisition(String), + #[error("credential acquisition failed: Vertex AI credentials: {0}")] + VertexTokenAcquisition(String), + #[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::>().join("; "))] + CredentialChain(Vec), + #[error("credential caller failed: credential caller returned an empty credential")] + EmptyCallerCredential, + #[error("credential caller failed: Azure AD token provider returned an empty token")] + EmptyAzureToken, + #[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")] + UnresolvedOidcReference, + #[error( + "Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params" + )] + MissingApiKey { provider: &'static str }, + #[error( + "Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter" + )] + MissingApiBase { + provider: &'static str, + environment_variable: &'static str, + }, + #[error("{0}")] + MissingCredential(#[from] MissingCredential), + #[error("{0}")] + Aws(#[from] AwsAuthError), + #[error("invalid authentication header")] + InvalidHeader, +} + +#[derive(Clone, Debug, Error, PartialEq, Eq)] +pub enum AuthConfigurationError { + #[error("credential header already exists")] + ExistingCredentialHeader, + #[error("credential plan is not allowed by the provider auth policy")] + DisallowedCredentialPlan, + #[error("credential cannot be empty")] + EmptyCredential, + #[error("invalid Azure credential selector")] + InvalidAzureSelector, + #[error("ClientSecretCredential requires tenant_id, client_id, and client_secret")] + MissingClientSecretFields, + #[error("WorkloadIdentityCredential requires tenant_id")] + MissingWorkloadTenant, + #[error("WorkloadIdentityCredential requires client_id")] + MissingWorkloadClient, + #[error("WorkloadIdentityCredential requires azure_federated_token_file")] + MissingWorkloadTokenFile, + #[error("credential reference requires a host credential resolver")] + MissingHostResolver, + #[error("caller credential plan requires provider-specific inputs")] + MissingCallerInputs, + #[error("credential header {0} already exists")] + DuplicateHeader(&'static str), + #[error("{0} must be a string or null")] + InvalidFieldType(String), + #[error("unsupported OIDC reference")] + UnsupportedOidcReference, + #[error("{0} cannot be empty")] + EmptyReference(String), + #[error("Azure credential initialization failed: {0}")] + AzureCredentialInitialization(String), + #[error("Azure authority must be an HTTPS origin without credentials, query, or fragment")] + InvalidAzureAuthority, + #[error("request-controlled Azure auth inputs cannot be combined with host credentials")] + MixedAzureCredentialSources, + #[error("request-controlled Azure credential references are not allowed")] + RequestAzureCredentialReference, + #[error("host credentials cannot be sent to a request-controlled Azure endpoint")] + RequestAzureCredentialDestination, + #[error("credentials cannot be sent to a request-controlled Vertex AI endpoint")] + RequestVertexCredentialDestination, + #[error( + "request-controlled Vertex credentials must use the canonical Google OAuth token endpoint" + )] + RequestVertexTokenEndpoint, +} + +#[derive(Clone, Debug, Error, PartialEq, Eq)] +pub enum MissingCredential { + #[error( + "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable" + )] + AnthropicApiKey, + #[error("Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable")] + AzureApiKey, + #[error( + "Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://.services.ai.azure.com/anthropic" + )] + AzureApiBase, + #[error( + "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable" + )] + OpenAiRealtimeApiKey, + #[error( + "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable" + )] + OpenAiResponsesApiKey, +} + +#[derive(Clone, Debug, Error, PartialEq, Eq)] +pub enum AwsAuthError { + #[error("AWS profile credentials failed: {0}")] + Profile(String), + #[error("AWS default credentials failed: {0}")] + DefaultChain(String), + #[error("AWS role credentials failed: {0}")] + AssumeRole(String), + #[error("AWS web identity credentials failed: {0}")] + WebIdentity(String), + #[error("AWS web identity expiration was invalid: {0}")] + WebIdentityExpiration(String), + #[error("AWS signing parameters failed: {0}")] + SigningParameters(String), + #[error("AWS signable request failed: {0}")] + SignableRequest(String), + #[error("AWS request signing failed: {0}")] + Signing(String), + #[error("AWS web identity response had no credentials")] + MissingWebIdentityCredentials, +} diff --git a/litellm-rust/crates/core/src/auth/http.rs b/litellm-rust/crates/core/src/auth/http.rs new file mode 100644 index 00000000000..83931311550 --- /dev/null +++ b/litellm-rust/crates/core/src/auth/http.rs @@ -0,0 +1,86 @@ +use crate::AuthError; +use crate::auth::error::AuthConfigurationError; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CredentialPlacement { + Bearer, + Header(&'static str), +} + +impl CredentialPlacement { + pub fn header_name(self) -> &'static str { + match self { + Self::Bearer => "Authorization", + Self::Header(name) => name, + } + } +} + +pub(crate) fn apply_credential( + headers: Vec<(String, String)>, + credential: &str, + placement: CredentialPlacement, +) -> Result, AuthError> { + if credential.trim().is_empty() { + return Err(AuthError::Configuration( + AuthConfigurationError::EmptyCredential, + )); + } + if headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case(placement.header_name())) + { + return Err(AuthError::Configuration( + AuthConfigurationError::DuplicateHeader(placement.header_name()), + )); + } + let value = match placement { + CredentialPlacement::Bearer => format!("Bearer {credential}"), + CredentialPlacement::Header(_) => credential.to_string(), + }; + Ok( + std::iter::once((placement.header_name().to_string(), value)) + .chain(headers) + .collect(), + ) +} + +/// How the upstream call is authenticated. API-key strategies are resolved in +/// `prepare`; SigV4 needs the serialized body, so the handler signs it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RequestAuth { + Header { name: &'static str, value: String }, + Bearer { token: String }, + AwsSigV4 { region: String }, +} + +#[cfg(test)] +mod tests { + use super::{CredentialPlacement, apply_credential}; + + #[test] + fn bearer_uses_authorization_header() { + let headers = apply_credential(Vec::new(), "key", CredentialPlacement::Bearer) + .expect("credential applies"); + + assert_eq!( + headers, + vec![("Authorization".to_string(), "Bearer key".to_string())] + ); + } + + #[test] + fn named_header_rejects_existing_value() { + let error = apply_credential( + vec![( + "ocp-apim-subscription-key".to_string(), + "caller-key".to_string(), + )], + "configured-key", + CredentialPlacement::Header("Ocp-Apim-Subscription-Key"), + ) + .expect_err("provider policy must handle existing credentials"); + + assert!(error.to_string().contains("already exists")); + } +} diff --git a/litellm-rust/crates/core/src/auth/mod.rs b/litellm-rust/crates/core/src/auth/mod.rs new file mode 100644 index 00000000000..35d9c676f65 --- /dev/null +++ b/litellm-rust/crates/core/src/auth/mod.rs @@ -0,0 +1,56 @@ +mod credential; +pub mod error; +pub(crate) mod vertex; +pub use error::AuthError; +pub(crate) mod http; +mod policy; +mod secret; +mod token; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum InputSource { + Request, + #[default] + Deployment, + Environment, +} + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct Sourced { + value: T, + source: InputSource, +} + +impl Sourced { + pub fn new(value: T, source: InputSource) -> Self { + Self { value, source } + } + + pub fn value(&self) -> &T { + &self.value + } + + pub fn source(&self) -> InputSource { + self.source + } + + pub fn into_value(self) -> T { + self.value + } + + pub fn map(self, map: impl FnOnce(T) -> U) -> Sourced { + Sourced::new(map(self.value), self.source) + } +} + +pub use credential::{ + CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, + CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, +}; +pub use http::{CredentialPlacement, RequestAuth}; +pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; +pub use secret::SecretValue; +pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/core/src/auth/policy.rs b/litellm-rust/crates/core/src/auth/policy.rs new file mode 100644 index 00000000000..b796dedf0d8 --- /dev/null +++ b/litellm-rust/crates/core/src/auth/policy.rs @@ -0,0 +1,114 @@ +use crate::AuthError; +use crate::auth::error::AuthConfigurationError; + +use super::http::apply_credential; +use super::{CredentialPlacement, ResolvedCredential}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CredentialPlanKind { + Static, + Entra, + Caller, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CredentialRule { + pub kind: CredentialPlanKind, + pub placement: CredentialPlacement, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExistingHeaderBehavior { + Preserve, + Reject, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ProviderAuthPolicy { + pub rules: &'static [CredentialRule], + pub accepted_existing_headers: &'static [&'static str], + pub existing_header_behavior: ExistingHeaderBehavior, + pub scope: Option<&'static str>, + pub audience: Option<&'static str>, +} + +impl ProviderAuthPolicy { + pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool { + headers.iter().any(|(name, _)| { + self.accepted_existing_headers + .iter() + .any(|accepted| name.eq_ignore_ascii_case(accepted)) + }) + } + + pub fn apply( + &self, + headers: Vec<(String, String)>, + kind: CredentialPlanKind, + credential: &ResolvedCredential, + ) -> Result, AuthError> { + if self.has_existing_credential(&headers) { + return match self.existing_header_behavior { + ExistingHeaderBehavior::Preserve => Ok(headers), + ExistingHeaderBehavior::Reject => Err(AuthError::Configuration( + AuthConfigurationError::ExistingCredentialHeader, + )), + }; + } + let rule = + self.rules + .iter() + .find(|rule| rule.kind == kind) + .ok_or(AuthError::Configuration( + AuthConfigurationError::DisallowedCredentialPlan, + ))?; + apply_credential(headers, credential.secret().expose(), rule.placement) + } +} + +#[cfg(test)] +mod tests { + use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; + use crate::auth::{CredentialPlacement, ResolvedCredential, SecretValue}; + + const RULES: &[CredentialRule] = &[CredentialRule { + kind: CredentialPlanKind::Static, + placement: CredentialPlacement::Header("x-api-key"), + }]; + const POLICY: ProviderAuthPolicy = ProviderAuthPolicy { + rules: RULES, + accepted_existing_headers: &["x-api-key"], + existing_header_behavior: ExistingHeaderBehavior::Preserve, + scope: None, + audience: None, + }; + + #[test] + fn rules_define_allowed_plans_and_credential_placement() { + let headers = POLICY + .apply( + Vec::new(), + CredentialPlanKind::Static, + &ResolvedCredential::Static(SecretValue::new("secret")), + ) + .unwrap(); + + assert_eq!( + headers, + vec![("x-api-key".to_string(), "secret".to_string())] + ); + } + + #[test] + fn unsupported_plan_is_rejected() { + let error = POLICY + .apply( + Vec::new(), + CredentialPlanKind::Entra, + &ResolvedCredential::Static(SecretValue::new("secret")), + ) + .unwrap_err(); + + assert!(error.to_string().contains("not allowed")); + } +} diff --git a/litellm-rust/crates/core/src/auth/secret.rs b/litellm-rust/crates/core/src/auth/secret.rs new file mode 100644 index 00000000000..3ecb0a835ee --- /dev/null +++ b/litellm-rust/crates/core/src/auth/secret.rs @@ -0,0 +1,41 @@ +use veil::Redact; + +#[derive(Redact, Clone)] +pub struct SecretValue(#[redact(with = "[REDACTED]")] String); + +impl SecretValue { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub fn expose(&self) -> &str { + &self.0 + } +} + +impl PartialEq for SecretValue { + fn eq(&self, other: &Self) -> bool { + subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into() + } +} + +impl Eq for SecretValue {} + +#[cfg(test)] +mod tests { + use super::SecretValue; + + #[test] + fn debug_redacts_plaintext() { + let debug = format!("{:?}", SecretValue::new("credential-value")); + + assert!(!debug.contains("credential-value")); + assert!(debug.contains("REDACTED")); + } + + #[test] + fn equality_compares_plaintext_values() { + assert_eq!(SecretValue::new("same"), SecretValue::new("same")); + assert_ne!(SecretValue::new("same"), SecretValue::new("different")); + } +} diff --git a/litellm-rust/crates/core/src/auth/token.rs b/litellm-rust/crates/core/src/auth/token.rs new file mode 100644 index 00000000000..cfc6b8f0d6b --- /dev/null +++ b/litellm-rust/crates/core/src/auth/token.rs @@ -0,0 +1,47 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::SystemTime; + +use veil::Redact; + +use crate::AuthError; + +use super::secret::SecretValue; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ResolvedCredential { + Static(SecretValue), + AccessToken { + token: SecretValue, + expires_on: Option, + }, +} + +impl ResolvedCredential { + pub fn secret(&self) -> &SecretValue { + match self { + Self::Static(secret) | Self::AccessToken { token: secret, .. } => secret, + } + } +} + +pub type TokenFuture<'a> = + Pin> + Send + 'a>>; + +pub trait TokenProvider: std::fmt::Debug + Send + Sync { + fn acquire(&self) -> TokenFuture<'_>; +} + +#[derive(Clone, Redact)] +pub struct TokenProviderHandle(#[redact(with = "[REDACTED]")] Arc); + +impl TokenProviderHandle { + pub fn new(caller: Arc) -> Self { + Self(caller) + } + + pub async fn acquire(&self) -> Result { + self.0.acquire().await + } +} diff --git a/litellm-rust/crates/core/src/auth/vertex.rs b/litellm-rust/crates/core/src/auth/vertex.rs new file mode 100644 index 00000000000..00a0a7ea7ee --- /dev/null +++ b/litellm-rust/crates/core/src/auth/vertex.rs @@ -0,0 +1,592 @@ +use std::collections::BTreeMap; +use std::future::Future; +use std::path::Path; +use std::pin::Pin; +use std::sync::Arc; + +use gcp_auth::{CustomServiceAccount, TokenProvider}; +use moka::future::Cache; +use serde_json::{Map, Value}; +use sha2::{Digest, Sha256}; + +use crate::auth::error::AuthConfigurationError; +use crate::auth::http::apply_credential; +use crate::auth::{AuthError, CredentialPlacement, InputSource, SecretValue, Sourced}; + +const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; +const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; +const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS"; +const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; +const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; +const VERTEXAI_CREDENTIALS_ENV: &str = "VERTEXAI_CREDENTIALS"; +const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; +const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; +const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; + +#[derive(Clone, Debug, Default)] +pub(crate) struct VertexConfig { + credentials: Option>, + project_id: Option, + location: Option, +} + +impl VertexConfig { + pub(crate) fn from_sourced_optional_params( + params: &Map, + sources: &BTreeMap, + ) -> Result { + Ok(Self { + credentials: optional_credentials( + params, + sources, + &["vertex_credentials", "vertex_ai_credentials"], + )?, + project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?, + location: optional_string(params, &["vertex_location", "vertex_ai_location"])?, + }) + } + + pub(crate) fn project_id(&self) -> Option<&str> { + self.project_id.as_deref() + } + + pub(crate) fn location(&self) -> Option<&str> { + self.location.as_deref() + } +} + +pub(crate) struct VertexEnvironment { + pub headers: Vec<(String, String)>, + pub project_id: String, +} + +struct VertexAccessToken { + token: String, + project_id: String, +} + +pub(crate) fn get_vertex_ai_project( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + config + .project_id() + .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, VERTEXAI_PROJECT_ENV)) +} + +pub(crate) fn get_vertex_ai_location( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + config + .location() + .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, VERTEXAI_LOCATION_ENV)) + .or_else(|| non_empty_env(env_lookup, VERTEX_LOCATION_ENV)) +} + +#[derive(Clone)] +pub(crate) struct VertexAuth { + providers: Cache>, + loader: Arc, +} + +impl Default for VertexAuth { + fn default() -> Self { + Self::new(Arc::new(GcpProviderLoader)) + } +} + +impl VertexAuth { + fn new(loader: Arc) -> Self { + Self { + providers: Cache::builder().max_capacity(64).build(), + loader, + } + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + pub(crate) async fn validate_environment( + &self, + headers: Vec<(String, String)>, + api_key: Option<&str>, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + let has_authorization = headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("Authorization")); + let static_token = api_key + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, VERTEX_AI_API_KEY_ENV)) + .or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV)); + let project_id = get_vertex_ai_project(config, env_lookup); + + if !has_authorization && static_token.is_none() { + let access = self.get_access_token(config, env_lookup).await?; + return Ok(VertexEnvironment { + headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, + project_id: project_id.unwrap_or(access.project_id), + }); + } + + let project_id = match project_id { + Some(project_id) => project_id, + None => { + self.load_provider(config, env_lookup) + .await? + .project_id() + .await? + } + }; + let headers = if has_authorization { + headers + } else { + apply_credential( + headers, + static_token.as_deref().expect("static token was checked"), + CredentialPlacement::Bearer, + )? + }; + Ok(VertexEnvironment { + headers, + project_id, + }) + } + + async fn get_access_token( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + let provider = self.load_provider(config, env_lookup).await?; + let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?; + Ok(VertexAccessToken { token, project_id }) + } + + async fn load_provider( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, AuthError> { + let source = credential_source(config, env_lookup); + let key = source.cache_key(); + self.providers + .try_get_with(key, self.loader.load(source)) + .await + .map_err(|error| (*error).clone()) + } +} + +trait VertexTokenSource: Send + Sync { + fn project_id(&self) -> VertexAuthFuture<'_, String>; + fn token(&self) -> VertexAuthFuture<'_, String>; +} + +trait VertexProviderLoader: Send + Sync { + fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc>; +} + +type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; + +struct GcpTokenSource(Arc); + +impl VertexTokenSource for GcpTokenSource { + fn project_id(&self) -> VertexAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .project_id() + .await + .map(|project| project.to_string()) + .map_err(auth_acquisition_error) + }) + } + + fn token(&self) -> VertexAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .token(&[CLOUD_PLATFORM_SCOPE]) + .await + .map(|token| token.as_str().to_string()) + .map_err(auth_acquisition_error) + }) + } +} + +struct GcpProviderLoader; + +impl VertexProviderLoader for GcpProviderLoader { + fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { + Box::pin(async move { + let provider: Arc = match source { + CredentialSource::Inline(configured) => Arc::new( + CustomServiceAccount::from_json(validate_request_credentials( + configured.expose(), + )?) + .map_err(auth_acquisition_error)?, + ), + CredentialSource::Trusted(configured) => { + let configured = configured.expose(); + let service_account = if Path::new(configured).is_file() { + CustomServiceAccount::from_file(configured) + } else { + CustomServiceAccount::from_json(configured) + } + .map_err(auth_acquisition_error)?; + Arc::new(service_account) + } + CredentialSource::ApplicationCredentials(path) => { + Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?) + } + CredentialSource::Adc => { + gcp_auth::provider().await.map_err(auth_acquisition_error)? + } + }; + Ok(Arc::new(GcpTokenSource(provider)) as Arc) + }) + } +} + +fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> { + let token_uri = serde_json::from_str::(configured) + .ok() + .and_then(|credentials| { + credentials + .get("token_uri") + .and_then(Value::as_str) + .map(str::to_string) + }); + if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { + return Err(AuthConfigurationError::RequestVertexTokenEndpoint.into()); + } + Ok(configured) +} + +#[derive(Clone, Debug)] +enum CredentialSource { + Inline(SecretValue), + Trusted(SecretValue), + ApplicationCredentials(String), + Adc, +} + +impl CredentialSource { + fn cache_key(&self) -> CredentialCacheKey { + match self { + Self::Inline(configured) => { + CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) + } + Self::Trusted(configured) => { + CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) + } + Self::ApplicationCredentials(path) => { + CredentialCacheKey::ApplicationCredentials(path.clone()) + } + Self::Adc => CredentialCacheKey::Adc, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +enum CredentialCacheKey { + Inline([u8; 32]), + Trusted([u8; 32]), + ApplicationCredentials(String), + Adc, +} + +fn credential_source( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> CredentialSource { + if let Some(configured) = config.credentials.clone() { + return match configured.source() { + InputSource::Request => CredentialSource::Inline(configured.into_value()), + InputSource::Deployment | InputSource::Environment => { + CredentialSource::Trusted(configured.into_value()) + } + }; + } + if let Some(configured) = non_empty_env(env_lookup, VERTEXAI_CREDENTIALS_ENV) { + return CredentialSource::Trusted(SecretValue::new(configured)); + } + non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) + .map(CredentialSource::ApplicationCredentials) + .unwrap_or(CredentialSource::Adc) +} + +fn optional_credentials( + params: &Map, + sources: &BTreeMap, + names: &[&str], +) -> Result>, AuthError> { + for name in names { + let source = source_for(sources, name); + match params.get(*name) { + None | Some(Value::Null) => continue, + Some(Value::String(value)) if value.trim().is_empty() => continue, + Some(Value::String(value)) => { + return Ok(Some(Sourced::new(SecretValue::new(value), source))); + } + Some(Value::Object(value)) if value.is_empty() => continue, + Some(Value::Object(value)) => { + return serde_json::to_string(value) + .map(SecretValue::new) + .map(|value| Sourced::new(value, source)) + .map(Some) + .map_err(|error| { + AuthError::Configuration(AuthConfigurationError::InvalidFieldType(format!( + "{}: {error}", + names[0] + ))) + }); + } + Some(_) => { + return Err(AuthError::Configuration( + AuthConfigurationError::InvalidFieldType(names[0].to_string()), + )); + } + } + } + Ok(None) +} + +fn source_for(sources: &BTreeMap, name: &str) -> InputSource { + sources.get(name).copied().unwrap_or_default() +} + +fn optional_string( + params: &Map, + names: &[&str], +) -> Result, AuthError> { + for name in names { + match params.get(*name) { + None | Some(Value::Null) => continue, + Some(Value::String(value)) if value.trim().is_empty() => continue, + Some(Value::String(value)) => return Ok(Some(value.clone())), + Some(_) => { + return Err(AuthError::Configuration( + AuthConfigurationError::InvalidFieldType(names[0].to_string()), + )); + } + } + } + Ok(None) +} + +fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option, name: &str) -> Option { + env_lookup(name) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +fn auth_acquisition_error(error: gcp_auth::Error) -> AuthError { + AuthError::VertexTokenAcquisition(error.to_string()) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use serde_json::json; + + use super::*; + + struct FakeProvider { + calls: Arc, + } + + impl VertexTokenSource for FakeProvider { + fn project_id(&self) -> VertexAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok("adc-project".into()) }) + } + + fn token(&self) -> VertexAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok("adc-token".into()) }) + } + } + + struct FakeLoader { + loads: Arc, + provider: Arc, + } + + impl VertexProviderLoader for FakeLoader { + fn load( + &self, + _source: CredentialSource, + ) -> VertexAuthFuture<'_, Arc> { + let loads = self.loads.clone(); + let provider = self.provider.clone(); + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok(provider) + }) + } + } + + fn config(value: Value) -> VertexConfig { + VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) + .unwrap() + } + + fn auth(calls: Arc, loads: Arc) -> VertexAuth { + let provider: Arc = Arc::new(FakeProvider { calls }); + VertexAuth::new(Arc::new(FakeLoader { loads, provider })) + } + + #[test] + fn config_is_typed_and_secrets_are_redacted() { + let config = config(json!({ + "vertex_credentials":{"private_key":"secret-key"}, + "vertex_project":"project-1", + "vertex_location":"europe-west4" + })); + assert_eq!(config.project_id(), Some("project-1")); + assert_eq!(config.location(), Some("europe-west4")); + assert!(!format!("{config:?}").contains("secret-key")); + assert!( + VertexConfig::from_sourced_optional_params( + json!({"vertex_credentials":true}).as_object().unwrap(), + &BTreeMap::new() + ) + .is_err() + ); + } + + #[test] + fn empty_primary_values_fall_back_to_python_aliases() { + let config = config(json!({ + "vertex_credentials": null, + "vertex_ai_credentials": "alias-credentials", + "vertex_project": " ", + "vertex_ai_project": "alias-project", + "vertex_location": null, + "vertex_ai_location": "alias-location" + })); + assert_eq!( + config.credentials.as_ref().unwrap().value().expose(), + "alias-credentials" + ); + assert_eq!(config.project_id(), Some("alias-project")); + assert_eq!(config.location(), Some("alias-location")); + } + + #[test] + fn project_and_location_prefer_input_then_environment() { + let configured = + config(json!({"vertex_project":"input-project","vertex_location":"input-location"})); + let env = |name: &str| Some(format!("env-{name}")); + assert_eq!( + get_vertex_ai_project(&configured, &env).as_deref(), + Some("input-project") + ); + assert_eq!( + get_vertex_ai_location(&configured, &env).as_deref(), + Some("input-location") + ); + let empty = VertexConfig::default(); + assert_eq!( + get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(), + Some("env-project") + ); + assert_eq!( + get_vertex_ai_location(&empty, &|name| (name == VERTEX_LOCATION_ENV) + .then(|| "fallback-location".into())) + .as_deref(), + Some("fallback-location") + ); + } + + #[test] + fn credential_discovery_prefers_input_then_environment_then_adc() { + let params = json!({"vertex_credentials":"input-json"}); + let sources = BTreeMap::from([("vertex_credentials".to_string(), InputSource::Request)]); + let configured = + VertexConfig::from_sourced_optional_params(params.as_object().unwrap(), &sources) + .unwrap(); + assert!( + matches!(credential_source(&configured, &|_| Some("environment-value".into())), CredentialSource::Inline(value) if value.expose() == "input-json") + ); + let empty = VertexConfig::default(); + assert!( + matches!(credential_source(&empty, &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())), CredentialSource::Trusted(value) if value.expose() == "environment-json") + ); + assert!( + matches!(credential_source(&empty, &|name| (name == GOOGLE_APPLICATION_CREDENTIALS_ENV).then(|| "adc.json".into())), CredentialSource::ApplicationCredentials(path) if path == "adc.json") + ); + assert!(matches!( + credential_source(&empty, &|_| None), + CredentialSource::Adc + )); + assert_ne!( + CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), + CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() + ); + } + + #[test] + fn request_credentials_require_canonical_token_endpoint() { + assert!( + validate_request_credentials(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#) + .is_ok() + ); + assert!(matches!( + validate_request_credentials(r#"{"token_uri":"http://127.0.0.1/token"}"#), + Err(AuthError::Configuration( + AuthConfigurationError::RequestVertexTokenEndpoint + )) + )); + assert!(matches!( + validate_request_credentials("{}"), + Err(AuthError::Configuration( + AuthConfigurationError::RequestVertexTokenEndpoint + )) + )); + } + + #[tokio::test] + async fn explicit_token_and_header_do_not_acquire_adc() { + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); + let configured = config(json!({"vertex_project":"project-1"})); + let explicit = auth + .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) + .await + .unwrap(); + assert_eq!(explicit.headers[0].1, "Bearer access-token"); + let existing = auth + .validate_environment( + vec![("authorization".into(), "Bearer existing".into())], + None, + &configured, + &|_| None, + ) + .await + .unwrap(); + assert_eq!(existing.headers[0].1, "Bearer existing"); + assert_eq!(loads.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn provider_is_reused_across_authentication_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(calls.clone(), loads.clone()); + for _ in 0..2 { + let environment = auth + .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None) + .await + .unwrap(); + assert_eq!(environment.project_id, "adc-project"); + assert_eq!(environment.headers[0].1, "Bearer adc-token"); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 4); + } +} diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index fc81f4fa029..9469d379462 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -43,3 +43,23 @@ pub const EMPTY_TEXT_PLACEHOLDER: &str = "[System: Empty message content sanitised to satisfy protocol]"; pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; + +pub(crate) const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10; + +pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600; +pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10; +pub(crate) const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024; +pub(crate) const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024; +pub(crate) const OCR_MAX_FETCH_REDIRECTS: usize = 10; +pub(crate) const OCR_POLL_TIMEOUT_SECS: u64 = 120; +pub(crate) const OCR_POLL_RETRY_SECS: u64 = 2; +pub(crate) const AZURE_DI_API_VERSION: &str = "2024-11-30"; +pub(crate) const AZURE_DI_SUBSCRIPTION_HEADER: &str = "Ocp-Apim-Subscription-Key"; +pub(crate) const AZURE_DI_DEFAULT_DPI: i64 = 96; +pub(crate) const AZURE_DI_DEFAULT_WIDTH: f64 = 8.5; +pub(crate) const AZURE_DI_DEFAULT_HEIGHT: f64 = 11.0; +pub(crate) const REDUCTO_API_BASE: &str = "https://platform.reducto.ai"; +pub(crate) const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY"; +pub(crate) const REDUCTO_ID_PREFIX: &str = "reducto://"; +pub(crate) const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr"; +pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1"; diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index db3fa2ec704..fa4a9d36e03 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -17,6 +17,22 @@ pub enum Error { InvalidRequest(String), #[error("{0}")] Auth(String), + #[error( + "Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params" + )] + MissingApiKey { provider: &'static str }, + #[error( + "invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID" + )] + MissingAzureAiCredentials, + #[error( + "invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID" + )] + MissingAzureDocumentIntelligenceCredentials, + #[error( + "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" + )] + MissingReductoApiKey, #[error("upstream request failed with status {status}: {body}")] Http { status: u16, body: String }, #[error("upstream network error: {0}")] @@ -36,6 +52,90 @@ pub enum Error { Unsupported(&'static str), } +#[derive(Debug, ThisError)] +pub(crate) enum MediaError { + #[error("media URL rejected by network policy")] + BlockedUrl, + #[error("media download is disabled")] + DownloadDisabled, + #[error("media download exceeds the maximum size")] + DownloadTooLarge, + #[error("too many redirects while fetching media")] + TooManyRedirects, + #[error("media redirect is missing a Location header")] + MissingRedirectLocation, + #[error("invalid media redirect")] + InvalidRedirect, + #[error("media download failed with status {0}")] + Http(u16), + #[error("media download timed out")] + Timeout, + #[error("{0}")] + Transport(#[from] TransportError), +} + +#[derive(Clone, Debug, ThisError, PartialEq, Eq)] +pub enum TransportError { + #[error("upstream request failed with status {status}: {body}")] + Http { status: u16, body: String }, + #[error("upstream network error: {0}")] + Network(String), + #[error("could not reach the provider: {0}")] + Connect(String), +} + +impl TransportError { + pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self { + let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder()); + let message = error.without_url().to_string(); + if before_dispatch { + Self::Connect(message) + } else { + Self::Network(message) + } + } +} + +impl From for TransportError { + fn from(error: reqwest::Error) -> Self { + Self::Network(error.without_url().to_string()) + } +} + +impl From for Error { + fn from(error: crate::ocr::error::OcrRequestError) -> Self { + match error { + crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field), + error => Self::InvalidRequest(error.to_string()), + } + } +} + +impl From for Error { + fn from(error: crate::ocr::error::OcrResponseError) -> Self { + Self::InvalidResponse(error.to_string()) + } +} + +impl From for Error { + fn from(error: TransportError) -> Self { + match error { + TransportError::Http { status, body } => Self::Http { status, body }, + TransportError::Network(message) => Self::Network(message), + TransportError::Connect(message) => Self::Connect(message), + } + } +} + +impl From for Error { + fn from(error: crate::AuthError) -> Self { + match error { + crate::AuthError::MissingApiKey { provider } => Self::MissingApiKey { provider }, + error => Self::Auth(error.to_string()), + } + } +} + pub fn json_type_name(value: &serde_json::Value) -> &'static str { match value { serde_json::Value::Null => "null", @@ -46,3 +146,61 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str { serde_json::Value::Object(_) => "object", } } + +#[cfg(test)] +mod transport_tests { + use super::*; + + #[test] + fn missing_auth_key_preserves_provider_in_public_error() { + assert_eq!( + Error::from(crate::AuthError::MissingApiKey { provider: "Vertex" }), + Error::MissingApiKey { provider: "Vertex" } + ); + } + + #[tokio::test] + async fn transport_errors_remove_urls_and_keep_dispatch_context() { + let error = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get("http://localhost:invalid/private?api_key=secret") + .send() + .await + .expect_err("invalid port"); + let error = TransportError::from_reqwest_before_dispatch(error); + assert!(matches!(error, TransportError::Connect(_))); + assert!(!error.to_string().contains("secret")); + assert!(!error.to_string().contains("private")); + } + + #[tokio::test] + async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() { + use std::time::Duration; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let address = listener.local_addr().expect("address"); + let request = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get(format!("http://{address}")) + .timeout(Duration::from_millis(200)) + .send(); + let (response, accepted) = tokio::join!( + request, + tokio::time::timeout(Duration::from_secs(2), listener.accept()) + ); + let _connection = accepted + .expect("accept deadline") + .expect("accepted connection"); + let error = response.expect_err("server does not respond"); + assert!(error.is_timeout()); + assert!(matches!( + TransportError::from_reqwest_before_dispatch(error), + TransportError::Network(_) + )); + } +} diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index cb472dd5a57..9299bb77ac8 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -1,10 +1,43 @@ -//! Header and upstream-body helpers shared by every route module. - use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; use crate::error::{Error, json_type_name}; +#[allow( + dead_code, + reason = "used by the OCR architecture in the next stacked PR" +)] +pub(crate) enum HeaderPolicy<'a> { + All, + Only(&'a [&'a str]), + Except(&'a [&'a str]), +} + +#[allow( + dead_code, + reason = "used by the OCR architecture in the next stacked PR" +)] +pub(crate) fn with_headers( + builder: reqwest::RequestBuilder, + headers: &[(String, String)], + policy: HeaderPolicy<'_>, +) -> reqwest::RequestBuilder { + headers + .iter() + .filter(|(name, _)| match policy { + HeaderPolicy::All => true, + HeaderPolicy::Only(names) => names + .iter() + .any(|allowed| name.eq_ignore_ascii_case(allowed)), + HeaderPolicy::Except(names) => !names + .iter() + .any(|excluded| name.eq_ignore_ascii_case(excluded)), + }) + .fold(builder, |builder, (name, value)| { + builder.header(name, value) + }) +} + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn http_request( request: reqwest::RequestBuilder, @@ -12,8 +45,6 @@ pub async fn http_request( request.send().await } -/// Bound an upstream error body before it crosses a host boundary, so provider -/// bodies stay data-minimized. pub fn truncate_error_body(body: &str) -> String { if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS { return body.to_string(); @@ -61,11 +92,77 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool { }) } +#[allow( + dead_code, + reason = "used by the OCR architecture in the next stacked PR" +)] +pub(crate) fn deserialize_optional_param<'de, D, T>( + deserializer: D, +) -> Result>, D::Error> +where + D: serde::Deserializer<'de>, + T: serde::Deserialize<'de>, +{ + as serde::Deserialize>::deserialize(deserializer).map(Some) +} + #[cfg(test)] mod tests { use super::*; use serde_json::json; + #[rstest::rstest] + #[case(HeaderPolicy::All, true, true)] + #[case(HeaderPolicy::Only(&["authorization"]), true, false)] + #[case(HeaderPolicy::Except(&["authorization"]), false, true)] + fn forwarding_policy_preserves_matching_headers_and_duplicates( + #[case] policy: HeaderPolicy<'_>, + #[case] auth: bool, + #[case] trace: bool, + ) { + let request = with_headers( + reqwest::Client::new().get("https://example.com"), + &[ + ("AuThOrIzAtIoN".into(), "Bearer token".into()), + ("X-Trace".into(), "first".into()), + ("x-trace".into(), "second".into()), + ], + policy, + ) + .build() + .unwrap(); + assert_eq!(request.headers().contains_key("authorization"), auth); + let traces: Vec<_> = request.headers().get_all("x-trace").iter().collect(); + if trace { + assert_eq!(traces, ["first", "second"]); + } else { + assert!(traces.is_empty()); + } + } + + #[test] + fn multipart_policy_leaves_content_headers_to_reqwest() { + let request = with_headers( + reqwest::Client::new() + .post("https://example.com") + .multipart(reqwest::multipart::Form::new().text("file", "abc")), + &[ + ("Content-Type".into(), "application/json".into()), + ("CONTENT-LENGTH".into(), "0".into()), + ], + HeaderPolicy::Except(&["content-type", "content-length"]), + ) + .build() + .unwrap(); + assert!( + request.headers()["content-type"] + .to_str() + .unwrap() + .starts_with("multipart/form-data; boundary=") + ); + assert_ne!(request.headers()["content-length"], "0"); + } + #[test] fn truncate_leaves_short_bodies_untouched() { assert_eq!(truncate_error_body("short"), "short"); diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index b93e084f57e..0b3573deab2 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,10 +1,12 @@ pub mod audio_transcription; +pub mod auth; pub mod caching; pub mod call_lifecycle; pub mod chat_completions; pub mod constants; pub mod error; pub mod http_utils; +mod media; pub mod messages; #[cfg(any(feature = "observability", test))] pub mod observability; @@ -14,5 +16,7 @@ pub mod realtime; pub mod responses; pub mod router; pub mod routing_utils; +mod url_utils; +pub use auth::AuthError; pub use error::Error; diff --git a/litellm-rust/crates/core/src/media.rs b/litellm-rust/crates/core/src/media.rs new file mode 100644 index 00000000000..5f9a43794c2 --- /dev/null +++ b/litellm-rust/crates/core/src/media.rs @@ -0,0 +1,528 @@ +use std::future::Future; +use std::io; +use std::net::{IpAddr, SocketAddr}; +use std::pin::Pin; +use std::sync::Arc; +use std::time::Duration; + +use reqwest::Url; +use reqwest::dns::{Addrs, Name, Resolve, Resolving}; + +use crate::constants::MEDIA_CONNECT_TIMEOUT_SECS; +use crate::error::{MediaError, TransportError}; + +#[derive(Clone)] +pub(crate) struct MediaFetcher { + client: reqwest::Client, + address_resolver: Arc, + allow_private_network: bool, +} + +type AddressResolution<'a> = Pin>> + Send + 'a>>; + +trait AddressResolver: Send + Sync { + fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a>; +} + +#[derive(Clone, Copy)] +pub(crate) struct DownloadPolicy { + pub(crate) timeout: Duration, + pub(crate) max_bytes: u64, + pub(crate) max_redirects: usize, +} + +#[derive(Debug)] +pub(crate) struct DownloadedMedia { + pub(crate) bytes: Vec, + pub(crate) content_type: String, +} + +impl MediaFetcher { + pub(crate) fn new() -> Result { + Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver)) + } + + fn with_resolvers( + transport_resolver: Arc, + address_resolver: Arc, + ) -> Result + where + R: Resolve + 'static, + { + let client = reqwest::Client::builder() + .connect_timeout(Duration::from_secs(MEDIA_CONNECT_TIMEOUT_SECS)) + .redirect(reqwest::redirect::Policy::none()) + .no_proxy() + .dns_resolver(transport_resolver) + .build()?; + Ok(Self { + client, + address_resolver, + allow_private_network: false, + }) + } + + #[cfg(test)] + pub(crate) fn for_test(client: reqwest::Client) -> Self { + Self { + client, + address_resolver: Arc::new(AllowPrivateResolver), + allow_private_network: true, + } + } + + pub(crate) async fn fetch( + &self, + url: Url, + policy: DownloadPolicy, + ) -> Result { + if policy.max_bytes == 0 { + return Err(MediaError::DownloadDisabled); + } + tokio::time::timeout(policy.timeout, self.fetch_before_deadline(url, policy)) + .await + .map_err(|_| MediaError::Timeout)? + } + + async fn fetch_before_deadline( + &self, + mut url: Url, + policy: DownloadPolicy, + ) -> Result { + let mut redirects_followed = 0; + loop { + self.validate_url(&url).await?; + let mut response = self + .client + .get(url.clone()) + .send() + .await + .map_err(TransportError::from)?; + if response.status().is_redirection() { + if redirects_followed == policy.max_redirects { + return Err(MediaError::TooManyRedirects); + } + let location = response + .headers() + .get(reqwest::header::LOCATION) + .and_then(|value| value.to_str().ok()) + .ok_or(MediaError::MissingRedirectLocation)?; + url = url + .join(location) + .map_err(|_| MediaError::InvalidRedirect)?; + redirects_followed += 1; + continue; + } + if !response.status().is_success() { + return Err(MediaError::Http(response.status().as_u16())); + } + enforce_download_size(response.content_length().unwrap_or(0), policy.max_bytes)?; + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.split(';').next()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or("application/octet-stream") + .to_string(); + let mut bytes = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(TransportError::from)? { + enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?; + bytes.extend_from_slice(&chunk); + } + return Ok(DownloadedMedia { + bytes, + content_type, + }); + } + } + + async fn validate_url(&self, url: &Url) -> Result<(), MediaError> { + if !matches!(url.scheme(), "http" | "https") + || !url.username().is_empty() + || url.password().is_some() + { + return Err(MediaError::BlockedUrl); + } + let host = url.host_str().ok_or(MediaError::BlockedUrl)?; + if self.allow_private_network { + return Ok(()); + } + if let Ok(ip) = host.parse::() { + return (!is_blocked_ip(ip)) + .then_some(()) + .ok_or(MediaError::BlockedUrl); + } + let port = url.port_or_known_default().ok_or(MediaError::BlockedUrl)?; + let addresses = self + .address_resolver + .resolve(host, port) + .await + .map_err(|error| TransportError::Network(error.to_string()))?; + validate_addresses(&addresses) + } +} + +fn enforce_download_size(length: u64, max_bytes: u64) -> Result<(), MediaError> { + if length > max_bytes { + return Err(MediaError::DownloadTooLarge); + } + Ok(()) +} + +fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), MediaError> { + if addresses.is_empty() || addresses.iter().any(|address| is_blocked_ip(address.ip())) { + return Err(MediaError::BlockedUrl); + } + Ok(()) +} + +fn is_blocked_ip(ip: IpAddr) -> bool { + match ip { + IpAddr::V4(ip) => { + let [first, second, third, _] = ip.octets(); + first == 0 + || first == 10 + || first == 127 + || (first == 100 && (64..=127).contains(&second)) + || (first == 169 && second == 254) + || (first == 172 && (16..=31).contains(&second)) + || (first == 192 && second == 0 && (third == 0 || third == 2)) + || (first == 192 && second == 168) + || (first == 192 && second == 88 && third == 99) + || (first == 198 && (second == 18 || second == 19)) + || (first == 198 && second == 51 && third == 100) + || (first == 203 && second == 0 && third == 113) + || first >= 224 + } + IpAddr::V6(ip) => { + let segments = ip.segments(); + ip.is_loopback() + || ip.is_unspecified() + || ip.is_multicast() + || (segments[0] & 0xfe00) == 0xfc00 + || (segments[0] & 0xffc0) == 0xfe80 + || (segments[0] & 0xffc0) == 0xfec0 + || (segments[0] == 0x2001 && segments[1] == 0x0db8) + || ip + .to_ipv4_mapped() + .or_else(|| ip.to_ipv4()) + .map(|ipv4| is_blocked_ip(IpAddr::V4(ipv4))) + .unwrap_or(false) + } + } +} + +#[derive(Default)] +struct PublicDnsResolver; + +struct SystemAddressResolver; + +impl AddressResolver for SystemAddressResolver { + fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> { + Box::pin(async move { + Ok(tokio::net::lookup_host((host, port)) + .await? + .collect::>()) + }) + } +} + +#[cfg(test)] +struct AllowPrivateResolver; + +#[cfg(test)] +impl AddressResolver for AllowPrivateResolver { + fn resolve<'a>(&'a self, _host: &'a str, port: u16) -> AddressResolution<'a> { + Box::pin(async move { Ok(vec![SocketAddr::from(([8, 8, 8, 8], port))]) }) + } +} + +impl Resolve for PublicDnsResolver { + fn resolve(&self, name: Name) -> Resolving { + let host = name.as_str().to_string(); + Box::pin(async move { + let addresses = tokio::net::lookup_host((host.as_str(), 0)) + .await + .map_err(|error| Box::new(error) as Box)? + .collect::>(); + validate_addresses(&addresses).map_err(|_| { + Box::new(io::Error::other("destination rejected by network policy")) + as Box + })?; + Ok(Box::new(addresses.into_iter()) as Addrs) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + async fn serve(response: &'static [u8]) -> (Url, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has address"); + let task = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let mut request = [0_u8; 1024]; + let bytes_read = socket.read(&mut request).await.expect("reads request"); + assert!(bytes_read > 0); + socket.write_all(response).await.expect("writes response"); + }); + ( + Url::parse(&format!("http://{address}/document")).expect("valid test URL"), + task, + ) + } + + async fn serve_named( + host: &str, + responses: Vec<&'static [u8]>, + ) -> (Url, tokio::task::JoinHandle>, SocketAddr) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has address"); + let task = tokio::spawn(async move { + let mut requests = Vec::with_capacity(responses.len()); + for response in responses { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let mut request = [0_u8; 4096]; + let bytes_read = socket.read(&mut request).await.expect("reads request"); + requests.push(String::from_utf8_lossy(&request[..bytes_read]).into_owned()); + socket.write_all(response).await.expect("writes response"); + } + requests + }); + ( + Url::parse(&format!("http://{host}:{}/document", address.port())) + .expect("valid test URL"), + task, + address, + ) + } + + struct LoopbackDnsResolver(SocketAddr); + + impl Resolve for LoopbackDnsResolver { + fn resolve(&self, _name: Name) -> Resolving { + let address = self.0; + Box::pin(async move { Ok(Box::new(vec![address].into_iter()) as Addrs) }) + } + } + + struct TestAddressResolver { + blocked_hosts: HashSet<&'static str>, + } + + impl AddressResolver for TestAddressResolver { + fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> { + let blocked = self.blocked_hosts.contains(host); + Box::pin(async move { + let ip = if blocked { + IpAddr::from([127, 0, 0, 1]) + } else { + IpAddr::from([8, 8, 8, 8]) + }; + Ok(vec![SocketAddr::new(ip, port)]) + }) + } + } + + fn policy_checked_fetcher( + address: SocketAddr, + blocked_hosts: HashSet<&'static str>, + ) -> MediaFetcher { + MediaFetcher::with_resolvers( + Arc::new(LoopbackDnsResolver(address)), + Arc::new(TestAddressResolver { blocked_hosts }), + ) + .expect("test fetcher builds") + } + + fn policy(max_bytes: u64, max_redirects: usize) -> DownloadPolicy { + DownloadPolicy { + timeout: Duration::from_secs(1), + max_bytes, + max_redirects, + } + } + + #[test] + fn blocks_non_public_addresses() { + for address in [ + "0.0.0.1", + "10.0.0.1", + "100.64.0.1", + "127.0.0.1", + "169.254.1.1", + "172.16.0.1", + "192.168.0.1", + "198.18.0.1", + "198.51.100.1", + "203.0.113.1", + "224.0.0.1", + "::1", + "fc00::1", + "fe80::1", + "2001:db8::1", + "::ffff:127.0.0.1", + ] { + assert!(is_blocked_ip(address.parse().expect("valid test address"))); + } + assert!(!is_blocked_ip( + "8.8.8.8".parse().expect("valid public address") + )); + } + + #[tokio::test] + async fn fetches_exact_limit_and_normalizes_content_type() { + let (url, server) = serve( + b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", + ) + .await; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test client builds"); + let media = MediaFetcher::for_test(client) + .fetch(url, policy(3, 0)) + .await + .expect("download succeeds at exact limit"); + server.await.expect("server completes"); + assert_eq!(media.bytes, b"abc"); + assert_eq!(media.content_type, "application/pdf"); + } + + #[tokio::test] + async fn rejects_declared_oversize_body() { + let (url, server) = serve( + b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", + ) + .await; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test client builds"); + let error = MediaFetcher::for_test(client) + .fetch(url, policy(2, 0)) + .await + .expect_err("oversize body is rejected"); + server.await.expect("server completes"); + assert!(matches!(error, MediaError::DownloadTooLarge)); + } + + #[tokio::test] + async fn rejects_streamed_oversize_body() { + let (url, server) = serve( + b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n2\r\ncd\r\n0\r\n\r\n", + ) + .await; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test client builds"); + let error = MediaFetcher::for_test(client) + .fetch(url, policy(3, 0)) + .await + .expect_err("stream crossing limit is rejected"); + server.await.expect("server completes"); + assert!(matches!(error, MediaError::DownloadTooLarge)); + } + + #[tokio::test] + async fn follows_allowed_redirects_and_revalidates_each_destination() { + let (url, server, address) = serve_named( + "public.test", + vec![ + b"HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\n\r\n", + b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok", + ], + ) + .await; + let media = policy_checked_fetcher(address, HashSet::new()) + .fetch(url, policy(2, 1)) + .await + .expect("redirected fetch succeeds"); + let requests = server.await.expect("server completes"); + assert_eq!(requests.len(), 2); + assert!(requests[1].starts_with("GET /final ")); + assert_eq!(media.bytes, b"ok"); + } + + #[tokio::test] + async fn blocks_redirected_private_destination_before_second_request() { + let (url, server, address) = serve_named( + "public.test", + vec![b"HTTP/1.1 302 Found\r\nLocation: http://blocked.test/document\r\nContent-Length: 0\r\n\r\n"], + ) + .await; + let error = policy_checked_fetcher(address, HashSet::from(["blocked.test"])) + .fetch(url, policy(10, 1)) + .await + .expect_err("private redirect is rejected"); + let requests = server.await.expect("server completes"); + assert_eq!(requests.len(), 1); + assert!(matches!(error, MediaError::BlockedUrl)); + } + + #[tokio::test] + async fn enforces_total_timeout() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has address"); + let server = tokio::spawn(async move { + let (_socket, _) = listener.accept().await.expect("accepts request"); + tokio::time::sleep(Duration::from_millis(100)).await; + }); + let url = Url::parse(&format!("http://public.test:{}/document", address.port())) + .expect("valid test URL"); + let error = policy_checked_fetcher(address, HashSet::new()) + .fetch( + url, + DownloadPolicy { + timeout: Duration::from_millis(20), + max_bytes: 10, + max_redirects: 0, + }, + ) + .await + .expect_err("fetch times out"); + server.await.expect("server completes"); + assert!(matches!(error, MediaError::Timeout)); + } + + #[tokio::test] + async fn document_client_does_not_send_ambient_credentials() { + let (url, server, address) = serve_named( + "public.test", + vec![b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"], + ) + .await; + policy_checked_fetcher(address, HashSet::new()) + .fetch(url, policy(2, 0)) + .await + .expect("fetch succeeds"); + let requests = server.await.expect("server completes"); + assert!(!requests[0].to_ascii_lowercase().contains("authorization:")); + assert!(!requests[0].to_ascii_lowercase().contains("api-key:")); + } + + #[tokio::test] + async fn rejects_url_credentials_before_network_access() { + let fetcher = MediaFetcher::new().expect("media fetcher builds"); + let url = + Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses"); + assert!(matches!( + fetcher.validate_url(&url).await, + Err(MediaError::BlockedUrl) + )); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs new file mode 100644 index 00000000000..71ca69ddc58 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs @@ -0,0 +1,214 @@ +use super::super::OcrAdapter; +use crate::Error; +use crate::auth::{InputSource, Sourced}; +use crate::constants::{AZURE_DI_API_VERSION, AZURE_DI_SUBSCRIPTION_HEADER}; +use crate::ocr::OcrClient; +use crate::ocr::codecs::document_intelligence::{ + self, AzureDocumentIntelligenceOperation, DocumentIntelligenceParams, +}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{credential_env, transform_request_body}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat}; +use crate::ocr::wire::DecodedOcrResponse; +use crate::providers::azure_ai::auth::AzureAuthInputs; +use crate::url_utils::ApiUrl; + +mod polling; + +const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; +const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; + +#[derive(Clone, Debug)] +pub(crate) struct AzureDocumentIntelligenceAdapter; + +impl OcrAdapter for AzureDocumentIntelligenceAdapter { + type ProviderResponse = AzureDocumentIntelligenceOperation; + const PROVIDER: OcrProvider = OcrProvider::AzureAi; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let params = map_ocr_params(request)?; + let config = AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)?; + let headers = validate_environment(&request.connection, &config, &credential_env).await?; + let endpoint = nonblank(request.connection.api_base.clone()) + .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) + .ok_or_else(|| Error::Auth("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into()))?; + let url = get_complete_url(&endpoint, &request.model, ¶ms)?; + let body = document_intelligence::transform_ocr_request(request.document.clone())?; + transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + document_intelligence::transform_ocr_response(&request.model, response) + } + + async fn read_response( + &self, + client: &OcrClient, + response: reqwest::Response, + url: &str, + headers: &[(String, String)], + request: &LiteLLMOcrRequest, + ) -> Result, OcrError> { + polling::read_operation_response( + client.polling_http(), + response, + url, + headers, + &request.connection, + request.response_format()? == OcrResponseFormat::Native, + ) + .await + } +} + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +fn map_ocr_params( + request: &LiteLLMOcrRequest, +) -> Result { + let params = document_intelligence::decode_input_params( + request.optional_params.clone(), + "optional_params", + )?; + let crate::ocr::prepare::ParsedProviderParams { + known: params, + extra_params: _extra_params, + } = params; + document_intelligence::map_ocr_params(params) +} + +fn get_complete_url( + endpoint: &str, + model: &str, + params: &DocumentIntelligenceParams, +) -> Result { + let model = format!("{}:analyze", model_id(model)?); + ApiUrl::parse(endpoint) + .and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model])) + .map(|url| { + url.append_query_pairs( + [("api-version", AZURE_DI_API_VERSION)] + .into_iter() + .chain(params.pages.iter().map(|pages| ("pages", pages.as_str()))) + .chain( + params + .features + .iter() + .map(|features| ("features", features.as_str())), + ), + ) + .into_string() + }) + .map_err(|_| OcrRequestError::RequestField { + path: "api_base".into(), + }) + .map_err(OcrError::from) +} + +async fn validate_environment( + connection: &OcrConnection, + config: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), +) -> Result, OcrError> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") + || crate::http_utils::has_header(&connection.extra_headers, AZURE_DI_SUBSCRIPTION_HEADER) + { + super::validate_destination(connection, connection.extra_headers_source)?; + return Ok(connection.extra_headers.clone()); + } + let key = nonblank(connection.api_key.clone()) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(env_lookup(AZURE_DI_API_KEY_ENV)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); + if let Some(key) = key { + super::validate_destination(connection, key.source())?; + return Ok( + std::iter::once((AZURE_DI_SUBSCRIPTION_HEADER.into(), key.into_value())) + .chain(connection.extra_headers.clone()) + .collect(), + ); + } + let token = super::resolve_entra(config, env_lookup) + .await? + .ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?; + super::validate_destination(connection, token.source())?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {}", token.value()))) + .chain(connection.extra_headers.clone()) + .collect(), + ) +} + +fn model_id(model: &str) -> Result<&str, OcrRequestError> { + let model = model.rsplit('/').next().unwrap_or(model); + if matches!(model, "." | "..") { + return Err(OcrRequestError::DotModel); + } + Ok(model) +} + +fn nonblank(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn request_endpoint_cannot_receive_environment_key() { + let connection = OcrConnection { + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let error = validate_environment(&connection, &Default::default(), &|name| { + (name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into()) + }) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains("request-controlled Azure endpoint") + ); + } + + #[tokio::test] + async fn request_endpoint_accepts_request_owned_key() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + api_key_source: InputSource::Request, + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let headers = validate_environment(&connection, &Default::default(), &|_| None) + .await + .unwrap(); + + assert_eq!( + headers[0], + (AZURE_DI_SUBSCRIPTION_HEADER.into(), "request-key".into()) + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs new file mode 100644 index 00000000000..1bddea0da4f --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs @@ -0,0 +1,100 @@ +use std::time::Duration; + +use reqwest::Url; +use tokio::time::Instant; + +use crate::constants::{AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS}; +use crate::ocr::client::read_json_response; +use crate::ocr::codecs::document_intelligence::{ + AzureDocumentIntelligenceOperation, OperationStatus, +}; +use crate::ocr::error::{OcrError, OcrPollingError, OcrResponseError}; +use crate::ocr::types::OcrConnection; +use crate::ocr::wire::DecodedOcrResponse; + +pub(super) async fn read_operation_response( + http_client: &reqwest::Client, + response: reqwest::Response, + original_url: &str, + headers: &[(String, String)], + connection: &OcrConnection, + native: bool, +) -> Result, OcrError> { + if response.status() != reqwest::StatusCode::ACCEPTED { + return read_json_response(response, native).await; + } + let location = response + .headers() + .get("operation-location") + .and_then(|value| value.to_str().ok()) + .ok_or(OcrPollingError::PollLocation)?; + let original = Url::parse(original_url).map_err(|_| OcrPollingError::PollOrigin)?; + let operation = Url::parse(location).map_err(|_| OcrPollingError::PollOrigin)?; + if original.origin() != operation.origin() + || !operation.username().is_empty() + || operation.password().is_some() + { + return Err(OcrPollingError::PollOrigin.into()); + } + poll_operation(http_client, operation, headers, connection, native).await +} + +async fn poll_operation( + http_client: &reqwest::Client, + url: Url, + headers: &[(String, String)], + connection: &OcrConnection, + native: bool, +) -> Result, OcrError> { + let deadline = Instant::now() + .checked_add(connection.poll_timeout) + .ok_or(OcrPollingError::PollTimeout)?; + loop { + let remaining = deadline + .checked_duration_since(Instant::now()) + .filter(|remaining| !remaining.is_zero()) + .ok_or(OcrPollingError::PollTimeout)?; + let builder = http_client + .get(url.clone()) + .timeout(remaining.min(connection.timeout)); + let builder = crate::http_utils::with_headers( + builder, + headers, + crate::http_utils::HeaderPolicy::Only(&[AZURE_DI_SUBSCRIPTION_HEADER, "authorization"]), + ); + let response = tokio::time::timeout_at(deadline, crate::http_utils::http_request(builder)) + .await + .map_err(|_| OcrPollingError::PollTimeout)? + .map_err(crate::error::TransportError::from)?; + let retry = response + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .unwrap_or(OCR_POLL_RETRY_SECS) + .max(1); + let decoded = tokio::time::timeout_at( + deadline, + read_json_response::(response, native), + ) + .await + .map_err(|_| OcrPollingError::PollTimeout)??; + match &decoded.data.status { + Some(OperationStatus::Succeeded) => return Ok(decoded), + Some(OperationStatus::Running | OperationStatus::NotStarted) => { + tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry))) + .await + .map_err(|_| OcrPollingError::PollTimeout)?; + } + status => { + return Err(OcrResponseError::OperationStatus( + status + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| "None".into()), + ) + .into()); + } + } + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs new file mode 100644 index 00000000000..3107494d39e --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs @@ -0,0 +1,217 @@ +use super::super::OcrAdapter; +use crate::Error; +use crate::auth::{InputSource, Sourced}; +use crate::constants::AZURE_AI_OCR_PATH; +use crate::ocr::OcrClient; +use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse}; +use crate::ocr::document::{inline_remote_document, validate_inline_document}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{ + _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, +}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection}; +use crate::providers::azure_ai::auth::AzureAuthInputs; +use crate::url_utils::ApiUrl; + +const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; +const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; + +#[derive(Clone, Debug)] +pub(crate) struct AzureMistralAdapter; + +impl OcrAdapter for AzureMistralAdapter { + type ProviderResponse = MistralOcrResponse; + const PROVIDER: OcrProvider = OcrProvider::AzureAi; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let ParsedProviderParams { + known: params, + extra_params: _extra_params, + } = _prepare_ocr_request::(request)?; + let config = AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)?; + let headers = validate_environment(&request.connection, &config, &credential_env).await?; + let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?; + let document = inline_remote_document( + client.document_fetcher(), + request.document.clone(), + &request.connection, + ) + .await?; + let body = mistral::transform_ocr_request(&request.model, document, ¶ms)?; + transform_request_body(client, request, &url, &headers, body, |body| { + validate_inline_document(&body.document) + }) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + mistral::transform_ocr_response(&request.model, response) + } +} + +fn get_complete_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + let base = nonblank(api_base.map(str::to_string)) + .or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV))) + .ok_or_else(|| Error::Auth( + "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter".into(), + ))?; + let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect(); + ApiUrl::parse(&base) + .and_then(|url| url.complete_path(&path)) + .map(|url| url.into_string()) + .map_err(|_| { + OcrRequestError::RequestField { + path: "api_base".into(), + } + .into() + }) +} + +async fn validate_environment( + connection: &OcrConnection, + config: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), +) -> Result, OcrError> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + super::validate_destination(connection, connection.extra_headers_source)?; + return Ok(connection.extra_headers.clone()); + } + let key = nonblank(connection.api_key.clone()) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(env_lookup(AZURE_AI_API_KEY_ENV)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); + if let Some(key) = key { + super::validate_destination(connection, key.source())?; + return Ok(bearer_headers(connection, key.value())); + } + let key = super::resolve_entra(config, env_lookup) + .await? + .ok_or(Error::MissingAzureAiCredentials)?; + super::validate_destination(connection, key.source())?; + Ok(bearer_headers(connection, key.value())) +} + +fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> { + std::iter::once(("Authorization".into(), format!("Bearer {key}"))) + .chain(connection.extra_headers.clone()) + .collect() +} + +fn nonblank(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn completes_azure_path_and_preserves_query() { + assert_eq!( + get_complete_url(Some("https://example.com/?tenant=a"), &|_| None).unwrap(), + "https://example.com/providers/mistral/azure/ocr?tenant=a" + ); + assert_eq!( + get_complete_url( + Some("https://example.com/providers/mistral/azure/ocr"), + &|_| None + ) + .unwrap(), + "https://example.com/providers/mistral/azure/ocr" + ); + } + + #[tokio::test] + async fn supplied_authorization_precedes_keys() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + extra_headers: vec![("authorization".into(), "Bearer prepared".into())], + ..Default::default() + }; + assert_eq!( + validate_environment(&connection, &Default::default(), &|_| { + Some("environment-key".into()) + }) + .await + .unwrap(), + connection.extra_headers + ); + } + + #[tokio::test] + async fn request_key_precedes_environment_key() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + ..Default::default() + }; + assert_eq!( + validate_environment(&connection, &Default::default(), &|_| { + Some("environment-key".into()) + }) + .await + .unwrap()[0], + ("Authorization".into(), "Bearer request-key".into()) + ); + } + + #[tokio::test] + async fn request_endpoint_cannot_receive_environment_key() { + let connection = OcrConnection { + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let error = validate_environment(&connection, &Default::default(), &|name| { + (name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into()) + }) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains("request-controlled Azure endpoint") + ); + } + + #[tokio::test] + async fn request_endpoint_accepts_request_owned_key() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + api_key_source: InputSource::Request, + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let headers = validate_environment(&connection, &Default::default(), &|_| None) + .await + .unwrap(); + + assert_eq!( + headers[0], + ("Authorization".into(), "Bearer request-key".into()) + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs new file mode 100644 index 00000000000..9c02a7471c9 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs @@ -0,0 +1,49 @@ +mod document_intelligence; +mod mistral; + +use std::sync::OnceLock; + +use crate::Error; +use crate::auth::error::AuthConfigurationError; +use crate::auth::{InputSource, Sourced}; +use crate::ocr::error::OcrError; +use crate::ocr::types::OcrConnection; +use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService}; + +pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter; +pub(crate) use mistral::AzureMistralAdapter; + +async fn resolve_entra( + config: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), +) -> Result>, Error> { + static SERVICE: OnceLock = OnceLock::new(); + SERVICE + .get_or_init(AzureAuthService::default) + .get_azure_ad_token(config, env_lookup) + .await + .map(|credential| { + credential.map(|credential| { + let source = credential.source(); + let value = credential.value().secret().expose().to_string(); + Sourced::new(value, source) + }) + }) + .map_err(Error::from) +} + +fn validate_destination( + connection: &OcrConnection, + credential_source: InputSource, +) -> Result<(), OcrError> { + if connection.api_base.is_some() + && connection.api_base_source == InputSource::Request + && credential_source != InputSource::Request + { + return Err(Error::from(crate::AuthError::Configuration( + AuthConfigurationError::RequestAzureCredentialDestination, + )) + .into()); + } + Ok(()) +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/mistral.rs new file mode 100644 index 00000000000..ea569ffb34f --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/mistral.rs @@ -0,0 +1,147 @@ +use super::OcrAdapter; +use crate::Error; +use crate::constants::MISTRAL_OCR_API_BASE; +use crate::ocr::OcrClient; +use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{ + _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, +}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection}; +use crate::url_utils::ApiUrl; + +const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY"; + +#[derive(Clone, Debug)] +pub(crate) struct MistralAdapter; + +impl OcrAdapter for MistralAdapter { + type ProviderResponse = MistralOcrResponse; + const PROVIDER: OcrProvider = OcrProvider::Mistral; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let ParsedProviderParams { + known: params, + extra_params: _extra_params, + } = _prepare_ocr_request::(request)?; + let headers = validate_environment(&request.connection, &credential_env)?; + let url = get_complete_url(request.connection.api_base.as_deref())?; + let body = + mistral::transform_ocr_request(&request.model, request.document.clone(), ¶ms)?; + transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + mistral::transform_ocr_response(&request.model, response) + } +} + +pub(crate) fn get_complete_url(api_base: Option<&str>) -> Result { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(MISTRAL_OCR_API_BASE); + ApiUrl::parse(base) + .and_then(|url| url.complete_path(&["v1", "ocr"])) + .map(|url| url.into_string()) + .map_err(|_| { + OcrRequestError::RequestField { + path: "api_base".into(), + } + .into() + }) +} + +fn validate_environment( + connection: &OcrConnection, + env_lookup: &(dyn Fn(&str) -> Option + Sync), +) -> Result, OcrError> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let api_key = connection + .api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .ok_or(Error::MissingApiKey { + provider: "Mistral", + })?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {api_key}"))) + .chain(connection.extra_headers.clone()) + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn complete_url_defaults_and_dedupes_v1() { + assert_eq!( + get_complete_url(None).unwrap(), + "https://api.mistral.ai/v1/ocr" + ); + assert_eq!( + get_complete_url(Some("https://example.com/v1?tenant=a")).unwrap(), + "https://example.com/v1/ocr?tenant=a" + ); + assert_eq!( + get_complete_url(Some("https://example.com/v1/ocr?tenant=a")).unwrap(), + "https://example.com/v1/ocr?tenant=a" + ); + } + + #[test] + fn environment_prefers_explicit_key_then_environment() { + let explicit = OcrConnection { + api_key: Some("explicit".into()), + ..OcrConnection::default() + }; + assert_eq!( + validate_environment(&explicit, &|_| Some("environment".into())).unwrap()[0], + ("Authorization".into(), "Bearer explicit".into()) + ); + + assert_eq!( + validate_environment(&OcrConnection::default(), &|_| Some("environment".into())) + .unwrap()[0], + ("Authorization".into(), "Bearer environment".into()) + ); + } + + #[test] + fn environment_preserves_forwarded_authorization() { + let connection = OcrConnection { + extra_headers: vec![("authorization".into(), "Bearer forwarded".into())], + ..OcrConnection::default() + }; + assert_eq!( + validate_environment(&connection, &|_| None).unwrap(), + connection.extra_headers + ); + } + + #[test] + fn environment_rejects_missing_key() { + assert!(matches!( + validate_environment(&OcrConnection::default(), &|_| None), + Err(OcrError::Public(Error::MissingApiKey { + provider: "Mistral" + })) + )); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/mod.rs new file mode 100644 index 00000000000..9171d11836c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/mod.rs @@ -0,0 +1,81 @@ +use std::future::Future; + +use serde::de::DeserializeOwned; + +use super::OcrClient; +use super::error::{OcrError, OcrResponseError}; +use super::registry::OcrProvider; +use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrResponseFormat}; +use super::wire::DecodedOcrResponse; + +mod azure; +mod mistral; +mod reducto; +mod vertex; + +pub(crate) use azure::{AzureDocumentIntelligenceAdapter, AzureMistralAdapter}; +pub(crate) use mistral::MistralAdapter; +pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter}; +pub(crate) use vertex::{VertexDeepSeekAdapter, VertexMistralAdapter}; + +/// Converts a complete LiteLLM OCR call to provider HTTP and normalizes its response. +pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static { + /// Provider JSON schema; direct and Vertex Mistral share `MistralOcrResponse`. + type ProviderResponse: DeserializeOwned + Send; + + const PROVIDER: OcrProvider; + + /// Prepares the complete provider HTTP request. + /// `request` contains the model, document, connection, and unmapped caller options. + /// `client` supplies reusable provider and document HTTP clients. + /// Returns the complete HTTP request, whereas Python returns body data. + fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> impl Future> + Send; + + /// Python: `transform_ocr_response`. + /// `request` supplies caller context, including the fallback model. + /// `response` is the decoded provider payload; the output is the shared LiteLLM schema. + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result; + + /// Decodes provider HTTP; adapters may override this to poll asynchronous operations. + /// Python performs that polling inside `async_transform_ocr_response`. + /// `client` is reused for polling; `response` is the initial HTTP response. + /// `url` and `headers` describe the submitted call; `request` supplies limits and format. + fn read_response( + &self, + _client: &OcrClient, + response: reqwest::Response, + _url: &str, + _headers: &[(String, String)], + request: &LiteLLMOcrRequest, + ) -> impl Future, OcrError>> + Send + { + let retain_native = request + .response_format() + .map(|format| format == OcrResponseFormat::Native); + async move { super::client::read_json_response(response, retain_native?).await } + } +} + +macro_rules! for_each_ocr_adapter { + ($callback:ident) => { + $callback! { + Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral; + AzureMistral, $crate::ocr::adapters::AzureMistralAdapter, $crate::ocr::adapters::AzureMistralAdapter, AzureAi; + AzureDocumentIntelligence, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, AzureAi; + ReductoLegacy, $crate::ocr::adapters::ReductoLegacyAdapter, $crate::ocr::adapters::ReductoLegacyAdapter, Reducto; + ReductoV3, $crate::ocr::adapters::ReductoV3Adapter, $crate::ocr::adapters::ReductoV3Adapter, Reducto; + VertexMistral, $crate::ocr::adapters::VertexMistralAdapter, $crate::ocr::adapters::VertexMistralAdapter, VertexAi; + VertexDeepSeek, $crate::ocr::adapters::VertexDeepSeekAdapter, $crate::ocr::adapters::VertexDeepSeekAdapter, VertexAi; + } + }; +} + +pub(crate) use for_each_ocr_adapter; diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs new file mode 100644 index 00000000000..062a0071a34 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs @@ -0,0 +1,45 @@ +use super::super::OcrAdapter; +use crate::ocr::OcrClient; +use crate::ocr::codecs::reducto::{self, ReductoLegacyParams, ReductoResponse}; +use crate::ocr::error::{OcrError, OcrResponseError}; +use crate::ocr::prepare::{ + _prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env, + guardrail_document, merge_extra_params, +}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; + +#[derive(Clone, Debug)] +pub(crate) struct ReductoLegacyAdapter; + +impl OcrAdapter for ReductoLegacyAdapter { + type ProviderResponse = ReductoResponse; + const PROVIDER: OcrProvider = OcrProvider::Reducto; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let ParsedProviderParams { + known: params, + extra_params, + } = _prepare_ocr_request::(request)?; + let headers = super::validate_environment(&request.connection, &credential_env)?; + let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?; + let document = guardrail_document(request, &url).await?; + let document = + super::prepare_document(client, document, &request.connection, &headers).await?; + let body = reducto::transform_legacy_ocr_request(&request.model, document, ¶ms)?; + let body = merge_extra_params(&body, extra_params)?; + build_http_request(client, request, &url, &headers, &body) + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + reducto::transform_ocr_response(&request.model, response) + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs new file mode 100644 index 00000000000..7621d0d326a --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs @@ -0,0 +1,148 @@ +mod legacy; +mod v3; + +use crate::Error; +use crate::constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX}; +use crate::ocr::document::InlineDocument; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::types::{OcrConnection, OcrDocument}; +use crate::url_utils::ApiUrl; + +pub(crate) use legacy::ReductoLegacyAdapter; +pub(crate) use v3::ReductoV3Adapter; + +pub(super) fn get_complete_url(api_base: Option<&str>, path: &str) -> Result { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(REDUCTO_API_BASE); + ApiUrl::parse(base) + .and_then(|url| url.complete_path(&[path])) + .map(|url| url.into_string()) + .map_err(|_| { + OcrRequestError::RequestField { + path: "api_base".into(), + } + .into() + }) +} + +pub(super) fn validate_environment( + connection: &OcrConnection, + env_lookup: &(dyn Fn(&str) -> Option + Sync), +) -> Result, OcrError> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let api_key = connection + .api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| { + env_lookup(REDUCTO_API_KEY_ENV) + .map(|key| key.trim().to_string()) + .filter(|key| !key.is_empty()) + }) + .ok_or(Error::MissingReductoApiKey)?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {api_key}"))) + .chain(connection.extra_headers.clone()) + .collect(), + ) +} + +pub(super) async fn prepare_document( + client: &crate::ocr::OcrClient, + document: OcrDocument, + connection: &OcrConnection, + headers: &[(String, String)], +) -> Result { + if document.source().starts_with(REDUCTO_ID_PREFIX) { + if document.source()[REDUCTO_ID_PREFIX.len()..] + .trim() + .is_empty() + { + return Err(OcrRequestError::RequestField { + path: "document file id".into(), + } + .into()); + } + return Ok(document); + } + let inline = InlineDocument::parse(document.source())?.ok_or(OcrRequestError::ReductoSource)?; + let mime = inline.mime_type().to_string(); + let bytes = inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?; + let part = reqwest::multipart::Part::bytes(bytes) + .file_name("document") + .mime_str(&mime) + .map_err(|_| OcrRequestError::InvalidDataUri)?; + let builder = client + .provider_http() + .post(get_complete_url(connection.api_base.as_deref(), "upload")?) + .multipart(reqwest::multipart::Form::new().part("file", part)) + .timeout(connection.timeout); + let builder = crate::http_utils::with_headers( + builder, + headers, + crate::http_utils::HeaderPolicy::Except(&["content-type", "content-length"]), + ); + let response = crate::http_utils::http_request(builder) + .await + .map_err(crate::error::TransportError::from)?; + let uploaded = crate::ocr::client::read_json_response::< + crate::ocr::codecs::reducto::ReductoUploadResponse, + >(response, false) + .await? + .data; + let file_id = uploaded + .file_id + .as_deref() + .map(str::trim) + .filter(|id| !id.is_empty()); + let Some(file_id) = file_id else { + return Err(OcrResponseError::ResponseField { + path: "file_id".into(), + } + .into()); + }; + Ok(document.with_source(file_id.to_string())) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn explicit_key_precedes_environment_key() { + let connection = OcrConnection { + api_key: Some("passed-key".into()), + ..Default::default() + }; + let headers = validate_environment(&connection, &|_| Some("env-key".into())).unwrap(); + assert_eq!(headers[0].1, "Bearer passed-key"); + } + + #[test] + fn blank_explicit_key_uses_environment_key() { + let connection = OcrConnection { + api_key: Some(" ".into()), + ..Default::default() + }; + let headers = validate_environment(&connection, &|_| Some(" env-key ".into())).unwrap(); + assert_eq!(headers[0].1, "Bearer env-key"); + } + + #[test] + fn existing_authorization_skips_key_lookup() { + let connection = OcrConnection { + extra_headers: vec![("authorization".into(), "Bearer existing".into())], + ..Default::default() + }; + assert_eq!( + validate_environment(&connection, &|_| None).unwrap(), + connection.extra_headers + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs new file mode 100644 index 00000000000..a49f8105e26 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs @@ -0,0 +1,45 @@ +use super::super::OcrAdapter; +use crate::ocr::OcrClient; +use crate::ocr::codecs::reducto::{self, ReductoResponse, ReductoV3Params}; +use crate::ocr::error::{OcrError, OcrResponseError}; +use crate::ocr::prepare::{ + _prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env, + guardrail_document, merge_extra_params, +}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; + +#[derive(Clone, Debug)] +pub(crate) struct ReductoV3Adapter; + +impl OcrAdapter for ReductoV3Adapter { + type ProviderResponse = ReductoResponse; + const PROVIDER: OcrProvider = OcrProvider::Reducto; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let ParsedProviderParams { + known: params, + extra_params, + } = _prepare_ocr_request::(request)?; + let headers = super::validate_environment(&request.connection, &credential_env)?; + let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?; + let document = guardrail_document(request, &url).await?; + let document = + super::prepare_document(client, document, &request.connection, &headers).await?; + let body = reducto::transform_v3_ocr_request(&request.model, document, ¶ms)?; + let body = merge_extra_params(&body, extra_params)?; + build_http_request(client, request, &url, &headers, &body) + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + reducto::transform_ocr_response(&request.model, response) + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs new file mode 100644 index 00000000000..ef188f8b9ac --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs @@ -0,0 +1,134 @@ +use super::super::OcrAdapter; +use super::validate_destination; +use crate::Error; +use crate::auth::vertex::{self, VertexConfig}; +use crate::ocr::OcrClient; +use crate::ocr::codecs::deepseek::{self, DeepSeekOcrParams, DeepSeekOcrResponse}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{ + _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, +}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; +use crate::url_utils::ApiUrl; +const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com"; +const MODEL_NAMESPACE: &str = "deepseek-ai"; +const DEFAULT_LOCATION: &str = "us-central1"; + +#[derive(Clone, Debug)] +pub(crate) struct VertexDeepSeekAdapter; + +impl OcrAdapter for VertexDeepSeekAdapter { + type ProviderResponse = DeepSeekOcrResponse; + const PROVIDER: OcrProvider = OcrProvider::VertexAi; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + validate_destination(&request.connection)?; + let ParsedProviderParams { + known: params, + extra_params: _extra_params, + } = _prepare_ocr_request::(request)?; + let config = VertexConfig::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)?; + let authentication = client + .vertex_auth() + .validate_environment( + request.connection.extra_headers.clone(), + request.connection.api_key.as_deref(), + &config, + &credential_env, + ) + .await + .map_err(Error::from)?; + let location = vertex::get_vertex_ai_location(&config, &credential_env) + .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); + let url = get_complete_url( + request.connection.api_base.as_deref(), + &authentication.project_id, + &location, + )?; + let document = request.document.clone(); + let body = + deepseek::transform_ocr_request(&provider_model(&request.model), document, ¶ms)?; + transform_request_body(client, request, &url, &authentication.headers, body, |_| { + Ok(()) + }) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + deepseek::transform_ocr_response(&request.model, response) + } +} + +fn provider_model(model: &str) -> String { + if model.starts_with(&format!("{MODEL_NAMESPACE}/")) { + model.to_string() + } else { + format!("{MODEL_NAMESPACE}/{model}") + } +} + +fn get_complete_url( + api_base: Option<&str>, + project: &str, + location: &str, +) -> Result { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(DEFAULT_API_BASE); + ApiUrl::parse(base) + .and_then(|url| { + url.complete_path(&[ + "v1", + "projects", + project, + "locations", + location, + "endpoints", + "openapi", + "chat", + "completions", + ]) + }) + .map(|url| url.into_string()) + .map_err(|_| { + OcrRequestError::RequestField { + path: "api_base".into(), + } + .into() + }) +} + +#[cfg(test)] +mod tests { + use super::{get_complete_url, provider_model}; + + #[test] + fn adapter_owns_model_namespace_and_endpoint() { + assert_eq!( + provider_model("deepseek-ocr-maas"), + "deepseek-ai/deepseek-ocr-maas" + ); + assert_eq!( + provider_model("deepseek-ai/deepseek-ocr-maas"), + "deepseek-ai/deepseek-ocr-maas" + ); + assert_eq!( + get_complete_url(None, "proj-1", "europe-west4").unwrap(), + "https://aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/endpoints/openapi/chat/completions" + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs new file mode 100644 index 00000000000..f3335bf497c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs @@ -0,0 +1,154 @@ +use super::super::OcrAdapter; +use super::validate_destination; +use crate::Error; +use crate::auth::vertex::{self, VertexConfig}; +use crate::ocr::OcrClient; +use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse}; +use crate::ocr::document::{inline_remote_document, validate_inline_document}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{ + _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, +}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; +use crate::url_utils::ApiUrl; +const DEFAULT_LOCATION: &str = "us-central1"; + +#[derive(Clone, Debug)] +pub(crate) struct VertexMistralAdapter; + +impl OcrAdapter for VertexMistralAdapter { + type ProviderResponse = MistralOcrResponse; + const PROVIDER: OcrProvider = OcrProvider::VertexAi; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + validate_destination(&request.connection)?; + let ParsedProviderParams { + known: params, + extra_params: _extra_params, + } = _prepare_ocr_request::(request)?; + let config = VertexConfig::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)?; + let authentication = client + .vertex_auth() + .validate_environment( + request.connection.extra_headers.clone(), + request.connection.api_key.as_deref(), + &config, + &credential_env, + ) + .await + .map_err(Error::from)?; + let location = vertex::get_vertex_ai_location(&config, &credential_env) + .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); + let url = get_complete_url( + request.connection.api_base.as_deref(), + &authentication.project_id, + &location, + &request.model, + )?; + let document = inline_remote_document( + client.document_fetcher(), + request.document.clone(), + &request.connection, + ) + .await?; + let body = mistral::transform_ocr_request(&request.model, document, ¶ms)?; + transform_request_body( + client, + request, + &url, + &authentication.headers, + body, + |body| validate_inline_document(&body.document), + ) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + mistral::transform_ocr_response(&request.model, response) + } +} + +fn get_complete_url( + api_base: Option<&str>, + project: &str, + location: &str, + model: &str, +) -> Result { + validate_location(location)?; + let default_base = format!("https://{location}-aiplatform.googleapis.com"); + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(&default_base); + let prediction = format!("{model}:rawPredict"); + ApiUrl::parse(base) + .and_then(|url| { + url.complete_path(&[ + "v1", + "projects", + project, + "locations", + location, + "publishers", + "mistralai", + "models", + &prediction, + ]) + }) + .map(|url| url.into_string()) + .map_err(|_| { + OcrRequestError::RequestField { + path: "api_base".into(), + } + .into() + }) +} + +fn validate_location(location: &str) -> Result<(), OcrError> { + let valid = !location.is_empty() + && location + .bytes() + .all(|value| value.is_ascii_lowercase() || value.is_ascii_digit() || value == b'-') + && location + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphanumeric) + && location + .as_bytes() + .last() + .is_some_and(u8::is_ascii_alphanumeric); + if valid { + return Ok(()); + } + Err(OcrRequestError::RequestField { + path: "vertex_location".into(), + } + .into()) +} + +#[cfg(test)] +mod tests { + use super::get_complete_url; + + #[test] + fn endpoint_uses_location_project_and_model() { + assert_eq!( + get_complete_url(None, "proj-1", "europe-west4", "mistral-ocr-maas").unwrap(), + "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + assert!(get_complete_url(None, "proj-1", "attacker.example/path", "model").is_err()); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs new file mode 100644 index 00000000000..270c41e647d --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs @@ -0,0 +1,21 @@ +mod deepseek; +mod mistral; + +use crate::Error; +use crate::auth::InputSource; +use crate::auth::error::AuthConfigurationError; +use crate::ocr::error::OcrError; +use crate::ocr::types::OcrConnection; + +pub(crate) use deepseek::VertexDeepSeekAdapter; +pub(crate) use mistral::VertexMistralAdapter; + +fn validate_destination(connection: &OcrConnection) -> Result<(), OcrError> { + if connection.api_base.is_some() && connection.api_base_source == InputSource::Request { + return Err(Error::from(crate::AuthError::Configuration( + AuthConfigurationError::RequestVertexCredentialDestination, + )) + .into()); + } + Ok(()) +} diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs new file mode 100644 index 00000000000..ab2d098d0bb --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -0,0 +1,111 @@ +use std::sync::OnceLock; +use std::time::Duration; + +use serde::de::DeserializeOwned; + +use super::error::OcrError; +use super::handler::perform_ocr_request; +use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; +use super::wire::{DecodedOcrResponse, decode_response}; +use crate::Error; +use crate::auth::vertex::VertexAuth; +use crate::constants::OCR_CONNECT_TIMEOUT_SECS; +use crate::error::TransportError; +use crate::media::MediaFetcher; + +#[derive(Clone)] +pub struct OcrClient { + provider_http: reqwest::Client, + polling_http: reqwest::Client, + document_fetcher: MediaFetcher, + vertex_auth: VertexAuth, +} + +impl OcrClient { + pub fn new(provider_http: reqwest::Client) -> Result { + let document_fetcher = MediaFetcher::new().map_err(TransportError::from)?; + Ok(Self { + provider_http, + polling_http: no_redirect_http()?, + document_fetcher, + vertex_auth: VertexAuth::default(), + }) + } + + #[tracing::instrument( + name = "ocr", + target = "litellm::function_trace", + level = "trace", + skip_all + )] + pub async fn perform(&self, request: LiteLLMOcrRequest) -> Result { + perform_ocr_request(self, request).await + } + + pub(crate) fn provider_http(&self) -> &reqwest::Client { + &self.provider_http + } + + pub(crate) fn polling_http(&self) -> &reqwest::Client { + &self.polling_http + } + + pub(crate) fn document_fetcher(&self) -> &MediaFetcher { + &self.document_fetcher + } + + pub(crate) fn vertex_auth(&self) -> &VertexAuth { + &self.vertex_auth + } + + #[cfg(test)] + pub(crate) fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self { + Self { + provider_http, + polling_http: no_redirect_http().expect("test polling client builds"), + document_fetcher: MediaFetcher::for_test(document_http), + vertex_auth: VertexAuth::default(), + } + } +} + +fn no_redirect_http() -> Result { + reqwest::Client::builder() + .connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS)) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(TransportError::from) +} + +pub async fn ocr(request: LiteLLMOcrRequest) -> Result { + static CLIENT: OnceLock> = OnceLock::new(); + let client = CLIENT + .get_or_init(|| { + reqwest::Client::builder() + .connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS)) + .build() + .map_err(TransportError::from) + .and_then(OcrClient::new) + }) + .clone()?; + client.perform(request).await +} + +pub async fn read_json_response( + response: reqwest::Response, + native: bool, +) -> Result, OcrError> { + let status = response.status(); + let bytes = response + .bytes() + .await + .map_err(crate::error::TransportError::from)?; + if !status.is_success() { + return Err(crate::error::TransportError::Http { + status: status.as_u16(), + body: crate::http_utils::truncate_error_body(&String::from_utf8_lossy(&bytes)), + } + .into()); + } + Ok(decode_response(&bytes, native)?) +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs new file mode 100644 index 00000000000..682b3addde7 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs @@ -0,0 +1,5 @@ +mod transformation; +mod types; + +pub(crate) use transformation::{transform_ocr_request, transform_ocr_response}; +pub(crate) use types::{DeepSeekOcrParams, DeepSeekOcrResponse}; diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs new file mode 100644 index 00000000000..98cfc0db78d --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs @@ -0,0 +1,98 @@ +use serde::de::IntoDeserializer; +use serde_json::{Value, json}; + +use super::types::*; +use crate::ocr::error::{OcrRequestError, OcrResponseError}; +use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub(crate) fn transform_ocr_request( + provider_model: &str, + document: OcrDocument, + params: &DeepSeekOcrParams, +) -> Result { + if document.source().is_empty() { + return Err(OcrRequestError::MissingField("document URL")); + } + Ok(DeepSeekOcrRequest { + model: provider_model.to_string(), + messages: vec![DeepSeekOcrMessage { + role: UserRole::User, + content: vec![document], + }], + params: params.clone(), + }) +} + +pub(crate) fn transform_ocr_response( + model: &str, + response: DeepSeekOcrResponse, +) -> Result { + let content = response + .choices + .into_iter() + .next() + .and_then(|choice| choice.message.content) + .ok_or(OcrResponseError::EmptyContent)?; + let decoded = decode_content(content)?; + let pages = match decoded.result.pages { + Some(pages) if !pages.is_empty() => pages + .into_iter() + .map(|page| serde_json::to_value(page).expect("DeepSeek page serializes")) + .collect(), + _ => vec![json!({ + "index":0, + "markdown":decoded.fallback_markdown, + "images":null + })], + }; + Ok(LiteLLMOcrResponse { + pages, + model: decoded.result.model.unwrap_or_else(|| model.to_string()), + document_annotation: decoded.result.document_annotation, + usage_info: decoded.result.usage_info.or(response.usage), + object: "ocr".into(), + extra_fields: decoded.result.extra_fields, + provider_native_response: None, + }) +} + +struct DecodedContent { + result: DeepSeekOcrResult, + fallback_markdown: String, +} + +fn decode_content(content: DeepSeekContent) -> Result { + let (result, fallback_markdown) = match content { + DeepSeekContent::Text(text) if text.is_empty() => { + return Err(OcrResponseError::EmptyContent); + } + DeepSeekContent::Text(text) => (decode_json_content(&text)?, text), + DeepSeekContent::Object(object) => { + let fallback = + serde_json::to_string(&object).map_err(|_| OcrResponseError::ResponseField { + path: "choices[0].message.content".into(), + })?; + (Some(object), fallback) + } + }; + Ok(DecodedContent { + result: result.unwrap_or_default(), + fallback_markdown, + }) +} + +fn decode_json_content(text: &str) -> Result, OcrResponseError> { + if !text.trim_start().starts_with('{') { + return Ok(None); + } + let value = match serde_json::from_str::(text) { + Ok(value) => value, + Err(_) => return Ok(None), + }; + serde_path_to_error::deserialize(value.into_deserializer()) + .map(Some) + .map_err(|error| OcrResponseError::ResponseField { + path: format!("choices[0].message.content.{}", error.path()), + }) +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs new file mode 100644 index 00000000000..0ce2d9913f7 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs @@ -0,0 +1,95 @@ +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub(crate) struct DeepSeekOcrParams { + #[serde(skip_serializing_if = "Option::is_none")] + pub stream: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub n: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stop: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum StopSequences { + One(String), + Many(Vec), +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct DeepSeekOcrRequest { + pub model: String, + pub messages: Vec, + #[serde(flatten)] + pub params: DeepSeekOcrParams, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct DeepSeekOcrMessage { + pub role: UserRole, + pub content: Vec, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum UserRole { + User, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DeepSeekOcrResponse { + #[serde(default)] + pub choices: Vec, + pub usage: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DeepSeekChoice { + pub message: DeepSeekResponseMessage, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DeepSeekResponseMessage { + pub content: Option, +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(untagged)] +pub(crate) enum DeepSeekContent { + Text(String), + Object(DeepSeekOcrResult), +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub(crate) struct DeepSeekOcrResult { + #[serde(skip_serializing_if = "Option::is_none")] + pub pages: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub usage_info: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub document_annotation: Option, + #[serde(flatten)] + pub extra_fields: Map, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct DeepSeekPage { + #[serde(default)] + pub index: i64, + #[serde(default)] + pub markdown: String, + pub images: Option, + pub dimensions: Option, + #[serde(flatten)] + pub extra_fields: Map, +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/mod.rs new file mode 100644 index 00000000000..8031f2124a3 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/mod.rs @@ -0,0 +1,9 @@ +mod params; +mod transformation; +mod types; + +pub(crate) use params::{decode_input_params, map_ocr_params}; +pub(crate) use transformation::{transform_ocr_request, transform_ocr_response}; +pub(crate) use types::{ + AzureDocumentIntelligenceOperation, DocumentIntelligenceParams, OperationStatus, +}; diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs new file mode 100644 index 00000000000..85d1dafa542 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs @@ -0,0 +1,195 @@ +use std::collections::BTreeSet; + +use serde_json::{Map, Value}; + +use super::types::{ + DocumentIntelligenceInputParams, DocumentIntelligenceParams, FeaturesInput, PagesInput, +}; +use crate::ocr::error::OcrRequestError; +use crate::ocr::prepare::ParsedProviderParams; + +pub(crate) fn decode_input_params( + params: Map, + prefix: &str, +) -> Result, OcrRequestError> { + if let Some(Value::Array(pages)) = params.get("pages") { + if pages.iter().any(Value::is_boolean) { + return Err(OcrRequestError::Pages("boolean page index".into())); + } + if pages + .iter() + .any(|page| page.is_number() && page.as_i64().is_none()) + { + return Err(OcrRequestError::Pages("page index is out of range".into())); + } + if !pages.iter().all(Value::is_i64) && !pages.iter().all(Value::is_string) { + return Err(OcrRequestError::Pages("mixed page element types".into())); + } + } + crate::ocr::wire::decode_request_value(Value::Object(params), prefix) +} + +pub(crate) fn map_ocr_params( + params: DocumentIntelligenceInputParams, +) -> Result { + Ok(DocumentIntelligenceParams { + pages: params.pages.map(normalize_pages).transpose()?.flatten(), + features: params + .features + .map(normalize_features) + .transpose()? + .flatten(), + }) +} + +fn normalize_pages(pages: PagesInput) -> Result, OcrRequestError> { + let normalized = match pages { + PagesInput::ZeroBasedIndices(indices) => { + if indices.is_empty() { + return Ok(None); + } + indices + .into_iter() + .map(|page| { + if page < 0 { + return Err(OcrRequestError::Pages("negative page index".into())); + } + page.checked_add(1) + .ok_or_else(|| OcrRequestError::Pages("page index is out of range".into())) + }) + .collect::, _>>()? + .into_iter() + .map(|page| page.to_string()) + .collect::>() + .join(",") + } + PagesInput::NativeTokens(tokens) => { + if tokens.is_empty() { + return Ok(None); + } + tokens + .iter() + .map(|token| token.trim()) + .collect::>() + .join(",") + } + PagesInput::NativeRange(range) => range + .split(',') + .map(str::trim) + .collect::>() + .join(","), + }; + if !normalized.split(',').all(valid_page_token) { + return Err(OcrRequestError::Pages("invalid native page range".into())); + } + Ok(Some(normalized)) +} + +fn valid_page_token(token: &str) -> bool { + let mut parts = token.split('-'); + let start = parts.next().unwrap_or_default(); + if start.is_empty() || !start.chars().all(|character| character.is_ascii_digit()) { + return false; + } + match parts.next() { + None => true, + Some(end) => { + !end.is_empty() + && end.chars().all(|character| character.is_ascii_digit()) + && parts.next().is_none() + } + } +} + +fn normalize_features(features: FeaturesInput) -> Result, OcrRequestError> { + let tokens = match features { + FeaturesInput::Names(names) => names, + FeaturesInput::CommaSeparated(names) => names.split(',').map(str::to_string).collect(), + }; + if tokens.is_empty() { + return Ok(None); + } + let normalized = tokens.iter().map(|token| token.trim()).collect::>(); + if !normalized.iter().all(|token| { + let Some((first, rest)) = token.as_bytes().split_first() else { + return false; + }; + first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric) + }) { + return Err(OcrRequestError::Features); + } + Ok(Some(normalized.join(","))) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + fn map(value: Value) -> Result { + let fields = value.as_object().unwrap().clone(); + map_ocr_params(decode_input_params(fields, "optional_params")?.known) + } + + #[test] + fn input_params_retain_unknown_fields() { + let parsed = decode_input_params( + json!({ + "pages": [0], + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + }) + .as_object() + .unwrap() + .clone(), + "optional_params", + ) + .unwrap(); + + assert_eq!( + parsed.known.pages, + Some(PagesInput::ZeroBasedIndices(vec![0])) + ); + assert_eq!(parsed.extra_params["future_ocr_option"], true); + assert_eq!( + parsed.extra_params["extra_body"], + json!({"provider_option": "value"}) + ); + assert_eq!( + serde_json::to_value(map_ocr_params(parsed.known).unwrap()).unwrap(), + json!({"pages": "1", "features": null}) + ); + } + + #[rstest] + #[case(json!(["keyValuePairs"]), "keyValuePairs")] + #[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")] + #[case(json!("keyValuePairs"), "keyValuePairs")] + #[case(json!("keyValuePairs,languages"), "keyValuePairs,languages")] + #[case(json!("keyValuePairs, languages"), "keyValuePairs,languages")] + fn feature_mapping_matches_python(#[case] input: Value, #[case] expected: &str) { + assert_eq!( + map(json!({"features": input})).unwrap().features.as_deref(), + Some(expected) + ); + } + + #[rstest] + #[case(json!("keyValuePairs&pages=9"))] + #[case(json!("key value pairs"))] + #[case(json!(""))] + #[case(json!([1, 2]))] + #[case(json!([["keyValuePairs"]]))] + #[case(json!({"feature":"keyValuePairs"}))] + #[case(json!(5))] + fn invalid_feature_mapping_matches_python(#[case] input: Value) { + assert!(map(json!({"features": input})).is_err()); + } + + #[test] + fn empty_feature_list_is_omitted() { + assert_eq!(map(json!({"features": []})).unwrap().features, None); + } +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs new file mode 100644 index 00000000000..2b848fcfb7a --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs @@ -0,0 +1,111 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +use serde_json::{Map, Value, json}; + +use super::types::*; +use crate::constants::{AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT, AZURE_DI_DEFAULT_WIDTH}; +use crate::ocr::document::InlineDocument; +use crate::ocr::error::{OcrRequestError, OcrResponseError}; +use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub(crate) fn transform_ocr_request( + document: OcrDocument, +) -> Result { + let source = document.source(); + if source.is_empty() { + return Err(OcrRequestError::MissingField("document URL")); + } + Ok(if let Some(document) = InlineDocument::parse(source)? { + DocumentIntelligenceRequest::Base64Source( + STANDARD.encode(document.decode(crate::constants::OCR_INLINE_MAX_BYTES)?), + ) + } else { + DocumentIntelligenceRequest::UrlSource(source.to_string()) + }) +} + +pub(crate) fn transform_ocr_response( + model: &str, + response: AzureDocumentIntelligenceOperation, +) -> Result { + if response.status != Some(OperationStatus::Succeeded) { + return Err(OcrResponseError::OperationStatus( + response + .status + .map(|status| status.to_string()) + .unwrap_or_else(|| "None".into()), + )); + } + let result = response.analyze_result.unwrap_or_default(); + let pages = result + .pages + .into_iter() + .map(normalize_page) + .collect::, _>>()?; + let pages_processed = pages.len(); + let mut extra_fields = Map::new(); + extra_fields.insert("content".into(), option_value(result.content)); + extra_fields.insert("tables".into(), option_value(result.tables)); + extra_fields.insert( + "key_value_pairs".into(), + option_value(result.key_value_pairs), + ); + Ok(LiteLLMOcrResponse { + pages, + model: model.into(), + document_annotation: None, + usage_info: Some(json!({"pages_processed":pages_processed})), + object: "ocr".into(), + extra_fields, + provider_native_response: None, + }) +} + +fn normalize_page(page: AzureDocumentIntelligencePage) -> Result { + let index = page + .page_number + .unwrap_or(1) + .checked_sub(1) + .ok_or(OcrResponseError::NumericRange("page.pageNumber"))?; + let scale = if page.unit.as_deref().unwrap_or("inch") == "inch" { + AZURE_DI_DEFAULT_DPI as f64 + } else { + 1.0 + }; + let width = pixel_dimension( + page.width.unwrap_or(AZURE_DI_DEFAULT_WIDTH), + scale, + "page.width", + )?; + let height = pixel_dimension( + page.height.unwrap_or(AZURE_DI_DEFAULT_HEIGHT), + scale, + "page.height", + )?; + let markdown = page + .lines + .iter() + .map(|line| line.content.as_deref().unwrap_or_default()) + .collect::>() + .join("\n"); + Ok(json!({ + "index":index, + "markdown":markdown, + "images":null, + "dimensions":{"width":width,"height":height,"dpi":AZURE_DI_DEFAULT_DPI} + })) +} + +fn pixel_dimension(value: f64, scale: f64, field: &'static str) -> Result { + let value = value * scale; + if !value.is_finite() || value < i64::MIN as f64 || value > i64::MAX as f64 { + return Err(OcrResponseError::NumericRange(field)); + } + Ok(value.trunc() as i64) +} + +fn option_value(value: Option) -> Value { + value + .and_then(|value| serde_json::to_value(value).ok()) + .unwrap_or(Value::Null) +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/types.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/types.rs new file mode 100644 index 00000000000..793f4547e99 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/types.rs @@ -0,0 +1,138 @@ +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::{Map, Value}; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum PagesInput { + ZeroBasedIndices(Vec), + NativeTokens(Vec), + NativeRange(String), +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum FeaturesInput { + Names(Vec), + CommaSeparated(String), +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub(crate) struct DocumentIntelligenceInputParams { + pub pages: Option, + pub features: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub(crate) struct DocumentIntelligenceParams { + pub pages: Option, + pub features: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) enum DocumentIntelligenceRequest { + #[serde(rename = "urlSource")] + UrlSource(String), + #[serde(rename = "base64Source")] + Base64Source(String), +} + +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum OperationStatus { + Succeeded, + Running, + NotStarted, + Failed, + Unknown(String), +} + +impl<'de> Deserialize<'de> for OperationStatus { + fn deserialize>(deserializer: D) -> Result { + Ok(match String::deserialize(deserializer)?.as_str() { + "succeeded" => Self::Succeeded, + "running" => Self::Running, + "notStarted" => Self::NotStarted, + "failed" => Self::Failed, + value => Self::Unknown(value.to_string()), + }) + } +} + +impl std::fmt::Display for OperationStatus { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(match self { + Self::Succeeded => "succeeded", + Self::Running => "running", + Self::NotStarted => "notStarted", + Self::Failed => "failed", + Self::Unknown(value) => value, + }) + } +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct AzureDocumentIntelligenceOperation { + pub status: Option, + #[serde(rename = "analyzeResult")] + pub analyze_result: Option, +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub(crate) struct AzureDocumentIntelligenceAnalyzeResult { + pub content: Option, + #[serde(default)] + pub pages: Vec, + pub tables: Option>>, + #[serde(rename = "keyValuePairs")] + pub key_value_pairs: Option>>, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct AzureDocumentIntelligencePage { + #[serde(rename = "pageNumber", default, deserialize_with = "optional_i64")] + pub page_number: Option, + #[serde(default, deserialize_with = "optional_f64")] + pub width: Option, + #[serde(default, deserialize_with = "optional_f64")] + pub height: Option, + pub unit: Option, + #[serde(default)] + pub lines: Vec, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct AzureDocumentIntelligenceLine { + pub content: Option, +} + +fn optional_i64<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + match Option::::deserialize(deserializer)? { + None | Some(Value::Null) => Ok(None), + Some(Value::Number(number)) => number + .as_i64() + .map(Some) + .ok_or_else(|| serde::de::Error::custom("expected an integer")), + Some(Value::String(value)) => value + .parse::() + .map(Some) + .map_err(|_| serde::de::Error::custom("expected an integer")), + Some(_) => Err(serde::de::Error::custom("expected an integer")), + } +} + +fn optional_f64<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + match Option::::deserialize(deserializer)? { + None | Some(Value::Null) => Ok(None), + Some(Value::Number(number)) => number + .as_f64() + .filter(|value| value.is_finite()) + .map(Some) + .ok_or_else(|| serde::de::Error::custom("expected a finite number")), + Some(Value::String(value)) => value + .parse::() + .ok() + .filter(|value| value.is_finite()) + .map(Some) + .ok_or_else(|| serde::de::Error::custom("expected a finite number")), + Some(_) => Err(serde::de::Error::custom("expected a number")), + } +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs new file mode 100644 index 00000000000..eea4254779e --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs @@ -0,0 +1,5 @@ +mod transformation; +mod types; + +pub(crate) use transformation::{transform_ocr_request, transform_ocr_response}; +pub(crate) use types::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse}; diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs new file mode 100644 index 00000000000..5bd7e555a1e --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs @@ -0,0 +1,228 @@ +use super::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse}; +use crate::ocr::error::{OcrRequestError, OcrResponseError}; +use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub(crate) fn transform_ocr_request( + model: &str, + document: OcrDocument, + params: &MistralOcrParams, +) -> Result { + Ok(MistralOcrRequest { + model: model.to_string(), + document, + params: params.clone(), + }) +} + +pub(crate) fn transform_ocr_response( + model: &str, + response: MistralOcrResponse, +) -> Result { + Ok(LiteLLMOcrResponse { + pages: response.pages, + model: response.model.unwrap_or_else(|| model.to_string()), + document_annotation: response.document_annotation, + usage_info: response.usage_info, + object: "ocr".to_string(), + extra_fields: response.extra_fields, + provider_native_response: None, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use serde_json::{Value, json}; + + fn mapped_params(value: Value) -> Value { + serde_json::to_value(serde_json::from_value::(value).unwrap()).unwrap() + } + + fn document() -> OcrDocument { + serde_json::from_value( + json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + ) + .unwrap() + } + + #[rstest] + fn extract_header_is_a_supported_ocr_param() { + assert_eq!( + mapped_params(json!({"extract_header":true}))["extract_header"], + true + ); + } + + #[rstest] + fn extract_footer_is_a_supported_ocr_param() { + assert_eq!( + mapped_params(json!({"extract_footer":false}))["extract_footer"], + false + ); + } + + #[rstest] + fn existing_ocr_params_remain_supported() { + let mapped = mapped_params(json!({ + "pages":[0,2], + "include_image_base64":true, + "image_limit":2, + "image_min_size":100, + "bbox_annotation_format":{"type":"json_schema"}, + "document_annotation_format":{"type":"json_schema"} + })); + assert_eq!(mapped["pages"], json!([0, 2])); + assert_eq!(mapped["include_image_base64"], true); + assert_eq!(mapped["image_limit"], 2); + assert_eq!(mapped["image_min_size"], 100); + assert_eq!(mapped["bbox_annotation_format"]["type"], "json_schema"); + assert_eq!(mapped["document_annotation_format"]["type"], "json_schema"); + } + + #[rstest] + fn map_ocr_params_forwards_extract_header() { + assert_eq!( + mapped_params(json!({"extract_header":true}))["extract_header"], + true + ); + } + + #[rstest] + fn map_ocr_params_forwards_extract_footer() { + assert_eq!( + mapped_params(json!({"extract_footer":true}))["extract_footer"], + true + ); + } + + #[rstest] + fn map_ocr_params_forwards_extract_header_and_footer() { + let mapped = mapped_params(json!({"extract_header":true,"extract_footer":false})); + assert_eq!(mapped["extract_header"], true); + assert_eq!(mapped["extract_footer"], false); + } + + #[rstest] + fn map_ocr_params_drops_unknown_params() { + let mapped = mapped_params(json!({"extract_header":true,"unsupported_param":"value"})); + assert_eq!(mapped["extract_header"], true); + assert!(mapped.get("unsupported_param").is_none()); + } + + #[rstest] + #[case("table_format", json!("html"))] + #[case("confidence_scores_granularity", json!("word"))] + #[case("document_annotation_prompt", json!("extract"))] + #[case("include_blocks", json!(true))] + #[case("id", json!("req-123"))] + fn new_ocr_params_are_supported(#[case] name: &str, #[case] value: Value) { + assert_eq!(mapped_params(json!({name:value.clone()}))[name], value); + } + + #[rstest] + #[case("table_format", json!("html"))] + #[case("confidence_scores_granularity", json!("word"))] + #[case("document_annotation_prompt", json!("extract"))] + #[case("include_blocks", json!(true))] + #[case("id", json!("req-123"))] + fn map_ocr_params_forwards_new_ocr_params(#[case] name: &str, #[case] value: Value) { + assert_eq!(mapped_params(json!({name:value.clone()}))[name], value); + } + + #[rstest] + #[case("pages", json!([0, 2]))] + #[case("include_image_base64", json!(true))] + #[case("image_limit", json!(2))] + #[case("image_min_size", json!(100))] + #[case("bbox_annotation_format", json!({"type":"json_schema"}))] + #[case("document_annotation_format", json!({"type":"json_schema"}))] + #[case("document_annotation_prompt", json!("extract"))] + #[case("extract_header", json!(true))] + #[case("extract_footer", json!(false))] + #[case("table_format", json!("html"))] + #[case("confidence_scores_granularity", json!("word"))] + #[case("include_blocks", json!(true))] + #[case("id", json!("req-123"))] + fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { + let params: MistralOcrParams = + serde_json::from_value(json!({name: value.clone()})).unwrap(); + let result = + serde_json::to_value(transform_ocr_request("model", document(), ¶ms).unwrap()) + .unwrap(); + assert_eq!(result["model"], "model"); + assert_eq!(result[name], value); + } + + #[rstest] + #[case("table_format", json!("html"))] + #[case("confidence_scores_granularity", json!("word"))] + #[case("document_annotation_prompt", json!("extract"))] + #[case("id", json!("req-123"))] + #[case("extract_header", json!(true))] + #[case("include_blocks", json!(true))] + #[case("pages", json!([0,1]))] + fn transform_ocr_request_includes_each_optional_param( + #[case] name: &str, + #[case] value: Value, + ) { + let params: MistralOcrParams = serde_json::from_value(json!({name:value.clone()})).unwrap(); + let result = serde_json::to_value( + transform_ocr_request("mistral-ocr-latest", document(), ¶ms).unwrap(), + ) + .unwrap(); + assert_eq!(result[name], value); + assert_eq!(result["model"], "mistral-ocr-latest"); + } + + #[rstest] + fn transform_ocr_request_includes_multiple_new_params() { + let params: MistralOcrParams = serde_json::from_value(json!({ + "table_format":"html", + "confidence_scores_granularity":"page", + "extract_header":true + })) + .unwrap(); + let result = serde_json::to_value( + transform_ocr_request("mistral-ocr-latest", document(), ¶ms).unwrap(), + ) + .unwrap(); + assert_eq!(result["table_format"], "html"); + assert_eq!(result["confidence_scores_granularity"], "page"); + assert_eq!(result["extract_header"], true); + } + + #[rstest] + fn transform_ocr_response_preserves_blocks_and_confidence_scores() { + let response: MistralOcrResponse = serde_json::from_value(json!({ + "pages":[{"index":0,"markdown":"hello","blocks":[{"type":"title"}],"confidence_scores":{"mean":0.99}}], + "model":"returned-model", + "usage_info":{"pages_processed":1} + })) + .unwrap(); + let result = transform_ocr_response("model", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["blocks"][0]["type"], "title"); + assert_eq!(result["pages"][0]["confidence_scores"]["mean"], 0.99); + } + + #[rstest] + fn transform_ocr_response_preserves_ocr4_page_fields() { + let page = json!({ + "index":0, + "markdown":"table page", + "tables":[{"rows":2,"cols":3}], + "hyperlinks":["https://example.com"], + "header":"header", + "footer":"footer" + }); + let response: MistralOcrResponse = + serde_json::from_value(json!({"pages":[page.clone()]})).unwrap(); + let result = transform_ocr_response("model", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0], page); + } +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs new file mode 100644 index 00000000000..0e601cd8319 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs @@ -0,0 +1,53 @@ +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +use crate::ocr::types::OcrDocument; + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub(crate) struct MistralOcrParams { + #[serde(skip_serializing_if = "Option::is_none")] + pub pages: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub include_image_base64: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub image_limit: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub image_min_size: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub bbox_annotation_format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub document_annotation_format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub document_annotation_prompt: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub extract_header: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub extract_footer: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub table_format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub confidence_scores_granularity: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub include_blocks: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub id: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct MistralOcrRequest { + pub model: String, + pub document: OcrDocument, + #[serde(flatten)] + pub params: MistralOcrParams, +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub(crate) struct MistralOcrResponse { + #[serde(default)] + pub pages: Vec, + pub model: Option, + pub document_annotation: Option, + pub usage_info: Option, + #[serde(flatten)] + pub extra_fields: Map, +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/mod.rs new file mode 100644 index 00000000000..7c752749901 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod deepseek; +pub(crate) mod document_intelligence; +pub(crate) mod mistral; +pub(crate) mod reducto; diff --git a/litellm-rust/crates/core/src/ocr/codecs/reducto/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/reducto/mod.rs new file mode 100644 index 00000000000..3fff40451c6 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/reducto/mod.rs @@ -0,0 +1,9 @@ +mod transformation; +mod types; + +pub(crate) use transformation::{ + transform_legacy_ocr_request, transform_ocr_response, transform_v3_ocr_request, +}; +pub(crate) use types::{ + ReductoLegacyParams, ReductoResponse, ReductoUploadResponse, ReductoV3Params, +}; diff --git a/litellm-rust/crates/core/src/ocr/codecs/reducto/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/reducto/transformation.rs new file mode 100644 index 00000000000..7073643f6b6 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/reducto/transformation.rs @@ -0,0 +1,115 @@ +use std::collections::BTreeMap; + +use serde_json::{Value, json}; + +use super::types::*; +use crate::ocr::error::{OcrRequestError, OcrResponseError}; +use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; + +#[tracing::instrument( + name = "transform_ocr_request", + target = "litellm::function_trace", + level = "trace", + skip_all +)] +pub(crate) fn transform_v3_ocr_request( + _model: &str, + document: OcrDocument, + params: &ReductoV3Params, +) -> Result { + Ok(ReductoV3Request { + input: document.source().to_string(), + params: params.clone(), + }) +} + +#[tracing::instrument( + name = "transform_ocr_request", + target = "litellm::function_trace", + level = "trace", + skip_all +)] +pub(crate) fn transform_legacy_ocr_request( + _model: &str, + document: OcrDocument, + params: &ReductoLegacyParams, +) -> Result { + Ok(ReductoLegacyRequest { + document_url: document.source().to_string(), + options: params.enhance.as_ref().map(|_| params.clone()), + }) +} + +pub(crate) fn transform_ocr_response( + model: &str, + response: ReductoResponse, +) -> Result { + let result = match response.result { + Some(result) => result.unwrap_or_default(), + None => ReductoResult { + chunks: response.chunks, + }, + }; + let usage = response.usage.unwrap_or_default(); + Ok(LiteLLMOcrResponse { + pages: build_pages(result.chunks.unwrap_or_default()), + model: model.to_string(), + document_annotation: None, + usage_info: Some(json!({ + "pages_processed": usage.num_pages, + "credits": usage.credits, + })), + object: "ocr".to_string(), + extra_fields: serde_json::Map::new(), + provider_native_response: None, + }) +} + +fn build_pages(chunks: Vec) -> Vec { + let blocks_by_page = chunks + .iter() + .flat_map(|chunk| chunk.blocks.iter().flatten()) + .filter_map(|block| block.bbox.as_ref()?.page.map(|page| (page, block))) + .fold( + BTreeMap::>::new(), + |mut pages, (page, block)| { + pages.entry(page).or_default().push(block); + pages + }, + ); + if blocks_by_page.is_empty() { + let markdown = join_content(chunks.iter().map(|chunk| chunk.content.as_deref())); + return if markdown.is_empty() { + Vec::new() + } else { + vec![page(0, markdown, None)] + }; + } + blocks_by_page + .into_iter() + .map(|(index, blocks)| { + let markdown = join_content(blocks.iter().map(|block| block.content.as_deref())); + page( + index.saturating_sub(1).max(0), + markdown, + Some(json!(blocks)), + ) + }) + .collect() +} + +fn join_content<'a>(content: impl Iterator>) -> String { + content + .flatten() + .filter(|text| !text.is_empty()) + .collect::>() + .join("\n\n") +} + +fn page(index: i64, markdown: String, blocks: Option) -> Value { + let mut result = json!({"index":index,"markdown":markdown,"images":null}); + if let (Value::Object(fields), Some(blocks)) = (&mut result, blocks) { + fields.insert("blocks".into(), blocks); + } + result +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/reducto/types.rs b/litellm-rust/crates/core/src/ocr/codecs/reducto/types.rs new file mode 100644 index 00000000000..c03720cc8ae --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/reducto/types.rs @@ -0,0 +1,128 @@ +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::{Map, Value}; + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub(crate) struct ReductoV3Params { + #[serde(skip_serializing_if = "Option::is_none")] + pub formatting: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub retrieval: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub settings: Option>, +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub(crate) struct ReductoLegacyParams { + #[serde(skip_serializing_if = "Option::is_none")] + pub enhance: Option>, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct ReductoV3Request { + pub input: String, + #[serde(flatten)] + pub params: ReductoV3Params, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct ReductoLegacyRequest { + pub document_url: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub options: Option, +} + +#[derive(Deserialize)] +pub(crate) struct ReductoUploadResponse { + pub file_id: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct ReductoResponse { + #[serde(default, deserialize_with = "present_nullable")] + pub result: Option>, + pub usage: Option, + #[serde(default)] + pub chunks: Option>, +} + +fn present_nullable<'de, D: Deserializer<'de>, T: Deserialize<'de>>( + deserializer: D, +) -> Result>, D::Error> { + Option::::deserialize(deserializer).map(Some) +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub(crate) struct ReductoResult { + pub chunks: Option>, +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub(crate) struct ReductoUsage { + #[serde(default, deserialize_with = "optional_i64")] + pub num_pages: Option, + #[serde(default, deserialize_with = "optional_f64")] + pub credits: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct ReductoChunk { + pub content: Option, + pub blocks: Option>, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct ReductoBlock { + #[serde(skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub bbox: Option, + #[serde(flatten)] + pub extra_fields: Map, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct ReductoBoundingBox { + #[serde(default, deserialize_with = "optional_i64")] + pub page: Option, + #[serde(flatten)] + pub extra_fields: Map, +} + +fn optional_i64<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + match Option::::deserialize(deserializer)? { + None | Some(Value::Null) => Ok(None), + Some(Value::Number(number)) => number + .as_i64() + .or_else(|| number.as_f64().and_then(checked_truncated_i64)) + .map(Some) + .ok_or_else(|| serde::de::Error::custom("expected an integer")), + Some(Value::String(value)) => value + .trim() + .parse::() + .map(Some) + .map_err(|_| serde::de::Error::custom("expected an integer")), + Some(Value::Bool(value)) => Ok(Some(i64::from(value))), + Some(_) => Ok(None), + } +} + +fn optional_f64<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + match Option::::deserialize(deserializer)? { + None | Some(Value::Null) => Ok(None), + Some(Value::Number(number)) => number + .as_f64() + .map(Some) + .ok_or_else(|| serde::de::Error::custom("expected a number")), + Some(Value::String(value)) => value + .trim() + .parse::() + .map(Some) + .map_err(|_| serde::de::Error::custom("expected a number")), + Some(_) => Ok(None), + } +} + +fn checked_truncated_i64(value: f64) -> Option { + (value.is_finite() && value >= i64::MIN as f64 && value <= i64::MAX as f64) + .then(|| value.trunc() as i64) +} diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs new file mode 100644 index 00000000000..e89b1c5c569 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -0,0 +1,212 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +use data_url::mime::Mime; +use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError}; +use reqwest::Url; + +use super::error::{OcrError, OcrRequestError, OcrResponseError}; +use super::types::{OcrConnection, OcrDocument}; +use crate::constants::OCR_MAX_FETCH_REDIRECTS; +use crate::error::{MediaError, TransportError}; +use crate::media::{DownloadPolicy, MediaFetcher}; + +pub(crate) struct InlineDocument<'a>(DataUrl<'a>); + +impl<'a> InlineDocument<'a> { + pub(crate) fn parse(source: &'a str) -> Result, OcrRequestError> { + match DataUrl::process(source) { + Ok(url) => Ok(Some(Self(url))), + Err(DataUrlError::NotADataUrl) => Ok(None), + Err(DataUrlError::NoComma) => Err(OcrRequestError::InvalidDataUri), + } + } + + pub(crate) fn mime_type(&self) -> &Mime { + self.0.mime_type() + } + + pub(crate) fn decode(&self, max_bytes: usize) -> Result, OcrRequestError> { + let mut body = Vec::new(); + self.0 + .decode(|bytes| { + if bytes.len() > max_bytes.saturating_sub(body.len()) { + return Err(OcrRequestError::InlineDocumentTooLarge); + } + body.extend_from_slice(bytes); + Ok(()) + }) + .map_err(|error| match error { + DecodeError::InvalidBase64(_) => OcrRequestError::InvalidDataUri, + DecodeError::WriteError(error) => error, + })?; + Ok(body) + } +} + +pub(crate) fn validate_inline_document(document: &OcrDocument) -> Result<(), OcrRequestError> { + let inline = + InlineDocument::parse(document.source())?.ok_or(OcrRequestError::InvalidDataUri)?; + inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?; + Ok(()) +} + +pub(crate) async fn inline_remote_document( + fetcher: &MediaFetcher, + document: OcrDocument, + connection: &OcrConnection, +) -> Result { + let source = document.source(); + if !source.starts_with("http://") && !source.starts_with("https://") { + validate_inline_document(&document)?; + return Ok(document); + } + let url = Url::parse(source).map_err(|_| OcrRequestError::RequestField { + path: "document URL".into(), + })?; + let downloaded = fetcher + .fetch( + url, + DownloadPolicy { + timeout: connection.timeout, + max_bytes: connection.max_download_bytes, + max_redirects: OCR_MAX_FETCH_REDIRECTS, + }, + ) + .await + .map_err(map_media_error)?; + let result = document.with_source(format!( + "data:{};base64,{}", + downloaded.content_type, + STANDARD.encode(downloaded.bytes) + )); + validate_inline_document(&result)?; + Ok(result) +} + +fn map_media_error(error: MediaError) -> OcrError { + match error { + MediaError::BlockedUrl => OcrRequestError::BlockedDocumentUrl.into(), + MediaError::DownloadDisabled => OcrRequestError::DownloadDisabled.into(), + MediaError::DownloadTooLarge => OcrRequestError::DownloadTooLarge.into(), + MediaError::TooManyRedirects => OcrRequestError::TooManyRedirects.into(), + MediaError::MissingRedirectLocation => OcrResponseError::MissingRedirectLocation.into(), + MediaError::InvalidRedirect => OcrResponseError::InvalidRedirect.into(), + MediaError::Http(status) => TransportError::Http { + status, + body: "OCR document download failed".into(), + } + .into(), + MediaError::Timeout => { + TransportError::Network("OCR document download timed out".into()).into() + } + MediaError::Transport(error) => error.into(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::Map; + + fn document(source: &str) -> OcrDocument { + OcrDocument::DocumentUrl { + document_url: source.into(), + extra_fields: Map::new(), + } + } + + #[test] + fn decodes_data_urls_and_limits_decoded_size() { + for (source, expected) in [ + ("data:application/pdf;base64,YWJj", b"abc".as_slice()), + ("DATA:application/pdf;BASE64,YWI", b"ab".as_slice()), + ("data:,a%20b%00%FF", b"a b\0\xff".as_slice()), + ] { + let inline = InlineDocument::parse(source).unwrap().unwrap(); + assert_eq!(inline.decode(expected.len()).unwrap(), expected); + assert_eq!( + inline.decode(expected.len() - 1), + Err(OcrRequestError::InlineDocumentTooLarge) + ); + } + } + + #[test] + fn preserves_mime_parameters_and_standard_default() { + let inline = InlineDocument::parse("data:application/pdf;version=1.7;base64,YQ==") + .unwrap() + .unwrap(); + assert!(inline.mime_type().matches("application", "pdf")); + assert_eq!(inline.mime_type().get_parameter("version"), Some("1.7")); + let default = InlineDocument::parse("data:,a").unwrap().unwrap(); + assert!(default.mime_type().matches("text", "plain")); + assert_eq!( + default.mime_type().get_parameter("charset"), + Some("US-ASCII") + ); + } + + #[test] + fn rejects_invalid_inline_documents() { + for source in [ + "https://example.com/document.pdf", + "data:application/pdf;base64", + "data:application/pdf;base64,INVALID!", + ] { + assert!(validate_inline_document(&document(source)).is_err()); + } + } + + #[tokio::test] + async fn remote_conversion_preserves_kind_and_isolates_provider_credentials() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0_u8; 2048]; + let count = socket.read(&mut request).await.unwrap(); + socket + .write_all(b"HTTP/1.1 200 OK\r\nContent-Type: image/png; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc") + .await + .unwrap(); + String::from_utf8_lossy(&request[..count]).into_owned() + }); + let mut provider_headers = reqwest::header::HeaderMap::new(); + provider_headers.insert( + reqwest::header::AUTHORIZATION, + reqwest::header::HeaderValue::from_static("Bearer provider-secret"), + ); + let provider_http = reqwest::Client::builder() + .default_headers(provider_headers) + .build() + .unwrap(); + let document_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .unwrap(); + let client = super::super::OcrClient::for_test(provider_http, document_http); + let converted = inline_remote_document( + client.document_fetcher(), + OcrDocument::ImageUrl { + image_url: format!("http://{address}/image"), + extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]), + }, + &OcrConnection::default(), + ) + .await + .unwrap(); + let request = server.await.unwrap(); + + assert_eq!( + converted, + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,YWJj".into(), + extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]), + } + ); + assert!(!request.to_ascii_lowercase().contains("authorization")); + assert!(!request.contains("provider-secret")); + } +} diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs new file mode 100644 index 00000000000..522d059ec48 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -0,0 +1,85 @@ +use thiserror::Error; + +use crate::error::TransportError; + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum OcrRequestError { + #[error("Invalid `req_format`. Expected 'native' or 'litellm'.")] + RequestFormat, + #[error("invalid OCR request field: {path}")] + RequestField { path: String }, + #[error("missing required field: {0}")] + MissingField(&'static str), + #[error("invalid OCR document data URI")] + InvalidDataUri, + #[error("Reducto requires a reducto:// id or a data URI")] + ReductoSource, + #[error("inline OCR document exceeds the size limit")] + InlineDocumentTooLarge, + #[error("OCR document URL is blocked by network policy")] + BlockedDocumentUrl, + #[error("OCR document downloads are disabled")] + DownloadDisabled, + #[error("OCR document download exceeds the size limit")] + DownloadTooLarge, + #[error("OCR document download exceeded the redirect limit")] + TooManyRedirects, + #[error("invalid OCR pages: {0}")] + Pages(String), + #[error("invalid OCR features")] + Features, + #[error("OCR model cannot be a dot segment")] + DotModel, +} + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum OcrResponseError { + #[error("invalid OCR response field: {path}")] + ResponseField { path: String }, + #[error("OCR response is missing non-empty content")] + EmptyContent, + #[error("OCR document redirect is missing a location")] + MissingRedirectLocation, + #[error("OCR document redirect location is invalid")] + InvalidRedirect, + #[error("OCR operation ended with status {0}")] + OperationStatus(String), + #[error("OCR response numeric value is out of range: {0}")] + NumericRange(&'static str), +} + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum OcrPollingError { + #[error("OCR accepted response is missing a valid operation-location")] + PollLocation, + #[error("OCR operation-location must use the submission origin without credentials")] + PollOrigin, + #[error("OCR polling timed out")] + PollTimeout, +} + +#[derive(Debug, Error)] +pub enum OcrError { + #[error("{0}")] + Request(#[from] OcrRequestError), + #[error("{0}")] + Response(#[from] OcrResponseError), + #[error("{0}")] + Transport(#[from] TransportError), + #[error("{0}")] + Polling(#[from] OcrPollingError), + #[error("{0}")] + Public(#[from] crate::Error), +} + +impl From for crate::Error { + fn from(error: OcrError) -> Self { + match error { + OcrError::Request(error) => error.into(), + OcrError::Response(error) => error.into(), + OcrError::Transport(error) => error.into(), + OcrError::Polling(error) => crate::Error::InvalidResponse(error.to_string()), + OcrError::Public(error) => error, + } + } +} diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs new file mode 100644 index 00000000000..0b04319d966 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -0,0 +1,72 @@ +use super::OcrClient; +use super::adapters::OcrAdapter; +use super::hooks::OcrLifecycleHooks; +use super::registry::OcrAdapterKind; +use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; +use crate::Error; +use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext}; + +pub(crate) async fn perform_ocr_request( + client: &OcrClient, + request: LiteLLMOcrRequest, +) -> Result { + let context = CallLifecycleContext::new( + "ocr", + request.model.clone(), + request.adapter.provider().as_str(), + request + .litellm_call_id + .clone() + .unwrap_or_else(|| format!("ocr-{:032x}", rand::random::())), + ); + let hooks = OcrLifecycleHooks { + hooks: request.hooks.clone(), + provider_name: context.custom_llm_provider.clone(), + }; + CallLifecycle::default().run(context, request, &hooks, |request| async move { + macro_rules! execute_selected_adapter { + ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { + match request.adapter { + $( OcrAdapterKind::$variant => execute_ocr_provider_call(client, &$instance, request).await, )+ + } + }; + } + super::adapters::for_each_ocr_adapter!(execute_selected_adapter) + }).await +} + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +async fn execute_ocr_provider_call( + client: &OcrClient, + adapter: &A, + request: LiteLLMOcrRequest, +) -> Result { + let provider_request = adapter.prepare_request(&request, client).await?; + let url = provider_request.url().to_string(); + let headers = provider_request + .headers() + .iter() + .map(|(name, value)| { + value + .to_str() + .map(|value| (name.to_string(), value.to_string())) + .map_err(|_| super::error::OcrRequestError::RequestField { + path: "headers".into(), + }) + }) + .collect::, _>>()?; + let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts( + client.provider_http().clone(), + provider_request, + )) + .await + .map_err(crate::error::TransportError::from)?; + let decoded = adapter + .read_response(client, response, &url, &headers, &request) + .await?; + let response = adapter.transform_ocr_response(&request, decoded.data)?; + Ok(LiteLLMOcrResponse { + provider_native_response: decoded.native, + ..response + }) +} diff --git a/litellm-rust/crates/core/src/ocr/hooks.rs b/litellm-rust/crates/core/src/ocr/hooks.rs new file mode 100644 index 00000000000..7dd3c6bf8b2 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/hooks.rs @@ -0,0 +1,146 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument}; +use crate::Error; +use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; +use serde::Serialize; +use serde_json::Value; + +pub type OcrHookFuture<'a, T> = Pin> + Send + 'a>>; +pub type OcrLogFuture<'a> = Pin + Send + 'a>>; + +#[derive(Clone, Debug, Serialize)] +pub struct OcrPreCallRequest { + pub model: String, + pub custom_llm_provider: String, + pub document: OcrDocument, + pub optional_params: Value, +} + +#[derive(Clone, Debug, Serialize)] +pub struct OcrDuringCallRequest { + pub model: String, + pub custom_llm_provider: String, + pub url: String, + pub body: Value, +} + +pub trait OcrHooks: Send + Sync { + fn has_guardrails(&self) -> bool { + false + } + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + Box::pin(async move { Ok(request) }) + } + fn during_call( + &self, + request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { Ok(request) }) + } + fn success<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a LiteLLMOcrResponse, + _timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async {}) + } + fn failure<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a Error, + _timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async {}) + } +} + +pub struct NoopOcrHooks; +impl OcrHooks for NoopOcrHooks {} + +pub(crate) struct OcrLifecycleHooks { + pub hooks: Arc, + pub provider_name: String, +} + +impl CallLifecycleHooks + for OcrLifecycleHooks +{ + type PreCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; + type DuringCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; + type SuccessFuture<'a> = OcrLogFuture<'a>; + type FailureFuture<'a> = OcrLogFuture<'a>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: LiteLLMOcrRequest, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { + if !self.hooks.has_guardrails() { + return Ok(request); + } + let changed = self + .hooks + .pre_call(OcrPreCallRequest { + model: request.model.clone(), + custom_llm_provider: self.provider_name.clone(), + document: request.document, + optional_params: Value::Object(request.optional_params), + }) + .await?; + let Value::Object(optional_params) = changed.optional_params else { + return Err(super::error::OcrRequestError::RequestField { + path: "guardrail.optional_params".into(), + } + .into()); + }; + Ok(LiteLLMOcrRequest { + document: changed.document, + optional_params, + ..request + }) + }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: LiteLLMOcrRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { Ok(request) }) + } + + #[tracing::instrument( + name = "success_callback", + target = "litellm::function_trace", + level = "trace", + skip_all + )] + fn async_log_success_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a LiteLLMOcrResponse, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + self.hooks.success(context, response, timing) + } + + #[tracing::instrument( + name = "failure_callback", + target = "litellm::function_trace", + level = "trace", + skip_all + )] + fn async_log_failure_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a Error, + timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + self.hooks.failure(context, error, timing) + } +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index ec2fbb969a6..1e975c3f521 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,2 +1,39 @@ -pub mod transformation; +mod adapters; +pub mod client; +mod codecs; +mod document; +pub mod error; +mod handler; +pub mod hooks; +mod prepare; +mod registry; pub mod types; +pub mod wire; + +pub use client::{OcrClient, ocr}; +pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument}; + +#[cfg(test)] +#[path = "../../tests/azure_ai_ocr.rs"] +mod azure_ai_tests; +#[cfg(test)] +#[path = "../../tests/azure_document_intelligence_ocr.rs"] +mod azure_document_intelligence_tests; +#[cfg(test)] +#[path = "../../tests/deepseek_ocr.rs"] +mod deepseek_tests; +#[cfg(test)] +#[path = "../../tests/reducto_ocr.rs"] +mod reducto_tests; +#[cfg(test)] +#[path = "../../tests/ocr/support.rs"] +pub(crate) mod test_support; +#[cfg(test)] +#[path = "../../tests/ocr.rs"] +pub(crate) mod tests; +#[cfg(test)] +#[path = "../../tests/vertex_ai_deepseek_ocr.rs"] +mod vertex_ai_deepseek_tests; +#[cfg(test)] +#[path = "../../tests/vertex_ai_ocr.rs"] +mod vertex_ai_tests; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs new file mode 100644 index 00000000000..bf6f924088c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -0,0 +1,197 @@ +use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use serde_json::{Map, Value}; + +use super::OcrClient; +use super::error::{OcrError, OcrRequestError}; +use super::hooks::OcrDuringCallRequest; +use super::types::{LiteLLMOcrRequest, OcrDocument}; + +#[derive(Debug, Deserialize)] +pub(crate) struct ParsedProviderParams { + #[serde(flatten)] + pub known: T, + #[serde(default, flatten)] + pub extra_params: Map, +} + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub(crate) fn _prepare_ocr_request( + request: &LiteLLMOcrRequest, +) -> Result, OcrRequestError> { + super::wire::decode_request_value( + Value::Object(request.optional_params.clone()), + "optional_params", + ) +} + +pub(crate) fn merge_extra_params( + body: &B, + extra_params: Map, +) -> Result { + let Value::Object(fields) = + serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { + path: "body".into(), + })? + else { + return Err(OcrRequestError::RequestField { + path: "body".into(), + }); + }; + let extra_body = extra_params + .get("extra_body") + .and_then(Value::as_object) + .cloned() + .unwrap_or_default() + .into_iter() + .collect::>(); + Ok(Value::Object( + fields + .into_iter() + .chain( + extra_params + .into_iter() + .filter(|(name, _)| name != "extra_body"), + ) + .chain(extra_body) + .collect(), + )) +} + +pub(crate) async fn transform_request_body( + client: &OcrClient, + request: &LiteLLMOcrRequest, + url: &str, + headers: &[(String, String)], + body: B, + validate: impl FnOnce(&B) -> Result<(), OcrRequestError>, +) -> Result +where + B: Serialize + DeserializeOwned, +{ + let body = if request.hooks.has_guardrails() { + let changed = request + .hooks + .during_call(OcrDuringCallRequest { + model: request.model.clone(), + custom_llm_provider: request.adapter.provider().as_str().into(), + url: url.into(), + body: serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { + path: "body".into(), + })?, + }) + .await?; + let body = OcrWireBody::::decode(changed.body)?; + validate(&body.body)?; + body + } else { + OcrWireBody { + body, + extra: Map::new(), + } + }; + build_http_request(client, request, url, headers, &body) +} + +pub(crate) fn build_http_request( + client: &OcrClient, + request: &LiteLLMOcrRequest, + url: &str, + headers: &[(String, String)], + body: &B, +) -> Result { + let builder = client + .provider_http() + .post(url) + .json(body) + .timeout(request.connection.timeout); + crate::http_utils::with_headers(builder, headers, crate::http_utils::HeaderPolicy::All) + .build() + .map_err(crate::error::TransportError::from) + .map_err(OcrError::from) +} + +pub(crate) async fn guardrail_document( + request: &LiteLLMOcrRequest, + url: &str, +) -> Result { + if !request.hooks.has_guardrails() { + return Ok(request.document.clone()); + } + let changed = request + .hooks + .during_call(OcrDuringCallRequest { + model: request.model.clone(), + custom_llm_provider: request.adapter.provider().as_str().into(), + url: url.into(), + body: serde_json::to_value(&request.document).map_err(|_| { + OcrRequestError::RequestField { + path: "document".into(), + } + })?, + }) + .await?; + super::wire::decode_request_value(changed.body, "guardrail.document").map_err(OcrError::from) +} + +#[derive(Serialize)] +struct OcrWireBody { + #[serde(flatten)] + body: B, + #[serde(flatten)] + extra: Map, +} + +impl OcrWireBody { + fn decode(value: Value) -> Result { + let body: B = super::wire::decode_request_value(value.clone(), "guardrail.body")?; + let Value::Object(fields) = value else { + return Err(OcrRequestError::RequestField { + path: "guardrail.body".into(), + }); + }; + let known = serde_json::to_value(&body).map_err(|_| OcrRequestError::RequestField { + path: "guardrail.body".into(), + })?; + let extra = fields + .into_iter() + .filter(|(key, _)| known.get(key).is_none()) + .collect(); + Ok(Self { body, extra }) + } +} + +pub(crate) fn credential_env(name: &str) -> Option { + std::env::var(name).ok() +} +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[derive(Debug, Deserialize, PartialEq)] + struct KnownParams { + pages: Option>, + } + + #[test] + fn parsed_provider_params_separates_known_and_extra_params() { + let parsed: ParsedProviderParams = super::super::wire::decode_request_value( + json!({ + "pages": [0, 2], + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + }), + "optional_params", + ) + .unwrap(); + + assert_eq!(parsed.known.pages, Some(vec![0, 2])); + assert_eq!(parsed.extra_params["future_ocr_option"], true); + assert_eq!( + parsed.extra_params["extra_body"], + json!({"provider_option": "value"}) + ); + assert_eq!(parsed.extra_params.len(), 2); + } +} diff --git a/litellm-rust/crates/core/src/ocr/registry.rs b/litellm-rust/crates/core/src/ocr/registry.rs new file mode 100644 index 00000000000..1b20a91143b --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/registry.rs @@ -0,0 +1,128 @@ +use super::adapters::OcrAdapter; +use crate::Error; +use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; + +macro_rules! define_adapter_types { + ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + pub(crate) enum OcrAdapterKind { + $( $variant, )+ + } + + impl OcrAdapterKind { + pub(crate) const fn provider(self) -> OcrProvider { + match self { + $( Self::$variant => <$adapter>::PROVIDER, )+ + } + } + } + }; +} + +super::adapters::for_each_ocr_adapter!(define_adapter_types); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum OcrProvider { + Mistral, + AzureAi, + Reducto, + VertexAi, +} + +impl OcrProvider { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Mistral => "mistral", + Self::AzureAi => "azure_ai", + Self::Reducto => "reducto", + Self::VertexAi => "vertex_ai", + } + } +} + +pub(crate) fn resolve_wire_adapter( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result<(String, OcrAdapterKind), Error> { + let provider = + get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider { + model, + custom_llm_provider: OcrProvider::Mistral.as_str(), + }); + let typed_provider = match provider.custom_llm_provider { + "mistral" => OcrProvider::Mistral, + "azure_ai" => OcrProvider::AzureAi, + "reducto" => OcrProvider::Reducto, + "vertex_ai" => OcrProvider::VertexAi, + value => return Err(Error::InvalidProvider(value.to_string())), + }; + let adapter = match typed_provider { + OcrProvider::Mistral => OcrAdapterKind::Mistral, + OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => { + OcrAdapterKind::AzureDocumentIntelligence + } + OcrProvider::AzureAi => OcrAdapterKind::AzureMistral, + OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-legacy") => { + OcrAdapterKind::ReductoLegacy + } + OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-v3") => { + OcrAdapterKind::ReductoV3 + } + OcrProvider::Reducto => { + return Err(Error::InvalidRequest(format!( + "unsupported Reducto OCR model: {}", + provider.model + ))); + } + OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => { + OcrAdapterKind::VertexDeepSeek + } + OcrProvider::VertexAi => OcrAdapterKind::VertexMistral, + }; + Ok((provider.model.to_string(), adapter)) +} + +fn is_document_intelligence_model(model: &str) -> bool { + let model = model.to_ascii_lowercase(); + model.contains("doc-intelligence") || model.contains("documentintelligence") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn provider_models_are_preserved_without_a_local_allowlist() { + let cases = [ + ("mistral/future-ocr-model", OcrAdapterKind::Mistral), + ("azure_ai/future-ocr-model", OcrAdapterKind::AzureMistral), + ]; + + for (qualified_model, expected_adapter) in cases { + let expected_model = qualified_model.split_once('/').unwrap().1; + let (model, adapter) = resolve_wire_adapter(qualified_model, None).unwrap(); + assert_eq!(model, expected_model); + assert_eq!(adapter, expected_adapter); + } + } + + #[test] + fn unknown_reducto_models_are_rejected() { + assert!(matches!( + resolve_wire_adapter("reducto/future-parse-model", None), + Err(Error::InvalidRequest(_)) + )); + } + + #[test] + fn known_protocol_models_still_select_specialized_adapters() { + let (model, adapter) = resolve_wire_adapter("reducto/parse-legacy", None).unwrap(); + assert_eq!(model, "parse-legacy"); + assert_eq!(adapter, OcrAdapterKind::ReductoLegacy); + + let (model, adapter) = + resolve_wire_adapter("azure_ai/doc-intelligence/prebuilt-layout", None).unwrap(); + assert_eq!(model, "doc-intelligence/prebuilt-layout"); + assert_eq!(adapter, OcrAdapterKind::AzureDocumentIntelligence); + } +} diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs deleted file mode 100644 index 62299faf9ed..00000000000 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ /dev/null @@ -1,107 +0,0 @@ -use crate::Error; -use serde_json::{Map, Value}; - -use super::types::{OcrRequestData, OcrResponseData}; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum OcrAuthStrategy { - Bearer, - Header(&'static str), -} - -impl OcrAuthStrategy { - pub fn header_name(self) -> &'static str { - match self { - Self::Bearer => "authorization", - Self::Header(header_name) => header_name, - } - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum OcrResponseHandling { - Json, - AzureDocumentIntelligencePoll, -} - -pub trait OcrProviderConfig: Sync { - fn supported_ocr_params(&self) -> &'static [&'static str]; - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn map_ocr_params(&self, non_default_params: &Map) -> Map { - let mut mapped_params = Map::new(); - for (param, value) in non_default_params { - if self.supported_ocr_params().contains(¶m.as_str()) { - mapped_params.insert(param.clone(), value.clone()); - } - } - mapped_params - } - - fn transform_ocr_request( - &self, - model: &str, - document: Value, - optional_params: Map, - ) -> Result; - - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result; - - fn transform_ocr_response_with_params( - &self, - model: &str, - response_json: Value, - _optional_params: &Map, - ) -> Result { - self.transform_ocr_response(model, response_json) - } - - fn complete_url( - &self, - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn validate_environment( - &self, - headers: Vec<(String, String)>, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result, Error> { - let strategy = self.auth_strategy(); - if crate::http_utils::has_header(&headers, strategy.header_name()) { - return Ok(headers); - } - let api_key = self.resolve_api_key(api_key, env_lookup)?; - let auth_header = match strategy { - OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), - OcrAuthStrategy::Header(name) => (name.to_string(), api_key), - }; - Ok(std::iter::once(auth_header).chain(headers).collect()) - } - - fn auth_strategy(&self) -> OcrAuthStrategy { - OcrAuthStrategy::Bearer - } - - fn requires_data_uri_document(&self) -> bool { - false - } - - fn response_handling(&self) -> OcrResponseHandling { - OcrResponseHandling::Json - } -} diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 71cdb232a87..06519f86c91 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,38 +1,190 @@ +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::Duration; + use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +use super::hooks::{NoopOcrHooks, OcrHooks}; +use super::registry::{OcrAdapterKind, resolve_wire_adapter}; +use crate::Error; +use crate::auth::InputSource; +use crate::constants::OCR_HTTP_TIMEOUT_SECS; + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct OcrRequestData { - pub data: Value, - pub files: Option, +#[serde(tag = "type")] +pub enum OcrDocument { + #[serde(rename = "document_url")] + DocumentUrl { + document_url: String, + #[serde(flatten)] + extra_fields: Map, + }, + #[serde(rename = "image_url")] + ImageUrl { + image_url: String, + #[serde(flatten)] + extra_fields: Map, + }, +} + +impl OcrDocument { + pub(crate) fn source(&self) -> &str { + match self { + Self::DocumentUrl { document_url, .. } => document_url, + Self::ImageUrl { image_url, .. } => image_url, + } + } + + pub(crate) fn with_source(self, source: String) -> Self { + match self { + Self::DocumentUrl { extra_fields, .. } => Self::DocumentUrl { + document_url: source, + extra_fields, + }, + Self::ImageUrl { extra_fields, .. } => Self::ImageUrl { + image_url: source, + extra_fields, + }, + } + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum OcrResponseFormat { + #[default] + Litellm, + Native, +} + +#[derive(Clone)] +pub struct OcrConnection { + pub api_key: Option, + pub api_key_source: InputSource, + pub api_base: Option, + pub api_base_source: InputSource, + pub extra_headers: Vec<(String, String)>, + pub extra_headers_source: InputSource, + pub timeout: Duration, + pub max_download_bytes: u64, + pub poll_timeout: Duration, +} + +impl Default for OcrConnection { + fn default() -> Self { + Self { + api_key: None, + api_key_source: InputSource::Deployment, + api_base: None, + api_base_source: InputSource::Deployment, + extra_headers: Vec::new(), + extra_headers_source: InputSource::Deployment, + timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS), + max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES, + poll_timeout: Duration::from_secs(crate::constants::OCR_POLL_TIMEOUT_SECS), + } + } +} + +pub struct LiteLLMOcrRequest { + pub model: String, + pub document: OcrDocument, + pub connection: OcrConnection, + pub hooks: Arc, + pub litellm_call_id: Option, + pub optional_params: Map, + pub input_sources: BTreeMap, + pub(crate) adapter: OcrAdapterKind, +} + +impl LiteLLMOcrRequest { + pub fn new( + model: String, + document: OcrDocument, + custom_llm_provider: Option<&str>, + optional_params: Map, + ) -> Result { + let (model, adapter_kind) = resolve_wire_adapter(&model, custom_llm_provider)?; + + Ok(Self { + model, + document, + connection: OcrConnection::default(), + hooks: Arc::new(NoopOcrHooks), + litellm_call_id: None, + optional_params, + input_sources: BTreeMap::new(), + adapter: adapter_kind, + }) + } + + pub(crate) fn response_format( + &self, + ) -> Result { + self.optional_params + .get("req_format") + .map(|value| { + serde_json::from_value(value.clone()) + .map_err(|_| super::error::OcrRequestError::RequestFormat) + }) + .transpose() + .map(|format| format.unwrap_or_default()) + } + + pub fn with_host_hooks( + self, + hooks: Arc, + litellm_call_id: Option, + ) -> Self { + Self { + hooks, + litellm_call_id, + ..self + } + } } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct OcrResponseData { +pub struct LiteLLMOcrResponse { pub pages: Vec, pub model: String, pub document_annotation: Option, pub usage_info: Option, pub object: String, + #[serde(flatten)] pub extra_fields: Map, + #[serde(skip_serializing_if = "Option::is_none")] pub provider_native_response: Option, } -impl OcrResponseData { +impl LiteLLMOcrResponse { pub fn into_json(self) -> Value { - let mut response = serde_json::json!({ - "pages": self.pages, - "model": self.model, - "document_annotation": self.document_annotation, - "usage_info": self.usage_info, - "object": self.object, - }); - if let Value::Object(object) = &mut response { - object.extend(self.extra_fields); - if let Some(native_response) = self.provider_native_response { - object.insert("provider_native_response".to_string(), native_response); - } - } - response + serde_json::to_value(self).expect("OCR response fields are JSON-compatible") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn response_serialization_flattens_extra_fields_and_omits_absent_native_response() { + let response = LiteLLMOcrResponse { + pages: vec![], + model: "model".into(), + document_annotation: None, + usage_info: None, + object: "ocr".into(), + extra_fields: json!({"provider_field":"kept"}) + .as_object() + .unwrap() + .clone(), + provider_native_response: None, + }; + let serialized = response.into_json(); + assert_eq!(serialized["provider_field"], "kept"); + assert!(serialized.get("provider_native_response").is_none()); } } diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs new file mode 100644 index 00000000000..34d0a7d7b86 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -0,0 +1,171 @@ +use crate::ocr::error::OcrRequestError; +use crate::ocr::error::OcrResponseError; +use std::collections::BTreeMap; +use std::time::Duration; + +use super::hooks::{OcrDuringCallRequest, OcrPreCallRequest}; +use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument}; +use crate::Error; +use crate::auth::InputSource; +use serde::{ + Deserialize, + de::{DeserializeOwned, IntoDeserializer}, +}; +use serde_json::{Map, Value}; + +#[derive(Debug)] +pub struct DecodedOcrResponse { + pub data: T, + pub native: Option, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct OcrWireRequest { + pub model: String, + pub document: Value, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub extra_headers: Option>, + #[serde(default)] + pub optional_params: Map, + #[serde(default)] + pub input_sources: BTreeMap, + pub timeout_seconds: Option, +} + +pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> bool { + super::registry::resolve_wire_adapter(model, custom_llm_provider).is_ok() +} + +pub fn decode_request(wire: OcrWireRequest) -> Result { + let api_key_source = source_for(&wire.input_sources, "api_key"); + let api_base_source = source_for(&wire.input_sources, "api_base"); + let extra_headers_source = source_for(&wire.input_sources, "extra_headers"); + let document = decode_request_value(wire.document, "document")?; + let headers = wire + .extra_headers + .unwrap_or_default() + .into_iter() + .map(|(name, value)| { + let value = value + .as_str() + .ok_or_else(|| OcrRequestError::RequestField { + path: format!("extra_headers.{name}"), + })?; + Ok((name, value.to_string())) + }) + .collect::, OcrRequestError>>()?; + let timeout = wire + .timeout_seconds + .map(|seconds| { + Duration::try_from_secs_f64(seconds).map_err(|_| OcrRequestError::RequestField { + path: "timeout_seconds".into(), + }) + }) + .transpose()?; + let defaults = OcrConnection::default(); + let request = LiteLLMOcrRequest::new( + wire.model, + document, + wire.custom_llm_provider.as_deref(), + wire.optional_params, + )?; + let connection = OcrConnection { + api_key: nonblank(wire.api_key), + api_key_source, + api_base: nonblank(wire.api_base), + api_base_source, + extra_headers: headers, + extra_headers_source, + timeout: timeout.unwrap_or(defaults.timeout), + max_download_bytes: defaults.max_download_bytes, + poll_timeout: defaults.poll_timeout, + }; + Ok(LiteLLMOcrRequest { + connection, + input_sources: wire.input_sources, + ..request + }) +} + +fn source_for(sources: &BTreeMap, name: &str) -> InputSource { + sources.get(name).copied().unwrap_or_default() +} + +fn nonblank(value: Option) -> Option { + value + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) +} +pub fn decode_request_value( + value: Value, + prefix: &str, +) -> Result { + serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| { + OcrRequestError::RequestField { + path: format!("{prefix}.{}", error.path()), + } + }) +} + +pub fn decode_response( + bytes: &[u8], + native: bool, +) -> Result, OcrResponseError> { + let mut deserializer = serde_json::Deserializer::from_slice(bytes); + let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| { + OcrResponseError::ResponseField { + path: error.path().to_string(), + } + })?; + deserializer + .end() + .map_err(|_| OcrResponseError::ResponseField { + path: "response".into(), + })?; + let native = if native { + Some( + serde_json::from_slice(bytes).map_err(|_| OcrResponseError::ResponseField { + path: "response".into(), + })?, + ) + } else { + None + }; + Ok(DecodedOcrResponse { data, native }) +} + +pub fn decode_pre_call_result( + original: OcrPreCallRequest, + value: Value, +) -> Result { + #[derive(Deserialize)] + struct Changed { + document: OcrDocument, + #[serde(default)] + optional_params: Map, + } + let changed: Changed = decode_request_value(value, "guardrail")?; + Ok(OcrPreCallRequest { + document: changed.document, + optional_params: Value::Object(changed.optional_params), + ..original + }) +} + +pub fn decode_during_call_result( + original: OcrDuringCallRequest, + value: Value, +) -> Result { + #[derive(Deserialize)] + struct Changed { + body: Value, + } + let changed: Changed = decode_request_value(value, "guardrail")?; + Ok(OcrDuringCallRequest { + body: changed.body, + ..original + }) +} diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index f31b961e78a..3ed00b7cc5f 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -1,3 +1,4 @@ +use crate::auth::error::MissingCredential; use crate::error::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; @@ -21,13 +22,7 @@ pub fn resolve_anthropic_api_key( non_empty(api_key) .map(str::to_string) .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - Error::Auth( - "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \ - environment variable" - .to_string(), - ) - }) + .ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AnthropicApiKey))) } pub fn complete_anthropic_url( diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs new file mode 100644 index 00000000000..297e4cc6502 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs @@ -0,0 +1,43 @@ +use std::future::Future; +use std::sync::Arc; + +use azure_core::credentials::TokenCredential; +use moka::future::Cache; + +use crate::AuthError; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(crate) struct AzureCredentialProviderCacheKey { + pub(crate) mechanism: &'static str, + pub(crate) authority: String, + pub(crate) tenant_id: String, + pub(crate) client_id: String, + pub(crate) scope: String, + pub(crate) secret_identity: String, +} + +pub(crate) struct AzureCredentialProviderCache { + entries: Cache>, +} + +impl AzureCredentialProviderCache { + pub(crate) fn new(capacity: u64) -> Self { + Self { + entries: Cache::builder().max_capacity(capacity).build(), + } + } + + pub(crate) async fn get_or_create( + &self, + key: AzureCredentialProviderCacheKey, + create: F, + ) -> Result, AuthError> + where + F: Future, AuthError>>, + { + self.entries + .try_get_with(key, create) + .await + .map_err(|error| (*error).clone()) + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs new file mode 100644 index 00000000000..33d007c1945 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs @@ -0,0 +1,7 @@ +mod credential_provider_cache; +mod native; +mod resolve; +mod types; + +pub(crate) use resolve::AzureAuthService; +pub(crate) use types::AzureAuthInputs; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs new file mode 100644 index 00000000000..b8f19818d16 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs @@ -0,0 +1,702 @@ +use crate::auth::error::AuthConfigurationError; +use std::sync::Arc; +use std::time::{Duration, UNIX_EPOCH}; + +use azure_core::cloud::{CloudConfiguration, CustomConfiguration}; +use azure_core::credentials::{Secret, TokenCredential}; +use azure_core::http::ClientOptions; +use azure_identity::{ + ClientAssertion, ClientAssertionCredential, ClientAssertionCredentialOptions, + ClientSecretCredential, ClientSecretCredentialOptions, DeveloperToolsCredential, + ManagedIdentityCredential, ManagedIdentityCredentialOptions, UserAssignedId, + WorkloadIdentityCredential, WorkloadIdentityCredentialOptions, +}; +use sha2::{Digest, Sha256}; + +use crate::AuthError; +use crate::auth::{InputSource, ResolvedCredential, SecretValue, Sourced}; + +use super::credential_provider_cache::{ + AzureCredentialProviderCache, AzureCredentialProviderCacheKey, +}; + +#[derive(Clone, Debug)] +pub(crate) enum NativeAzureRequest { + ClientSecret { + tenant_id: Sourced, + client_id: Sourced, + client_secret: Sourced, + scope: Sourced, + authority: Option>, + }, + ClientAssertion { + tenant_id: Sourced, + client_id: Sourced, + assertion: Sourced, + assertion_identity: String, + scope: Sourced, + authority: Option>, + }, + WorkloadIdentity { + tenant_id: Sourced, + client_id: Sourced, + token_file_path: Sourced, + scope: Sourced, + authority: Option>, + }, + ManagedIdentity { + client_id: Option>, + scope: Sourced, + selection_source: InputSource, + }, + DeveloperTools { + scope: Sourced, + selection_source: InputSource, + }, +} + +#[derive(Clone, Debug)] +pub(crate) struct ValidatedAzureRequest { + request: NativeAzureRequest, + credential_source: InputSource, +} + +impl ValidatedAzureRequest { + pub(crate) fn new(request: NativeAzureRequest) -> Result { + validate_authority(&request)?; + let credential_source = validate_sources(&request)?; + Ok(Self { + request, + credential_source, + }) + } + + pub(crate) fn credential_source(&self) -> InputSource { + self.credential_source + } + + #[cfg(test)] + pub(super) fn kind(&self) -> &'static str { + match self.request { + NativeAzureRequest::ClientSecret { .. } => "client-secret", + NativeAzureRequest::ClientAssertion { .. } => "client-assertion", + NativeAzureRequest::WorkloadIdentity { .. } => "workload-identity", + NativeAzureRequest::ManagedIdentity { .. } => "managed-identity", + NativeAzureRequest::DeveloperTools { .. } => "developer-tools", + } + } +} + +pub(crate) struct NativeAzureTokenAcquirer { + cache: AzureCredentialProviderCache, + transport: Option, +} + +impl Default for NativeAzureTokenAcquirer { + fn default() -> Self { + Self::new(64) + } +} + +impl NativeAzureTokenAcquirer { + pub(crate) fn new(cache_capacity: u64) -> Self { + Self { + cache: AzureCredentialProviderCache::new(cache_capacity), + transport: None, + } + } + + #[cfg(test)] + pub(super) fn with_transport( + cache_capacity: u64, + transport: azure_core::http::Transport, + ) -> Self { + Self { + cache: AzureCredentialProviderCache::new(cache_capacity), + transport: Some(transport), + } + } + + pub(crate) async fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> Result { + let scope = request.request.scope().to_string(); + let key = request.request.cache_key(); + let transport = self.transport.clone(); + let credential = self + .cache + .get_or_create( + key, + async move { build_credential(request.request, transport) }, + ) + .await?; + let token = credential + .get_token(&[scope.as_str()], None) + .await + .map_err(|error| AuthError::AzureTokenAcquisition(error.to_string()))?; + let expires_on = u64::try_from(token.expires_on.unix_timestamp()) + .ok() + .map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds)); + + Ok(ResolvedCredential::AccessToken { + token: SecretValue::new(token.token.secret()), + expires_on, + }) + } +} + +impl NativeAzureRequest { + fn scope(&self) -> &str { + match self { + Self::ClientSecret { scope, .. } + | Self::ClientAssertion { scope, .. } + | Self::WorkloadIdentity { scope, .. } + | Self::ManagedIdentity { scope, .. } + | Self::DeveloperTools { scope, .. } => scope.value(), + } + } + + fn cache_key(&self) -> AzureCredentialProviderCacheKey { + match self { + Self::ClientSecret { + tenant_id, + client_id, + client_secret, + scope, + authority, + } => AzureCredentialProviderCacheKey { + mechanism: "client-secret", + authority: authority + .as_ref() + .map(|value| value.value().clone()) + .unwrap_or_default(), + tenant_id: tenant_id.value().clone(), + client_id: client_id.value().clone(), + scope: scope.value().clone(), + secret_identity: secret_digest(client_secret.value().expose()), + }, + Self::ClientAssertion { + tenant_id, + client_id, + assertion, + assertion_identity, + scope, + authority, + } => AzureCredentialProviderCacheKey { + mechanism: "client-assertion", + authority: authority + .as_ref() + .map(|value| value.value().clone()) + .unwrap_or_default(), + tenant_id: tenant_id.value().clone(), + client_id: client_id.value().clone(), + scope: scope.value().clone(), + secret_identity: format!( + "{assertion_identity}:{}", + secret_digest(assertion.value().expose()) + ), + }, + Self::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + scope, + authority, + } => AzureCredentialProviderCacheKey { + mechanism: "workload-identity", + authority: authority + .as_ref() + .map(|value| value.value().clone()) + .unwrap_or_default(), + tenant_id: tenant_id.value().clone(), + client_id: client_id.value().clone(), + scope: scope.value().clone(), + secret_identity: token_file_path.value().clone(), + }, + Self::ManagedIdentity { + client_id, scope, .. + } => AzureCredentialProviderCacheKey { + mechanism: "managed-identity", + authority: String::new(), + tenant_id: String::new(), + client_id: client_id + .as_ref() + .map(|value| value.value().clone()) + .unwrap_or_default(), + scope: scope.value().clone(), + secret_identity: String::new(), + }, + Self::DeveloperTools { scope, .. } => AzureCredentialProviderCacheKey { + mechanism: "developer-tools", + authority: String::new(), + tenant_id: String::new(), + client_id: String::new(), + scope: scope.value().clone(), + secret_identity: String::new(), + }, + } + } +} + +fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> { + let authority = match request { + NativeAzureRequest::ClientSecret { authority, .. } + | NativeAzureRequest::ClientAssertion { authority, .. } + | NativeAzureRequest::WorkloadIdentity { authority, .. } => authority.as_ref(), + NativeAzureRequest::ManagedIdentity { .. } | NativeAzureRequest::DeveloperTools { .. } => { + None + } + }; + let Some(authority) = authority else { + return Ok(()); + }; + let url = url::Url::parse(authority.value()) + .map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureAuthority))?; + if url.scheme() != "https" + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || !matches!(url.path(), "" | "/") + { + return Err(AuthError::Configuration( + AuthConfigurationError::InvalidAzureAuthority, + )); + } + Ok(()) +} + +fn validate_sources(request: &NativeAzureRequest) -> Result { + match request { + NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret, + scope, + authority, + } => { + let identity_sources = [ + tenant_id.source(), + client_id.source(), + client_secret.source(), + ]; + let request_identity = identity_sources.contains(&InputSource::Request); + if request_identity + && !identity_sources + .iter() + .all(|source| *source == InputSource::Request) + { + return mixed_sources(); + } + if !request_identity && is_request_controlled(scope, authority.as_ref()) { + return mixed_sources(); + } + Ok(if request_identity { + InputSource::Request + } else { + trusted_source(&identity_sources) + }) + } + NativeAzureRequest::ClientAssertion { + tenant_id, + client_id, + assertion, + scope, + authority, + .. + } => trusted_only(&[ + tenant_id.source(), + client_id.source(), + assertion.source(), + scope.source(), + authority + .as_ref() + .map(Sourced::source) + .unwrap_or(InputSource::Environment), + ]), + NativeAzureRequest::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + scope, + authority, + } => trusted_only(&[ + tenant_id.source(), + client_id.source(), + token_file_path.source(), + scope.source(), + authority + .as_ref() + .map(Sourced::source) + .unwrap_or(InputSource::Environment), + ]), + NativeAzureRequest::ManagedIdentity { + client_id, + scope, + selection_source, + } => trusted_only(&[ + client_id + .as_ref() + .map(Sourced::source) + .unwrap_or(InputSource::Environment), + scope.source(), + *selection_source, + ]), + NativeAzureRequest::DeveloperTools { + scope, + selection_source, + } => trusted_only(&[scope.source(), *selection_source]), + } +} + +fn is_request_controlled(value: &Sourced, optional: Option<&Sourced>) -> bool { + value.source() == InputSource::Request + || optional.is_some_and(|value| value.source() == InputSource::Request) +} + +fn trusted_only(sources: &[InputSource]) -> Result { + if sources.contains(&InputSource::Request) { + return mixed_sources(); + } + Ok(trusted_source(sources)) +} + +fn trusted_source(sources: &[InputSource]) -> InputSource { + if sources.contains(&InputSource::Deployment) { + InputSource::Deployment + } else { + InputSource::Environment + } +} + +fn mixed_sources() -> Result { + Err(AuthError::Configuration( + AuthConfigurationError::MixedAzureCredentialSources, + )) +} + +fn build_credential( + request: NativeAzureRequest, + transport: Option, +) -> Result, AuthError> { + match request { + NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret, + authority, + .. + } => ClientSecretCredential::new( + tenant_id.value(), + client_id.into_value(), + Secret::new(client_secret.value().expose().to_string()), + Some(ClientSecretCredentialOptions { + client_options: client_options(authority.map(Sourced::into_value), transport), + }), + ) + .map(|credential| credential as Arc), + NativeAzureRequest::ClientAssertion { + tenant_id, + client_id, + assertion, + authority, + .. + } => ClientAssertionCredential::new( + tenant_id.into_value(), + client_id.into_value(), + StaticAssertion(assertion.into_value()), + Some(ClientAssertionCredentialOptions { + client_options: client_options(authority.map(Sourced::into_value), transport), + }), + ) + .map(|credential| credential as Arc), + NativeAzureRequest::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + authority, + .. + } => WorkloadIdentityCredential::new(Some(WorkloadIdentityCredentialOptions { + credential_options: azure_identity::ClientAssertionCredentialOptions { + client_options: client_options(authority.map(Sourced::into_value), transport), + }, + client_id: Some(client_id.into_value()), + tenant_id: Some(tenant_id.into_value()), + token_file_path: Some(token_file_path.into_value().into()), + })) + .map(|credential| credential as Arc), + NativeAzureRequest::ManagedIdentity { client_id, .. } => { + ManagedIdentityCredential::new(Some(ManagedIdentityCredentialOptions { + user_assigned_id: client_id + .map(Sourced::into_value) + .map(UserAssignedId::ClientId), + client_options: client_options(None, transport), + })) + .map(|credential| credential as Arc) + } + NativeAzureRequest::DeveloperTools { .. } => DeveloperToolsCredential::new(None) + .map(|credential| credential as Arc), + } + .map_err(|error| { + AuthError::Configuration(AuthConfigurationError::AzureCredentialInitialization( + error.to_string(), + )) + }) +} + +fn client_options( + authority: Option, + transport: Option, +) -> ClientOptions { + let cloud = authority.map(|authority_host| { + let mut custom = CustomConfiguration::default(); + custom.authority_host = authority_host; + Arc::new(CloudConfiguration::from(custom)) + }); + ClientOptions { + cloud, + transport, + ..Default::default() + } +} + +fn secret_digest(secret: &str) -> String { + format!("{:x}", Sha256::digest(secret.as_bytes())) +} + +#[derive(Debug)] +struct StaticAssertion(SecretValue); + +impl ClientAssertion for StaticAssertion { + fn secret<'life0, 'life1, 'async_trait>( + &'life0 self, + _options: Option>, + ) -> std::pin::Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { Ok(self.0.expose().to_string()) }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use azure_core::http::headers::Headers; + use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport}; + use azure_core::{Bytes, Result}; + + use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; + use crate::auth::{InputSource, SecretValue, Sourced}; + + fn deployment(value: T) -> Sourced { + Sourced::new(value, InputSource::Deployment) + } + + fn sourced_client_secret( + credential_source: InputSource, + authority_source: InputSource, + authority: &str, + ) -> NativeAzureRequest { + NativeAzureRequest::ClientSecret { + tenant_id: Sourced::new("tenant".to_string(), credential_source), + client_id: Sourced::new("client".to_string(), credential_source), + client_secret: Sourced::new(SecretValue::new("secret"), credential_source), + scope: Sourced::new("scope".to_string(), InputSource::Environment), + authority: Some(Sourced::new(authority.to_string(), authority_source)), + } + } + + fn client_secret_request( + tenant: &str, + client: &str, + secret: &str, + scope: &str, + authority: &str, + ) -> ValidatedAzureRequest { + ValidatedAzureRequest::new(NativeAzureRequest::ClientSecret { + tenant_id: deployment(tenant.to_string()), + client_id: deployment(client.to_string()), + client_secret: deployment(SecretValue::new(secret)), + scope: deployment(scope.to_string()), + authority: Some(deployment(authority.to_string())), + }) + .unwrap() + } + + #[derive(Debug, Default)] + struct RecordingTokenClient { + requests: Mutex>, + } + + impl HttpClient for RecordingTokenClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + request: &'life1 Request, + ) -> std::pin::Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let body = Bytes::from(request.body()); + self.requests.lock().unwrap().push(( + request.url().to_string(), + String::from_utf8(body.to_vec()).unwrap(), + )); + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + Headers::new(), + r#"{"token_type":"Bearer","expires_in":3600,"ext_expires_in":3600,"access_token":"native-token"}"#, + )) + }) + } + } + + #[tokio::test] + async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() { + let transport = Arc::new(RecordingTokenClient::default()); + let acquirer = + NativeAzureTokenAcquirer::with_transport(4, Transport::new(transport.clone())); + let request = client_secret_request( + "tenant", + "client", + "secret", + "https://service.test/.default", + "https://login.test", + ); + + let first = acquirer.acquire(request.clone()).await.unwrap(); + let second = acquirer.acquire(request).await.unwrap(); + + assert_eq!(first.secret().expose(), "native-token"); + assert_eq!(second.secret().expose(), "native-token"); + let requests = transport.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].0, "https://login.test/tenant/oauth2/v2.0/token"); + assert!(requests[0].1.contains("client_id=client")); + assert!(requests[0].1.contains("client_secret=secret")); + assert!( + requests[0] + .1 + .contains("scope=https%3A%2F%2Fservice.test%2F.default") + ); + } + + #[tokio::test] + async fn credential_provider_cache_isolates_every_client_secret_identity_field() { + let transport = Arc::new(RecordingTokenClient::default()); + let acquirer = + NativeAzureTokenAcquirer::with_transport(16, Transport::new(transport.clone())); + let request = client_secret_request; + let base = request("tenant", "client", "secret", "scope", "https://login.test"); + let variants = [ + base.clone(), + request( + "other-tenant", + "client", + "secret", + "scope", + "https://login.test", + ), + request( + "tenant", + "other-client", + "secret", + "scope", + "https://login.test", + ), + request( + "tenant", + "client", + "other-secret", + "scope", + "https://login.test", + ), + request( + "tenant", + "client", + "secret", + "other-scope", + "https://login.test", + ), + request( + "tenant", + "client", + "secret", + "scope", + "https://other-login.test", + ), + ]; + + acquirer.acquire(base.clone()).await.unwrap(); + acquirer.acquire(base).await.unwrap(); + for request in variants.into_iter().skip(1) { + acquirer.acquire(request).await.unwrap(); + } + + assert_eq!(transport.requests.lock().unwrap().len(), 6); + } + + #[test] + fn request_authority_requires_request_owned_client_secret_identity() { + let error = ValidatedAzureRequest::new(sourced_client_secret( + InputSource::Deployment, + InputSource::Request, + "https://login.example", + )) + .unwrap_err(); + + assert!(matches!( + error, + crate::AuthError::Configuration( + crate::auth::error::AuthConfigurationError::MixedAzureCredentialSources + ) + )); + } + + #[test] + fn request_owned_client_secret_identity_can_select_custom_authority() { + let request = ValidatedAzureRequest::new(sourced_client_secret( + InputSource::Request, + InputSource::Request, + "https://login.example", + )) + .unwrap(); + + assert_eq!(request.credential_source(), InputSource::Request); + } + + #[test] + fn authority_is_restricted_to_an_https_origin() { + for authority in [ + "http://login.example", + "https://user@login.example", + "https://login.example/tenant", + "https://login.example?target=other", + ] { + let error = ValidatedAzureRequest::new(sourced_client_secret( + InputSource::Deployment, + InputSource::Deployment, + authority, + )) + .unwrap_err(); + assert!(matches!( + error, + crate::AuthError::Configuration( + crate::auth::error::AuthConfigurationError::InvalidAzureAuthority + ) + )); + } + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs new file mode 100644 index 00000000000..025dd4f8740 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs @@ -0,0 +1,683 @@ +use crate::AuthError; +use crate::auth::error::AuthConfigurationError; +use crate::auth::{ + CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, + SecretValue, Sourced, TokenProviderHandle, +}; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; +use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE}; + +const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; +const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID"; +const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID"; +const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET"; +const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE"; +const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST"; +const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL"; +const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; + +#[derive(Clone, Debug)] +pub(crate) enum AzureCredentialPlan { + Supplied(Sourced), + Caller(TokenProviderHandle), + Oidc { + reference: Sourced, + tenant_id: Sourced, + client_id: Sourced, + scope: Sourced, + authority: Option>, + }, + Native(ValidatedAzureRequest), + Chain(Vec), + Missing, +} + +/// Rust counterpart to Python's `get_azure_ad_token`, not `BaseAzureLLM`. +pub(crate) struct AzureAuthService { + native: Arc, +} + +trait AzureTokenAcquirer: Send + Sync { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> Pin> + Send + '_>>; +} + +impl AzureTokenAcquirer for NativeAzureTokenAcquirer { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> Pin> + Send + '_>> { + Box::pin(NativeAzureTokenAcquirer::acquire(self, request)) + } +} + +impl Default for AzureAuthService { + fn default() -> Self { + Self { + native: Arc::new(NativeAzureTokenAcquirer::default()), + } + } +} + +impl AzureAuthService { + #[cfg(test)] + fn with_acquirer(native: Arc) -> Self { + Self { native } + } + + pub(crate) async fn get_azure_ad_token( + &self, + inputs: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result>, AuthError> { + match select_auth_plan(inputs, env_lookup)? { + AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)), + AzureCredentialPlan::Caller(caller) => { + let credential = caller.acquire().await?; + if credential.secret().expose().is_empty() { + return Err(AuthError::EmptyAzureToken); + } + Ok(Some(Sourced::new(credential, InputSource::Deployment))) + } + AzureCredentialPlan::Oidc { + reference, + tenant_id, + client_id, + scope, + authority, + } => { + let assertion = resolve_reference(inputs, env_lookup, reference.value()) + .await? + .ok_or(AuthError::UnresolvedOidcReference)?; + let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion { + tenant_id, + client_id, + assertion: Sourced::new(assertion, reference.source()), + assertion_identity: format!("{:?}", reference.value()), + scope, + authority, + })?; + let source = request.credential_source(); + self.native + .acquire(request) + .await + .map(|credential| Sourced::new(credential, source)) + .map(Some) + } + AzureCredentialPlan::Native(request) => { + let source = request.credential_source(); + self.native + .acquire(request) + .await + .map(|credential| Some(Sourced::new(credential, source))) + } + AzureCredentialPlan::Chain(requests) => { + let mut failures = Vec::new(); + for request in requests { + let source = request.credential_source(); + match self.native.acquire(request).await { + Ok(credential) => return Ok(Some(Sourced::new(credential, source))), + Err(error) => failures.push(error), + } + } + Err(AuthError::CredentialChain(failures)) + } + AzureCredentialPlan::Missing => Ok(None), + } + } +} + +pub(crate) fn select_auth_plan( + inputs: &AzureAuthInputs, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup); + let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup); + let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup); + let client_secret = + configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup); + let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup) + .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment)); + let authority = configured_string( + &inputs.azure_authority_host, + AZURE_AUTHORITY_HOST_ENV, + env_lookup, + ); + let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup) + .map(|value| { + value + .value() + .parse::() + .map(|selector| Sourced::new(selector, value.source())) + }) + .transpose() + .map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureSelector))?; + let federated_token_file = configured_string( + &inputs.federated_token_file, + AZURE_FEDERATED_TOKEN_FILE_ENV, + env_lookup, + ); + + if inputs.azure_ad_token_provider.is_none() + && let (Some(tenant_id), Some(client_id), Some(client_secret)) = + (tenant_id.clone(), client_id.clone(), client_secret) + { + return Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( + NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret, + scope, + authority, + }, + )?)); + } + + if let (Some(reference), Some(tenant_id), Some(client_id)) = ( + oidc_reference(&token)?, + tenant_id.clone(), + client_id.clone(), + ) { + return Ok(AzureCredentialPlan::Oidc { + reference, + tenant_id, + client_id, + scope, + authority, + }); + } + + if let Some(caller) = &inputs.azure_ad_token_provider { + return Ok(AzureCredentialPlan::Caller(caller.clone())); + } + + if let Some(token) = token { + return Ok(AzureCredentialPlan::Supplied(token.map(|token| { + ResolvedCredential::AccessToken { + token, + expires_on: None, + } + }))); + } + + if !*inputs.enable_azure_ad_token_refresh.value() && selector.is_none() { + return Ok(AzureCredentialPlan::Missing); + } + + select_native_plan( + selector, + tenant_id, + client_id, + federated_token_file, + scope, + authority, + inputs.enable_azure_ad_token_refresh.source(), + ) +} + +fn select_native_plan( + selector: Option>, + tenant_id: Option>, + client_id: Option>, + federated_token_file: Option>, + scope: Sourced, + authority: Option>, + refresh_source: InputSource, +) -> Result { + let selected = selector.unwrap_or_else(|| { + Sourced::new( + { + if federated_token_file.is_some() { + AzureCredentialType::DefaultAzureCredential + } else if client_id.is_some() { + AzureCredentialType::ManagedIdentityCredential + } else { + AzureCredentialType::DefaultAzureCredential + } + }, + refresh_source, + ) + }); + let selection_source = selected.source(); + + match selected.into_value() { + AzureCredentialType::ClientSecretCredential => Err(AuthError::Configuration( + AuthConfigurationError::MissingClientSecretFields, + )), + AzureCredentialType::WorkloadIdentityCredential => { + Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( + workload_request(tenant_id, client_id, federated_token_file, scope, authority)?, + )?)) + } + AzureCredentialType::ManagedIdentityCredential => Ok(AzureCredentialPlan::Native( + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id, + scope, + selection_source, + })?, + )), + AzureCredentialType::DefaultAzureCredential => { + let workload = match (tenant_id, client_id.clone(), federated_token_file) { + (Some(tenant_id), Some(client_id), Some(token_file_path)) => { + Some(NativeAzureRequest::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + scope: scope.clone(), + authority, + }) + } + _ => None, + }; + Ok(AzureCredentialPlan::Chain( + workload + .into_iter() + .chain(std::iter::once(NativeAzureRequest::ManagedIdentity { + client_id, + scope: scope.clone(), + selection_source, + })) + .chain(std::iter::once(NativeAzureRequest::DeveloperTools { + scope, + selection_source, + })) + .map(ValidatedAzureRequest::new) + .collect::, _>>()?, + )) + } + AzureCredentialType::DeploymentIdentityCredential => { + let workload = match (tenant_id, client_id.clone(), federated_token_file) { + (Some(tenant_id), Some(client_id), Some(token_file_path)) => { + Some(NativeAzureRequest::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + scope: scope.clone(), + authority, + }) + } + _ => None, + }; + let user_assigned = client_id.map(|client_id| NativeAzureRequest::ManagedIdentity { + client_id: Some(client_id), + scope: scope.clone(), + selection_source, + }); + Ok(AzureCredentialPlan::Chain( + workload + .into_iter() + .chain(user_assigned) + .chain(std::iter::once(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope, + selection_source, + })) + .map(ValidatedAzureRequest::new) + .collect::, _>>()?, + )) + } + } +} + +fn workload_request( + tenant_id: Option>, + client_id: Option>, + token_file_path: Option>, + scope: Sourced, + authority: Option>, +) -> Result { + Ok(NativeAzureRequest::WorkloadIdentity { + tenant_id: tenant_id.ok_or(AuthError::Configuration( + AuthConfigurationError::MissingWorkloadTenant, + ))?, + client_id: client_id.ok_or(AuthError::Configuration( + AuthConfigurationError::MissingWorkloadClient, + ))?, + token_file_path: token_file_path.ok_or(AuthError::Configuration( + AuthConfigurationError::MissingWorkloadTokenFile, + ))?, + scope, + authority, + }) +} + +fn configured_string( + configured: &ConfigValue, + environment_name: &str, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option> { + configured + .as_value() + .filter(|value| !value.value().is_empty()) + .cloned() + .or_else(|| { + env_lookup(environment_name) + .filter(|value| !value.is_empty()) + .map(|value| Sourced::new(value, InputSource::Environment)) + }) +} + +fn configured_secret( + configured: &ConfigValue, + environment_name: &str, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option> { + configured + .as_value() + .filter(|value| !value.value().expose().is_empty()) + .cloned() + .or_else(|| { + env_lookup(environment_name) + .filter(|value| !value.is_empty()) + .map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment)) + }) +} + +async fn resolve_reference( + inputs: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + reference: &CredentialRef, +) -> Result, AuthError> { + let lookup = match reference { + CredentialRef::Explicit(secret) => return Ok(Some(secret.clone())), + CredentialRef::Env(name) => env_lookup(name) + .filter(|value| !value.is_empty()) + .map(SecretValue::new) + .map_or(CredentialLookup::Missing, CredentialLookup::Found), + CredentialRef::None => return Ok(None), + CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => { + let resolver = inputs + .credential_resolver + .as_ref() + .ok_or(AuthError::Configuration( + AuthConfigurationError::MissingHostResolver, + ))?; + resolver.resolve(reference).await? + } + }; + Ok(match lookup { + CredentialLookup::Found(secret) => Some(secret), + CredentialLookup::Missing | CredentialLookup::Declined => None, + }) +} + +fn oidc_reference( + token: &Option>, +) -> Result>, AuthError> { + let Some(token) = token.as_ref() else { + return Ok(None); + }; + let value = token.value().expose(); + if token.source() == InputSource::Request && value.starts_with("oidc/") { + return Err(AuthError::Configuration( + AuthConfigurationError::RequestAzureCredentialReference, + )); + } + if let Some(name) = value.strip_prefix("oidc/env/") { + return non_empty_reference(name, "OIDC environment reference") + .map(CredentialRef::Env) + .map(|reference| Sourced::new(reference, token.source())) + .map(Some); + } + if let Some(name) = value.strip_prefix("oidc/env_path/") { + return non_empty_reference(name, "OIDC environment path reference") + .map(|name| CredentialRef::File(CredentialFileRef::EnvironmentVariable(name))) + .map(|reference| Sourced::new(reference, token.source())) + .map(Some); + } + if let Some(path) = value.strip_prefix("oidc/file/") { + let path = non_empty_reference(path, "OIDC file reference")?; + return Ok(Some(Sourced::new( + CredentialRef::File(CredentialFileRef::Path(path.into())), + token.source(), + ))); + } + if value.starts_with("oidc/") { + return Err(AuthError::Configuration( + AuthConfigurationError::UnsupportedOidcReference, + )); + } + Ok(None) +} + +fn non_empty_reference(value: &str, kind: &str) -> Result { + if value.is_empty() { + return Err(AuthError::Configuration( + AuthConfigurationError::EmptyReference(kind.to_string()), + )); + } + Ok(value.to_string()) +} + +#[cfg(test)] +mod tests { + use std::future::Future; + use std::sync::{Arc, Mutex}; + + use serde_json::json; + + use super::{ + AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, oidc_reference, + resolve_reference, select_auth_plan, + }; + use crate::AuthError; + use crate::auth::ResolvedCredential; + use crate::auth::{ + CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, + CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced, + }; + use crate::providers::azure_ai::auth::native::ValidatedAzureRequest; + use crate::providers::azure_ai::auth::types::AzureAuthInputs; + + #[derive(Debug)] + struct FileResolver; + + struct ChainAcquirer { + requests: Mutex>, + succeed_on: Option<&'static str>, + } + + impl AzureTokenAcquirer for ChainAcquirer { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> std::pin::Pin< + Box> + Send + '_>, + > { + let kind = request.kind(); + self.requests.lock().unwrap().push(kind); + Box::pin(async move { + if self.succeed_on == Some(kind) { + Ok(ResolvedCredential::AccessToken { + token: SecretValue::new("chain-token"), + expires_on: None, + }) + } else { + Err(AuthError::AzureTokenAcquisition(format!("{kind} failed"))) + } + }) + } + } + + impl CredentialResolver for FileResolver { + fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { + Box::pin(async move { + Ok(match reference { + CredentialRef::File(CredentialFileRef::Path(path)) + if path == std::path::Path::new("/run/secrets/assertion") => + { + CredentialLookup::Found(SecretValue::new("rotated-assertion")) + } + _ => CredentialLookup::Declined, + }) + }) + } + } + + #[test] + fn null_and_empty_values_fall_back_to_environment() { + let params = json!({"tenant_id": null, "client_id": "", "client_secret": null}); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + let plan = select_auth_plan(&inputs, &|name| match name { + "AZURE_TENANT_ID" => Some("tenant".to_string()), + "AZURE_CLIENT_ID" => Some("client".to_string()), + "AZURE_CLIENT_SECRET" => Some("secret".to_string()), + _ => None, + }) + .unwrap(); + + assert!(matches!(plan, AzureCredentialPlan::Native(_))); + } + + #[test] + fn supplied_token_does_not_require_refresh() { + let params = json!({"azure_ad_token": "token"}); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + + assert!(matches!( + select_auth_plan(&inputs, &|_| None).unwrap(), + AzureCredentialPlan::Supplied(_) + )); + } + + #[test] + fn oidc_reference_is_deferred() { + let params = json!({ + "azure_ad_token": "oidc/env/ASSERTION", + "tenant_id": "tenant", + "client_id": "client" + }); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + + assert!(matches!( + select_auth_plan(&inputs, &|_| None).unwrap(), + AzureCredentialPlan::Oidc { + reference, + .. + } if reference.value() == &CredentialRef::Env("ASSERTION".to_string()) + )); + } + + #[test] + fn oidc_file_location_is_typed_before_resolution() { + assert_eq!( + oidc_reference(&Some(Sourced::new( + SecretValue::new("oidc/file//run/secrets/assertion"), + InputSource::Deployment, + ))) + .unwrap() + .map(Sourced::into_value), + Some(CredentialRef::File(CredentialFileRef::Path( + "/run/secrets/assertion".into() + ))) + ); + } + + #[test] + fn unsupported_oidc_reference_is_rejected_during_plan_creation() { + let error = oidc_reference(&Some(Sourced::new( + SecretValue::new("oidc/vault/assertion"), + InputSource::Deployment, + ))) + .expect_err("unsupported backend must fail validation"); + + assert!(error.to_string().contains("unsupported OIDC reference")); + } + + #[test] + fn request_oidc_reference_is_rejected_before_lookup() { + let params = json!({ + "azure_ad_token": "oidc/env/ASSERTION", + "tenant_id": "tenant", + "client_id": "client" + }); + let sources = std::collections::BTreeMap::from([ + ("azure_ad_token".to_string(), InputSource::Request), + ("tenant_id".to_string(), InputSource::Request), + ("client_id".to_string(), InputSource::Request), + ]); + let inputs = + AzureAuthInputs::from_sourced_optional_params(params.as_object().unwrap(), &sources) + .unwrap(); + + let error = select_auth_plan(&inputs, &|name| { + assert_ne!(name, "ASSERTION"); + None + }) + .unwrap_err(); + + assert!(matches!( + error, + AuthError::Configuration( + crate::auth::error::AuthConfigurationError::RequestAzureCredentialReference + ) + )); + } + + #[tokio::test] + async fn host_resolver_owns_file_access() { + let inputs = AzureAuthInputs { + credential_resolver: Some(CredentialResolverHandle::new(Arc::new(FileResolver))), + ..AzureAuthInputs::default() + }; + let reference = + CredentialRef::File(CredentialFileRef::Path("/run/secrets/assertion".into())); + + let resolved = resolve_reference(&inputs, &|_| None, &reference) + .await + .unwrap(); + + assert_eq!(resolved, Some(SecretValue::new("rotated-assertion"))); + } + + #[tokio::test] + async fn default_chain_uses_declared_order_and_stops_after_success() { + let acquirer = Arc::new(ChainAcquirer { + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let inputs = AzureAuthInputs { + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + + let credential = service + .get_azure_ad_token(&inputs, &|_| None) + .await + .unwrap() + .unwrap(); + + assert_eq!(credential.value().secret().expose(), "chain-token"); + assert_eq!( + *acquirer.requests.lock().unwrap(), + ["managed-identity", "developer-tools"] + ); + } + + #[tokio::test] + async fn chain_reports_each_acquisition_failure() { + let acquirer = Arc::new(ChainAcquirer { + requests: Mutex::new(Vec::new()), + succeed_on: None, + }); + let service = AzureAuthService::with_acquirer(acquirer); + let inputs = AzureAuthInputs { + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + + let error = service + .get_azure_ad_token(&inputs, &|_| None) + .await + .unwrap_err(); + + assert!(matches!(error, AuthError::CredentialChain(errors) if errors.len() == 2)); + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs new file mode 100644 index 00000000000..f15d526d945 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs @@ -0,0 +1,195 @@ +use crate::auth::error::AuthConfigurationError; +use serde_json::{Map, Value}; +use std::collections::BTreeMap; +use strum::EnumString; + +use crate::AuthError; +use crate::auth::{ + CredentialResolverHandle, InputSource, SecretValue, Sourced, TokenProviderHandle, +}; + +pub const DEFAULT_AZURE_SCOPE: &str = "https://cognitiveservices.azure.com/.default"; + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub enum ConfigValue { + #[default] + Absent, + ExplicitNone(InputSource), + Value(Sourced), +} + +impl ConfigValue { + pub fn as_value(&self) -> Option<&Sourced> { + match self { + Self::Value(value) => Some(value), + Self::Absent | Self::ExplicitNone(_) => None, + } + } +} + +#[derive(Clone, Copy, Debug, EnumString, PartialEq, Eq, Hash)] +#[allow(clippy::enum_variant_names)] +pub enum AzureCredentialType { + ClientSecretCredential, + ManagedIdentityCredential, + DefaultAzureCredential, + DeploymentIdentityCredential, + WorkloadIdentityCredential, +} + +#[derive(Clone, Debug, Default)] +pub struct AzureAuthInputs { + pub azure_ad_token: ConfigValue, + pub azure_ad_token_provider: Option, + pub credential_resolver: Option, + pub tenant_id: ConfigValue, + pub client_id: ConfigValue, + pub client_secret: ConfigValue, + pub azure_scope: ConfigValue, + pub azure_authority_host: ConfigValue, + pub azure_credential: ConfigValue, + pub federated_token_file: ConfigValue, + pub enable_azure_ad_token_refresh: Sourced, +} + +impl AzureAuthInputs { + #[cfg(test)] + pub fn from_optional_params(params: &Map) -> Result { + Self::from_sourced_optional_params(params, &BTreeMap::new()) + } + + pub fn from_sourced_optional_params( + params: &Map, + sources: &BTreeMap, + ) -> Result { + Ok(Self { + azure_ad_token: secret_config(params, sources, "azure_ad_token")?, + azure_ad_token_provider: None, + credential_resolver: None, + tenant_id: string_config(params, sources, "tenant_id")?, + client_id: string_config(params, sources, "client_id")?, + client_secret: secret_config(params, sources, "client_secret")?, + azure_scope: string_config(params, sources, "azure_scope")?, + azure_authority_host: string_config(params, sources, "azure_authority_host")?, + azure_credential: string_config(params, sources, "azure_credential")?, + federated_token_file: string_config(params, sources, "azure_federated_token_file")?, + enable_azure_ad_token_refresh: Sourced::new( + params + .get("enable_azure_ad_token_refresh") + .and_then(Value::as_bool) + .unwrap_or(false), + source_for(sources, "enable_azure_ad_token_refresh"), + ), + }) + } +} + +fn string_config( + params: &Map, + sources: &BTreeMap, + name: &str, +) -> Result, AuthError> { + let source = source_for(sources, name); + match params.get(name) { + None => Ok(ConfigValue::Absent), + Some(Value::Null) => Ok(ConfigValue::ExplicitNone(source)), + Some(Value::String(value)) => Ok(ConfigValue::Value(Sourced::new(value.clone(), source))), + Some(_) => Err(AuthError::Configuration( + AuthConfigurationError::InvalidFieldType(name.to_string()), + )), + } +} + +fn secret_config( + params: &Map, + sources: &BTreeMap, + name: &str, +) -> Result, AuthError> { + Ok(match string_config(params, sources, name)? { + ConfigValue::Absent => ConfigValue::Absent, + ConfigValue::ExplicitNone(source) => ConfigValue::ExplicitNone(source), + ConfigValue::Value(value) => ConfigValue::Value(value.map(SecretValue::new)), + }) +} + +fn source_for(sources: &BTreeMap, name: &str) -> InputSource { + sources.get(name).copied().unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use std::collections::BTreeMap; + + use super::{AzureAuthInputs, AzureCredentialType, ConfigValue}; + use crate::auth::{InputSource, Sourced}; + + #[test] + fn selector_parsing_is_exact() { + assert_eq!( + "ClientSecretCredential".parse::(), + Ok(AzureCredentialType::ClientSecretCredential) + ); + assert!( + "clientsecretcredential" + .parse::() + .is_err() + ); + } + + #[test] + fn defaults_preserve_absence() { + let inputs = AzureAuthInputs::default(); + + assert_eq!(inputs.tenant_id, ConfigValue::Absent); + assert_eq!(inputs.azure_ad_token, ConfigValue::Absent); + } + + #[test] + fn parsing_distinguishes_null_empty_and_absent() { + let params = json!({"tenant_id": null, "client_id": ""}); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + + assert_eq!( + inputs.tenant_id, + ConfigValue::ExplicitNone(InputSource::Deployment) + ); + assert_eq!( + inputs.client_id, + ConfigValue::Value(Sourced::new(String::new(), InputSource::Deployment)) + ); + assert_eq!(inputs.client_secret, ConfigValue::Absent); + } + + #[test] + fn parsing_preserves_trusted_input_sources() { + let params = json!({"tenant_id": "tenant", "client_secret": null}); + let sources = BTreeMap::from([ + ("tenant_id".to_string(), InputSource::Request), + ("client_secret".to_string(), InputSource::Request), + ]); + let inputs = + AzureAuthInputs::from_sourced_optional_params(params.as_object().unwrap(), &sources) + .unwrap(); + + assert_eq!( + inputs.tenant_id, + ConfigValue::Value(Sourced::new("tenant".to_string(), InputSource::Request)) + ); + assert_eq!( + inputs.client_secret, + ConfigValue::ExplicitNone(InputSource::Request) + ); + } + + #[test] + fn debug_does_not_expose_secrets() { + let params = json!({"azure_ad_token": "token-value", "client_secret": "secret-value"}); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + let debug = format!("{inputs:?}"); + + assert!(!debug.contains("token-value")); + assert!(!debug.contains("secret-value")); + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index b8ca10461fb..585b34f393f 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -1,3 +1,4 @@ +use crate::auth::error::MissingCredential; use crate::error::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use crate::messages::types::{ @@ -32,12 +33,7 @@ pub fn resolve_azure_api_key( non_empty(api_key) .map(str::to_string) .or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - Error::Auth( - "Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable" - .to_string(), - ) - }) + .ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiKey))) } pub fn complete_azure_anthropic_url( @@ -47,13 +43,7 @@ pub fn complete_azure_anthropic_url( let api_base = non_empty(api_base) .map(str::to_string) .or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - Error::Auth( - "Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. \ - Expected format: https://.services.ai.azure.com/anthropic" - .to_string(), - ) - })?; + .ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiBase)))?; let api_base = api_base.trim_end_matches('/'); diff --git a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs index 5d13fa93e00..4f41d1d6abb 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs @@ -1,2 +1,2 @@ +pub(crate) mod auth; pub mod messages; -pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs deleted file mode 100644 index f239b6921fa..00000000000 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs deleted file mode 100644 index d15c032f0bc..00000000000 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ /dev/null @@ -1,1381 +0,0 @@ -use std::collections::BTreeSet; - -use crate::error::{Error, json_type_name}; -use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; -use crate::ocr::types::{OcrRequestData, OcrResponseData}; -use serde_json::{Map, Value, json}; - -use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; - -const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; -const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; -const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; -const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; -const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30"; -const AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: i64 = 96; - -const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = - &["pages", "features", "req_format"]; - -pub struct AzureAiOcrConfig; -pub struct AzureDocumentIntelligenceOcrConfig; - -pub const AZURE_AI_OCR_CONFIG: AzureAiOcrConfig = AzureAiOcrConfig; -pub const AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG: AzureDocumentIntelligenceOcrConfig = - AzureDocumentIntelligenceOcrConfig; - -fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -fn resolve_value( - explicit: Option<&str>, - env_name: &str, - env_lookup: &dyn Fn(&str) -> Option, - missing_message: &str, -) -> Result { - non_empty(explicit) - .map(str::to_string) - .or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| Error::Auth(missing_message.to_string())) -} - -pub fn resolve_azure_ai_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - resolve_value( - api_key, - AZURE_AI_API_KEY_ENV, - env_lookup, - "Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params", - ) -} - -pub fn resolve_azure_ai_api_base( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - resolve_value( - api_base, - AZURE_AI_API_BASE_ENV, - env_lookup, - "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter", - ) -} - -pub fn complete_azure_ai_url( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - let base = resolve_azure_ai_api_base(api_base, env_lookup)?; - Ok(format!( - "{}/providers/mistral/azure/ocr", - base.trim_end_matches('/') - )) -} - -pub fn resolve_document_intelligence_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - resolve_value( - api_key, - AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, - env_lookup, - "Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter", - ) -} - -pub fn resolve_document_intelligence_endpoint( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - resolve_value( - api_base, - AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV, - env_lookup, - "Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter", - ) -} - -fn prepend_auth_header( - headers: Vec<(String, String)>, - name: &str, - value: String, -) -> Vec<(String, String)> { - std::iter::once((name.to_string(), value)) - .chain(headers) - .collect() -} - -pub fn validate_azure_ai_environment( - headers: Vec<(String, String)>, - api_key: Option<&str>, - azure_ad_token: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result, Error> { - if crate::http_utils::has_header(&headers, "Authorization") - || crate::http_utils::has_header(&headers, "Api-Key") - { - return Ok(headers); - } - if let Ok(api_key) = resolve_azure_ai_api_key(api_key, env_lookup) { - return Ok(prepend_auth_header(headers, "Api-Key", api_key)); - } - non_empty(azure_ad_token) - .map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}"))) - .ok_or_else(|| { - Error::Auth( - "Missing Azure AI credentials - set AZURE_AI_API_KEY or provide azure_ad_token" - .to_string(), - ) - }) -} - -pub fn validate_document_intelligence_environment( - headers: Vec<(String, String)>, - api_key: Option<&str>, - azure_ad_token: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result, Error> { - if crate::http_utils::has_header(&headers, "Authorization") - || crate::http_utils::has_header(&headers, "Ocp-Apim-Subscription-Key") - { - return Ok(headers); - } - if let Ok(api_key) = resolve_document_intelligence_api_key(api_key, env_lookup) { - return Ok(prepend_auth_header( - headers, - "Ocp-Apim-Subscription-Key", - api_key, - )); - } - non_empty(azure_ad_token) - .map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}"))) - .ok_or_else(|| { - Error::Auth( - "Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or provide azure_ad_token" - .to_string(), - ) - }) -} - -fn encode_model_id(model: &str) -> Result { - let model_id = model.rsplit('/').next().unwrap_or(model); - if matches!(model_id, "." | "..") { - return Err(Error::InvalidRequest( - "model_id cannot be a dot path segment".to_string(), - )); - } - Ok(model_id - .bytes() - .flat_map(|byte| match byte { - b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => { - vec![byte as char] - } - _ => format!("%{byte:02X}").chars().collect(), - }) - .collect()) -} - -fn pages_token_is_valid(token: &str) -> bool { - let mut parts = token.split('-'); - let Some(start) = parts.next() else { - return false; - }; - if start.is_empty() || !start.chars().all(|ch| ch.is_ascii_digit()) { - return false; - } - match parts.next() { - None => true, - Some(end) => { - !end.is_empty() && end.chars().all(|ch| ch.is_ascii_digit()) && parts.next().is_none() - } - } -} - -fn normalize_pages_param(pages: &Value) -> Result, Error> { - match pages { - Value::String(value) => { - let normalized = value - .split(',') - .map(str::trim) - .collect::>() - .join(","); - if normalized.split(',').all(pages_token_is_valid) { - Ok(Some(normalized)) - } else { - Err(Error::InvalidRequest(format!( - "Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'." - ))) - } - } - Value::Array(values) => { - if values.is_empty() { - return Ok(None); - } - if values.iter().any(Value::is_boolean) { - return Err(Error::InvalidRequest( - "`pages` must be integers, not booleans".to_string(), - )); - } - if values.iter().all(Value::is_i64) { - let mut pages = BTreeSet::new(); - for value in values { - let page = value.as_i64().expect("checked is_i64"); - if page < 0 { - return Err(Error::InvalidRequest( - "`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(), - )); - } - pages.insert(page + 1); - } - return Ok(Some( - pages - .into_iter() - .map(|page| page.to_string()) - .collect::>() - .join(","), - )); - } - if values.iter().all(Value::is_string) { - let normalized = values - .iter() - .filter_map(Value::as_str) - .map(str::trim) - .collect::>() - .join(","); - if normalized.split(',').all(pages_token_is_valid) { - return Ok(Some(normalized)); - } - return Err(Error::InvalidRequest(format!( - "Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'." - ))); - } - Err(Error::InvalidRequest( - "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." - .to_string(), - )) - } - _ => Err(Error::InvalidRequest( - "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." - .to_string(), - )), - } -} - -fn feature_token_is_valid(token: &str) -> bool { - let Some((first, rest)) = token.as_bytes().split_first() else { - return false; - }; - first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric) -} - -fn invalid_features_error(features: &Value) -> Error { - Error::InvalidRequest(format!( - "Invalid `features` for Azure Document Intelligence: {features:?}. Expected a list of feature names or a comma-separated string like 'keyValuePairs' or 'keyValuePairs,languages'." - )) -} - -fn normalize_features_param(features: &Value) -> Result, Error> { - let normalized = match features { - Value::String(value) => value - .split(',') - .map(str::trim) - .collect::>() - .join(","), - Value::Array(values) if values.is_empty() => return Ok(None), - Value::Array(values) => values - .iter() - .map(Value::as_str) - .collect::>>() - .ok_or_else(|| invalid_features_error(features))? - .into_iter() - .map(str::trim) - .collect::>() - .join(","), - _ => return Err(invalid_features_error(features)), - }; - - if normalized.split(',').all(feature_token_is_valid) { - Ok(Some(normalized)) - } else { - Err(invalid_features_error(features)) - } -} - -fn normalize_req_format(req_format: &Value) -> Result { - match req_format.as_str() { - Some(value @ ("native" | "litellm")) => Ok(value.to_string()), - _ => Err(Error::InvalidRequest(format!( - "Invalid `req_format` for Azure Document Intelligence: {req_format:?}. Expected 'native' or 'litellm'." - ))), - } -} - -pub fn map_document_intelligence_ocr_params( - non_default_params: &Map, -) -> Result, Error> { - let mut mapped = Map::new(); - if let Some(pages) = non_default_params.get("pages") - && let Some(normalized) = normalize_pages_param(pages)? - { - mapped.insert("pages".to_string(), Value::String(normalized)); - } - if let Some(features) = non_default_params.get("features") - && let Some(normalized) = normalize_features_param(features)? - { - mapped.insert("features".to_string(), Value::String(normalized)); - } - if let Some(req_format) = non_default_params.get("req_format") { - mapped.insert( - "req_format".to_string(), - Value::String(normalize_req_format(req_format)?), - ); - } - Ok(mapped) -} - -pub fn complete_document_intelligence_url( - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?; - let mut url = format!( - "{}/documentintelligence/documentModels/{}:analyze?api-version={}", - endpoint.trim_end_matches('/'), - encode_model_id(model)?, - AZURE_DOCUMENT_INTELLIGENCE_API_VERSION - ); - - if let Some(pages) = optional_params.get("pages") - && let Some(normalized) = normalize_pages_param(pages)? - { - url.push_str("&pages="); - url.push_str(&normalized); - } - - if let Some(features) = optional_params.get("features") - && let Some(normalized) = normalize_features_param(features)? - { - url.push_str("&features="); - url.push_str(&normalized); - } - - if let Some(req_format) = optional_params.get("req_format") { - normalize_req_format(req_format)?; - } - - Ok(url) -} - -fn document_url_from_mistral_document(document: &Value) -> Result<&str, Error> { - let object = document.as_object().ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(document), - })?; - let doc_type = object - .get("type") - .and_then(Value::as_str) - .ok_or(Error::MissingField("document.type"))?; - let field_name = match doc_type { - "document_url" => "document_url", - "image_url" => "image_url", - other => { - return Err(Error::InvalidRequest(format!( - "Invalid document type: {other}. Must be 'document_url' or 'image_url'" - ))); - } - }; - object - .get(field_name) - .and_then(Value::as_str) - .filter(|value| !value.is_empty()) - .ok_or(Error::MissingField(field_name)) -} - -fn extract_base64_from_data_uri(data_uri: &str) -> &str { - data_uri - .split_once(',') - .map(|(_, data)| data) - .unwrap_or(data_uri) -} - -fn page_markdown(page: &Map) -> String { - page.get("lines") - .and_then(Value::as_array) - .map(|lines| { - lines - .iter() - .filter_map(|line| line.get("content").and_then(Value::as_str)) - .collect::>() - .join("\n") - }) - .unwrap_or_default() -} - -fn page_dimensions(page: &Map) -> Value { - let width = page.get("width").and_then(Value::as_f64).unwrap_or(8.5); - let height = page.get("height").and_then(Value::as_f64).unwrap_or(11.0); - let unit = page.get("unit").and_then(Value::as_str).unwrap_or("inch"); - let (width, height) = if unit == "inch" { - ( - (width * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64, - (height * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64, - ) - } else { - (width as i64, height as i64) - }; - json!({ - "width": width, - "height": height, - "dpi": AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI, - }) -} - -fn transform_document_intelligence_response( - model: &str, - response_json: Value, - preserve_native_response: bool, -) -> Result { - let response = response_json - .as_object() - .ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(&response_json), - })?; - let status = response - .get("status") - .and_then(Value::as_str) - .ok_or(Error::MissingField("status"))?; - if status != "succeeded" { - return Err(Error::InvalidResponse(format!( - "Azure Document Intelligence analysis failed with status: {status}" - ))); - } - - let analyze_result = response.get("analyzeResult").and_then(Value::as_object); - let azure_pages = analyze_result - .and_then(|result| result.get("pages")) - .and_then(Value::as_array) - .cloned() - .unwrap_or_default(); - let pages = azure_pages - .iter() - .filter_map(Value::as_object) - .map(|page| { - let page_number = page.get("pageNumber").and_then(Value::as_i64).unwrap_or(1); - json!({ - "index": page_number - 1, - "markdown": page_markdown(page), - "dimensions": page_dimensions(page), - }) - }) - .collect::>(); - let extra_fields = ["content", "tables", "keyValuePairs"] - .into_iter() - .map(|field| { - ( - field.to_string(), - analyze_result - .and_then(|result| result.get(field)) - .cloned() - .unwrap_or(Value::Null), - ) - }) - .collect(); - - Ok(OcrResponseData { - usage_info: Some(json!({ - "pages_processed": pages.len(), - "doc_size_bytes": null, - })), - pages, - model: model.to_string(), - document_annotation: None, - object: "ocr".to_string(), - extra_fields, - provider_native_response: preserve_native_response.then_some(response_json), - }) -} - -impl OcrProviderConfig for AzureAiOcrConfig { - fn supported_ocr_params(&self) -> &'static [&'static str] { - MISTRAL_OCR_CONFIG.supported_ocr_params() - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn transform_ocr_request( - &self, - model: &str, - document: Value, - optional_params: Map, - ) -> Result { - MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) - } - - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result { - MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - _optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - complete_azure_ai_url(api_base, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_azure_ai_api_key(api_key, env_lookup) - } - - fn requires_data_uri_document(&self) -> bool { - true - } -} - -impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn supported_ocr_params(&self) -> &'static [&'static str] { - AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn map_ocr_params(&self, non_default_params: &Map) -> Map { - map_document_intelligence_ocr_params(non_default_params).unwrap_or_else(|_| { - non_default_params - .iter() - .filter(|(name, _)| { - AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS.contains(&name.as_str()) - }) - .map(|(name, value)| (name.clone(), value.clone())) - .collect() - }) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn transform_ocr_request( - &self, - _model: &str, - document: Value, - _optional_params: Map, - ) -> Result { - let document_url = document_url_from_mistral_document(&document)?; - let mut data = Map::new(); - if document_url.starts_with("data:") { - data.insert( - "base64Source".to_string(), - Value::String(extract_base64_from_data_uri(document_url).to_string()), - ); - } else { - data.insert( - "urlSource".to_string(), - Value::String(document_url.to_string()), - ); - } - Ok(OcrRequestData { - data: Value::Object(data), - files: None, - }) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result { - transform_document_intelligence_response(model, response_json, false) - } - - fn transform_ocr_response_with_params( - &self, - model: &str, - response_json: Value, - optional_params: &Map, - ) -> Result { - transform_document_intelligence_response( - model, - response_json, - optional_params.get("req_format").and_then(Value::as_str) == Some("native"), - ) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn complete_url( - &self, - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - complete_document_intelligence_url(api_base, model, optional_params, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_document_intelligence_api_key(api_key, env_lookup) - } - - fn auth_strategy(&self) -> OcrAuthStrategy { - OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key") - } - - fn response_handling(&self) -> OcrResponseHandling { - OcrResponseHandling::AzureDocumentIntelligencePoll - } -} - -#[cfg(test)] -mod tests { - use super::*; - use rstest::{fixture, rstest}; - - const ENDPOINT: &str = "https://example.cognitiveservices.azure.com"; - - #[fixture] - fn document_intelligence_config() -> AzureDocumentIntelligenceOcrConfig { - AzureDocumentIntelligenceOcrConfig - } - - fn header_value<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { - headers - .iter() - .find(|(header_name, _)| header_name.eq_ignore_ascii_case(name)) - .map(|(_, value)| value.as_str()) - } - - #[fixture] - fn native_operation() -> Value { - json!({ - "status": "succeeded", - "createdDateTime": "2026-07-02T00:00:00Z", - "lastUpdatedDateTime": "2026-07-02T00:00:05Z", - "analyzeResult": { - "content": "Invoice\nInvoice No: INV-12345\nTotal: $100.00", - "pages": [{ - "pageNumber": 1, - "width": 8.5, - "height": 11, - "unit": "inch", - "angle": 0.13, - "lines": [ - {"content": "Invoice"}, - {"content": "Invoice No: INV-12345"}, - {"content": "Total: $100.00"} - ], - "words": [{"content": "Invoice", "confidence": 0.994}] - }], - "tables": [ - { - "rowCount": 2, - "columnCount": 2, - "cells": [ - {"kind": "columnHeader", "rowIndex": 0, "columnIndex": 0, "content": "Item"}, - {"kind": "columnHeader", "rowIndex": 0, "columnIndex": 1, "content": "Price"}, - {"rowIndex": 1, "columnIndex": 0, "content": "Widget"}, - {"rowIndex": 1, "columnIndex": 1, "content": "$100.00"} - ] - }, - { - "rowCount": 1, - "columnCount": 1, - "cells": [{"rowIndex": 0, "columnIndex": 0, "content": "Totals"}] - } - ], - "keyValuePairs": [ - { - "key": {"content": "Invoice No"}, - "value": {"content": "INV-12345"}, - "confidence": 0.98 - }, - { - "key": {"content": "Total"}, - "value": {"content": "$100.00"}, - "confidence": 0.95 - } - ], - "paragraphs": [{"content": "Invoice"}] - } - }) - } - - fn assert_native_fields_preserved(response: &OcrResponseData, operation: &Value) { - let analyze_result = &operation["analyzeResult"]; - - assert_eq!(response.extra_fields["content"], analyze_result["content"]); - assert_eq!(response.extra_fields["tables"], analyze_result["tables"]); - assert_eq!( - response.extra_fields["keyValuePairs"], - analyze_result["keyValuePairs"] - ); - assert_eq!(response.object, "ocr"); - assert_eq!( - response.usage_info, - Some(json!({"pages_processed": 1, "doc_size_bytes": null})) - ); - assert_eq!(response.pages[0]["index"], 0); - assert_eq!( - response.pages[0]["markdown"], - "Invoice\nInvoice No: INV-12345\nTotal: $100.00" - ); - assert_eq!( - response.pages[0]["dimensions"], - json!({"width": 816, "height": 1056, "dpi": 96}) - ); - } - - #[test] - fn azure_ai_reuses_mistral_body_transform() { - let body = AZURE_AI_OCR_CONFIG - .transform_ocr_request( - "pixtral-12b-2409", - json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc"}), - serde_json::Map::from_iter([("include_image_base64".to_string(), json!(true))]), - ) - .expect("request transforms") - .data; - - assert_eq!(body["model"], "pixtral-12b-2409"); - assert_eq!(body["include_image_base64"], true); - assert_eq!( - body["document"]["document_url"], - "data:application/pdf;base64,abc" - ); - } - - #[test] - fn document_intelligence_url_normalizes_zero_based_pages() { - let params = serde_json::Map::from_iter([("pages".to_string(), json!([2, 0, 2]))]); - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com/"), - "azure_ai/doc-intelligence/prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,3" - ); - } - - #[test] - fn document_intelligence_url_normalizes_features() { - let params = serde_json::Map::from_iter([( - "features".to_string(), - json!("keyValuePairs, languages"), - )]); - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com"), - "prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&features=keyValuePairs,languages" - ); - } - - #[test] - fn document_intelligence_url_combines_pages_and_feature_list() { - let params = serde_json::Map::from_iter([ - ("pages".to_string(), json!([0, 1, 2])), - ( - "features".to_string(), - json!([" keyValuePairs ", "languages"]), - ), - ]); - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com"), - "prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,2,3&features=keyValuePairs,languages" - ); - } - - #[test] - fn document_intelligence_url_omits_empty_feature_list() { - let params = serde_json::Map::from_iter([("features".to_string(), json!([]))]); - assert!( - map_document_intelligence_ocr_params(¶ms) - .expect("empty features map") - .is_empty() - ); - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com"), - "prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30" - ); - } - - #[rstest] - #[case::query_injection(json!("keyValuePairs&pages=9"))] - #[case::spaces(json!("key value pairs"))] - #[case::empty_string(json!(""))] - #[case::integer_list(json!([1, 2]))] - #[case::nested_list(json!([["keyValuePairs"]]))] - #[case::object(json!({"feature": "keyValuePairs"}))] - #[case::number(json!(5))] - fn document_intelligence_mapping_rejects_invalid_features(#[case] features: Value) { - let params = serde_json::Map::from_iter([("features".to_string(), features)]); - let error = - map_document_intelligence_ocr_params(¶ms).expect_err("invalid features must fail"); - - assert!(matches!( - error, - Error::InvalidRequest(message) if message.contains("Invalid `features`") - )); - } - - #[rstest] - #[case::single_list(json!(["keyValuePairs"]), "keyValuePairs")] - #[case::multiple_list( - json!(["keyValuePairs", "languages"]), - "keyValuePairs,languages" - )] - #[case::single_string(json!("keyValuePairs"), "keyValuePairs")] - #[case::comma_separated(json!("keyValuePairs,languages"), "keyValuePairs,languages")] - #[case::spaces(json!("keyValuePairs, languages"), "keyValuePairs,languages")] - fn document_intelligence_maps_features(#[case] features: Value, #[case] expected: &str) { - let params = Map::from_iter([ - ("features".to_string(), features), - ("unsupported".to_string(), json!(true)), - ]); - - assert_eq!( - AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.map_ocr_params(¶ms), - Map::from_iter([("features".to_string(), json!(expected))]) - ); - } - - #[test] - fn document_intelligence_request_uses_base64_source_for_data_uri() { - let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_request( - "prebuilt-read", - json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc123"}), - Map::new(), - ) - .expect("request transforms") - .data; - - assert_eq!(body, json!({"base64Source": "abc123"})); - } - - #[rstest] - fn document_intelligence_response_normalizes_pages(native_operation: Value) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response("prebuilt-layout", native_operation.clone()) - .expect("response transforms"); - - assert_native_fields_preserved(&response, &native_operation); - } - - #[test] - fn azure_document_intelligence_model_id_is_encoded() { - let url = complete_document_intelligence_url( - Some(ENDPOINT), - "prebuilt-layout?x=1#frag", - &Map::new(), - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout%3Fx%3D1%23frag:analyze?api-version=2024-11-30" - ); - } - - #[test] - fn azure_document_intelligence_dot_segment_model_id_is_rejected() { - let error = complete_document_intelligence_url( - Some(ENDPOINT), - "azure_ai/doc-intelligence/..", - &Map::new(), - &|_| None, - ) - .expect_err("dot segment must fail"); - - assert_eq!( - error, - Error::InvalidRequest("model_id cannot be a dot path segment".to_string()) - ); - } - - #[rstest] - fn document_intelligence_async_response_preserves_normalized_fields(native_operation: Value) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response( - "azure_ai/doc-intelligence/prebuilt-layout", - native_operation.clone(), - ) - .expect("response transforms"); - - assert_native_fields_preserved(&response, &native_operation); - } - - #[test] - fn document_intelligence_response_tolerates_missing_native_fields() { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response( - "azure_ai/doc-intelligence/prebuilt-read", - json!({ - "status": "succeeded", - "analyzeResult": { - "pages": [{ - "pageNumber": 1, - "width": 8.5, - "height": 11, - "unit": "inch", - "lines": [{"content": "hello"}] - }] - } - }), - ) - .expect("missing optional fields are allowed"); - - assert_eq!(response.pages[0]["markdown"], "hello"); - assert_eq!(response.extra_fields["content"], Value::Null); - assert_eq!(response.extra_fields["tables"], Value::Null); - assert_eq!(response.extra_fields["keyValuePairs"], Value::Null); - } - - #[test] - fn document_intelligence_non_succeeded_status_is_rejected() { - let error = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response( - "azure_ai/doc-intelligence/prebuilt-layout", - json!({"status": "failed"}), - ) - .expect_err("failed status must fail"); - - assert_eq!( - error, - Error::InvalidResponse( - "Azure Document Intelligence analysis failed with status: failed".to_string() - ) - ); - } - - #[test] - fn document_intelligence_supported_params_include_features() { - assert_eq!( - AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.supported_ocr_params(), - &["pages", "features", "req_format"] - ); - } - - #[rstest] - fn document_intelligence_native_format_carries_raw_operation(native_operation: Value) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response_with_params( - "azure_ai/doc-intelligence/prebuilt-layout", - native_operation.clone(), - &Map::from_iter([("req_format".to_string(), json!("native"))]), - ) - .expect("native response transforms"); - - assert_eq!( - response.provider_native_response, - Some(native_operation.clone()) - ); - assert_native_fields_preserved(&response, &native_operation); - } - - #[rstest] - fn document_intelligence_async_native_format_carries_raw_operation(native_operation: Value) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response_with_params( - "azure_ai/doc-intelligence/prebuilt-layout", - native_operation.clone(), - &Map::from_iter([("req_format".to_string(), json!("native"))]), - ) - .expect("native response transforms"); - - assert_eq!( - response.provider_native_response, - Some(native_operation.clone()) - ); - assert_native_fields_preserved(&response, &native_operation); - } - - #[rstest] - #[case::default(Map::new())] - #[case::litellm(Map::from_iter([("req_format".to_string(), json!("litellm"))]))] - fn document_intelligence_default_format_omits_raw_operation( - #[case] optional_params: Map, - native_operation: Value, - ) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response_with_params( - "azure_ai/doc-intelligence/prebuilt-layout", - native_operation.clone(), - &optional_params, - ) - .expect("response transforms"); - - assert_eq!(response.provider_native_response, None); - assert_native_fields_preserved(&response, &native_operation); - } - - #[rstest] - #[case::native("native")] - #[case::litellm("litellm")] - fn document_intelligence_maps_req_format(#[case] req_format: &str) { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "req_format".to_string(), - json!(req_format), - )])) - .expect("req_format maps"); - - assert_eq!( - mapped, - Map::from_iter([("req_format".to_string(), json!(req_format))]) - ); - } - - #[test] - fn document_intelligence_rejects_unknown_req_format() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "req_format".to_string(), - json!("azure"), - )])) - .expect_err("unknown req_format must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("Invalid `req_format`")) - ); - } - - #[test] - fn document_intelligence_url_omits_req_format() { - let url = complete_document_intelligence_url( - Some(ENDPOINT), - "prebuilt-layout", - &Map::from_iter([("req_format".to_string(), json!("native"))]), - &|_| None, - ) - .expect("url builds"); - - assert!(!url.contains("req_format")); - } - - #[test] - fn document_intelligence_validate_environment_uses_subscription_key() { - let headers = - validate_document_intelligence_environment(Vec::new(), Some("my-key"), None, &|_| None) - .expect("api key authenticates"); - - assert_eq!( - header_value(&headers, "Ocp-Apim-Subscription-Key"), - Some("my-key") - ); - } - - #[test] - fn document_intelligence_validate_environment_falls_back_to_entra_token() { - let headers = validate_document_intelligence_environment( - Vec::new(), - None, - Some("entra-token"), - &|_| None, - ) - .expect("Entra token authenticates"); - - assert_eq!( - header_value(&headers, "Authorization"), - Some("Bearer entra-token") - ); - assert_eq!(header_value(&headers, "Ocp-Apim-Subscription-Key"), None); - } - - #[test] - fn document_intelligence_supported_params_include_pages_features_and_req_format() { - assert_eq!( - AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.supported_ocr_params(), - &["pages", "features", "req_format"] - ); - } - - #[test] - fn document_intelligence_maps_zero_based_page_list() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([0, 1, 2]), - )])) - .expect("pages map"); - - assert_eq!( - mapped, - Map::from_iter([("pages".to_string(), json!("1,2,3"))]) - ); - } - - #[test] - fn document_intelligence_page_mapping_dedupes_and_sorts() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([2, 0, 0, 1]), - )])) - .expect("pages map"); - - assert_eq!(mapped["pages"], "1,2,3"); - } - - #[test] - fn document_intelligence_page_mapping_omits_empty_list() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([]), - )])) - .expect("empty pages map"); - - assert!(mapped.is_empty()); - } - - #[test] - fn document_intelligence_page_mapping_accepts_native_range() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!("3-9"), - )])) - .expect("range maps"); - - assert_eq!(mapped["pages"], "3-9"); - } - - #[test] - fn document_intelligence_page_mapping_strips_spaces() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!("1-3, 5"), - )])) - .expect("range maps"); - - assert_eq!(mapped["pages"], "1-3,5"); - } - - #[test] - fn document_intelligence_page_mapping_accepts_string_tokens() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!(["1", "3-5"]), - )])) - .expect("tokens map"); - - assert_eq!(mapped["pages"], "1,3-5"); - } - - #[test] - fn document_intelligence_page_mapping_rejects_invalid_string() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!("a,b"), - )])) - .expect_err("invalid pages must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("Invalid `pages` string")) - ); - } - - #[test] - fn document_intelligence_page_mapping_rejects_negative_index() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([-1]), - )])) - .expect_err("negative pages must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("must be >= 0")) - ); - } - - #[test] - fn document_intelligence_page_mapping_rejects_bool_list() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([true, false]), - )])) - .expect_err("boolean pages must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("integers, not booleans")) - ); - } - - #[test] - fn document_intelligence_page_mapping_rejects_unsupported_type() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!(5), - )])) - .expect_err("unsupported pages must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("Mistral-style")) - ); - } - - #[test] - fn document_intelligence_url_appends_pages_query() { - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com/"), - "azure_ai/doc-intelligence/prebuilt-layout", - &Map::from_iter([("pages".to_string(), json!("1-3,5"))]), - &|_| None, - ) - .expect("url builds"); - - assert!(url.contains("api-version=2024-11-30")); - assert!(url.contains("pages=1-3,5")); - assert!(url.contains("/documentintelligence/documentModels/prebuilt-layout:analyze")); - } - - #[test] - fn document_intelligence_url_has_no_pages_when_params_are_empty() { - let url = complete_document_intelligence_url( - Some(ENDPOINT), - "prebuilt-layout", - &Map::new(), - &|_| None, - ) - .expect("url builds"); - - assert!(!url.contains("pages=")); - } - - #[rstest] - fn document_intelligence_request_keeps_pages_out_of_body( - document_intelligence_config: AzureDocumentIntelligenceOcrConfig, - ) { - let request = document_intelligence_config - .transform_ocr_request( - "prebuilt-layout", - json!({"type": "document_url", "document_url": "https://example.com/x.pdf"}), - Map::from_iter([("pages".to_string(), json!("1,2,3"))]), - ) - .expect("request transforms"); - - assert_eq!( - request.data, - json!({"urlSource": "https://example.com/x.pdf"}) - ); - } - - #[test] - fn document_intelligence_mistral_pages_flow_to_query_only() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([2, 3, 4, 5, 6, 7, 8]), - )])) - .expect("pages map"); - let url = - complete_document_intelligence_url(Some(ENDPOINT), "prebuilt-layout", &mapped, &|_| { - None - }) - .expect("url builds"); - let request = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_request( - "prebuilt-layout", - json!({"type": "document_url", "document_url": "https://example.com/x.pdf"}), - mapped, - ) - .expect("request transforms"); - - assert!(url.contains("pages=3,4,5,6,7,8,9")); - assert_eq!( - request.data, - json!({"urlSource": "https://example.com/x.pdf"}) - ); - } - - #[test] - fn document_intelligence_endpoint_ignores_generic_azure_ai_base() { - let resolved = resolve_document_intelligence_endpoint(None, &|name| match name { - AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV => Some(ENDPOINT.to_string()), - AZURE_AI_API_BASE_ENV => Some("https://generic.example.com".to_string()), - _ => None, - }) - .expect("endpoint resolves"); - - assert_eq!(resolved, ENDPOINT); - } - - #[test] - fn document_intelligence_endpoint_honors_explicit_api_base() { - let resolved = resolve_document_intelligence_endpoint( - Some("https://my-di.cognitiveservices.azure.com"), - &|name| match name { - AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV => Some(ENDPOINT.to_string()), - AZURE_AI_API_BASE_ENV => Some("https://generic.example.com".to_string()), - _ => None, - }, - ) - .expect("endpoint resolves"); - - assert_eq!(resolved, "https://my-di.cognitiveservices.azure.com"); - } - - #[test] - fn azure_ai_mistral_ocr_uses_generic_api_base() { - let resolved = resolve_azure_ai_api_base(None, &|name| match name { - AZURE_AI_API_BASE_ENV => Some("https://generic-azure-ai.example.com".to_string()), - AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV => Some(ENDPOINT.to_string()), - _ => None, - }) - .expect("api base resolves"); - - assert_eq!(resolved, "https://generic-azure-ai.example.com"); - } - - #[test] - fn azure_ai_ocr_authenticates_with_entra_token() { - let headers = - validate_azure_ai_environment(Vec::new(), None, Some("entra-token"), &|_| None) - .expect("Entra token authenticates"); - - assert_eq!( - header_value(&headers, "Authorization"), - Some("Bearer entra-token") - ); - } -} diff --git a/litellm-rust/crates/core/src/providers/mistral/mod.rs b/litellm-rust/crates/core/src/providers/mistral/mod.rs deleted file mode 100644 index 3621ff6a2fd..00000000000 --- a/litellm-rust/crates/core/src/providers/mistral/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs deleted file mode 100644 index f239b6921fa..00000000000 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs deleted file mode 100644 index 11e8fe7db18..00000000000 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ /dev/null @@ -1,433 +0,0 @@ -use crate::error::{Error, json_type_name}; -use crate::ocr::transformation::OcrProviderConfig; -use crate::ocr::types::{OcrRequestData, OcrResponseData}; -use serde_json::{Map, Value}; - -const SUPPORTED_OCR_PARAMS: &[&str] = &[ - "pages", - "include_image_base64", - "image_limit", - "image_min_size", - "bbox_annotation_format", - "document_annotation_format", - "document_annotation_prompt", - "extract_header", - "extract_footer", - "table_format", - "confidence_scores_granularity", - "include_blocks", - "id", -]; - -/// Default Mistral API base, used when the caller does not override `api_base`. -pub const MISTRAL_DEFAULT_API_BASE: &str = "https://api.mistral.ai/v1"; - -/// Environment variable holding the Mistral API key. -pub const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY"; - -/// Error message raised when no Mistral API key can be resolved. -pub const MISSING_KEY_MESSAGE: &str = "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params"; - -/// Build the complete OCR endpoint URL, de-duplicating a trailing `/v1`. -/// -/// Blank/whitespace `api_base` is treated as absent (guard at resolution time). -pub fn complete_url(api_base: Option<&str>) -> String { - let base = api_base - .map(str::trim) - .filter(|base| !base.is_empty()) - .unwrap_or(MISTRAL_DEFAULT_API_BASE) - .trim_end_matches('/'); - - if base.ends_with("/v1") { - format!("{base}/ocr") - } else { - format!("{base}/v1/ocr") - } -} - -/// Resolve the Mistral API key from the explicit param or the environment. -/// -/// Blank/whitespace values are treated as absent. Returns `Error::Auth` -/// when no usable key is available. -/// -/// Note: the env fallback only reads the process environment. Secret-manager -/// backends (AWS/Azure/GCP/Vault) are resolved on the Python side and passed in -/// via `api_key`; this fallback is a last resort for direct/standalone use. -pub fn resolve_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - api_key - .map(str::trim) - .filter(|key| !key.is_empty()) - .map(str::to_string) - .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) -} - -pub struct MistralOcrConfig; - -pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig; - -impl OcrProviderConfig for MistralOcrConfig { - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn supported_ocr_params(&self) -> &'static [&'static str] { - SUPPORTED_OCR_PARAMS - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn transform_ocr_request( - &self, - model: &str, - document: Value, - optional_params: Map, - ) -> Result { - if !document.is_object() { - return Err(Error::InvalidType { - expected: "object", - actual: json_type_name(&document), - }); - } - - let mut data = Map::new(); - data.insert("model".to_string(), Value::String(model.to_string())); - data.insert("document".to_string(), document); - for (param, value) in optional_params { - data.insert(param, value); - } - - Ok(OcrRequestData { - data: Value::Object(data), - files: None, - }) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result { - let response_object = response_json - .as_object() - .ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(&response_json), - })?; - - let pages = response_object - .get("pages") - .and_then(Value::as_array) - .cloned() - .unwrap_or_default(); - let model = response_object - .get("model") - .and_then(Value::as_str) - .unwrap_or(model) - .to_string(); - let document_annotation = response_object.get("document_annotation").cloned(); - let usage_info = response_object.get("usage_info").cloned(); - - Ok(OcrResponseData { - pages, - model, - document_annotation, - usage_info, - object: "ocr".to_string(), - extra_fields: Map::new(), - provider_native_response: None, - }) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - _optional_params: &Map, - _env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(complete_url(api_base)) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_api_key(api_key, env_lookup) - } -} - -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub fn supported_ocr_params() -> &'static [&'static str] { - MISTRAL_OCR_CONFIG.supported_ocr_params() -} - -pub fn map_ocr_params(non_default_params: &Map) -> Map { - MISTRAL_OCR_CONFIG.map_ocr_params(non_default_params) -} - -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub fn transform_ocr_request( - model: &str, - document: Value, - optional_params: Map, -) -> Result { - MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) -} - -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub fn transform_ocr_response(model: &str, response_json: Value) -> Result { - MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn extract_header_is_a_supported_ocr_param() { - assert!(supported_ocr_params().contains(&"extract_header")); - } - - #[test] - fn extract_footer_is_a_supported_ocr_param() { - assert!(supported_ocr_params().contains(&"extract_footer")); - } - - #[test] - fn existing_ocr_params_remain_supported() { - for param in [ - "pages", - "include_image_base64", - "image_limit", - "image_min_size", - "bbox_annotation_format", - "document_annotation_format", - ] { - assert!(supported_ocr_params().contains(¶m)); - } - } - - #[test] - fn map_ocr_params_forwards_extract_header() { - let params = json!({"extract_header": true}); - assert_eq!( - map_ocr_params(params.as_object().unwrap()), - params.as_object().unwrap().clone() - ); - } - - #[test] - fn map_ocr_params_forwards_extract_footer() { - let params = json!({"extract_footer": true}); - assert_eq!( - map_ocr_params(params.as_object().unwrap()), - params.as_object().unwrap().clone() - ); - } - - #[test] - fn map_ocr_params_forwards_extract_header_and_footer() { - let params = json!({"extract_header": true, "extract_footer": false}); - assert_eq!( - map_ocr_params(params.as_object().unwrap()), - params.as_object().unwrap().clone() - ); - } - - #[test] - fn map_ocr_params_drops_unknown_params() { - let params = json!({"extract_header": true, "unsupported_param": "value"}); - let mapped = map_ocr_params(params.as_object().unwrap()); - assert_eq!(mapped.get("extract_header"), Some(&json!(true))); - assert!(!mapped.contains_key("unsupported_param")); - } - - #[test] - fn new_ocr_params_are_supported() { - for param in [ - "table_format", - "confidence_scores_granularity", - "document_annotation_prompt", - "include_blocks", - "id", - ] { - assert!(supported_ocr_params().contains(¶m)); - } - } - - #[test] - fn map_ocr_params_forwards_new_ocr_params() { - for (param, value) in [ - ("table_format", json!("html")), - ("confidence_scores_granularity", json!("word")), - ( - "document_annotation_prompt", - json!("Extract all invoice line items"), - ), - ("include_blocks", json!(true)), - ("id", json!("req-123")), - ] { - let params = json!({param: value}); - assert_eq!( - map_ocr_params(params.as_object().unwrap()), - params.as_object().unwrap().clone() - ); - } - } - - #[test] - fn transform_ocr_request_includes_each_optional_param() { - let document = json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }); - for (param, value) in [ - ("table_format", json!("html")), - ("confidence_scores_granularity", json!("word")), - ( - "document_annotation_prompt", - json!("Extract all invoice line items"), - ), - ("id", json!("req-123")), - ("extract_header", json!(true)), - ("include_blocks", json!(true)), - ("pages", json!([0, 1])), - ] { - let result = transform_ocr_request( - "mistral-ocr-latest", - document.clone(), - json!({param: value}).as_object().unwrap().clone(), - ) - .expect("request should transform"); - assert_eq!(result.data.get(param), Some(&value)); - assert_eq!(result.data.get("model"), Some(&json!("mistral-ocr-latest"))); - assert_eq!(result.data.get("document"), Some(&document)); - assert_eq!(result.files, None); - } - } - - #[test] - fn transform_ocr_request_includes_multiple_new_params() { - let document = json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }); - let optional_params = json!({ - "table_format": "html", - "confidence_scores_granularity": "page", - "extract_header": true - }) - .as_object() - .unwrap() - .clone(); - let result = transform_ocr_request("mistral-ocr-latest", document, optional_params) - .expect("request should transform"); - assert_eq!(result.data.get("table_format"), Some(&json!("html"))); - assert_eq!( - result.data.get("confidence_scores_granularity"), - Some(&json!("page")) - ); - assert_eq!(result.data.get("extract_header"), Some(&json!(true))); - } - - #[test] - fn transform_ocr_response_preserves_blocks_and_confidence_scores() { - let blocks = json!([{"type": "title", "content": "Invoice"}]); - let confidence_scores = json!({"page": 0.98}); - let response = json!({ - "pages": [{"index": 0, "markdown": "# Invoice", "blocks": blocks, "confidence_scores": confidence_scores}], - "model": "mistral-ocr-4-0", - "usage_info": {"pages_processed": 1} - }); - let result = - transform_ocr_response("mistral-ocr-4-0", response).expect("response should transform"); - assert_eq!(result.pages[0].get("blocks"), Some(&blocks)); - assert_eq!( - result.pages[0].get("confidence_scores"), - Some(&confidence_scores) - ); - } - - #[test] - fn transform_ocr_response_preserves_ocr4_page_fields() { - let response = json!({ - "pages": [{"index": 0, "markdown": "table page", "tables": [{"rows": 2, "cols": 3}], "hyperlinks": ["https://example.com"], "header": "Acme Corp", "footer": "Page 1"}], - "model": "mistral-ocr-4-0", - "usage_info": {"pages_processed": 1} - }); - let result = transform_ocr_response("mistral-ocr-4-0", response.clone()) - .expect("response should transform"); - assert_eq!(result.pages[0], response["pages"][0]); - } - - #[test] - fn transform_ocr_request_rejects_non_object_document() { - let err = transform_ocr_request("mistral-ocr-latest", json!("bad"), Map::new()) - .expect_err("string document should be rejected"); - - assert_eq!( - err, - Error::InvalidType { - expected: "object", - actual: "string", - } - ); - } - - #[test] - fn transform_ocr_response_normalizes_mistral_json() { - let response = json!({ - "pages": [{"index": 0, "markdown": "hello"}], - "model": "mistral-ocr-2505-completion", - "document_annotation": null, - "usage_info": {"pages_processed": 1} - }); - - let result = transform_ocr_response("mistral-ocr-latest", response) - .expect("response should transform"); - - assert_eq!(result.pages, vec![json!({"index": 0, "markdown": "hello"})]); - assert_eq!(result.model, "mistral-ocr-2505-completion"); - assert_eq!(result.document_annotation, Some(Value::Null)); - assert_eq!(result.usage_info, Some(json!({"pages_processed": 1}))); - assert_eq!(result.object, "ocr"); - } - - #[test] - fn complete_url_defaults_and_dedupes_v1() { - assert_eq!(complete_url(None), "https://api.mistral.ai/v1/ocr"); - assert_eq!(complete_url(Some(" ")), "https://api.mistral.ai/v1/ocr"); - assert_eq!( - complete_url(Some("https://proxy.internal")), - "https://proxy.internal/v1/ocr" - ); - assert_eq!( - complete_url(Some("https://proxy.internal/v1/")), - "https://proxy.internal/v1/ocr" - ); - } - - #[test] - fn resolve_api_key_prefers_param_then_env() { - let no_env = |_: &str| None; - assert_eq!( - resolve_api_key(Some("sk-param"), &no_env).unwrap(), - "sk-param" - ); - - let with_env = |key: &str| (key == MISTRAL_API_KEY_ENV).then(|| "sk-env".to_string()); - assert_eq!(resolve_api_key(None, &with_env).unwrap(), "sk-env"); - // Blank param falls through to the environment. - assert_eq!(resolve_api_key(Some(" "), &with_env).unwrap(), "sk-env"); - } - - #[test] - fn resolve_api_key_errors_when_absent() { - let err = resolve_api_key(None, &|_| None).expect_err("missing key should error"); - assert_eq!(err, Error::Auth(MISSING_KEY_MESSAGE.to_string())); - } -} diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index c0c2c69831b..1aeb75063d6 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -2,7 +2,4 @@ pub mod anthropic; pub mod azure_ai; #[cfg(feature = "bedrock-auth")] pub mod bedrock; -pub mod mistral; pub mod openai; -pub mod reducto; -pub mod vertex_ai; diff --git a/litellm-rust/crates/core/src/providers/reducto/mod.rs b/litellm-rust/crates/core/src/providers/reducto/mod.rs deleted file mode 100644 index 3621ff6a2fd..00000000000 --- a/litellm-rust/crates/core/src/providers/reducto/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/reducto/ocr/mod.rs b/litellm-rust/crates/core/src/providers/reducto/ocr/mod.rs deleted file mode 100644 index 8acee8f770c..00000000000 --- a/litellm-rust/crates/core/src/providers/reducto/ocr/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -pub mod transformation; - -#[cfg(test)] -mod tests; diff --git a/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs b/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs deleted file mode 100644 index 2b66d058b5d..00000000000 --- a/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs +++ /dev/null @@ -1,202 +0,0 @@ -use rstest::{fixture, rstest}; -use serde_json::{Value, json}; - -use super::transformation::*; -use crate::ocr::transformation::OcrProviderConfig; - -#[fixture] -fn parse_response() -> Value { - json!({ - "job_id": "job_123", - "usage": {"num_pages": 3, "credits": 3}, - "result": { - "chunks": [ - { - "content": "Page 1 block A", - "blocks": [{ - "content": "Page 1 block A", - "bbox": {"page": 1}, - "kind": "text", - }], - }, - { - "content": "Page 2 block A", - "blocks": [{ - "content": "Page 2 block A", - "bbox": {"page": 2}, - "kind": "table", - }], - }, - { - "content": "Page 1 block B", - "blocks": [{ - "content": "Page 1 block B", - "bbox": {"page": 1}, - "kind": "text", - }], - }, - { - "content": "Page 3 block A", - "blocks": [{ - "content": "Page 3 block A", - "bbox": {"page": 3}, - "kind": "figure", - }], - }, - ], - }, - }) -} - -#[rstest] -fn test_parse_v3_file_upload_and_response_mapping(parse_response: Value) { - let source = classify_document_source("data:application/pdf;base64,JVBERi0xLjQ=") - .expect("PDF data URI should be valid"); - let upload = build_upload_request( - source, - "Bearer test-key", - Some("https://platform.reducto.ai"), - ) - .expect("data URI should require upload"); - assert_eq!(upload.url, "https://platform.reducto.ai/upload"); - assert_eq!(upload.authorization, "Bearer test-key"); - assert_eq!(upload.file_name, "document"); - assert_eq!(upload.mime_type, "application/pdf"); - assert_eq!(upload.bytes, b"%PDF-1.4"); - - let optional_params = json!({ - "formatting": {"table_output_format": "html"}, - "retrieval": {"chunk_mode": "section"}, - "settings": {"ocr_system": "standard"}, - }) - .as_object() - .expect("params should be an object") - .clone(); - let request = build_parse_v3_request("reducto://uploaded.pdf", optional_params); - assert_eq!( - request.data, - json!({ - "input": "reducto://uploaded.pdf", - "formatting": {"table_output_format": "html"}, - "retrieval": {"chunk_mode": "section"}, - "settings": {"ocr_system": "standard"}, - }) - ); - - let transformed = transform_reducto_response("parse-v3", parse_response.clone()) - .expect("response should transform"); - assert_eq!( - transformed.usage_info, - Some(json!({"pages_processed": 3, "credits": 3})) - ); - assert_eq!(transformed.pages.len(), 3); - assert_eq!( - transformed.pages[0], - json!({ - "index": 0, - "markdown": "Page 1 block A\n\nPage 1 block B", - "blocks": [ - {"content": "Page 1 block A", "bbox": {"page": 1}, "kind": "text"}, - {"content": "Page 1 block B", "bbox": {"page": 1}, "kind": "text"}, - ], - }) - ); - assert_eq!(transformed.pages[1]["markdown"], "Page 2 block A"); - assert_eq!(transformed.pages[2]["markdown"], "Page 3 block A"); - assert_eq!(transformed.provider_native_response, Some(parse_response)); -} - -#[rstest] -fn test_parse_v3_reducto_id_passthrough_skips_upload(parse_response: Value) { - let document = json!({ - "type": "document_url", - "document_url": "reducto://already-uploaded.pdf", - }); - let source = extract_document_source(&document).expect("Reducto ID should be valid"); - assert!(build_upload_request(source.clone(), "Bearer test-key", None).is_none()); - assert_eq!( - source, - ReductoDocumentSource::FileId("reducto://already-uploaded.pdf".to_string()) - ); - - let request = REDUCTO_PARSE_V3_CONFIG - .transform_ocr_request( - "parse-v3", - document, - json!({"retrieval": {"chunk_mode": "section"}}) - .as_object() - .expect("params should be object") - .clone(), - ) - .expect("direct ID should transform"); - assert_eq!(request.data["input"], "reducto://already-uploaded.pdf"); - assert_eq!(request.data["retrieval"]["chunk_mode"], "section"); - - let response = REDUCTO_PARSE_V3_CONFIG - .transform_ocr_response("parse-v3", parse_response) - .expect("response should transform"); - assert!( - response.pages[0]["markdown"] - .as_str() - .expect("markdown should be string") - .starts_with("Page 1 block A") - ); -} - -#[rstest] -fn test_parse_legacy_wraps_enhance_under_options() { - let request = build_parse_legacy_request( - "reducto://legacy.pdf", - json!({"enhance": {"agentic": [{"type": "table"}]}}) - .as_object() - .expect("params should be object"), - ); - assert_eq!( - request.data, - json!({ - "document_url": "reducto://legacy.pdf", - "options": {"enhance": {"agentic": [{"type": "table"}]}}, - }) - ); -} - -#[rstest] -fn test_parse_v3_image_data_uri_upload_uses_image_mime() { - let source = classify_document_source("data:image/png;base64,iVBORw0KGgo=") - .expect("PNG data URI should be valid"); - let upload = build_upload_request( - source, - "Bearer programmatic-key", - Some("https://custom.reducto.test/"), - ) - .expect("data URI should require upload"); - assert_eq!(upload.url, "https://custom.reducto.test/upload"); - assert_eq!(upload.authorization, "Bearer programmatic-key"); - assert_eq!(upload.mime_type, "image/png"); - assert_eq!(upload.bytes, b"\x89PNG\r\n\x1a\n"); -} - -#[rstest] -#[case::http("http://example.com/document.pdf")] -#[case::https("https://example.com/document.pdf")] -fn test_parse_v3_rejects_plain_http_urls(#[case] source: &str) { - let error = classify_document_source(source).expect_err("plain URL should be rejected"); - assert!(error.to_string().contains("upload the file first")); -} - -#[rstest] -fn test_parse_v3_uses_programmatic_api_key_over_env() { - let key = resolve_api_key(Some("passed-key"), &|_| Some("env-reducto-key".to_string())) - .expect("explicit key should resolve"); - assert_eq!(key, "passed-key"); - - let headers = REDUCTO_PARSE_V3_CONFIG - .validate_environment(Vec::new(), Some("passed-key"), &|_| { - Some("env-reducto-key".to_string()) - }) - .expect("headers should validate"); - assert_eq!( - headers, - vec![("Authorization".to_string(), "Bearer passed-key".to_string())] - ); -} diff --git a/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs deleted file mode 100644 index b8507541e18..00000000000 --- a/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs +++ /dev/null @@ -1,407 +0,0 @@ -use std::collections::BTreeMap; - -use base64::Engine; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use serde_json::{Map, Value, json}; - -use crate::error::{Error, json_type_name}; -use crate::ocr::transformation::OcrProviderConfig; -use crate::ocr::types::{OcrRequestData, OcrResponseData}; - -pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai"; -pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY"; -pub const REDUCTO_ID_PREFIX: &str = "reducto://"; - -const PARSE_V3_SUPPORTED_OCR_PARAMS: &[&str] = &["formatting", "retrieval", "settings"]; -const PARSE_LEGACY_SUPPORTED_OCR_PARAMS: &[&str] = &["enhance"]; -const MISSING_KEY_MESSAGE: &str = "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"; -const DATA_URI_UPLOAD_REQUIRED: &str = - "Reducto data URI upload must complete before OCR request transformation"; - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum ReductoDocumentSource { - FileId(String), - Upload { bytes: Vec, mime_type: String }, -} - -#[derive(Clone, PartialEq, Eq)] -pub struct ReductoUploadRequest { - pub url: String, - pub authorization: String, - pub file_name: &'static str, - pub bytes: Vec, - pub mime_type: String, -} - -pub struct ReductoParseV3Config; -pub struct ReductoParseLegacyConfig; - -pub const REDUCTO_PARSE_V3_CONFIG: ReductoParseV3Config = ReductoParseV3Config; -pub const REDUCTO_PARSE_LEGACY_CONFIG: ReductoParseLegacyConfig = ReductoParseLegacyConfig; - -pub fn config_for_model(model: &str) -> Option<&'static dyn OcrProviderConfig> { - match model { - "parse-v3" => Some(&REDUCTO_PARSE_V3_CONFIG), - "parse-legacy" => Some(&REDUCTO_PARSE_LEGACY_CONFIG), - _ => None, - } -} - -pub fn normalize_api_base(api_base: Option<&str>) -> String { - api_base - .map(str::trim) - .filter(|base| !base.is_empty()) - .unwrap_or(REDUCTO_API_BASE) - .trim_end_matches('/') - .to_string() -} - -pub fn parse_url(api_base: Option<&str>) -> String { - format!("{}/parse", normalize_api_base(api_base)) -} - -pub fn upload_url(api_base: Option<&str>) -> String { - format!("{}/upload", normalize_api_base(api_base)) -} - -pub fn resolve_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - api_key - .map(str::trim) - .filter(|key| !key.is_empty()) - .map(str::to_string) - .or_else(|| { - env_lookup(REDUCTO_API_KEY_ENV) - .map(|key| key.trim().to_string()) - .filter(|key| !key.is_empty()) - }) - .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) -} - -pub fn extract_document_source(document: &Value) -> Result { - let document = document.as_object().ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(document), - })?; - let source = document - .get("document_url") - .and_then(Value::as_str) - .filter(|source| !source.is_empty()) - .or_else(|| document.get("image_url").and_then(Value::as_str)) - .ok_or_else(|| { - Error::InvalidRequest( - "Reducto expected OCR preprocessing to produce document_url or image_url" - .to_string(), - ) - })?; - classify_document_source(source) -} - -pub fn classify_document_source(source: &str) -> Result { - if source.starts_with(REDUCTO_ID_PREFIX) { - return Ok(ReductoDocumentSource::FileId(source.to_string())); - } - if source.starts_with("http://") || source.starts_with("https://") { - return Err(Error::InvalidRequest( - "Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first." - .to_string(), - )); - } - if !source.starts_with("data:") { - return Err(Error::InvalidRequest( - "Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing." - .to_string(), - )); - } - - let (header, encoded) = source - .split_once(',') - .ok_or_else(|| Error::InvalidRequest("Invalid Reducto data URI provided.".to_string()))?; - if !header.split(';').any(|part| part == "base64") { - return Err(Error::InvalidRequest( - "Reducto only supports base64-encoded data URIs.".to_string(), - )); - } - - let mime_type = header - .strip_prefix("data:") - .and_then(|header| header.split(';').next()) - .filter(|mime| !mime.is_empty()) - .unwrap_or("application/octet-stream") - .to_string(); - let bytes = BASE64_STANDARD.decode(encoded).map_err(|_| { - Error::InvalidRequest("Invalid Reducto base64 payload provided.".to_string()) - })?; - - Ok(ReductoDocumentSource::Upload { bytes, mime_type }) -} - -pub fn build_upload_request( - source: ReductoDocumentSource, - authorization: &str, - api_base: Option<&str>, -) -> Option { - let ReductoDocumentSource::Upload { bytes, mime_type } = source else { - return None; - }; - - Some(ReductoUploadRequest { - url: upload_url(api_base), - authorization: authorization.to_string(), - file_name: "document", - bytes, - mime_type, - }) -} - -pub fn extract_upload_file_id(response_json: &Value) -> Result<&str, Error> { - response_json - .as_object() - .and_then(|response| response.get("file_id")) - .and_then(Value::as_str) - .filter(|file_id| !file_id.is_empty()) - .ok_or_else(|| { - Error::InvalidResponse(format!( - "Reducto /upload returned 200 without a file_id; got payload={response_json}" - )) - }) -} - -pub fn build_parse_v3_request( - file_id: &str, - optional_params: Map, -) -> OcrRequestData { - let data = std::iter::once(("input".to_string(), Value::String(file_id.to_string()))) - .chain(optional_params) - .collect(); - OcrRequestData { - data: Value::Object(data), - files: None, - } -} - -pub fn build_parse_legacy_request( - file_id: &str, - optional_params: &Map, -) -> OcrRequestData { - let options = optional_params - .get("enhance") - .filter(|enhance| !enhance.is_null()) - .map(|enhance| json!({"options": {"enhance": enhance}})); - let data = match options { - Some(Value::Object(options)) => std::iter::once(( - "document_url".to_string(), - Value::String(file_id.to_string()), - )) - .chain(options) - .collect(), - _ => Map::from_iter([( - "document_url".to_string(), - Value::String(file_id.to_string()), - )]), - }; - OcrRequestData { - data: Value::Object(data), - files: None, - } -} - -fn source_file_id(document: &Value) -> Result { - match extract_document_source(document)? { - ReductoDocumentSource::FileId(file_id) => Ok(file_id), - ReductoDocumentSource::Upload { .. } => Err(Error::Unsupported(DATA_URI_UPLOAD_REQUIRED)), - } -} - -fn page_number(block: &Map) -> Option { - let page = block.get("bbox")?.as_object()?.get("page")?; - page.as_i64() - .or_else(|| page.as_u64().and_then(|page| i64::try_from(page).ok())) - .or_else(|| page.as_str().and_then(|page| page.parse().ok())) -} - -fn chunks(result: &Map) -> &[Value] { - result - .get("chunks") - .and_then(Value::as_array) - .map(Vec::as_slice) - .unwrap_or_default() -} - -fn build_pages(result: &Map) -> Vec { - let blocks_by_page = chunks(result) - .iter() - .filter_map(Value::as_object) - .filter_map(|chunk| chunk.get("blocks").and_then(Value::as_array)) - .flatten() - .filter_map(|block| block.as_object().map(|object| (block, object))) - .filter_map(|(block, object)| page_number(object).map(|page| (page, block.clone()))) - .fold( - BTreeMap::>::new(), - |mut pages, (page, block)| { - pages.entry(page).or_default().push(block); - pages - }, - ); - - if blocks_by_page.is_empty() { - let markdown = chunks(result) - .iter() - .filter_map(Value::as_object) - .filter_map(|chunk| chunk.get("content").and_then(Value::as_str)) - .filter(|content| !content.is_empty()) - .collect::>() - .join("\n\n"); - return if markdown.is_empty() { - Vec::new() - } else { - vec![json!({"index": 0, "markdown": markdown})] - }; - } - - blocks_by_page - .into_iter() - .map(|(page, blocks)| { - let markdown = blocks - .iter() - .filter_map(Value::as_object) - .filter_map(|block| block.get("content").and_then(Value::as_str)) - .filter(|content| !content.is_empty()) - .collect::>() - .join("\n\n"); - json!({ - "index": page.saturating_sub(1).max(0), - "markdown": markdown, - "blocks": blocks, - }) - }) - .collect() -} - -pub fn transform_reducto_response( - model: &str, - response_json: Value, -) -> Result { - let response = response_json - .as_object() - .ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(&response_json), - })?; - let empty_result = Map::new(); - let result = match response.get("result") { - Some(Value::Object(result)) => result, - Some(Value::Null) => &empty_result, - Some(_) => { - return Err(Error::InvalidResponse( - "Reducto result must be an object".to_string(), - )); - } - None => response, - }; - let usage = response - .get("usage") - .and_then(Value::as_object) - .cloned() - .unwrap_or_default(); - let usage_info = Some(json!({ - "pages_processed": usage.get("num_pages").cloned().unwrap_or(Value::Null), - "credits": usage.get("credits").cloned().unwrap_or(Value::Null), - })); - - Ok(OcrResponseData { - pages: build_pages(result), - model: model.to_string(), - document_annotation: None, - usage_info, - object: "ocr".to_string(), - extra_fields: Map::new(), - provider_native_response: Some(response_json), - }) -} - -impl OcrProviderConfig for ReductoParseV3Config { - fn supported_ocr_params(&self) -> &'static [&'static str] { - PARSE_V3_SUPPORTED_OCR_PARAMS - } - - fn transform_ocr_request( - &self, - _model: &str, - document: Value, - optional_params: Map, - ) -> Result { - let file_id = source_file_id(&document)?; - Ok(build_parse_v3_request(&file_id, optional_params)) - } - - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result { - transform_reducto_response(model, response_json) - } - - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - _optional_params: &Map, - _env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(parse_url(api_base)) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_api_key(api_key, env_lookup) - } -} - -impl OcrProviderConfig for ReductoParseLegacyConfig { - fn supported_ocr_params(&self) -> &'static [&'static str] { - PARSE_LEGACY_SUPPORTED_OCR_PARAMS - } - - fn transform_ocr_request( - &self, - _model: &str, - document: Value, - optional_params: Map, - ) -> Result { - let file_id = source_file_id(&document)?; - Ok(build_parse_legacy_request(&file_id, &optional_params)) - } - - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result { - transform_reducto_response(model, response_json) - } - - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - _optional_params: &Map, - _env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(parse_url(api_base)) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_api_key(api_key, env_lookup) - } -} diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs b/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs deleted file mode 100644 index 3621ff6a2fd..00000000000 --- a/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs deleted file mode 100644 index f239b6921fa..00000000000 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs deleted file mode 100644 index c324de8cb45..00000000000 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ /dev/null @@ -1,467 +0,0 @@ -use crate::error::{Error, json_type_name}; -use crate::ocr::transformation::OcrProviderConfig; -use crate::ocr::types::{OcrRequestData, OcrResponseData}; -use serde_json::{Map, Value, json}; - -use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; - -const VERTEX_DEFAULT_LOCATION: &str = "us-central1"; -const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com"; -const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; -const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; -const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; -const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; -const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; - -#[rustfmt::skip] -const DEEPSEEK_SUPPORTED_OCR_PARAMS: &[&str] = &[ - "stream", - "temperature", - "max_tokens", - "top_p", - "n", - "stop", -]; - -pub struct VertexAiOcrConfig; -pub struct VertexAiDeepSeekOcrConfig; - -pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig; -pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig; - -fn string_param<'a>(params: &'a Map, keys: &[&str]) -> Option<&'a str> { - keys.iter() - .find_map(|key| params.get(*key).and_then(Value::as_str)) - .map(str::trim) - .filter(|value| !value.is_empty()) -} - -pub fn is_deepseek_model(model: &str) -> bool { - model.to_ascii_lowercase().contains("deepseek") -} - -pub fn resolve_vertex_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - api_key - .map(str::trim) - .filter(|key| !key.is_empty()) - .map(str::to_string) - .or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .ok_or_else(|| { - Error::Auth( - "Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers" - .to_string(), - ) - }) -} - -fn vertex_project( - params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - string_param(params, &["vertex_project", "vertex_ai_project"]) - .map(str::to_string) - .or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - Error::InvalidRequest( - "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" - .to_string(), - ) - }) -} - -fn vertex_location( - params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> String { - string_param(params, &["vertex_location", "vertex_ai_location"]) - .map(str::to_string) - .or_else(|| env_lookup(VERTEXAI_LOCATION_ENV).filter(|value| !value.trim().is_empty())) - .or_else(|| env_lookup(VERTEX_LOCATION_ENV).filter(|value| !value.trim().is_empty())) - .unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string()) -} - -fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String { - api_base - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com")) - .trim_end_matches('/') - .to_string() -} - -pub fn complete_vertex_mistral_url( - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - let project = vertex_project(optional_params, env_lookup)?; - let location = vertex_location(optional_params, env_lookup); - let base = vertex_mistral_api_base(api_base, &location); - Ok(format!( - "{base}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict" - )) -} - -pub fn complete_vertex_deepseek_url( - api_base: Option<&str>, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - let project = vertex_project(optional_params, env_lookup)?; - let location = vertex_location(optional_params, env_lookup); - let base = api_base - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or(VERTEX_DEFAULT_DEEPSEEK_API_BASE) - .trim_end_matches('/'); - Ok(format!( - "{base}/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions" - )) -} - -fn document_content_item(document: &Value) -> Result { - let object = document.as_object().ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(document), - })?; - let doc_type = object - .get("type") - .and_then(Value::as_str) - .ok_or(Error::MissingField("document.type"))?; - let url_field = match doc_type { - "image_url" => "image_url", - "document_url" => "document_url", - other => { - return Err(Error::InvalidRequest(format!( - "Unsupported document type: {other}. Expected 'image_url' or 'document_url'" - ))); - } - }; - let url = object - .get(url_field) - .and_then(Value::as_str) - .filter(|value| !value.is_empty()) - .ok_or(Error::MissingField(url_field))?; - - Ok(json!({ - "type": "image_url", - "image_url": url, - })) -} - -fn deepseek_model_name(model: &str) -> String { - if model.starts_with("deepseek-ai/") { - model.to_string() - } else { - format!("deepseek-ai/{model}") - } -} - -fn first_choice_content(response: &Value) -> Result { - response - .get("choices") - .and_then(Value::as_array) - .and_then(|choices| choices.first()) - .and_then(|choice| choice.get("message")) - .and_then(|message| message.get("content")) - .cloned() - .filter(|content| match content { - Value::String(value) => !value.is_empty(), - Value::Object(_) => true, - _ => false, - }) - .ok_or_else(|| Error::InvalidResponse("No content in DeepSeek OCR response".to_string())) -} - -fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> Value { - match content { - Value::String(content) => { - if content.trim_start().starts_with('{') { - serde_json::from_str(&content).unwrap_or_else(|_| { - json!({ - "pages": [{"index": 0, "markdown": content}], - "model": model, - "usage_info": usage.unwrap_or_else(|| json!({})), - }) - }) - } else { - json!({ - "pages": [{"index": 0, "markdown": content}], - "model": model, - "usage_info": usage.unwrap_or_else(|| json!({})), - }) - } - } - Value::Object(_) => content, - other => json!({ - "pages": [{"index": 0, "markdown": other.to_string()}], - "model": model, - "usage_info": usage.unwrap_or_else(|| json!({})), - }), - } -} - -impl OcrProviderConfig for VertexAiOcrConfig { - fn supported_ocr_params(&self) -> &'static [&'static str] { - MISTRAL_OCR_CONFIG.supported_ocr_params() - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn transform_ocr_request( - &self, - model: &str, - document: Value, - optional_params: Map, - ) -> Result { - MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) - } - - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result { - MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn complete_url( - &self, - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - complete_vertex_mistral_url(api_base, model, optional_params, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_vertex_api_key(api_key, env_lookup) - } - - fn requires_data_uri_document(&self) -> bool { - true - } -} - -impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn supported_ocr_params(&self) -> &'static [&'static str] { - DEEPSEEK_SUPPORTED_OCR_PARAMS - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn map_ocr_params(&self, non_default_params: &Map) -> Map { - non_default_params - .iter() - .filter(|(name, _)| DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&name.as_str())) - .map(|(name, value)| (name.clone(), value.clone())) - .collect() - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn transform_ocr_request( - &self, - model: &str, - document: Value, - optional_params: Map, - ) -> Result { - let mut data = Map::new(); - data.insert( - "model".to_string(), - Value::String(deepseek_model_name(model)), - ); - data.insert( - "messages".to_string(), - json!([{"role": "user", "content": [document_content_item(&document)?]}]), - ); - for (key, value) in optional_params { - if DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&key.as_str()) { - data.insert(key, value); - } - } - Ok(OcrRequestData { - data: Value::Object(data), - files: None, - }) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result { - let response = response_json - .as_object() - .ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(&response_json), - })?; - let usage = response.get("usage").cloned(); - let content = first_choice_content(&response_json)?; - let mut ocr_data = ocr_data_from_content(content.clone(), usage.clone(), model); - - if !ocr_data.get("pages").is_some_and(Value::is_array) { - ocr_data = json!({ - "pages": [{ - "index": 0, - "markdown": match content { - Value::String(value) => value, - other => other.to_string(), - } - }], - "model": ocr_data.get("model").and_then(Value::as_str).unwrap_or(model), - "usage_info": ocr_data.get("usage_info").cloned().or(usage).unwrap_or_else(|| json!({})), - }); - } - - let object = ocr_data.as_object().ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(&ocr_data), - })?; - let pages = object - .get("pages") - .and_then(Value::as_array) - .cloned() - .unwrap_or_default(); - let usage_info = object - .get("usage_info") - .cloned() - .or_else(|| response.get("usage").cloned()); - Ok(OcrResponseData { - pages, - model: object - .get("model") - .and_then(Value::as_str) - .unwrap_or(model) - .to_string(), - document_annotation: object.get("document_annotation").cloned(), - usage_info, - object: "ocr".to_string(), - extra_fields: Map::new(), - provider_native_response: None, - }) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - complete_vertex_deepseek_url(api_base, optional_params, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_vertex_api_key(api_key, env_lookup) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use rstest::rstest; - - #[test] - fn vertex_mistral_url_uses_project_location_and_model() { - let params = Map::from_iter([ - ("vertex_project".to_string(), json!("proj-1")), - ("vertex_location".to_string(), json!("europe-west4")), - ]); - - let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", ¶ms, &|_| None) - .expect("url builds"); - - assert_eq!( - url, - "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - } - - #[test] - fn vertex_mistral_reuses_mistral_body_transform() { - let body = VERTEX_AI_OCR_CONFIG - .transform_ocr_request( - "mistral-ocr-maas", - json!({"type": "image_url", "image_url": "data:image/png;base64,abc"}), - Map::new(), - ) - .expect("request transforms") - .data; - - assert_eq!(body["model"], "mistral-ocr-maas"); - assert_eq!(body["document"]["image_url"], "data:image/png;base64,abc"); - } - - #[test] - fn vertex_deepseek_request_uses_ocr_endpoint_shape() { - let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG - .transform_ocr_request( - "deepseek-ocr-maas", - json!({"type": "document_url", "document_url": "gs://bucket/doc.pdf"}), - Map::from_iter([("temperature".to_string(), json!(0.1))]), - ) - .expect("request transforms") - .data; - - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type": "image_url", "image_url": "gs://bucket/doc.pdf"}) - ); - } - - #[rstest] - #[case::bare_model("deepseek-ocr-maas")] - #[case::namespaced_model("deepseek-ai/deepseek-ocr-maas")] - fn vertex_deepseek_request_uses_single_provider_namespace(#[case] model: &str) { - let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG - .transform_ocr_request( - model, - json!({"type": "image_url", "image_url": "data:image/png;base64,AA=="}), - Map::new(), - ) - .expect("request transforms") - .data; - - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - } - - #[test] - fn vertex_deepseek_response_wraps_markdown_content() { - let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG - .transform_ocr_response( - "deepseek-ocr-maas", - json!({ - "choices": [{"message": {"content": "# OCR text"}}], - "usage": {"prompt_tokens": 1} - }), - ) - .expect("response transforms"); - - assert_eq!( - response.pages, - vec![json!({"index": 0, "markdown": "# OCR text"})] - ); - assert_eq!(response.model, "deepseek-ocr-maas"); - assert_eq!(response.usage_info, Some(json!({"prompt_tokens": 1}))); - } -} diff --git a/litellm-rust/crates/core/src/url_utils.rs b/litellm-rust/crates/core/src/url_utils.rs new file mode 100644 index 00000000000..1150f93a5c7 --- /dev/null +++ b/litellm-rust/crates/core/src/url_utils.rs @@ -0,0 +1,118 @@ +use std::marker::PhantomData; + +use thiserror::Error; +use url::Url; + +#[derive(Debug, Error)] +pub(crate) enum ApiUrlError { + #[error("invalid URL: {0}")] + Parse(#[from] url::ParseError), + #[error("URL cannot be used as a base")] + CannotBeBase, +} + +pub(crate) struct Base; +pub(crate) struct Complete; + +pub(crate) struct ApiUrl { + url: Url, + state: PhantomData, +} + +impl ApiUrl { + pub(crate) fn parse(value: &str) -> Result { + Ok(Self { + url: Url::parse(value.trim())?, + state: PhantomData, + }) + } + + pub(crate) fn complete_path( + mut self, + target: &[&str], + ) -> Result, ApiUrlError> { + let existing: Vec = self + .url + .path_segments() + .ok_or(ApiUrlError::CannotBeBase)? + .filter(|segment| !segment.is_empty()) + .map(str::to_string) + .collect(); + let overlap = (0..=existing.len().min(target.len())) + .rev() + .find(|&length| { + existing[existing.len() - length..] + .iter() + .map(String::as_str) + .eq(target[..length].iter().copied()) + }) + .unwrap_or(0); + self.url + .path_segments_mut() + .map_err(|()| ApiUrlError::CannotBeBase)? + .pop_if_empty() + .extend(target[overlap..].iter().copied()); + Ok(ApiUrl { + url: self.url, + state: PhantomData, + }) + } +} + +impl ApiUrl { + pub(crate) fn append_query_pairs<'a>( + mut self, + pairs: impl IntoIterator, + ) -> Self { + self.url.query_pairs_mut().extend_pairs(pairs); + self + } + + pub(crate) fn into_string(self) -> String { + self.url.into() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn completion_appends_only_the_missing_path_suffix() { + for (base, expected) in [ + ("https://example.test", "https://example.test/v1/ocr"), + ("https://example.test/v1", "https://example.test/v1/ocr"), + ("https://example.test/v1/ocr", "https://example.test/v1/ocr"), + ] { + let actual = ApiUrl::parse(base) + .and_then(|url| url.complete_path(&["v1", "ocr"])) + .map(|url| url.into_string()) + .expect("url builds"); + assert_eq!(actual, expected); + } + } + + #[test] + fn completion_places_paths_before_queries() { + let actual = ApiUrl::parse("https://example.test/v1?tenant=a") + .and_then(|url| url.complete_path(&["v1", "ocr"])) + .map(|url| url.into_string()) + .expect("url builds"); + assert_eq!(actual, "https://example.test/v1/ocr?tenant=a"); + } + + #[test] + fn appended_query_pairs_are_encoded() { + let actual = ApiUrl::parse("https://example.test") + .and_then(|url| url.complete_path(&["analyze"])) + .map(|url| { + url.append_query_pairs([("model", "name with spaces")]) + .into_string() + }) + .expect("url builds"); + assert_eq!( + actual, + "https://example.test/analyze?model=name+with+spaces" + ); + } +} diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs new file mode 100644 index 00000000000..d7d532cfef1 --- /dev/null +++ b/litellm-rust/crates/core/tests/azure_ai_ocr.rs @@ -0,0 +1,97 @@ +use std::sync::Arc; + +use serde_json::{Value, json}; + +use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks}; +use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + +#[tokio::test] +async fn facade_executes_azure_mistral_with_prepared_auth() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"include_image_base64":true}), + ); + request.connection.api_key = None; + request.connection.extra_headers = vec![( + "Authorization".into(), + "Bearer python-prepared-token".into(), + )]; + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0]["markdown"], "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer python-prepared-token\r\n") + ); + let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "include_image_base64":true + }) + ); +} + +#[tokio::test] +async fn facade_acquires_supplied_entra_token_for_final_request() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"azure_ad_token":"rust-owned-token"}), + ); + request.connection.api_key = None; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer rust-owned-token\r\n") + ); +} + +struct ReplaceBodyDocument; + +impl OcrHooks for ReplaceBodyDocument { + fn has_guardrails(&self) -> bool { + true + } + + fn during_call( + &self, + mut request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { + request.body["document"] = json!({ + "type":"document_url", + "document_url":"https://example.com/not-inline.pdf" + }); + Ok(request) + }) + } +} + +#[tokio::test] +async fn rejects_non_inline_body_after_guardrails() { + let mut request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); + request.hooks = Arc::new(ReplaceBodyDocument); + let error = perform_ocr(request).await.unwrap_err(); + assert!(error.to_string().contains("data URI")); +} diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs new file mode 100644 index 00000000000..e4c81dea5a7 --- /dev/null +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -0,0 +1,395 @@ +use serde_json::{Value, json}; + +use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; +use super::wire::{OcrWireRequest, decode_request}; + +fn query_value(url: &str, key: &str) -> Option { + url::Url::parse(url) + .unwrap() + .query_pairs() + .find_map(|(name, value)| (name == key).then(|| value.into_owned())) +} + +#[tokio::test] +async fn facade_maps_pages_features_and_url_document() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[]} + }))]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), + ); + request.document = serde_json::from_value(json!({ + "type":"document_url", + "document_url":"https://example.com/document.pdf" + })) + .unwrap(); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let target = request.split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); + assert_eq!( + query_value(&url, "features").as_deref(), + Some("keyValuePairs,languages") + ); + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({"urlSource":"https://example.com/document.pdf"}) + ); +} + +#[tokio::test] +async fn rejects_invalid_pages_features_and_format() { + for options in [ + json!({"pages":[true]}), + json!({"pages":[1,"2"]}), + json!({"pages":[-1]}), + json!({"pages":"1&&features=bad"}), + json!({"features":"languages&pages=1"}), + json!({"req_format":"azure"}), + ] { + let result = decode_request(OcrWireRequest { + model: "azure_ai/doc-intelligence/prebuilt-read".into(), + document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + api_key: Some("key".into()), + api_base: Some("http://127.0.0.1:1".into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: None, + }); + let rejected = match result { + Ok(request) => perform_ocr(request).await.is_err(), + Err(_) => true, + }; + assert!(rejected, "accepted {options}"); + } +} + +#[tokio::test] +async fn inline_document_decodes_to_base64_source() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); +} + +#[tokio::test] +async fn immediate_response_normalizes_pages_and_preserves_native() { + let operation = json!({ + "status":"succeeded", + "operationExtension":42, + "analyzeResult":{ + "content":"A\n\nB", + "tables":[{"cells":[]}], + "keyValuePairs":[{"key":{"content":"A"}}], + "pages":[{ + "pageNumber":"2", + "width":"8.5", + "height":11, + "unit":"inch", + "lines":[{"content":"A"},{"content":null},{"content":"B"}] + }] + } + }); + let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; + let result = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(result.pages[0]["index"], 1); + assert_eq!(result.pages[0]["markdown"], "A\n\nB"); + assert_eq!( + result.pages[0]["dimensions"], + json!({"width":816,"height":1056,"dpi":96}) + ); + assert_eq!(result.usage_info, Some(json!({"pages_processed":1}))); + assert_eq!(result.provider_native_response, Some(operation)); +} + +#[tokio::test] +async fn accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "0".into())], + body: json!({"status":"running"}), + }, + MockResponse::json(operation.clone()), + ]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + ); + request + .connection + .extra_headers + .push(("X-Trace".into(), "initial-only".into())); + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(result.provider_native_response, Some(operation)); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); + for poll in &requests[1..] { + assert!(!poll.to_ascii_lowercase().contains("x-trace:")); + assert!( + poll.to_ascii_lowercase() + .contains("ocp-apim-subscription-key: test-key") + ); + } +} + +#[tokio::test] +async fn polling_forwards_bearer_credentials() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + request.connection.api_key = None; + request.connection.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert!( + requests[1] + .to_ascii_lowercase() + .contains("authorization: bearer token") + ); +} + +#[tokio::test] +async fn polling_does_not_follow_redirects() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 302, + headers: vec![("Location", "{base}/redirected".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + + assert!(error.to_string().contains("status 302"), "{error}"); + assert_eq!(seen.lock().unwrap().len(), 2); + server.abort(); +} + +#[tokio::test] +async fn polling_rejects_terminal_failure() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"failed"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("status failed")); +} + +#[tokio::test] +async fn malformed_provider_pages_report_response_paths() { + for (analysis, path) in [ + (json!({"pages":null}), "pages"), + (json!({"pages":[null]}), "pages[0]"), + (json!({"pages":[{"lines":null}]}), "lines"), + (json!({"pages":[{"width":"bad"}]}), "width"), + ] { + let (base, _, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":analysis + }))]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains(path), "{error}"); + } +} + +#[tokio::test] +async fn rejects_missing_invalid_and_cross_origin_operation_locations() { + for headers in [ + Vec::new(), + vec![("Operation-Location", "/relative".into())], + vec![("Operation-Location", "http://example.com/operation".into())], + vec![( + "Operation-Location", + "http://user:password@127.0.0.1/operation".into(), + )], + ] { + let (base, _, server) = mock_server(vec![MockResponse { + status: 202, + headers, + body: json!({}), + }]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("operation-location")); + } +} + +#[tokio::test] +async fn polling_deadline_bounds_retry_delay() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "9999".into())], + body: json!({"status":"notStarted"}), + }, + ]) + .await; + let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + request.connection.poll_timeout = std::time::Duration::from_millis(100); + + let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request)) + .await + .unwrap() + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("timed out")); +} + +#[tokio::test] +async fn model_id_is_encoded_and_dot_segments_are_rejected() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + perform_ocr(wire_request( + "azure_ai/doc-intelligence/a ?#é", + &base, + json!({}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); + + for model in [ + "azure_ai/doc-intelligence/.", + "azure_ai/doc-intelligence/..", + ] { + let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) + .await + .unwrap_err(); + assert!(error.to_string().contains("dot segment")); + } +} + +#[tokio::test] +async fn pre_call_guardrail_receives_caller_pages_before_mapping() { + use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest}; + use std::sync::Arc; + + struct RewritePages; + impl OcrHooks for RewritePages { + fn has_guardrails(&self) -> bool { + true + } + + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + Box::pin(async move { + assert_eq!(request.optional_params["pages"], json!([0, 2])); + Ok(OcrPreCallRequest { + optional_params: json!({"pages": [1]}), + ..request + }) + }) + } + } + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"status": "succeeded"}))]).await; + let request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages": [0, 2]}), + ) + .with_host_hooks(Arc::new(RewritePages), None); + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + let target = requests[0].split_whitespace().nth(1).unwrap(); + assert_eq!( + query_value(&format!("{base}{target}"), "pages").as_deref(), + Some("2") + ); + assert_eq!(requests.len(), 1); +} diff --git a/litellm-rust/crates/core/tests/deepseek_ocr.rs b/litellm-rust/crates/core/tests/deepseek_ocr.rs new file mode 100644 index 00000000000..875fc9e3dc6 --- /dev/null +++ b/litellm-rust/crates/core/tests/deepseek_ocr.rs @@ -0,0 +1,95 @@ +use rstest::rstest; +use serde_json::{Value, json}; + +use crate::ocr::codecs::deepseek::{ + DeepSeekOcrParams, DeepSeekOcrResponse, transform_ocr_request, transform_ocr_response, +}; +use crate::ocr::types::OcrDocument; + +fn document() -> OcrDocument { + serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap() +} + +#[rstest] +#[case("stream", json!(true))] +#[case("temperature", json!(0.1))] +#[case("max_tokens", json!(1024))] +#[case("top_p", json!(0.9))] +#[case("n", json!(2))] +#[case("stop", json!("done"))] +#[case("stop", json!(["done", "stop"]))] +fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { + let params: DeepSeekOcrParams = + serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); + let result = serde_json::to_value( + transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms).unwrap(), + ) + .unwrap(); + assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/a.png"}) + ); + assert_eq!(result[name], value); + assert!(result.get("ignored").is_none()); +} + +#[rstest] +#[case(json!("# hello"), "# hello")] +#[case(json!("{broken"), "{broken")] +#[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] +#[case(json!({"pages":[]}), "{\"pages\":[]}")] +#[case(json!({}), "{}")] +#[case(json!("[]"), "[]")] +#[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] +#[case(json!({"pages":[{"markdown":"object"}]}), "object")] +fn response_codec_handles_text_json_and_objects(#[case] content: Value, #[case] expected: &str) { + let response: DeepSeekOcrResponse = serde_json::from_value( + json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), + ) + .unwrap(); + let result = transform_ocr_response("model", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["markdown"], expected); + assert_eq!(result["pages"][0]["index"], 0); + assert_eq!(result["usage_info"]["prompt_tokens"], 1); +} + +#[test] +fn structured_result_maps_pages_usage_model_and_annotation() { + let response: DeepSeekOcrResponse = serde_json::from_value(json!({ + "choices":[{"message":{"content":{ + "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], + "model":"provider-model", + "usage_info":{"pages_processed":1}, + "document_annotation":{"language":"en"}, + "future":"kept" + }}}] + })) + .unwrap(); + let result = transform_ocr_response("requested", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["index"], 2); + assert_eq!(result["pages"][0]["images"][0]["id"], "one"); + assert_eq!(result["model"], "provider-model"); + assert_eq!(result["usage_info"]["pages_processed"], 1); + assert_eq!(result["document_annotation"]["language"], "en"); + assert_eq!(result["future"], "kept"); +} + +#[test] +fn response_codec_rejects_missing_empty_and_malformed_content() { + for value in [ + json!({"choices":[]}), + json!({"choices":[{"message":{"content":""}}]}), + json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), + json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), + ] { + let result = serde_json::from_value::(value) + .map_err(|_| ()) + .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); + assert!(result.is_err()); + } +} diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs new file mode 100644 index 00000000000..cecd8869741 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -0,0 +1,229 @@ +use std::sync::{Arc, Mutex}; + +use serde_json::{Value, json}; + +use super::OcrClient; +use super::hooks::{OcrHookFuture, OcrHooks, OcrLogFuture, OcrPreCallRequest}; +use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; +use super::wire::{OcrWireRequest, decode_request}; +use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; + +#[test] +fn request_boundary_selects_mistral_and_rejects_unknown_providers() { + let request = OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some("key".into()), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: json!({"extract_header":true,"unknown":42}) + .as_object() + .unwrap() + .clone(), + input_sources: Default::default(), + timeout_seconds: None, + }; + assert!(decode_request(request).is_ok()); + assert!( + decode_request(OcrWireRequest { + model: "model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some("key".into()), + api_base: None, + custom_llm_provider: Some("unknown".into()), + extra_headers: None, + optional_params: serde_json::Map::new(), + input_sources: Default::default(), + timeout_seconds: None, + }) + .is_err() + ); +} + +#[tokio::test] +async fn facade_executes_direct_mistral_once() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let result = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"extract_header":true,"unknown":"ignored"}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0]["markdown"], "hello"); + assert_eq!(result.pages[0]["custom"], "preserved"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /v1/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "extract_header":true + }) + ); +} + +#[tokio::test] +async fn facade_retains_native_response_when_requested() { + let provider_response = json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1}, + "provider_only":"preserved" + }); + let (base, _, server) = mock_server(vec![MockResponse::json(provider_response.clone())]).await; + let response = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + + server.await.unwrap(); + assert_eq!(response.provider_native_response, Some(provider_response)); +} + +#[tokio::test] +async fn facade_uses_the_injected_http_client() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut default_headers = reqwest::header::HeaderMap::new(); + default_headers.insert( + "x-transport-owner", + reqwest::header::HeaderValue::from_static("host"), + ); + let provider_http = reqwest::Client::builder() + .default_headers(default_headers) + .build() + .unwrap(); + OcrClient::new(provider_http) + .unwrap() + .perform(wire_request("mistral/model", &base, json!({}))) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host")); +} + +struct RecordingHooks { + events: Arc>>, + block: bool, +} + +impl OcrHooks for RecordingHooks { + fn has_guardrails(&self) -> bool { + true + } + + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + Box::pin(async move { + self.events.lock().unwrap().push("pre"); + if self.block { + return Err(crate::Error::InvalidRequest("blocked".into())); + } + Ok(request) + }) + } + + fn during_call( + &self, + request: super::hooks::OcrDuringCallRequest, + ) -> OcrHookFuture<'_, super::hooks::OcrDuringCallRequest> { + Box::pin(async move { + self.events.lock().unwrap().push("during"); + Ok(request) + }) + } + + fn success<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a super::LiteLLMOcrResponse, + _timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("success"); + }) + } + + fn failure<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a crate::Error, + _timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("failure"); + }) + } +} + +#[tokio::test] +async fn lifecycle_orders_hooks_and_emits_one_success() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let events = Arc::new(Mutex::new(Vec::new())); + let request = wire_request("mistral/model", &base, json!({})); + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(RecordingHooks { + events: events.clone(), + block: false, + }), + ..request + }; + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(*events.lock().unwrap(), ["pre", "during", "success"]); + assert_eq!(seen.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { + let events = Arc::new(Mutex::new(Vec::new())); + let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(RecordingHooks { + events: events.clone(), + block: true, + }), + ..request + }; + let error = perform_ocr(request).await.unwrap_err(); + assert!(matches!(error, crate::Error::InvalidRequest(_))); + assert_eq!(*events.lock().unwrap(), ["pre", "failure"]); +} + +#[tokio::test] +async fn upstream_failure_emits_one_terminal_failure() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 500, + headers: vec![], + body: json!({"error":"failed"}), + }]) + .await; + let events = Arc::new(Mutex::new(Vec::new())); + let request = wire_request("mistral/model", &base, json!({})); + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(RecordingHooks { + events: events.clone(), + block: false, + }), + ..request + }; + assert!(perform_ocr(request).await.is_err()); + server.await.unwrap(); + assert_eq!(*events.lock().unwrap(), ["pre", "during", "failure"]); + assert_eq!(seen.lock().unwrap().len(), 1); +} diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs new file mode 100644 index 00000000000..a2e67dffc7d --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -0,0 +1,111 @@ +use std::sync::{Arc, Mutex}; + +use serde_json::{Value, json}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; + +use crate::ocr::wire::{OcrWireRequest, decode_request}; +use crate::ocr::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient}; + +pub(crate) fn ocr_client() -> OcrClient { + let document_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test document client builds"); + OcrClient::for_test(reqwest::Client::new(), document_http) +} + +pub(crate) async fn perform_ocr( + request: LiteLLMOcrRequest, +) -> Result { + ocr_client().perform(request).await +} + +pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { + decode_request(OcrWireRequest { + model: model.into(), + document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + api_key: Some("test-key".into()), + api_base: Some(base.into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap() +} + +pub(crate) struct MockResponse { + pub status: u16, + pub headers: Vec<(&'static str, String)>, + pub body: Value, +} + +impl MockResponse { + pub fn json(body: Value) -> Self { + Self { + status: 200, + headers: vec![], + body, + } + } +} + +pub(crate) async fn mock_server( + responses: Vec, +) -> (String, Arc>>, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(Mutex::new(Vec::new())); + let seen = requests.clone(); + let server_base = base.clone(); + let task = tokio::spawn(async move { + for response in responses { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + let mut buffer = [0u8; 4096]; + let header_end = loop { + let n = socket.read(&mut buffer).await.unwrap(); + assert!(n > 0); + bytes.extend_from_slice(&buffer[..n]); + if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { + break index + 4; + } + }; + let length = String::from_utf8_lossy(&bytes[..header_end]) + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + .unwrap_or(0); + while bytes.len() < header_end + length { + let n = socket.read(&mut buffer).await.unwrap(); + assert!(n > 0); + bytes.extend_from_slice(&buffer[..n]); + } + seen.lock() + .unwrap() + .push(String::from_utf8_lossy(&bytes).into_owned()); + let body = serde_json::to_vec(&response.body).unwrap(); + let headers = response + .headers + .into_iter() + .map(|(name, value)| { + format!("{name}: {}\r\n", value.replace("{base}", &server_base)) + }) + .collect::(); + let head = format!( + "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n", + response.status, + body.len(), + headers + ); + socket.write_all(head.as_bytes()).await.unwrap(); + socket.write_all(&body).await.unwrap(); + } + }); + (base, requests, task) +} diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs new file mode 100644 index 00000000000..8e86e4713ef --- /dev/null +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -0,0 +1,232 @@ +use std::sync::Arc; + +use rstest::rstest; +use serde_json::{Value, json}; + +use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks}; +use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + +fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() +} + +#[rstest] +#[case( + "reducto/parse-v3", + json!({ + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://already.pdf", + json!({ + "input":"reducto://already.pdf", + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "provider_option":"value" + }) +)] +#[case( + "reducto/parse-legacy", + json!({ + "enhance":{"agentic":[{"type":"table"}]}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url":"reducto://legacy.pdf", + "options":{"enhance":{"agentic":[{"type":"table"}]}}, + "future_ocr_option":true, + "provider_option":"value" + }) +)] +#[tokio::test] +async fn request_mapping_matches_python( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, +) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[]} + }))]) + .await; + let mut request = wire_request(model, &base, options); + request.document = request.document.with_source(source.into()); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!(request_body(&requests[0]), expected); +} + +#[rstest] +#[case("parse-v3")] +#[case("parse-legacy")] +#[tokio::test] +async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), + ]) + .await; + let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); + request.connection.extra_headers = vec![ + ("Content-Type".into(), "application/json".into()), + ("X-Trace".into(), "upload-test".into()), + ]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0]["markdown"], "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("content-type: multipart/form-data; boundary=") + ); + assert!(requests[0].contains("x-trace: upload-test")); + assert!(requests[0].contains("application/pdf")); + assert!(requests[0].contains("abc")); + assert!(requests[1].starts_with("POST /parse ")); +} + +#[rstest] +#[case(json!({"file_id":""}))] +#[case(json!({}))] +#[case(json!({"file_id":null}))] +#[tokio::test] +async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { + let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; + let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("file_id")); + assert_eq!(seen.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn upload_failure_stops_before_parse() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 503, + headers: vec![], + body: json!({"error":"unavailable"}), + }]) + .await; + assert!( + perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .is_err() + ); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 1); +} + +#[rstest] +#[case("https://example.com/a.pdf")] +#[case("reducto://")] +#[case("data:application/pdf;base64")] +#[case("data:application/pdf;base64,INVALID!")] +#[tokio::test] +async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { + let mut request = wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})); + request.document = request.document.with_source(source.into()); + assert!(perform_ocr(request).await.is_err()); +} + +#[test] +fn response_normalization_groups_blocks_and_distinguishes_null_result() { + use crate::ocr::codecs::reducto::{ReductoResponse, transform_ocr_response}; + + let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"chunks":[ + {"blocks":[{"content":"B","bbox":{"page":2},"kind":"table"}]}, + {"blocks":[{"content":"A","bbox":{"page":1},"kind":"text"},{"content":"C","bbox":{"page":1}}]} + ]}}); + let response: ReductoResponse = serde_json::from_value(raw).unwrap(); + let normalized = transform_ocr_response("parse-v3", response) + .unwrap() + .into_json(); + assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); + assert_eq!(normalized["pages"][1]["markdown"], "B"); + assert_eq!(normalized["pages"][1]["blocks"][0]["kind"], "table"); + assert_eq!(normalized["usage_info"]["pages_processed"], 2); + assert_eq!(normalized["usage_info"]["credits"], 3.0); + + let missing: ReductoResponse = + serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); + let missing = transform_ocr_response("parse-v3", missing).unwrap(); + assert_eq!(missing.pages[0]["markdown"], "text"); + let null: ReductoResponse = serde_json::from_value( + json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), + ) + .unwrap(); + let null = transform_ocr_response("parse-v3", null).unwrap(); + assert!(null.pages.is_empty()); +} + +#[tokio::test] +async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { + let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); + let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; + let mut request = wire_request("reducto/parse-v3", &base, json!({})); + request.document = request.document.with_source("reducto://ready.pdf".into()); + request.connection.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.provider_native_response, None); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer existing") + ); +} + +struct RewriteDocument; + +impl OcrHooks for RewriteDocument { + fn has_guardrails(&self) -> bool { + true + } + + fn during_call( + &self, + request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { + assert_eq!( + request.body["document_url"], + "data:application/pdf;base64,YWJj" + ); + Ok(OcrDuringCallRequest { + body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), + ..request + }) + }) + } +} + +#[tokio::test] +async fn guardrail_rewrites_document_before_upload() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; + let mut request = wire_request("reducto/parse-v3", &base, json!({})); + request.hooks = Arc::new(RewriteDocument); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert!(requests[0].contains("reducto://guarded.pdf")); +} diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs new file mode 100644 index 00000000000..6d3061d8f5d --- /dev/null +++ b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs @@ -0,0 +1,83 @@ +use serde_json::{Value, json}; + +use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; +use crate::auth::InputSource; + +fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() +} + +#[tokio::test] +async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "choices":[{"message":{"content":"recognized"}}], + "usage":{"prompt_tokens":1} + }))]) + .await; + let mut request = wire_request( + "vertex_ai/deepseek-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "temperature":0.1, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + ); + request.document = request + .document + .with_source("gs://bucket/document.pdf".into()); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0]["markdown"], "recognized"); + assert_eq!(response.usage_info.unwrap()["prompt_tokens"], 1); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + let body = request_body(&requests[0]); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert!(body.get("future_ocr_option").is_none()); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type":"document_url","document_url":"gs://bucket/document.pdf"}) + ); +} + +#[test] +fn host_registration_selects_deepseek_without_affecting_mistral() { + assert!(crate::ocr::wire::is_supported_request( + "deepseek-ocr-maas", + Some("vertex_ai") + )); + assert!(crate::ocr::wire::is_supported_request( + "mistral-ocr-maas", + Some("vertex_ai") + )); +} + +#[tokio::test] +async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let mut request = wire_request( + "vertex_ai/deepseek-ocr-maas", + "https://caller.example", + json!({"vertex_project":"project-1"}), + ); + request.connection.api_base_source = InputSource::Request; + + let error = perform_ocr(request).await.unwrap_err(); + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); +} diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs new file mode 100644 index 00000000000..96a19dd62b4 --- /dev/null +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -0,0 +1,161 @@ +use serde_json::{Value, json}; + +use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; +use crate::auth::InputSource; + +fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() +} + +#[tokio::test] +async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/mistral-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "extract_footer":true + }), + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0]["markdown"], "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + assert_eq!( + request_body(&requests[0]), + json!({ + "model":"mistral-ocr-maas", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "extract_footer":true + }) + ); +} + +#[tokio::test] +async fn supplied_authorization_is_forwarded_without_a_static_token() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "vertex_ai/model", + &base, + json!({"vertex_project":"project-1"}), + ); + request.connection.api_key = None; + request.connection.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer supplied") + ); +} + +#[tokio::test] +async fn invalid_credentials_fail_before_provider_http() { + let request = wire_request( + "vertex_ai/model", + "http://127.0.0.1:1", + json!({"vertex_credentials": true}), + ); + let error = perform_ocr(request).await.unwrap_err(); + assert!(error.to_string().contains("vertex_credentials")); +} + +#[tokio::test] +async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let mut request = wire_request( + "vertex_ai/mistral-ocr-maas", + "https://caller.example", + json!({"vertex_project":"project-1"}), + ); + request.connection.api_base_source = InputSource::Request; + + let error = perform_ocr(request).await.unwrap_err(); + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); +} + +#[tokio::test] +async fn adapters_build_complete_requests_and_share_mistral_normalization() { + use std::time::Duration; + + use crate::ocr::adapters::{MistralAdapter, OcrAdapter, VertexMistralAdapter}; + use crate::ocr::test_support::ocr_client; + + let client = ocr_client(); + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "ignored" + }); + let direct = wire_request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + options.clone(), + ); + let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); + let direct_http = MistralAdapter + .prepare_request(&direct, &client) + .await + .unwrap(); + let vertex_http = VertexMistralAdapter + .prepare_request(&vertex, &client) + .await + .unwrap(); + assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr"); + assert_eq!( + vertex_http.url().as_str(), + "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + for http in [&direct_http, &vertex_http] { + assert_eq!(http.method(), reqwest::Method::POST); + assert_eq!(http.headers()["authorization"], "Bearer test-key"); + assert_eq!(http.headers()["content-type"], "application/json"); + assert_eq!(http.timeout(), Some(&Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!( + body, + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "pages": [0, 2], + "include_image_base64": true + }) + ); + } + let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); + let direct_response = MistralAdapter + .transform_ocr_response(&direct, serde_json::from_value(payload.clone()).unwrap()) + .unwrap() + .into_json(); + let vertex_response = VertexMistralAdapter + .transform_ocr_response(&vertex, serde_json::from_value(payload).unwrap()) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert_eq!(direct_response["extra"], "preserved"); +} diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs index fc0ab2b62a3..e7739fe7312 100644 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs @@ -1,6 +1,7 @@ -//! Enforcement: the litellm-rust workspace has exactly five crates. +//! Enforcement: the litellm-rust workspace has exactly six crates. //! -//! `core` (the Rust SDK), `config` (the config-loading boundary), +//! `core` (the Rust SDK), `token-counter` (standalone input token counting), +//! `config` (the config-loading boundary), //! `ai-gateway` (the HTTP/WebSocket host), //! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the //! PyO3 cdylib). Adding or removing a crate must be a @@ -20,6 +21,7 @@ use std::path::{Path, PathBuf}; /// workspace legitimately gains or loses a crate. const EXPECTED_MEMBERS: &[&str] = &[ "crates/core", + "crates/token-counter", "crates/config", "crates/ai-gateway", "crates/python-interop", @@ -29,6 +31,7 @@ const EXPECTED_MEMBERS: &[&str] = &[ /// The crate subdirectory names that must exist under `crates/`. const EXPECTED_CRATE_DIRS: &[&str] = &[ "core", + "token-counter", "config", "ai-gateway", "python-interop", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index bda09a7d840..337a1e8e5ac 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -24,16 +24,17 @@ trace-parity = [ futures-util.workspace = true tracing = { workspace = true, optional = true } litellm-core = { workspace = true, features = ["bedrock-auth"] } +litellm-token-counter.workspace = true litellm-ai-gateway = { workspace = true, default-features = false } litellm-python-interop.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true serde.workspace = true serde_json.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["sync"] } [dev-dependencies] -criterion = "0.8.2" +criterion.workspace = true tokio-tungstenite.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/constants.rs b/litellm-rust/crates/python-bridge/src/constants.rs new file mode 100644 index 00000000000..d5cf5749820 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/constants.rs @@ -0,0 +1,2 @@ +/// Concurrent token-count encodes allowed when the core count is unavailable. +pub(crate) const TOKEN_COUNT_FALLBACK_PARALLELISM: usize = 1; diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 76c298abf89..e1f458ea0bc 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -41,6 +41,10 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { | Error::InvalidRequest(_) | Error::InvalidType { .. } | Error::MissingField(_) + | Error::MissingApiKey { .. } + | Error::MissingAzureAiCredentials + | Error::MissingAzureDocumentIntelligenceCredentials + | Error::MissingReductoApiKey | Error::Routing(_) // Nothing reached the provider, so serving it on Python cannot double // bill and is the only way the caller gets an answer at all. diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index f3648158cf6..b57197b9ddf 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -3,7 +3,6 @@ use std::panic::AssertUnwindSafe; use std::time::Duration; use futures_util::FutureExt; -use litellm_core::error::Error; use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil}; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; @@ -11,14 +10,15 @@ use serde::Serialize; use tokio::runtime::{Handle, Runtime}; use tokio::time::{self, MissedTickBehavior}; -pub(crate) fn run_sync( +pub(crate) fn run_sync( py: Python<'_>, future: F, - map_error: fn(Error) -> PyErr, + map_error: fn(E) -> PyErr, ) -> PyResult> where T: Serialize + Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { run_sync_on( py, @@ -28,15 +28,16 @@ where ) } -fn run_sync_on( +fn run_sync_on( py: Python<'_>, runtime: &Runtime, future: F, - map_error: fn(Error) -> PyErr, + map_error: fn(E) -> PyErr, ) -> PyResult> where T: Serialize + Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { if Handle::try_current().is_ok() { return Err(PyRuntimeError::new_err( @@ -49,14 +50,15 @@ where Pythonized(result).into_pyobject(py).map(Bound::unbind) } -pub(crate) fn run_async( +pub(crate) fn run_async( py: Python<'_>, future: F, - map_error: fn(Error) -> PyErr, + map_error: fn(E) -> PyErr, ) -> PyResult> where T: Serialize + Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { pyo3_async_runtimes::tokio::future_into_py(py, async move { let result = catch_future_panic(future).await?; @@ -65,7 +67,7 @@ where }) } -fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { +fn map_core_result(result: Result, map_error: fn(E) -> PyErr) -> PyResult { match result { Ok(value) => Ok(value), Err(error) => Err( @@ -75,9 +77,9 @@ fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) - } } -async fn catch_future_panic(future: F) -> PyResult> +async fn catch_future_panic(future: F) -> PyResult> where - F: Future>, + F: Future>, { AssertUnwindSafe(future) .catch_unwind() @@ -85,9 +87,9 @@ where .map_err(panic_to_pyerr) } -async fn wait_for_sync_result(future: F) -> PyResult> +async fn wait_for_sync_result(future: F) -> PyResult> where - F: Future>, + F: Future>, { let future = catch_future_panic(future); tokio::pin!(future); @@ -114,6 +116,7 @@ mod tests { use std::thread; use std::time::Instant; + use litellm_core::error::Error; use pyo3::panic::PanicException; use pyo3::types::{PyDict, PyModule}; use serde::Serializer; @@ -237,7 +240,7 @@ mod tests { let error = runtime.block_on(async { Python::attach(|py| { - run_sync::(py, async { Ok(true) }, runtime_error) + run_sync::(py, async { Ok(true) }, runtime_error) .expect_err("sync route should reject a nested Tokio runtime") }) }); @@ -273,7 +276,7 @@ mod tests { fn sync_runner_maps_a_panicked_future() { Python::initialize(); Python::attach(|py| { - let error = run_sync::( + let error = run_sync::( py, poll_fn(|_| -> Poll> { panic!("route future panicked") }), runtime_error, @@ -289,7 +292,7 @@ mod tests { fn sync_runner_maps_a_panicked_error_mapper() { Python::initialize(); Python::attach(|py| { - let error = run_sync::( + let error = run_sync::( py, async { Err(Error::InvalidRequest("invalid".to_string())) }, panicking_error_mapper, diff --git a/litellm-rust/crates/python-bridge/src/function_trace.rs b/litellm-rust/crates/python-bridge/src/function_trace.rs index ea7d9f4993e..bc3c962f7a3 100644 --- a/litellm-rust/crates/python-bridge/src/function_trace.rs +++ b/litellm-rust/crates/python-bridge/src/function_trace.rs @@ -1,3 +1,4 @@ +use std::fmt::Display; use std::future::Future; use litellm_core::observability::{FunctionTrace, FunctionTraceEvent}; @@ -6,17 +7,32 @@ use tracing::instrument::WithSubscriber; #[derive(Serialize)] pub(crate) struct TracedResponse { - response: T, + #[serde(skip_serializing_if = "Option::is_none")] + response: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, trace: Vec, } pub(crate) async fn capture( future: impl Future>, -) -> Result, E> { +) -> Result, E> +where + E: Display, +{ let trace = FunctionTrace::default(); - let response = future.with_subscriber(trace.dispatcher()).await?; - Ok(TracedResponse { - response, - trace: trace.events(), + let result = future.with_subscriber(trace.dispatcher()).await; + let events = trace.events(); + Ok(match result { + Ok(response) => TracedResponse { + response: Some(response), + error: None, + trace: events, + }, + Err(error) => TracedResponse { + response: None, + error: Some(error.to_string()), + trace: events, + }, }) } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 384f0be5a1b..cf0450a1b30 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,3 +1,4 @@ +mod constants; mod diagnostics; mod errors; mod execution; @@ -5,6 +6,7 @@ mod execution; mod function_trace; mod marshal; mod routes; +mod token_counter; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use pyo3::prelude::*; @@ -71,6 +73,7 @@ mod _native { super::errors::register(module)?; super::routes::register(module)?; module.add_class::()?; + super::token_counter::register(module)?; super::diagnostics::register(module) } } @@ -106,6 +109,7 @@ mod tests { "chat_completions", "achat_completions", "ResponsesWebSocketConnection", + "TokenCounter", "gil_stats", ]; diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 3285da14d5f..97313651011 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -225,7 +225,7 @@ mod tests { ( "ocr", "aocr", - "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)", + "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, input_sources=None, timeout_seconds=None)", ), ( "transcription", @@ -486,10 +486,11 @@ asyncio.run(exercise()) let code = CString::new( r#" result = routes.echo("traced") -assert result == { - "response": "traced", - "trace": [{"function": "execute_echo", "depth": 0}], -} +assert result["response"] == "traced", result +assert [event["function"] for event in result["trace"]] == ["execute_echo"], result +failure = routes.echo("error") +assert failure["error"] == "invalid request: synthetic error", failure +assert [event["function"] for event in failure["trace"]] == ["execute_echo"], failure "#, ) .expect("Python source should not contain null bytes"); diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index cc2f8e43cea..c5def64c2f1 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -2,6 +2,7 @@ use litellm_core::Error; use std::future::Future; use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; +use litellm_core::ocr::wire::{OcrWireRequest, decode_request, is_supported_request}; use pyo3::prelude::*; use serde_json::Value; @@ -21,6 +22,12 @@ fn prepare_ocr( timeout_seconds: inputs.timeout_seconds, })?; let optional_params = object_or_empty("optional_params", inputs.optional_params)?; + let input_sources = inputs + .input_sources + .map(serde_json::from_value) + .transpose() + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))? + .unwrap_or_default(); Ok(async move { let RouteOptions { @@ -31,6 +38,22 @@ fn prepare_ocr( extra_headers, timeout, } = options; + if is_supported_request(&model, custom_llm_provider.as_deref()) { + let request = decode_request(OcrWireRequest { + model, + document, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + input_sources, + timeout_seconds: timeout.map(|value| value.as_secs_f64()), + })?; + return litellm_core::ocr::ocr(request) + .await + .map(|response| response.into_json()); + } run_ocr(OcrRequest { model: &model, document, @@ -66,8 +89,29 @@ bridge_route! { extra_headers: Option, #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + input_sources: Option, timeout_seconds: Option, }, prepare = prepare_ocr, errors = ocr_error_to_pyerr, } + +#[cfg(test)] +mod tests { + use litellm_core::ocr::wire::is_supported_request; + + #[test] + fn native_activation_includes_migrated_providers() { + assert!(is_supported_request("model", Some("mistral"))); + assert!(is_supported_request("pixtral-12b", Some("azure_ai"))); + assert!(is_supported_request( + "documentintelligence/prebuilt-read", + Some("azure_ai") + )); + assert!(is_supported_request("parse-v3", Some("reducto"))); + assert!(is_supported_request("parse-legacy", Some("reducto"))); + assert!(is_supported_request("mistral-ocr", Some("vertex_ai"))); + assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); + } +} diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs new file mode 100644 index 00000000000..ee82c170d55 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -0,0 +1,87 @@ +use std::num::NonZero; +use std::sync::Arc; +use std::thread::available_parallelism; + +use litellm_python_interop::release_gil; +use litellm_token_counter::{ + CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, +}; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyAny; +use tokio::sync::Semaphore; + +use crate::constants::TOKEN_COUNT_FALLBACK_PARALLELISM; +use crate::errors::RustBridgeDeclined; +use crate::execution::run_async; + +/// Counts the input tokens of a raw request body off the Python event loop with +/// the GIL released. Python owns which requests get here and what to do with +/// the count. At most one encode per core runs at a time; the rest wait in the +/// async task, where a cancelled Python awaiter drops them before any blocking +/// work is scheduled. +#[pyclass(frozen)] +struct TokenCounter { + inner: Arc, + encode_slots: Arc, +} + +#[pymethods] +impl TokenCounter { + #[new] + fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { + let inner = release_gil(py, || CoreTokenCounter::from_json(tokenizer_json)) + .map_err(token_count_error_to_pyerr)?; + Ok(Self { + inner: Arc::new(inner), + encode_slots: Arc::new(Semaphore::new(encode_parallelism())), + }) + } + + fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult> { + let counter = Arc::clone(&self.inner); + let encode_slots = Arc::clone(&self.encode_slots); + let body = body.to_vec(); + run_async( + py, + async move { + let _slot = encode_slots + .acquire_owned() + .await + .map_err(|error| Error::Task(error.to_string()))?; + tokio::task::spawn_blocking(move || count_body(&counter, &body)) + .await + .map_err(|error| Error::Task(error.to_string()))? + }, + token_count_error_to_pyerr, + ) + } +} + +fn encode_parallelism() -> usize { + available_parallelism().map_or(TOKEN_COUNT_FALLBACK_PARALLELISM, NonZero::get) +} + +fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result { + let request = CountableRequest::parse(body)?; + counter.count_request(&request) +} + +fn token_count_error_to_pyerr(error: Error) -> PyErr { + let message = error.to_string(); + match error { + Error::Load(_) => PyValueError::new_err(message), + Error::RequestParse(_) + | Error::MissingInput + | Error::FloatText + | Error::ContentBlock + | Error::ArrayItems + | Error::JsonSerialization(_) + | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message), + Error::Encode(_) | Error::Task(_) => PyRuntimeError::new_err(message), + } +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::() +} diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml new file mode 100644 index 00000000000..61c9cf6e991 --- /dev/null +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "litellm-token-counter" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +indexmap = { version = "2.14.0", features = ["serde"] } +itoa = "1.0" +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +unicode-normalization-alignments = "0.1.12" + +[dev-dependencies] +criterion.workspace = true +rand.workspace = true +rstest.workspace = true + +[[bench]] +name = "token_counter" +harness = false + +[[bench]] +name = "allocations" +harness = false diff --git a/litellm-rust/crates/token-counter/benches/allocations.rs b/litellm-rust/crates/token-counter/benches/allocations.rs new file mode 100644 index 00000000000..343f815749d --- /dev/null +++ b/litellm-rust/crates/token-counter/benches/allocations.rs @@ -0,0 +1,128 @@ +use std::alloc::{GlobalAlloc, Layout, System}; +use std::hint::black_box; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use litellm_token_counter::{CountableRequest, TokenCounter}; + +struct CountingAllocator; + +static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0); +static BYTES: AtomicUsize = AtomicUsize::new(0); + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + BYTES.fetch_add(layout.size(), Ordering::Relaxed); + // SAFETY: This allocator delegates the unchanged layout to `System`. + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) { + // SAFETY: The pointer and layout came from the delegated `System` allocation. + unsafe { System.dealloc(pointer, layout) } + } + + unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + BYTES.fetch_add(size, Ordering::Relaxed); + // SAFETY: The pointer and layout came from `System`; the new size is unchanged. + unsafe { System.realloc(pointer, layout, size) } + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +#[derive(Clone, Copy)] +struct AllocationCount { + allocations: usize, + bytes: usize, +} + +impl AllocationCount { + fn assert_max(self, label: &str, maximum: Self) { + eprintln!( + "{label}: {} allocations, {} bytes", + self.allocations, self.bytes + ); + assert!( + self.allocations <= maximum.allocations, + "{label} allocation count exceeded {}", + maximum.allocations + ); + assert!( + self.bytes <= maximum.bytes, + "{label} allocated bytes exceeded {}", + maximum.bytes + ); + } +} + +const TOKENIZER_JSON: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" +)); +const OBJECT_BODY: &[u8] = br#"{"model":"claude-sonnet-4-5","input":{"text":"caf\u00e9","n":3,"ok":true,"list":[1,"a",{"z":[]}]}}"#; +const INTEGER_BODY: &[u8] = br#"{"input":[-9223372036854775808,0,18446744073709551615]}"#; + +fn measure(operation: impl FnOnce()) -> AllocationCount { + ALLOCATIONS.store(0, Ordering::Relaxed); + BYTES.store(0, Ordering::Relaxed); + operation(); + AllocationCount { + allocations: ALLOCATIONS.load(Ordering::Relaxed), + bytes: BYTES.load(Ordering::Relaxed), + } +} + +fn main() { + measure(|| { + black_box(CountableRequest::parse(OBJECT_BODY).expect("object request parses")); + }) + .assert_max( + "parse object request", + AllocationCount { + allocations: 16, + bytes: 1_900, + }, + ); + + let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("tokenizer loads"); + let object = CountableRequest::parse(OBJECT_BODY).expect("object request parses"); + counter + .count_request(&object) + .expect("object warmup succeeds"); + measure(|| { + black_box( + counter + .count_request(black_box(&object)) + .expect("object counts"), + ); + }) + .assert_max( + "count object request", + AllocationCount { + allocations: 74, + bytes: 2_200, + }, + ); + + let integers = CountableRequest::parse(INTEGER_BODY).expect("integer request parses"); + counter + .count_request(&integers) + .expect("integer warmup succeeds"); + measure(|| { + black_box( + counter + .count_request(black_box(&integers)) + .expect("integers count"), + ); + }) + .assert_max( + "count integer list", + AllocationCount { + allocations: 26, + bytes: 1_050, + }, + ); +} diff --git a/litellm-rust/crates/token-counter/benches/token_counter.rs b/litellm-rust/crates/token-counter/benches/token_counter.rs new file mode 100644 index 00000000000..a7c3177b88a --- /dev/null +++ b/litellm-rust/crates/token-counter/benches/token_counter.rs @@ -0,0 +1,100 @@ +use std::hint::black_box; +use std::time::Duration; + +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use litellm_token_counter::TokenCounter; +use tokenizers::Tokenizer; + +const TOKENIZER_JSON: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" +)); +const FULL_CONTEXT_TOKENS: usize = 1_000_000; +const LARGE_PROMPT_UNIT: &str = "The quick brown fox jumps over the lazy dog. 0123456789\n"; + +fn full_context_input(tokenizer: &Tokenizer) -> String { + let unit_tokens = tokenizer + .encode_fast(LARGE_PROMPT_UNIT, false) + .expect("reference tokenizer should encode") + .len(); + let input = LARGE_PROMPT_UNIT.repeat(FULL_CONTEXT_TOKENS.div_ceil(unit_tokens)); + let actual_tokens = tokenizer + .encode_fast(input.as_str(), true) + .expect("reference tokenizer should encode") + .len(); + assert!((FULL_CONTEXT_TOKENS..FULL_CONTEXT_TOKENS + unit_tokens).contains(&actual_tokens)); + input +} + +fn inputs(tokenizer: &Tokenizer) -> Vec<(&'static str, String)> { + vec![ + ( + "ascii_chat", + "User: Summarize the benefits of statistical benchmarking.\nAssistant:".repeat(8), + ), + ( + "unicode_nfkc", + "Ａ quick café résumé: مرحبا 世界 🙂 fi Ⅳ.\n".repeat(16), + ), + ("large_prompt", LARGE_PROMPT_UNIT.repeat(256)), + ("full_context_1m_tokens", full_context_input(tokenizer)), + ] +} + +fn token_counter(c: &mut Criterion) { + let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("token counter should load"); + let tokenizer = TOKENIZER_JSON + .parse::() + .expect("reference tokenizer should load"); + let mut group = c.benchmark_group("anthropic_token_counter"); + + for (name, input) in inputs(&tokenizer) { + let expected = tokenizer + .encode_fast(input.as_str(), true) + .expect("reference tokenizer should encode") + .len(); + let actual = counter + .count_text(input.as_str()) + .expect("benchmark path should count"); + assert_eq!( + actual, expected, + "benchmark paths should produce the same count" + ); + group.throughput(Throughput::Bytes(input.len() as u64)); + group.bench_with_input( + BenchmarkId::new("byte_level_fast_path", name), + &input, + |b, input| { + b.iter(|| { + counter + .count_text(black_box(input.as_str())) + .expect("fast path should count") + }) + }, + ); + group.bench_with_input( + BenchmarkId::new("full_encoder", name), + &input, + |b, input| { + b.iter(|| { + tokenizer + .encode_fast(black_box(input.as_str()), true) + .expect("reference tokenizer should encode") + .len() + }) + }, + ); + } + + group.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(20) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(4)); + targets = token_counter +} +criterion_main!(benches); diff --git a/litellm-rust/crates/token-counter/src/byte_level.rs b/litellm-rust/crates/token-counter/src/byte_level.rs new file mode 100644 index 00000000000..2b435eb39bb --- /dev/null +++ b/litellm-rust/crates/token-counter/src/byte_level.rs @@ -0,0 +1,650 @@ +//! Exact token counting for a supported tokenizer configuration: optional +//! NFKC normalization, `ByteLevel` pre-tokenization with the GPT-2 split regex, +//! and no post-processing. A scanner reproduces the regex's piece boundaries +//! and hands each piece to the tokenizer's model. Unsupported configurations +//! and added-token inputs fall back to the full encoder. + +use std::borrow::Cow; +use std::iter; + +use tokenizers::normalizers::NormalizerWrapper; +use tokenizers::pre_tokenizers::PreTokenizerWrapper; +use tokenizers::{Model, Tokenizer}; +use unicode_normalization_alignments::{IsNormalized, UnicodeNormalization, is_nfkc_quick}; + +use super::unicode_classes::UnicodeClasses; + +const CONTRACTIONS: [&str; 7] = ["'s", "'t", "'re", "'ve", "'m", "'ll", "'d"]; + +pub(super) struct ByteLevelCounter { + nfkc: bool, + normalized_added_tokens: Vec, + unicode_classes: &'static UnicodeClasses, +} + +impl ByteLevelCounter { + pub(super) fn detect(tokenizer: &Tokenizer) -> Option { + let nfkc = match tokenizer.get_normalizer() { + None => false, + Some(NormalizerWrapper::NFKC(_)) => true, + Some(_) => return None, + }; + let Some(PreTokenizerWrapper::ByteLevel(byte_level)) = tokenizer.get_pre_tokenizer() else { + return None; + }; + let plain = !byte_level.add_prefix_space + && byte_level.use_regex + && tokenizer.get_post_processor().is_none() + && tokenizer.get_truncation().is_none() + && tokenizer.get_padding().is_none(); + if !plain { + return None; + } + let vocabulary = tokenizer.get_added_vocabulary(); + let normalized_added_tokens = vocabulary + .get_vocab() + .iter() + .filter_map(|(original, id)| { + vocabulary + .simple_id_to_token(*id) + .filter(|normalized| normalized != original) + }) + .collect(); + Some(Self { + nfkc, + normalized_added_tokens, + unicode_classes: UnicodeClasses::get()?, + }) + } + + /// `None` when the text contains an added token or the model rejects a + /// piece; the caller then runs the full encoder. + pub(super) fn count(&self, tokenizer: &Tokenizer, text: &str) -> Option { + let normalized = self.normalize(text); + let added_tokens = tokenizer.get_added_vocabulary().get_vocab(); + if added_tokens + .keys() + .chain(self.normalized_added_tokens.iter()) + .any(|token| text.contains(token.as_str()) || normalized.contains(token.as_str())) + { + return None; + } + let model = tokenizer.get_model(); + let mapped: String = normalized.bytes().map(byte_char).collect(); + pieces(&normalized, self.unicode_classes) + .try_fold((0, 0), |(start, total), piece| { + let end = start + mapped_len(piece); + let tokens = model.tokenize(&mapped[start..end]).ok()?; + Some((end, total + tokens.len())) + }) + .map(|(_, total)| total) + } + + /// Same crate and Unicode tables as `NormalizedString::nfkc`, so the + /// result is what the full encoder would have tokenized. + fn normalize<'a>(&self, text: &'a str) -> Cow<'a, str> { + if !self.nfkc || text.is_ascii() || is_nfkc_quick(text.chars()) == IsNormalized::Yes { + return Cow::Borrowed(text); + } + Cow::Owned(text.nfkc().map(|(character, _)| character).collect()) + } +} + +/// GPT-2 `bytes_to_unicode`: printable Latin-1 bytes map to themselves, the +/// rest to U+0100 onwards in byte order. +fn byte_char(byte: u8) -> char { + let code = match byte { + 0x21..=0x7E | 0xA1..=0xAC | 0xAE..=0xFF => u32::from(byte), + 0x00..=0x20 => 0x100 + u32::from(byte), + 0x7F..=0xA0 => 0x121 + u32::from(byte - 0x7F), + 0xAD => 0x143, + }; + char::from_u32(code).unwrap_or(char::REPLACEMENT_CHARACTER) +} + +fn mapped_len(piece: &str) -> usize { + piece.len() + + piece + .bytes() + .filter(|byte| !byte.is_ascii_graphic()) + .count() +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Class { + Letter, + Number, + Space, + Other, +} + +fn class(character: char, unicode_classes: &UnicodeClasses) -> Class { + match character { + 'A'..='Z' | 'a'..='z' => Class::Letter, + '0'..='9' => Class::Number, + '\t'..='\r' | ' ' => Class::Space, + _ if character.is_ascii() => Class::Other, + _ if unicode_classes.is_letter(character) => Class::Letter, + _ if unicode_classes.is_number(character) => Class::Number, + _ if unicode_classes.is_space(character) => Class::Space, + _ => Class::Other, + } +} + +/// The regex matches every character, so the pieces tile the text. +fn pieces<'a>( + text: &'a str, + unicode_classes: &'static UnicodeClasses, +) -> impl Iterator { + iter::successors(split_piece(text, unicode_classes), move |(_, rest)| { + split_piece(rest, unicode_classes) + }) + .map(|(piece, _)| piece) +} + +fn split_piece<'a>(text: &'a str, unicode_classes: &UnicodeClasses) -> Option<(&'a str, &'a str)> { + let first = text.chars().next()?; + Some(text.split_at(piece_len(text, first, unicode_classes))) +} + +fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize { + if let Some(contraction) = CONTRACTIONS.iter().find(|word| text.starts_with(**word)) { + return contraction.len(); + } + let first_class = class(first, unicode_classes); + if first_class != Class::Space { + return run_len(text, first_class, unicode_classes); + } + if first != ' ' { + return space_run_len(text, unicode_classes); + } + let after_space = &text[1..]; + match after_space + .chars() + .next() + .map(|character| class(character, unicode_classes)) + { + None | Some(Class::Space) => space_run_len(text, unicode_classes), + Some(run_class) => 1 + run_len(after_space, run_class, unicode_classes), + } +} + +fn run_len(text: &str, run_class: Class, unicode_classes: &UnicodeClasses) -> usize { + text.char_indices() + .find(|(_, character)| class(*character, unicode_classes) != run_class) + .map_or(text.len(), |(index, _)| index) +} + +/// `\s+(?!\S)|\s+`: whitespace followed by a non-space leaves its last +/// character to start the next piece (` ?` on the following alternatives). +fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize { + let run = run_len(text, Class::Space, unicode_classes); + if run == text.len() { + return run; + } + let last = text[..run].chars().next_back().map_or(0, char::len_utf8); + match run - last { + 0 => run, + shorter => shorter, + } +} + +#[cfg(test)] +mod tests { + use rand::rngs::StdRng; + use rand::seq::SliceRandom; + use rand::{Rng, SeedableRng}; + use rstest::{fixture, rstest}; + use tokenizers::normalizers::NFKC; + use tokenizers::pre_tokenizers::byte_level::ByteLevel; + use tokenizers::utils::SysRegex; + use tokenizers::{ + NormalizedString, Normalizer, OffsetReferential, OffsetType, PreTokenizedString, + PreTokenizer, + }; + + use super::*; + + #[fixture] + fn anthropic_tokenizer() -> Tokenizer { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" + ); + std::fs::read_to_string(path) + .expect("anthropic tokenizer json is in the repo") + .parse() + .expect("anthropic tokenizer loads") + } + + fn reference_count(tokenizer: &Tokenizer, text: &str) -> usize { + tokenizer.encode_fast(text, true).expect("encode").len() + } + + fn byte_level_counter(nfkc: bool) -> ByteLevelCounter { + ByteLevelCounter { + nfkc, + normalized_added_tokens: Vec::new(), + unicode_classes: UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"), + } + } + + const ALPHABET: &[&str] = &[ + "a", + "Z", + "e", + "s", + "t", + "d", + "m", + "'", + "'s", + "'re", + "'ll", + "'S", + "0", + "9", + " ", + " ", + "\t", + "\n", + "\r\n", + "\u{b}", + ".", + ",", + "!", + "-", + "(", + "\"", + "\u{a0}", + "\u{85}", + "\u{2028}", + "\u{3000}", + "\u{200b}", + "\u{200d}", + "é", + "e\u{301}", + "ß", + "漢", + "字", + "ع", + "३", + "½", + "Ⅳ", + "🙂", + "👍🏽", + "Ａ", + "fi", + "㍿", + "㋿", + "ꟲ", + "𐞁", + "a\u{30a}", + "\u{1e0b}\u{323}", + "<", + ">", + "EOT", + "", + "", + ]; + + fn random_text(rng: &mut StdRng) -> String { + let pieces = rng.gen_range(0..40); + (0..pieces) + .map(|_| *ALPHABET.choose(rng).expect("alphabet is not empty")) + .collect() + } + + #[rstest] + #[case::plain_text("Hello, how are you today?", true)] + #[case::added_token("stop here", false)] + #[case::normalized_added_token("stop ＜ＥＯＴ＞ here", false)] + fn anthropic_tokenizer_takes_the_fast_path( + anthropic_tokenizer: Tokenizer, + #[case] text: &str, + #[case] supported: bool, + ) { + let fast = + ByteLevelCounter::detect(&anthropic_tokenizer).expect("anthropic shape is supported"); + assert!(fast.nfkc); + let count = fast.count(&anthropic_tokenizer, text); + if supported { + assert_eq!(count, Some(reference_count(&anthropic_tokenizer, text))); + } else { + assert_eq!(count, None); + } + } + + #[rstest] + fn counts_match_the_full_encoder(anthropic_tokenizer: Tokenizer) { + let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); + let mut rng = StdRng::seed_from_u64(2026); + for _ in 0..4000 { + let text = random_text(&mut rng).replace('<', "("); + let expected = reference_count(&anthropic_tokenizer, &text); + assert_eq!( + fast.count(&anthropic_tokenizer, &text), + Some(expected), + "text {text:?}" + ); + } + } + + #[rstest] + fn nfkc_matches_the_tokenizer_normalizer_for_every_scalar_value() { + let fast = byte_level_counter(true); + let mut text = String::new(); + for character in (0..=0x10FFFFu32).filter_map(char::from_u32) { + text.clear(); + text.push(character); + let mut expected = NormalizedString::from(text.as_str()); + NFKC.normalize(&mut expected).expect("nfkc"); + assert_eq!( + fast.normalize(&text), + expected.get(), + "U+{:04X}", + u32::from(character) + ); + } + } + + #[rstest] + fn nfkc_matches_the_tokenizer_normalizer_on_random_texts() { + let fast = byte_level_counter(true); + let mut rng = StdRng::seed_from_u64(11); + for _ in 0..4000 { + let text = random_text(&mut rng); + let mut expected = NormalizedString::from(text.as_str()); + NFKC.normalize(&mut expected).expect("nfkc"); + assert_eq!(fast.normalize(&text), expected.get(), "text {text:?}"); + } + } + + #[rstest] + fn pieces_match_the_byte_level_pre_tokenizer() { + let byte_level = ByteLevel::new(false, true, true); + let mut rng = StdRng::seed_from_u64(7); + for _ in 0..4000 { + let text = byte_level_counter(true) + .normalize(&random_text(&mut rng)) + .into_owned(); + let mut pre_tokenized = PreTokenizedString::from(text.as_str()); + byte_level + .pre_tokenize(&mut pre_tokenized) + .expect("pre-tokenize"); + let expected: Vec<(String, (usize, usize))> = pre_tokenized + .get_splits(OffsetReferential::Original, OffsetType::Byte) + .into_iter() + .map(|(mapped, offsets, _)| (mapped.to_string(), offsets)) + .collect(); + let actual: Vec<(String, (usize, usize))> = pieces( + &text, + UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"), + ) + .map(|piece| { + let start = piece.as_ptr() as usize - text.as_ptr() as usize; + let mapped: String = piece.bytes().map(byte_char).collect(); + (mapped, (start, start + piece.len())) + }) + .collect(); + assert_eq!(actual, expected, "text {text:?}"); + } + } + + #[rstest] + fn byte_chars_match_the_byte_level_alphabet() { + let byte_level = ByteLevel::new(false, false, false); + let characters: Vec = (0..=0x10FFFFu32).filter_map(char::from_u32).collect(); + for chunk in characters.chunks(1024) { + let text: String = chunk.iter().collect(); + let mut pre_tokenized = PreTokenizedString::from(text.as_str()); + byte_level + .pre_tokenize(&mut pre_tokenized) + .expect("pre-tokenize"); + let expected: String = pre_tokenized + .get_splits(OffsetReferential::Original, OffsetType::Byte) + .into_iter() + .map(|(mapped, _, _)| mapped) + .collect(); + let actual: String = text.bytes().map(byte_char).collect(); + assert_eq!(actual.len(), mapped_len(&text)); + assert_eq!( + actual, + expected, + "chunk starting at U+{:04X}", + u32::from(chunk[0]) + ); + } + } + + #[rstest] + fn classes_match_oniguruma() { + let unicode_classes = UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"); + let letter = SysRegex::new(r"\p{L}").expect("regex"); + let number = SysRegex::new(r"\p{N}").expect("regex"); + let space = SysRegex::new(r"\s").expect("regex"); + let whole = + |regex: &SysRegex, text: &str| regex.find_iter(text).next() == Some((0, text.len())); + let mut text = String::new(); + for character in (0..=0x10FFFFu32).filter_map(char::from_u32) { + text.clear(); + text.push(character); + let expected = if whole(&letter, &text) { + Class::Letter + } else if whole(&number, &text) { + Class::Number + } else if whole(&space, &text) { + Class::Space + } else { + Class::Other + }; + assert_eq!( + class(character, unicode_classes), + expected, + "U+{:04X}", + u32::from(character) + ); + } + } + + #[rstest] + #[case("prefix")] + #[case("regex")] + #[case("normalizer")] + #[case("no_pre_tokenizer")] + #[case("other_pre_tokenizer")] + #[case("post_processor")] + #[case("truncation")] + #[case("padding")] + fn other_tokenizer_shapes_are_declined( + mut anthropic_tokenizer: Tokenizer, + #[case] shape: &str, + ) { + use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams}; + match shape { + "prefix" => { + anthropic_tokenizer.with_pre_tokenizer(Some(ByteLevel::new(true, true, true))); + } + "regex" => { + anthropic_tokenizer.with_pre_tokenizer(Some(ByteLevel::new(false, true, false))); + } + "normalizer" => { + anthropic_tokenizer + .with_normalizer(Some(tokenizers::normalizers::Lowercase)) + .expect("normalizer"); + } + "no_pre_tokenizer" => { + anthropic_tokenizer.with_pre_tokenizer(None::); + } + "other_pre_tokenizer" => { + anthropic_tokenizer + .with_pre_tokenizer(Some(tokenizers::pre_tokenizers::whitespace::Whitespace)); + } + "post_processor" => { + anthropic_tokenizer.with_post_processor(Some(ByteLevel::default())); + } + "truncation" => { + anthropic_tokenizer + .with_truncation(Some(TruncationParams { + max_length: 2, + ..Default::default() + })) + .expect("truncation"); + } + "padding" => { + anthropic_tokenizer.with_padding(Some(PaddingParams { + strategy: PaddingStrategy::Fixed(32), + ..Default::default() + })); + } + _ => unreachable!(), + } + assert!(ByteLevelCounter::detect(&anthropic_tokenizer).is_none()); + let counter = crate::TokenCounter::from_json( + &anthropic_tokenizer.to_string(false).expect("serialize"), + ) + .expect("load"); + for text in ["", "Hello WORLD! ＡＢ fi Ⅳ", " stop"] { + assert_eq!( + counter.count_text(text).expect("count"), + reference_count(&anthropic_tokenizer, text) + ); + } + } + + #[rstest] + #[case(false)] + #[case(true)] + fn arbitrary_unicode_and_long_inputs_use_fast_path( + mut anthropic_tokenizer: Tokenizer, + #[case] nfkc: bool, + ) { + if !nfkc { + anthropic_tokenizer + .with_normalizer(None::) + .expect("normalizer"); + } + let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); + let mut rng = StdRng::seed_from_u64(314159); + for _ in 0..1000 { + let text: String = (0..64) + .filter_map(|_| char::from_u32(rng.gen_range(0..=0x10ffff))) + .collect(); + assert_eq!( + fast.count(&anthropic_tokenizer, &text), + Some(reference_count(&anthropic_tokenizer, &text)), + "text {text:?}" + ); + } + for text in [ + "", + "'s't're've'm'll'd'S'RE", + " a \t\r\n b\u{85}\u{a0}c ", + "\0é漢🙂", + "a\u{30a}\u{301}", + "ＡfiⅣ", + ] { + let text = text.repeat(2048); + assert_eq!( + fast.count(&anthropic_tokenizer, &text), + Some(reference_count(&anthropic_tokenizer, &text)) + ); + } + } + + #[rstest] + #[case(false, false, false, false)] + #[case(true, false, false, false)] + #[case(false, true, false, false)] + #[case(false, false, true, false)] + #[case(false, false, false, true)] + fn added_token_options_fall_back( + mut anthropic_tokenizer: Tokenizer, + #[case] special: bool, + #[case] single_word: bool, + #[case] lstrip: bool, + #[case] rstrip: bool, + ) { + anthropic_tokenizer + .add_tokens([tokenizers::AddedToken::from("custom token", special) + .single_word(single_word) + .lstrip(lstrip) + .rstrip(rstrip)]) + .expect("add token"); + let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); + let counter = crate::TokenCounter::from_json( + &anthropic_tokenizer.to_string(false).expect("serialize"), + ) + .expect("load"); + for text in [ + "custom token", + "a custom token b", + "acustom tokenb", + " custom token ", + ] { + assert_eq!(fast.count(&anthropic_tokenizer, text), None); + assert_eq!( + counter.count_text(text).expect("count"), + reference_count(&anthropic_tokenizer, text) + ); + } + } + + #[test] + fn model_errors_reach_public_caller() { + let mut tokenizer = Tokenizer::new(tokenizers::models::wordpiece::WordPiece::default()); + tokenizer.with_pre_tokenizer(Some(ByteLevel::new(false, true, true))); + let fast = ByteLevelCounter::detect(&tokenizer).expect("supported"); + assert_eq!(fast.count(&tokenizer, "hello"), None); + assert!(tokenizer.encode_fast("hello", true).is_err()); + let counter = + crate::TokenCounter::from_json(&tokenizer.to_string(false).expect("serialize")) + .expect("load"); + assert!(matches!( + counter.count_text("hello"), + Err(crate::Error::Encode(_)) + )); + } + + #[rstest] + fn shared_counter_matches_encoder_across_threads(anthropic_tokenizer: Tokenizer) { + let counter = crate::TokenCounter::from_json( + &anthropic_tokenizer.to_string(false).expect("serialize"), + ) + .expect("load"); + let inputs = [ + "hello world", + "ＡＢ fi\n漢字🙂", + " stop", + "\t 're \r\n", + ]; + let expected = inputs.map(|text| reference_count(&anthropic_tokenizer, text)); + std::thread::scope(|scope| { + for _ in 0..8 { + let counter = &counter; + scope.spawn(move || { + for _ in 0..100 { + for (text, count) in inputs.iter().zip(expected) { + assert_eq!(counter.count_text(text).expect("count"), count); + } + } + }); + } + }); + } + + #[rstest] + fn normalized_added_token_spelling_declines_fast_path(mut anthropic_tokenizer: Tokenizer) { + anthropic_tokenizer + .add_tokens([tokenizers::AddedToken::from("ＡＢＣＤ　ＥＦＧＨ", false)]) + .expect("add token"); + let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); + assert_eq!(reference_count(&anthropic_tokenizer, "ABCD EFGH"), 1); + assert_eq!(fast.count(&anthropic_tokenizer, "ABCD EFGH"), None); + let counter = crate::TokenCounter::from_json( + &anthropic_tokenizer.to_string(false).expect("serialize"), + ) + .expect("load"); + assert_eq!(counter.count_text("ABCD EFGH").expect("count"), 1); + } +} diff --git a/litellm-rust/crates/token-counter/src/counter.rs b/litellm-rust/crates/token-counter/src/counter.rs new file mode 100644 index 00000000000..a3943515a64 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/counter.rs @@ -0,0 +1,194 @@ +use serde::Serialize; + +use crate::Error; +use crate::byte_level::ByteLevelCounter; +use crate::python_json; +use crate::tools::format_function_definitions; +use crate::types::{ + ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice, + ToolDefinition, +}; + +const TOKENS_PER_MESSAGE: usize = 3; +const TOKENS_PER_NAME: usize = 1; +const REPLY_PRIMING_TOKENS: usize = 3; +const TOOL_DEFINITIONS_TOKENS: usize = 9; +const TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT: usize = 4; +const TOOL_CHOICE_NONE_TOKENS: usize = 1; +const NAMED_TOOL_CHOICE_TOKENS: usize = 7; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct InputTokenCount { + pub model: Option, + pub input_tokens: usize, +} + +/// A loaded HuggingFace tokenizer plus the message accounting Python applies on +/// top of it. Encoding is CPU-bound and synchronous; hosts run it off their +/// event loop. +pub struct TokenCounter { + tokenizer: tokenizers::Tokenizer, + byte_level: Option, +} + +impl TokenCounter { + /// Load a HuggingFace `tokenizer.json` document. The host reads the file. + pub fn from_json(tokenizer_json: &str) -> Result { + let tokenizer = tokenizer_json + .parse::() + .map_err(Error::Load)?; + let byte_level = ByteLevelCounter::detect(&tokenizer); + Ok(Self { + tokenizer, + byte_level, + }) + } + + pub fn count_text(&self, text: &str) -> Result { + if let Some(count) = self + .byte_level + .as_ref() + .and_then(|counter| counter.count(&self.tokenizer, text)) + { + return Ok(count); + } + self.tokenizer + .encode_fast(text, true) + .map(|encoding| encoding.len()) + .map_err(Error::Encode) + } + + /// Mirrors the host's key precedence: `messages`, then `prompt`, then + /// `input`, then `query` plus `documents`. + pub fn count_request(&self, request: &CountableRequest) -> Result { + let input_tokens = if let Some(messages) = &request.messages { + self.count_messages(request, messages)? + } else if let Some(prompt) = &request.prompt { + self.count_text_value(prompt)? + } else if let Some(input) = &request.input { + self.count_text_value(input)? + } else if request.query.is_some() || request.documents.is_some() { + self.count_optional_text_value(request.query.as_ref())? + + self.count_optional_text_value(request.documents.as_ref())? + } else { + return Err(Error::MissingInput); + }; + Ok(InputTokenCount { + model: request.model.clone(), + input_tokens, + }) + } + + fn count_messages( + &self, + request: &CountableRequest, + messages: &[Message], + ) -> Result { + let message_tokens = messages + .iter() + .map(|message| self.count_message(message)) + .sum::>()?; + let includes_system_message = messages + .iter() + .any(|message| message.role.as_deref() == Some("system")); + let extra_tokens = self.count_extra( + request.tools.as_deref().unwrap_or_default(), + request.tool_choice.as_ref(), + includes_system_message, + )?; + Ok(message_tokens + extra_tokens) + } + + fn count_optional_text_value(&self, value: Option<&TextValue>) -> Result { + value.map_or(Ok(0), |value| self.count_text_value(value)) + } + + /// `str()` for scalars, `json.dumps()` for objects, lists flattened, nulls + /// skipped. Floats are declined because Python's `repr` and Rust's float + /// formatting disagree on exponents. + fn count_text_value(&self, value: &TextValue) -> Result { + match value { + TextValue::Null => Ok(0), + TextValue::Bool(true) => self.count_text("True"), + TextValue::Bool(false) => self.count_text("False"), + TextValue::Number(number) => match (number.as_i64(), number.as_u64()) { + (Some(number), _) => self.count_text(itoa::Buffer::new().format(number)), + (_, Some(number)) => self.count_text(itoa::Buffer::new().format(number)), + _ => Err(Error::FloatText), + }, + TextValue::Text(text) => self.count_text(text), + TextValue::List(items) => items + .iter() + .map(|item| self.count_text_value(item)) + .sum::>(), + TextValue::Object(_) => self.count_text(&python_json::dumps(value)?), + } + } + + fn count_message(&self, message: &Message) -> Result { + let role_tokens = match &message.role { + Some(role) => self.count_text(role)?, + None => 0, + }; + let name_tokens = match &message.name { + Some(name) => self.count_text(name)? + TOKENS_PER_NAME, + None => 0, + }; + let content_tokens = match &message.content { + Some(MessageContent::Text(text)) => self.count_text(text)?, + Some(MessageContent::Blocks(items)) => items + .iter() + .map(|item| self.count_content_item(item)) + .sum::>()?, + None => 0, + }; + Ok(TOKENS_PER_MESSAGE + role_tokens + name_tokens + content_tokens) + } + + fn count_content_item(&self, item: &ContentItem) -> Result { + match item { + ContentItem::Text(text) => self.count_text(text), + ContentItem::Block(ContentBlock::Text { text }) => self.count_text(text), + ContentItem::Block(ContentBlock::Thinking { thinking }) => { + if thinking.is_empty() { + return Ok(0); + } + self.count_text(thinking) + } + ContentItem::Block(ContentBlock::ToolReference { tool_name }) => { + match tool_name.as_deref().filter(|name| !name.is_empty()) { + Some(name) => self.count_text(name), + None => Ok(0), + } + } + ContentItem::Block(ContentBlock::Unsupported) => Err(Error::ContentBlock), + } + } + + fn count_extra( + &self, + tools: &[ToolDefinition], + tool_choice: Option<&ToolChoice>, + includes_system_message: bool, + ) -> Result { + let tool_tokens = if tools.is_empty() { + 0 + } else { + let definitions = self.count_text(&format_function_definitions(tools)?)?; + let discount = if includes_system_message { + TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT + } else { + 0 + }; + definitions + TOOL_DEFINITIONS_TOKENS - discount + }; + let choice_tokens = match tool_choice { + Some(ToolChoice::Mode(mode)) if mode == "none" => TOOL_CHOICE_NONE_TOKENS, + Some(ToolChoice::Mode(_)) | None => 0, + Some(ToolChoice::Named(named)) => { + NAMED_TOOL_CHOICE_TOKENS + self.count_text(&named.function.name)? + } + }; + Ok(REPLY_PRIMING_TOKENS + tool_tokens + choice_tokens) + } +} diff --git a/litellm-rust/crates/token-counter/src/error.rs b/litellm-rust/crates/token-counter/src/error.rs new file mode 100644 index 00000000000..dc590093f64 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/error.rs @@ -0,0 +1,31 @@ +use std::string::FromUtf8Error; + +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("unsupported by the rust token counter: request body could not be parsed: {0}")] + RequestParse(#[source] serde_json::Error), + #[error("unsupported by the rust token counter: request has no countable input")] + MissingInput, + #[error( + "unsupported by the rust token counter: float text values are counted by the python path" + )] + FloatText, + #[error( + "unsupported by the rust token counter: content block type is counted by the python path" + )] + ContentBlock, + #[error("unsupported by the rust token counter: array parameter without items")] + ArrayItems, + #[error("unsupported by the rust token counter: text value could not be serialized: {0}")] + JsonSerialization(#[source] serde_json::Error), + #[error("unsupported by the rust token counter: serialized text value is not UTF-8: {0}")] + JsonUtf8(#[source] FromUtf8Error), + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), + #[error("token counting task failed: {0}")] + Task(String), +} diff --git a/litellm-rust/crates/token-counter/src/lib.rs b/litellm-rust/crates/token-counter/src/lib.rs new file mode 100644 index 00000000000..eb602b3cade --- /dev/null +++ b/litellm-rust/crates/token-counter/src/lib.rs @@ -0,0 +1,17 @@ +//! Input token counting for a request body, mirroring `litellm.token_counter` +//! for the shapes it can count exactly. Everything else is declined so the host +//! keeps its own counter as the reference. + +#![forbid(unsafe_code)] + +mod byte_level; +mod counter; +mod error; +mod python_json; +mod tools; +mod types; +mod unicode_classes; + +pub use counter::{InputTokenCount, TokenCounter}; +pub use error::Error; +pub use types::CountableRequest; diff --git a/litellm-rust/crates/token-counter/src/python_json.rs b/litellm-rust/crates/token-counter/src/python_json.rs new file mode 100644 index 00000000000..00a55ad300c --- /dev/null +++ b/litellm-rust/crates/token-counter/src/python_json.rs @@ -0,0 +1,155 @@ +//! `json.dumps(value)` with Python's default arguments: `", "` and `": "` +//! separators, `ensure_ascii=True`, and keys in insertion order. + +use std::io::{self, Write}; + +use serde::Serialize; +use serde_json::ser::{Formatter, Serializer}; + +use super::Error; +use super::types::TextValue; + +pub(super) fn dumps(value: &TextValue) -> Result { + let mut output = Vec::with_capacity(serialized_len(value)?); + value + .serialize(&mut Serializer::with_formatter( + &mut output, + PythonFormatter, + )) + .map_err(Error::JsonSerialization)?; + debug_assert_eq!(output.len(), output.capacity()); + String::from_utf8(output).map_err(Error::JsonUtf8) +} + +fn serialized_len(value: &TextValue) -> Result { + match value { + TextValue::Null => Ok(4), + TextValue::Bool(true) => Ok(4), + TextValue::Bool(false) => Ok(5), + TextValue::Number(number) => match (number.as_i64(), number.as_u64()) { + (Some(number), _) => Ok(unsigned_len(number.unsigned_abs()) + usize::from(number < 0)), + (_, Some(number)) => Ok(unsigned_len(number)), + _ => Err(Error::FloatText), + }, + TextValue::Text(text) => Ok(quoted_len(text)), + TextValue::List(items) => items + .iter() + .try_fold(2 + items.len().saturating_sub(1) * 2, |len, item| { + Ok(len + serialized_len(item)?) + }), + TextValue::Object(entries) => entries.iter().try_fold( + 2 + entries.len().saturating_sub(1) * 2, + |len, (key, value)| Ok(len + quoted_len(key) + 2 + serialized_len(value)?), + ), + } +} + +fn unsigned_len(number: u64) -> usize { + if number == 0 { + 1 + } else { + number.ilog10() as usize + 1 + } +} + +fn quoted_len(value: &str) -> usize { + value.chars().fold(2, |len, character| { + len + match character { + '"' | '\\' | '\u{0008}' | '\u{000c}' | '\n' | '\r' | '\t' => 2, + '\u{0000}'..='\u{001f}' => 6, + ' '..='~' => 1, + _ => character.len_utf16() * 6, + } + }) +} + +struct PythonFormatter; + +impl Formatter for PythonFormatter { + fn begin_array_value(&mut self, writer: &mut W, first: bool) -> io::Result<()> + where + W: ?Sized + Write, + { + if !first { + writer.write_all(b", ")?; + } + Ok(()) + } + + fn begin_object_key(&mut self, writer: &mut W, first: bool) -> io::Result<()> + where + W: ?Sized + Write, + { + if !first { + writer.write_all(b", ")?; + } + Ok(()) + } + + fn begin_object_value(&mut self, writer: &mut W) -> io::Result<()> + where + W: ?Sized + Write, + { + writer.write_all(b": ") + } + + fn write_string_fragment(&mut self, writer: &mut W, fragment: &str) -> io::Result<()> + where + W: ?Sized + Write, + { + for character in fragment.chars() { + if (' '..='~').contains(&character) { + write!(writer, "{character}")?; + continue; + } + let mut units = [0u16; 2]; + for unit in character.encode_utf16(&mut units) { + write!(writer, "\\u{unit:04x}")?; + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::dumps; + + #[rstest] + #[case::null("null", "null")] + #[case::boolean("true", "true")] + #[case::signed_integer("-3", "-3")] + #[case::unsigned_integer("18446744073709551615", "18446744073709551615")] + #[case::empty_array("[]", "[]")] + #[case::empty_object("{}", "{}")] + #[case::nested( + r#"{"first":1,"second":{"ok":true,"none":null},"third":[false,2]}"#, + r#"{"first": 1, "second": {"ok": true, "none": null}, "third": [false, 2]}"# + )] + #[case::string_escaping( + r#""caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~ /""#, + r#""caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~ /""# + )] + #[case::short_control_escapes(r#""\b\f\r""#, r#""\b\f\r""#)] + fn matches_python_json_dumps(#[case] input: &str, #[case] expected: &str) { + let value = serde_json::from_str(input).expect("fixture parses"); + + assert_eq!(dumps(&value).expect("fixture dumps"), expected); + } + + #[rstest] + #[case::top_level("1.5")] + #[case::array("[1,2.5]")] + #[case::object(r#"{"nested":{"value":-0.25}}"#)] + fn rejects_floats(#[case] input: &str) { + let value = serde_json::from_str(input).expect("fixture parses"); + let error = dumps(&value).expect_err("floats are declined"); + + assert_eq!( + error.to_string(), + "unsupported by the rust token counter: float text values are counted by the python path" + ); + } +} diff --git a/litellm-rust/crates/token-counter/src/tools.rs b/litellm-rust/crates/token-counter/src/tools.rs new file mode 100644 index 00000000000..22150a4e4b1 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/tools.rs @@ -0,0 +1,304 @@ +//! Renders tool definitions the way `litellm.token_counter` does before +//! tokenizing them (the TypeScript-like namespace OpenAI appears to use). + +use std::fmt::Write; + +use super::Error; +use super::types::{EnumValue, Schema, SchemaType, ToolDefinition}; + +pub(super) fn format_function_definitions(tools: &[ToolDefinition]) -> Result { + ToolFormatter::new(tools.len()).format(tools) +} + +struct ToolFormatter { + output: String, +} + +impl ToolFormatter { + fn new(tool_count: usize) -> Self { + Self { + output: String::with_capacity(tool_count.saturating_mul(128).saturating_add(48)), + } + } + + fn format(mut self, tools: &[ToolDefinition]) -> Result { + self.output.push_str("namespace functions {\n\n"); + + for tool in tools { + self.write_function(tool)?; + } + + self.output.push_str("} // namespace functions"); + Ok(self.output) + } + + fn write_function(&mut self, tool: &ToolDefinition) -> Result<(), Error> { + let (name, description, parameters) = resolve_function(tool); + let Some(name) = name.filter(|name| !name.is_empty()) else { + return Ok(()); + }; + + if let Some(description) = description.filter(|description| !description.is_empty()) { + self.output.push_str("// "); + self.output.push_str(description); + self.output.push('\n'); + } + + match parameters.filter(|parameters| { + parameters + .properties + .as_ref() + .is_some_and(|properties| !properties.is_empty()) + }) { + Some(parameters) => { + self.output.push_str("type "); + self.output.push_str(name); + self.output.push_str(" = (_: {\n"); + self.write_object_parameters(parameters, 0)?; + self.output.push_str("\n}) => any;\n\n"); + } + _ => { + self.output.push_str("type "); + self.output.push_str(name); + self.output.push_str(" = () => any;\n\n"); + } + } + + Ok(()) + } + + fn write_object_parameters(&mut self, parameters: &Schema, indent: usize) -> Result<(), Error> { + let Some(properties) = parameters + .properties + .as_ref() + .filter(|properties| !properties.is_empty()) + else { + return Ok(()); + }; + let required = parameters.required.as_deref().unwrap_or_default(); + for (index, (key, props)) in properties.iter().enumerate() { + if index > 0 { + self.output.push('\n'); + } + if let Some(description) = props + .description + .as_deref() + .filter(|description| !description.is_empty()) + { + self.write_indent(indent); + self.output.push_str("// "); + self.output.push_str(description); + self.output.push('\n'); + } + + self.write_indent(indent); + self.output.push_str(key); + if !required.iter().any(|required| required == key) { + self.output.push('?'); + } + self.output.push_str(": "); + self.write_type(props, indent)?; + self.output.push(','); + } + + Ok(()) + } + + fn write_type(&mut self, props: &Schema, indent: usize) -> Result<(), Error> { + let Some(SchemaType::Name(schema_type)) = &props.schema_type else { + self.output.push_str("any"); + return Ok(()); + }; + + match schema_type.as_str() { + "string" | "integer" | "number" => match &props.enum_values { + Some(values) => self.write_enum(values), + None if schema_type == "string" => self.output.push_str("string"), + None => self.output.push_str("number"), + }, + "array" => { + let items = props.items.as_deref().ok_or(Error::ArrayItems)?; + self.write_type(items, indent)?; + self.output.push_str("[]"); + } + "object" => { + self.output.push_str("{\n"); + self.write_object_parameters(props, indent + 2)?; + self.output.push_str("\n}"); + } + "boolean" => self.output.push_str("boolean"), + "null" => self.output.push_str("null"), + _ => self.output.push_str("any"), + } + + Ok(()) + } + + fn write_enum(&mut self, values: &[EnumValue]) { + for (index, value) in values.iter().enumerate() { + if index > 0 { + self.output.push_str(" | "); + } + self.output.push('"'); + match value { + EnumValue::Text(text) => self.output.push_str(text), + EnumValue::Integer(number) => { + write!(self.output, "{number}").expect("writing to a String cannot fail"); + } + } + self.output.push('"'); + } + } + + fn write_indent(&mut self, indent: usize) { + for _ in 0..indent { + self.output.push(' '); + } + } +} + +fn resolve_function(tool: &ToolDefinition) -> (Option<&str>, Option<&str>, Option<&Schema>) { + match &tool.function { + Some(function) => ( + function.name.as_deref(), + function.description.as_deref(), + function.parameters.as_ref(), + ), + None => ( + tool.name.as_deref(), + tool.description.as_deref(), + tool.input_schema.as_ref().or(tool.parameters.as_ref()), + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn parse_tools(json: &str) -> Vec { + serde_json::from_str(json).expect("tool fixture parses") + } + + #[test] + fn empty_tool_list_renders_an_empty_namespace() { + assert_eq!( + format_function_definitions(&[]).expect("empty tool list renders"), + "namespace functions {\n\n} // namespace functions" + ); + } + + #[test] + fn unnamed_tools_are_skipped_and_function_shape_takes_precedence() { + let tools = parse_tools( + r#"[ + {}, + {"name":""}, + {"name":"ignored","function":{"description":"missing a function name"}}, + {"name":"ping","description":""} + ]"#, + ); + + assert_eq!( + format_function_definitions(&tools).expect("tools render"), + "namespace functions {\n\ntype ping = () => any;\n\n} // namespace functions" + ); + } + + #[test] + fn empty_or_missing_properties_render_a_no_argument_function() { + let tools = parse_tools( + r#"[ + {"name":"missing","input_schema":{"type":"object"}}, + {"name":"empty","input_schema":{"type":"object","properties":{}}} + ]"#, + ); + + assert_eq!( + format_function_definitions(&tools).expect("tools render"), + concat!( + "namespace functions {\n\n", + "type missing = () => any;\n\n", + "type empty = () => any;\n\n", + "} // namespace functions" + ) + ); + } + + #[test] + fn anthropic_parameters_render_all_supported_types() { + let tools = parse_tools( + r#"[{ + "name":"inspect", + "description":"Inspect a value", + "input_schema":{ + "type":"object", + "properties":{ + "text":{"type":"string"}, + "count":{"type":"integer","description":"Number of attempts"}, + "ratio":{"type":"number"}, + "enabled":{"type":"boolean"}, + "nothing":{"type":"null"}, + "unknown":{"type":"custom"}, + "union":{"type":["string","null"]}, + "labels":{"type":"array","items":{"type":"string"}}, + "config":{"type":"object","properties":{"retries":{"type":"integer"}},"required":["retries"]}, + "mode":{"type":"string","enum":["fast",2]} + }, + "required":["text"] + } + }]"#, + ); + + assert_eq!( + format_function_definitions(&tools).expect("tool renders"), + concat!( + "namespace functions {\n\n", + "// Inspect a value\n", + "type inspect = (_: {\n", + "text: string,\n", + "// Number of attempts\n", + "count?: number,\n", + "ratio?: number,\n", + "enabled?: boolean,\n", + "nothing?: null,\n", + "unknown?: any,\n", + "union?: any,\n", + "labels?: string[],\n", + "config?: {\n", + " retries: number,\n", + "},\n", + "mode?: \"fast\" | \"2\",\n", + "}) => any;\n\n", + "} // namespace functions" + ) + ); + } + + #[test] + fn input_schema_takes_precedence_over_legacy_parameters() { + let tools = parse_tools( + r#"[{ + "name":"choose", + "input_schema":{"type":"object","properties":{"current":{"type":"string"}}}, + "parameters":{"type":"object","properties":{"legacy":{"type":"string"}}} + }]"#, + ); + + let rendered = format_function_definitions(&tools).expect("tool renders"); + assert!(rendered.contains("current?: string,")); + assert!(!rendered.contains("legacy")); + } + + #[test] + fn array_without_items_returns_an_error() { + let tools = parse_tools( + r#"[{"name":"broken","parameters":{"type":"object","properties":{"values":{"type":"array"}}}}]"#, + ); + + assert!(matches!( + format_function_definitions(&tools), + Err(Error::ArrayItems) + )); + } +} diff --git a/litellm-rust/crates/token-counter/src/types.rs b/litellm-rust/crates/token-counter/src/types.rs new file mode 100644 index 00000000000..d25554beaac --- /dev/null +++ b/litellm-rust/crates/token-counter/src/types.rs @@ -0,0 +1,238 @@ +use std::fmt; + +use indexmap::IndexMap; +use serde::de::{MapAccess, SeqAccess, Visitor}; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::Number; + +use super::Error; + +/// The parts of a request body the host's budget counter reads. Chat and +/// Anthropic Messages bodies carry `messages`; completions carry `prompt`; +/// Responses and embeddings carry `input`; rerank carries `query` and +/// `documents`. The host checks key presence, not nullness, so an explicit +/// `null` is kept distinct from an absent key. Anything outside this shape is +/// declined so the host can fall back to its own counter instead of silently +/// miscounting. +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub struct CountableRequest { + pub(crate) model: Option, + #[serde(default, deserialize_with = "present_messages")] + pub(crate) messages: Option>, + pub(crate) tools: Option>, + pub(crate) tool_choice: Option, + #[serde(default, deserialize_with = "present_text")] + pub(crate) prompt: Option, + #[serde(default, deserialize_with = "present_text")] + pub(crate) input: Option, + #[serde(default, deserialize_with = "present_text")] + pub(crate) query: Option, + #[serde(default, deserialize_with = "present_text")] + pub(crate) documents: Option, +} + +impl CountableRequest { + pub fn parse(body: &[u8]) -> Result { + serde_json::from_slice(body).map_err(Error::RequestParse) + } +} + +fn present_messages<'de, D: Deserializer<'de>>( + deserializer: D, +) -> Result>, D::Error> { + Option::>::deserialize(deserializer) + .map(|messages| Some(messages.unwrap_or_default())) +} + +fn present_text<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + TextValue::deserialize(deserializer).map(Some) +} + +/// Free-form JSON the host counts as text: strings and integers via `str()`, +/// objects via `json.dumps()`, lists flattened. Objects keep document order so +/// the dumped text matches Python byte for byte. +#[derive(Clone, Debug, PartialEq, Serialize)] +#[serde(untagged)] +pub(crate) enum TextValue { + Null, + Bool(bool), + Number(Number), + Text(String), + List(Vec), + Object(IndexMap), +} + +impl<'de> Deserialize<'de> for TextValue { + fn deserialize>(deserializer: D) -> Result { + deserializer.deserialize_any(TextValueVisitor) + } +} + +struct TextValueVisitor; + +impl<'de> Visitor<'de> for TextValueVisitor { + type Value = TextValue; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a JSON value") + } + + fn visit_unit(self) -> Result { + Ok(TextValue::Null) + } + + fn visit_none(self) -> Result { + Ok(TextValue::Null) + } + + fn visit_bool(self, value: bool) -> Result { + Ok(TextValue::Bool(value)) + } + + fn visit_i64(self, value: i64) -> Result { + Ok(TextValue::Number(value.into())) + } + + fn visit_u64(self, value: u64) -> Result { + Ok(TextValue::Number(value.into())) + } + + fn visit_f64(self, value: f64) -> Result + where + E: serde::de::Error, + { + Number::from_f64(value) + .map(TextValue::Number) + .ok_or_else(|| E::custom("non-finite JSON number")) + } + + fn visit_str(self, value: &str) -> Result { + Ok(TextValue::Text(value.to_owned())) + } + + fn visit_string(self, value: String) -> Result { + Ok(TextValue::Text(value)) + } + + fn visit_seq(self, mut sequence: A) -> Result + where + A: SeqAccess<'de>, + { + let mut items = Vec::with_capacity(sequence.size_hint().unwrap_or(0)); + while let Some(item) = sequence.next_element()? { + items.push(item); + } + Ok(TextValue::List(items)) + } + + fn visit_map(self, mut map: A) -> Result + where + A: MapAccess<'de>, + { + let mut entries = IndexMap::with_capacity(map.size_hint().unwrap_or(0)); + while let Some((key, value)) = map.next_entry()? { + entries.insert(key, value); + } + Ok(TextValue::Object(entries)) + } +} + +/// Python counts every string-valued key of a message, so any key beyond these +/// makes the shape unsupported rather than silently uncounted. +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(deny_unknown_fields)] +pub(crate) struct Message { + pub(crate) role: Option, + pub(crate) name: Option, + pub(crate) content: Option, +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum MessageContent { + Text(String), + Blocks(Vec), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum ContentItem { + Text(String), + Block(ContentBlock), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(tag = "type")] +pub(crate) enum ContentBlock { + #[serde(rename = "text")] + Text { text: String }, + #[serde(rename = "thinking")] + Thinking { thinking: String }, + #[serde(rename = "tool_reference")] + ToolReference { tool_name: Option }, + /// Images, documents, files and tool use/result blocks price through + /// Python-only helpers, so they stay on the Python counter. + #[serde(other)] + Unsupported, +} + +/// Either the OpenAI `{"type": "function", "function": {...}}` shape or the +/// Anthropic `{"name", "description", "input_schema"}` shape. +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub(crate) struct ToolDefinition { + pub(crate) function: Option, + pub(crate) name: Option, + pub(crate) description: Option, + pub(crate) input_schema: Option, + pub(crate) parameters: Option, +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub(crate) struct FunctionDefinition { + pub(crate) name: Option, + pub(crate) description: Option, + pub(crate) parameters: Option, +} + +#[derive(Clone, Debug, Default, Deserialize, PartialEq)] +pub(crate) struct Schema { + #[serde(rename = "type")] + pub(crate) schema_type: Option, + pub(crate) description: Option, + #[serde(rename = "enum")] + pub(crate) enum_values: Option>, + pub(crate) items: Option>, + pub(crate) properties: Option>, + pub(crate) required: Option>, +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum SchemaType { + Name(String), + Union(Vec), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum EnumValue { + Text(String), + Integer(i64), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum ToolChoice { + Mode(String), + Named(NamedToolChoice), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub(crate) struct NamedToolChoice { + pub(crate) function: NamedFunction, +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub(crate) struct NamedFunction { + pub(crate) name: String, +} diff --git a/litellm-rust/crates/token-counter/src/unicode_classes.rs b/litellm-rust/crates/token-counter/src/unicode_classes.rs new file mode 100644 index 00000000000..405cee34949 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/unicode_classes.rs @@ -0,0 +1,73 @@ +use std::cmp::Ordering; +use std::sync::LazyLock; + +use tokenizers::utils::SysRegex; + +struct Ranges(Box<[(u32, u32)]>); + +pub(super) struct UnicodeClasses { + letters: Ranges, + numbers: Ranges, + spaces: Ranges, +} + +static CLASSES: LazyLock> = LazyLock::new(|| { + let scalars: String = (0..=u32::from(char::MAX)) + .filter_map(char::from_u32) + .collect(); + Some(UnicodeClasses { + letters: Ranges::load(r"\p{L}+", &scalars)?, + numbers: Ranges::load(r"\p{N}+", &scalars)?, + spaces: Ranges::load(r"\s+", &scalars)?, + }) +}); + +impl Ranges { + fn load(pattern: &str, scalars: &str) -> Option { + let regex = SysRegex::new(pattern).ok()?; + let ranges = regex + .find_iter(scalars) + .map(|(start, end)| { + let matched = scalars.get(start..end)?; + Some(( + u32::from(matched.chars().next()?), + u32::from(matched.chars().next_back()?), + )) + }) + .collect::>>()?; + Some(Self(ranges)) + } + + fn contains(&self, character: char) -> bool { + let code = u32::from(character); + self.0 + .binary_search_by(|(low, high)| { + if *high < code { + Ordering::Less + } else if *low > code { + Ordering::Greater + } else { + Ordering::Equal + } + }) + .is_ok() + } +} + +impl UnicodeClasses { + pub(super) fn get() -> Option<&'static Self> { + CLASSES.as_ref() + } + + pub(super) fn is_letter(&self, character: char) -> bool { + self.letters.contains(character) + } + + pub(super) fn is_number(&self, character: char) -> bool { + self.numbers.contains(character) + } + + pub(super) fn is_space(&self, character: char) -> bool { + self.spaces.contains(character) + } +} diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs new file mode 100644 index 00000000000..0e71a5c03cf --- /dev/null +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -0,0 +1,176 @@ +use rstest::rstest; + +use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter}; + +/// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)` +/// so this test also guards Python parity. +fn counter() -> TokenCounter { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" + ); + let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo"); + TokenCounter::from_json(&json).expect("anthropic tokenizer loads") +} + +const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#; + +const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[ + {"role":"system","content":"You are a terse assistant."}, + {"role":"user","name":"alice","content":[ + {"type":"text","text":"Summarise this paragraph about ships and harbours."}, + "plain string item", + {"type":"thinking","thinking":"pondering"}, + {"type":"tool_reference","tool_name":"get_weather"}]}, + {"role":"assistant","content":[{"type":"text","text":"Sure.","cache_control":{"type":"ephemeral"}}]}]}"#; + +const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}], + "tools":[ + {"type":"function","function":{"name":"get_weather","description":"Get weather","parameters":{ + "type":"object", + "properties":{ + "location":{"type":"string","description":"City name"}, + "unit":{"type":"string","enum":["celsius","fahrenheit"]}, + "days":{"type":"integer"}, + "tags":{"type":"array","items":{"type":"string"}}, + "opts":{"type":"object","properties":{"verbose":{"type":"boolean"},"level":{"type":"integer","enum":[1,2]}},"required":["verbose"]}, + "anything":{}}, + "required":["location"]}}}, + {"type":"function","function":{"name":"noop"}}], + "tool_choice":{"type":"function","function":{"name":"get_weather"}}}"#; + +const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5", + "messages":[{"role":"system","content":"sys"},{"role":"user","content":"weather?"}], + "tools":[{"name":"get_weather","description":"Get weather","input_schema":{ + "type":"object","properties":{"location":{"type":["string","null"]}},"required":["location"]}}], + "tool_choice":"none"}"#; + +const COMPLETIONS_PROMPT: &str = + r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#; + +const COMPLETIONS_PROMPT_LIST: &str = + r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#; + +const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[ + {"role":"user","content":[{"type":"input_text","text":"Summarise caf\u00e9 menus, na\u00efve \u2014 ok? \"quoted\"\n"}]}, + {"role":"assistant","content":"Sure."}],"instructions":"be terse"}"#; + +const EMBEDDINGS_TOKEN_IDS: &str = + r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#; + +const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", + "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#; + +/// Expected counts are pinned from +/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`. +#[rstest] +#[case::text_only(SIMPLE, 14)] +#[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)] +#[case::openai_tools_named_choice(TOOLS_OPENAI, 123)] +#[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)] +#[case::completions_prompt(COMPLETIONS_PROMPT, 7)] +#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)] +#[case::responses_input_items(RESPONSES_INPUT, 62)] +#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)] +#[case::rerank_query_and_documents(RERANK, 41)] +fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) { + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let count = counter().count_request(&request).expect("fixture counts"); + assert_eq!( + count, + InputTokenCount { + model: Some("claude-sonnet-4-5".to_string()), + input_tokens: expected, + } + ); +} + +#[rstest] +#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)] +#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] +#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] +#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] +fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let count = counter().count_request(&request).expect("fixture counts"); + assert_eq!(count.input_tokens, expected); +} + +#[rstest] +#[case::not_json(b"not json" as &[u8])] +#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)] +#[case::message_with_tool_calls( + br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"# +)] +#[case::dict_content( + br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"# +)] +#[case::float_enum( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"# +)] +#[case::anthropic_tool_choice_without_function( + br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"# +)] +fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { + assert!(matches!( + CountableRequest::parse(body), + Err(Error::RequestParse(_)) + )); +} + +#[rstest] +#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] +#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] +#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] +#[case::image_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# +)] +#[case::tool_result_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# +)] +#[case::array_without_items( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# +)] +fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { + let request = CountableRequest::parse(body).expect("shape parses"); + assert!(matches!( + counter().count_request(&request), + Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) + )); +} + +#[test] +fn tool_choice_and_system_discount_change_the_count() { + let counter = counter(); + let count = |body: &str| { + counter + .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) + .expect("counts") + .input_tokens + }; + let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#), + base + 1 + ); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"#), + base + ); + let with_tools = count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#, + ); + let with_tools_and_system = count( + r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#, + ); + assert_eq!(with_tools - with_tools_and_system, 4); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#), + base + ); +} + +#[test] +fn loading_a_bad_tokenizer_is_a_load_error() { + assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_)))); +} diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fd..1dfd146a00e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -45,8 +45,11 @@ from typing import ( TYPE_CHECKING, Union, ) +from collections.abc import Mapping from litellm.types.integrations.datadog import DatadogInitParams from litellm.types.integrations.newrelic import NewRelicInitParams +from litellm.litellm_core_utils.core_helpers import drop_params_env_flag +from litellm.types.integrations.pointfive import PointFiveInitParams from litellm._logging import ( set_verbose, _turn_on_debug, @@ -153,6 +156,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "smtp_email", "deepeval", "s3_v2", + "pointfive", "aws_sqs", "vector_store_pre_call_hook", "dotprompt", @@ -238,7 +242,7 @@ token: Optional[str] = ( ) telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults -drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) +drop_params = drop_params_env_flag(os.environ, verbose_logger) modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) use_chat_completions_url_for_anthropic_messages: bool = bool( os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) @@ -325,6 +329,9 @@ ssl_certificate: Optional[str] = None user_url_validation: bool = True user_url_allowed_hosts: List[str] = [] provider_url_destination_allowed_hosts: List[str] = [] +#: "override" (default) or "additive": whether a key or team destination replaces +#: the operator's exporter for that backend or exports alongside it. +otel_tenant_destination_mode: str | None = None ssl_ecdh_curve: Optional[str] = None # Set to 'X25519' to disable PQC and improve performance disable_streaming_logging: bool = False disable_token_counter: bool = False @@ -435,6 +442,7 @@ s3_audit_callback_params: Optional[Dict] = None datadog_llm_observability_params: Optional[Union[DatadogLLMObsInitParams, Dict]] = None datadog_params: Optional[Union[DatadogInitParams, Dict]] = None newrelic_params: Optional[Union[NewRelicInitParams, Dict]] = None +pointfive_params: Optional[Union[PointFiveInitParams, Mapping[str, object]]] = None aws_sqs_callback_params: Optional[Dict] = None generic_logger_headers: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None @@ -471,6 +479,7 @@ prometheus_metrics_config: Optional[List] = None prometheus_exclude_metrics: Optional[List[str]] = None prometheus_exclude_labels: Optional[List[str]] = None prometheus_emit_stream_label: bool = False +prometheus_emit_input_sequence_length_label: bool = False prometheus_deployment_and_latency_caller_identity: Literal[ "api_key_alias", "user_email", @@ -495,6 +504,7 @@ public_model_groups: Optional[List[str]] = None public_agent_groups: Optional[List[str]] = None agent_search_embedding_model: Optional[str] = None mcp_tool_search: Optional[Mapping[str, object]] = None +skill_search_embedding_model: Optional[str] = None # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]]) # New format: { "displayName": { "url": "...", "index": 0 } } # Old format: { "displayName": "url" } (for backward compatibility) @@ -541,7 +551,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None #### PII MASKING #### output_parse_pii: bool = False ############################################# -from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map +from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete model_cost = get_model_cost_map(url=model_cost_map_url) cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount @@ -1360,6 +1370,7 @@ from .exceptions import ( InvalidRequestError, BadRequestError, ImageFetchError, + VectorStoreSearchError, NotFoundError, PermissionDeniedError, RateLimitError, @@ -1461,9 +1472,11 @@ from .vector_stores.vector_store_registry import ( VectorStoreRegistry, VectorStoreIndexRegistry, ) +from .types.vector_stores import VectorStoreSearchFailureMode vector_store_registry: Optional[VectorStoreRegistry] = None vector_store_index_registry: Optional[VectorStoreIndexRegistry] = None +vector_store_search_failure_mode: VectorStoreSearchFailureMode = "annotate" ### RAG ### from . import rag @@ -2001,6 +2014,9 @@ if TYPE_CHECKING: from .llms.hosted_vllm.responses.transformation import ( HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig, ) + from .llms.fireworks_ai.responses.transformation import ( + FireworksAIResponsesAPIConfig as FireworksAIResponsesAPIConfig, + ) from .llms.github_copilot.chat.transformation import ( GithubCopilotConfig as GithubCopilotConfig, ) @@ -2397,3 +2413,5 @@ def __getattr__(name: str) -> Any: # ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time + +mark_litellm_import_complete() diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index f856fe0f2b3..8132008731f 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -6,9 +6,33 @@ be settable from user input. Context variables are scoped to the current asyncio task and cannot be injected via HTTP request bodies. """ +from collections.abc import Generator +from contextlib import contextmanager from contextvars import ContextVar +from datetime import datetime, timezone from typing import Final # When True, suppresses async logging and billing for internal sub-calls # (e.g., emulated file-search steps that make nested LLM calls). is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", default=False) + +# One request prices its totals, its per-token-type lines and the rates it reports on +# separate code paths. Each reads the clock for off-peak pricing, so without a pinned +# moment they can land on either side of a window boundary and disagree with each other. +_billing_time: Final[ContextVar[datetime | None]] = ContextVar("billing_time", default=None) + + +@contextmanager +def pinned_billing_time(moment: datetime) -> Generator[None]: + """Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read.""" + token: Final = _billing_time.set(moment) + try: + yield + finally: + _billing_time.reset(token) + + +def current_billing_time() -> datetime: + """The pinned billing moment, or now in UTC outside a pinned block.""" + pinned: Final = _billing_time.get() + return pinned if pinned is not None else datetime.now(timezone.utc) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index e9199e1ec80..dc323c8cc15 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -237,6 +237,7 @@ LLM_CONFIG_NAMES: Final = ( "XAIResponsesAPIConfig", "LiteLLMProxyResponsesAPIConfig", "HostedVLLMResponsesAPIConfig", + "FireworksAIResponsesAPIConfig", "VolcEngineResponsesAPIConfig", "PerplexityResponsesConfig", "DatabricksResponsesAPIConfig", @@ -957,6 +958,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.hosted_vllm.responses.transformation", "HostedVLLMResponsesAPIConfig", ), + "FireworksAIResponsesAPIConfig": ( + ".llms.fireworks_ai.responses.transformation", + "FireworksAIResponsesAPIConfig", + ), "VolcEngineResponsesAPIConfig": ( ".llms.volcengine.responses.transformation", "VolcEngineResponsesAPIConfig", diff --git a/litellm/_redis.py b/litellm/_redis.py index 3e68d50cf16..c5acdcb038b 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -24,6 +24,7 @@ from redis.credentials import CredentialProvider from litellm import get_secret, get_secret_str from litellm._redis_credential_provider import ( AzureADCredentialProvider, + ElastiCacheIAMCredentialProvider, GCPIAMCredentialProvider, _generate_gcp_iam_access_token, ) @@ -38,6 +39,14 @@ from ._logging import verbose_logger AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default" +_AWS_IAM_KWARG_NAMES: Final = ( + "aws_iam_auth", + "aws_iam_user_name", + "aws_iam_cache_name", + "aws_iam_region", + "aws_iam_serverless", +) + def _unwrapped_init_args(cls: type) -> frozenset[str]: """Every parameter on a single class's own ``__init__``, decorator-unwrapped. @@ -75,6 +84,7 @@ def _get_redis_kwargs(): "azure_client_id", "azure_tenant_id", "azure_client_secret", + *_AWS_IAM_KWARG_NAMES, } available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args @@ -270,6 +280,42 @@ def _redis_kwargs_from_environment(): return return_dict +def _coerces_to_true(value: object | None) -> bool: + return _str_to_bool(value) if isinstance(value, str) else bool(value) + + +def _uses_tls(redis_kwargs: Mapping[str, object]) -> bool: + if redis_kwargs.get("startup_nodes") is not None: + return _coerces_to_true(redis_kwargs.get("ssl")) + url: Final = redis_kwargs.get("url") + if isinstance(url, str): + return urlsplit(url).scheme.lower() == "rediss" + return _coerces_to_true(redis_kwargs.get("ssl")) + + +def _build_elasticache_iam_provider(redis_kwargs: Mapping[str, object]) -> ElastiCacheIAMCredentialProvider: + user_name: Final = redis_kwargs.get("aws_iam_user_name") + cache_name: Final = redis_kwargs.get("aws_iam_cache_name") + region: Final = ( + redis_kwargs.get("aws_iam_region") or get_secret_str("AWS_REGION") or get_secret_str("AWS_DEFAULT_REGION") + ) + required_settings: Final = ( + ("aws_iam_user_name", user_name), + ("aws_iam_cache_name", cache_name), + ("aws_iam_region", region), + ) + missing_settings: Final = tuple(name for name, value in required_settings if not value) + if missing_settings: + raise ValueError("AWS ElastiCache IAM Redis authentication requires: " + ", ".join(missing_settings)) + + return ElastiCacheIAMCredentialProvider( + user_name=str(user_name), + cache_name=str(cache_name), + region=str(region), + is_serverless=_coerces_to_true(redis_kwargs.get("aws_iam_serverless")), + ) + + def create_gcp_iam_redis_connect_func( service_account: str, ssl_ca_certs: str | None = None, @@ -540,6 +586,7 @@ def _get_redis_client_logic(**env_overrides): _azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN") _azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true" + _aws_iam_enabled: Final = _coerces_to_true(redis_kwargs.get("aws_iam_auth")) if _azure_ad_enabled and _gcp_service_account is not None: verbose_logger.warning( @@ -567,6 +614,22 @@ def _get_redis_client_logic(**env_overrides): # credentials via inspection or logging. redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True + if _aws_iam_enabled and _gcp_service_account is not None: + verbose_logger.warning( + "Both GCP IAM (gcp_service_account) and AWS ElastiCache IAM (aws_iam_auth) are configured " + "for Redis. Using GCP IAM. Remove one to avoid misconfiguration." + ) + elif _aws_iam_enabled and _azure_ad_enabled: + verbose_logger.warning( + "Both Azure AD (azure_redis_ad_token) and AWS ElastiCache IAM (aws_iam_auth) are configured " + "for Redis. Using Azure AD. Remove one to avoid misconfiguration." + ) + elif _aws_iam_enabled: + if not _uses_tls(redis_kwargs): + raise ValueError("AWS ElastiCache IAM Redis authentication requires TLS") + verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.") + redis_kwargs["credential_provider"] = _build_elasticache_iam_provider(redis_kwargs) + redis_kwargs.pop("gcp_service_account", None) redis_kwargs.pop("gcp_ssl_ca_certs", None) @@ -575,6 +638,8 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs.pop("azure_client_id", None) redis_kwargs.pop("azure_tenant_id", None) redis_kwargs.pop("azure_client_secret", None) + for aws_iam_key in _AWS_IAM_KWARG_NAMES: + redis_kwargs.pop(aws_iam_key, None) if redis_kwargs.get("credential_provider") is not None: redis_kwargs.pop("redis_connect_func", None) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index ba0398789a6..7d90f944657 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -1,10 +1,17 @@ +from __future__ import annotations + import asyncio import threading import time -from typing import Final, Protocol +from collections.abc import Callable +from typing import TYPE_CHECKING, Final, Protocol +from urllib.parse import urlencode from redis.credentials import CredentialProvider +if TYPE_CHECKING: + from botocore.credentials import Credentials + # Azure AD scope for Redis Cache for Azure. AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default" @@ -117,6 +124,82 @@ class GCPIAMCredentialProvider(CredentialProvider): return (token,) +_ELASTICACHE_SERVICE_NAME: Final = "elasticache" +_ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900 +_ELASTICACHE_SERVERLESS_RESOURCE_TYPE: Final = "ServerlessCache" + + +class ElastiCacheIAMCredentialProvider(CredentialProvider): + def __init__( + self, + user_name: str, + cache_name: str, + region: str, + is_serverless: bool = False, + credentials_resolver: Callable[[], Credentials | None] | None = None, + token_lifetime_seconds: int = _ELASTICACHE_TOKEN_TTL_SECONDS, + ) -> None: + self._user_name = user_name + self._cache_name = cache_name.lower() + self._region = region + self._is_serverless = is_serverless + self._credentials_resolver = credentials_resolver or self._resolve_credentials + self._credentials: Credentials | None = None + self._token_lifetime_seconds = token_lifetime_seconds + + @staticmethod + def _resolve_credentials() -> Credentials | None: + try: + import botocore.session + except ImportError as e: + raise ImportError( + "botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3" + ) from e + + return botocore.session.get_session().get_credentials() + + def _get_credentials(self) -> tuple[str, str]: + credentials: Final = self._credentials if self._credentials is not None else self._credentials_resolver() + if credentials is None: + raise RuntimeError("Unable to resolve AWS credentials for ElastiCache IAM Redis authentication") + self._credentials = credentials + + frozen_credentials: Final = credentials.get_frozen_credentials() + + try: + from botocore.auth import SigV4QueryAuth + from botocore.awsrequest import AWSRequest + except ImportError as e: + raise ImportError( + "botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3" + ) from e + + query: Final = urlencode( + ( + ("Action", "connect"), + ("User", self._user_name), + *((("ResourceType", _ELASTICACHE_SERVERLESS_RESOURCE_TYPE),) if self._is_serverless else ()), + ) + ) + request: Final = AWSRequest(method="GET", url=f"https://{self._cache_name}/?{query}") + SigV4QueryAuth( + frozen_credentials, + _ELASTICACHE_SERVICE_NAME, + self._region, + expires=self._token_lifetime_seconds, + ).add_auth(request) + signed_url: Final = request.url + if signed_url is None: + raise RuntimeError("Unable to generate AWS ElastiCache IAM credentials") + return self._user_name, signed_url.removeprefix("https://") + + def get_credentials(self) -> tuple[str, str]: + return self._get_credentials() + + async def get_credentials_async(self) -> tuple[str, str]: + return await asyncio.to_thread(self._get_credentials) + + class AzureADCredentialProvider(CredentialProvider): """ redis.credentials.CredentialProvider implementation that supplies Azure AD diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index a4e6fa50901..306a8871b12 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -13,6 +13,7 @@ from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( BedrockAgentCoreA2ATransformation, ) +from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -45,7 +46,8 @@ class BedrockAgentCoreA2AHandler: Returns: A2A JSON-RPC response dict from the AgentCore agent """ - url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + url, headers, body = await run_aws_signing( + BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, litellm_params=litellm_params, @@ -91,7 +93,8 @@ class BedrockAgentCoreA2AHandler: Yields: A2A streaming response events from the AgentCore agent """ - url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + url, headers, body = await run_aws_signing( + BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, litellm_params=litellm_params, diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 97be5f77d79..87f8fd3946e 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -10,7 +10,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.llms.openai import Batch -from litellm.types.utils import CallTypes, ModelInfo, Usage +from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -23,6 +23,29 @@ class BatchCostUsageResult: models: list[str] successful_requests: int failed_requests: int + prompt_cost: float = 0.0 + completion_cost: float = 0.0 + + +_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( @@ -130,7 +153,8 @@ class _LineOutcome(Enum): @dataclass(frozen=True, slots=True) class _BatchOutputLineStats: - cost: float + prompt_cost: float + completion_cost: float prompt_tokens: int completion_tokens: int total_tokens: int @@ -193,15 +217,16 @@ def _compute_output_line_stats( raw_model: Final = response_body.get("model") response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None completion_details: Final = usage.completion_tokens_details + line_prompt_cost, line_completion_cost = _output_line_cost( + usage=usage, + custom_llm_provider=custom_llm_provider, + model_name=model_name, + response_model=response_model, + model_info=model_info, + ) return _BatchOutputLineStats( - cost=_output_line_cost( - response_body=response_body, - usage=usage, - custom_llm_provider=custom_llm_provider, - model_name=model_name, - response_model=response_model, - model_info=model_info, - ), + prompt_cost=line_prompt_cost, + completion_cost=line_completion_cost, prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, total_tokens=usage.total_tokens, @@ -213,31 +238,24 @@ def _compute_output_line_stats( def _output_line_cost( - response_body: Mapping[str, object], usage: Usage, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, response_model: str | None, model_info: ModelInfo | None, -) -> float: +) -> tuple[float, float]: + """(prompt_cost, completion_cost) for one output line, priced at batch rates.""" from litellm.cost_calculator import batch_cost_calculator - if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"): - return litellm.completion_cost( - completion_response=response_body, - custom_llm_provider=custom_llm_provider, - call_type=CallTypes.aretrieve_batch.value, - ) cost_model: Final = ( model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or "" ) - prompt_cost, completion_cost = batch_cost_calculator( + return batch_cost_calculator( usage=usage, model=cost_model, custom_llm_provider=custom_llm_provider, model_info=model_info, ) - return prompt_cost + completion_cost def _aggregate_batch_cost_usage_models( @@ -270,7 +288,9 @@ def _aggregate_batch_cost_usage_models( **cache_token_params, ) batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model] - total_cost: Final = sum((stats.cost for stats in line_stats), 0.0) + total_prompt_cost: Final = sum((stats.prompt_cost for stats in line_stats), 0.0) + total_completion_cost: Final = sum((stats.completion_cost for stats in line_stats), 0.0) + total_cost: Final = total_prompt_cost + total_completion_cost verbose_logger.debug( "batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d", total_cost, @@ -285,6 +305,8 @@ def _aggregate_batch_cost_usage_models( models=batch_models, successful_requests=successful_requests, failed_requests=failed_requests, + prompt_cost=total_prompt_cost, + completion_cost=total_completion_cost, ) @@ -309,7 +331,8 @@ def calculate_vertex_ai_batch_cost_and_usage( """ from litellm.cost_calculator import batch_cost_calculator - total_cost = 0.0 + total_prompt_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below + total_completion_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 @@ -341,7 +364,8 @@ def calculate_vertex_ai_batch_cost_and_usage( model=actual_model_name, custom_llm_provider="vertex_ai", ) - total_cost += p_cost + c_cost + total_prompt_cost += p_cost + total_completion_cost += c_cost except Exception as e: verbose_logger.debug("vertex_ai batch cost calculation error for line: %s", str(e)) @@ -349,6 +373,7 @@ def calculate_vertex_ai_batch_cost_and_usage( completion_tokens += _completion total_tokens += _total + total_cost: Final = total_prompt_cost + total_completion_cost verbose_logger.info( "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d", total_cost, @@ -369,6 +394,8 @@ def calculate_vertex_ai_batch_cost_and_usage( models=[actual_model_name], successful_requests=successful_requests, failed_requests=failed_requests, + prompt_cost=total_prompt_cost, + completion_cost=total_completion_cost, ) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index c8360a81c7a..77a4fdebf16 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -319,6 +319,7 @@ def create_batch( timeout=timeout, max_retries=optional_params.max_retries, create_batch_data=_create_batch_request, + custom_endpoint=optional_params.get("custom_endpoint"), ) else: raise litellm.exceptions.BadRequestError( diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 884d095793c..d6dd2a073af 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -10,6 +10,7 @@ import ast import hashlib import json +import logging import time import traceback from collections.abc import Mapping @@ -32,7 +33,7 @@ from .dual_cache import DualCache # noqa: F401 from .gcs_cache import GCSCache from .in_memory_cache import InMemoryCache from .qdrant_semantic_cache import QdrantSemanticCache -from .redis_cache import RedisCache +from .redis_cache import RedisCache, log_redis_failure from .redis_cluster_cache import RedisClusterCache from .redis_semantic_cache import RedisSemanticCache from .s3_cache import S3Cache @@ -678,7 +679,7 @@ class Cache: cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) self.cache.set_cache(cache_key, cached_data, **kwargs) except Exception as e: - verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e) async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs): """ @@ -697,7 +698,7 @@ class Cache: else: await self.cache.async_set_cache(cache_key, cached_data, **kwargs) except Exception as e: - verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e) def _convert_to_cached_embedding( self, @@ -876,7 +877,7 @@ class Cache: else: await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs) except Exception as e: - verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e) def should_use_cache(self, **kwargs): """ diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 0de88eacaa5..139dcf058d2 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -1217,7 +1217,9 @@ class LLMCachingHandler: } if litellm.cache is not None: - litellm_params["preset_cache_key"] = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs) + litellm_params["preset_cache_key"] = ( + self.preset_cache_key or litellm.cache._get_preset_cache_key_from_kwargs(**kwargs) + ) else: litellm_params["preset_cache_key"] = None diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index ec17cc1d809..81e2af45686 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -8,8 +8,8 @@ Has 4 primary methods: - async_get_cache """ +import logging import time -import traceback from collections.abc import Sequence from threading import Lock from typing import TYPE_CHECKING, Any, Final @@ -22,8 +22,8 @@ from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from .base_cache import BaseCache -from .in_memory_cache import InMemoryCache -from .redis_cache import RedisCache +from .in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache +from .redis_cache import RedisCache, RedisCircuitBreakerOpenError, log_redis_failure if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -83,6 +83,9 @@ class DualCache(BaseCache): if default_redis_ttl is not None: self.default_redis_ttl = default_redis_ttl + def update_in_memory_max_size(self, max_size: int | None) -> None: + self.in_memory_cache.max_size_in_memory = DEFAULT_MAX_SIZE_IN_MEMORY if max_size is None else max_size + def attach_redis_cache( self, redis_cache: RedisCache | None = None, @@ -177,8 +180,10 @@ class DualCache(BaseCache): print_verbose(f"get cache: cache result: {result}") return result - except Exception: - verbose_logger.error(traceback.format_exc()) + except Exception as e: + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Cache: exception in get_cache", e, with_traceback=True + ) def batch_get_cache( self, @@ -204,9 +209,12 @@ class DualCache(BaseCache): redis_result: Final = self.redis_cache.batch_get_cache( key_list=sublist_keys, parent_otel_span=parent_otel_span ) - except Exception: + except Exception as e: # Do not throttle subsequent callers if the Redis read fails. self._rollback_redis_batch_key_reservations(previous_access_times) + if isinstance(e, RedisCircuitBreakerOpenError): + verbose_logger.debug("LiteLLM Cache: batch_get_cache served from memory only: %s", e) + return result raise if self.in_memory_cache is not None: @@ -217,8 +225,10 @@ class DualCache(BaseCache): return list( # mutable-ok: public list contract redis_result.get(key) if value is None else value for key, value in zip(keys, result) ) - except Exception: - verbose_logger.error(traceback.format_exc()) + except Exception as e: + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Cache: exception in batch_get_cache", e, with_traceback=True + ) async def async_get_cache( self, @@ -250,8 +260,10 @@ class DualCache(BaseCache): print_verbose(f"get cache: cache result: {result}") return result - except Exception: - verbose_logger.error(traceback.format_exc()) + except Exception as e: + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async_get_cache", e, with_traceback=True + ) def _reserve_redis_batch_keys( self, @@ -319,9 +331,12 @@ class DualCache(BaseCache): redis_result: Final = await self.redis_cache.async_batch_get_cache( sublist_keys, parent_otel_span=parent_otel_span ) - except Exception: + except Exception as e: # Do not throttle subsequent callers if the Redis read fails. self._rollback_redis_batch_key_reservations(previous_access_times) + if isinstance(e, RedisCircuitBreakerOpenError): + verbose_logger.debug("LiteLLM Cache: async_batch_get_cache served from memory only: %s", e) + return result raise # Short-circuit if redis_result is None or contains only None values @@ -339,8 +354,14 @@ class DualCache(BaseCache): await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs)) return result - except Exception: - verbose_logger.error(traceback.format_exc()) + except Exception as e: + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Cache: exception in async_batch_get_cache", + e, + with_traceback=True, + ) async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): print_verbose(f"async set cache: cache key: {key}; local_only: {local_only}; value: {value}") @@ -353,10 +374,14 @@ class DualCache(BaseCache): if self.redis_cache is not None and local_only is False: await self.redis_cache.async_set_cache(key, value, **kwargs) except Exception as e: - verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e) + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async add_cache", e, with_traceback=True + ) # async_batch_set_cache - async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs): + async def async_set_cache_pipeline( + self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs + ): """ Batch write values to the cache """ @@ -372,7 +397,9 @@ class DualCache(BaseCache): cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs ) except Exception as e: - verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e) + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async add_cache", e, with_traceback=True + ) async def async_increment_cache( self, @@ -410,8 +437,10 @@ class DualCache(BaseCache): return result except Exception as e: - verbose_logger.warning( - "Redis async_increment_cache failed, falling back to in-memory result: %s", + log_redis_failure( + verbose_logger, + logging.WARNING, + "Redis async_increment_cache failed, falling back to in-memory result", e, ) return result @@ -439,8 +468,10 @@ class DualCache(BaseCache): return result except Exception as e: - verbose_logger.warning( - "Redis async_increment_cache_pipeline failed, falling back to in-memory result: %s", + log_redis_failure( + verbose_logger, + logging.WARNING, + "Redis async_increment_cache_pipeline failed, falling back to in-memory result", e, ) return result diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 38a9966f9f9..4058a3d72dd 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -24,11 +24,13 @@ from litellm.constants import MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB from .base_cache import BaseCache +DEFAULT_MAX_SIZE_IN_MEMORY: Final = 200 + class InMemoryCache(BaseCache): def __init__( self, - max_size_in_memory: int | None = 200, + max_size_in_memory: int | None = DEFAULT_MAX_SIZE_IN_MEMORY, default_ttl: int | None = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute max_size_per_item: int | None = 1024, # 1MB = 1024KB @@ -37,7 +39,7 @@ class InMemoryCache(BaseCache): max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default """ self.max_size_in_memory = ( - max_size_in_memory if max_size_in_memory is not None else 200 + max_size_in_memory if max_size_in_memory is not None else DEFAULT_MAX_SIZE_IN_MEMORY ) # set an upper bound of 200 items in-memory self.default_ttl = default_ttl or 600 self.max_size_per_item = max_size_per_item or MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB # 1MB = 1024KB diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index aaee7188d86..2c36995c4f8 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -14,12 +14,16 @@ import functools import hashlib import inspect import json +import logging import time from collections.abc import Awaitable, Callable, Sequence from contextvars import ContextVar +from dataclasses import dataclass from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( @@ -80,11 +84,29 @@ class _AsyncRedisCommands(Protocol): def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... + def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... + _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( {"", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"} ) +_INCREMENT_WITH_FLOOR_LUA: Final = ( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " + "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end " + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " + "return count" +) + +_LUA_COUNT: Final = TypeAdapter(int) +_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) + + +def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]: + return _OPTIONAL_COUNTS.validate_python( + tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values) + ) + def _get_call_stack_info(num_frames: int = 2) -> str: """ @@ -175,8 +197,14 @@ class RedisCircuitBreaker: self._timeout_streak_started_at: float | None = None self._opened_at: float | None = None self._state = self.CLOSED + self._generation = 0 _breaker_metrics().record_state_change(None, self._state) + @property + def generation(self) -> int: + """Counts state transitions, so a call can tell whether the breaker moved while it ran.""" + return self._generation + def is_open(self) -> bool: """Returns True if Redis calls should be skipped.""" if not self.enabled: @@ -229,7 +257,7 @@ class RedisCircuitBreaker: self._set_state(self.OPEN) def record_success(self) -> None: - if not self.enabled: + if not self.enabled or self._state == self.OPEN: return if self._state == self.HALF_OPEN: verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered") @@ -245,6 +273,7 @@ class RedisCircuitBreaker: _breaker_metrics().record_transition(state) _breaker_metrics().record_state_change(self._state, state) self._state = state + self._generation += 1 _RedisCallResult = TypeVar("_RedisCallResult") @@ -371,21 +400,46 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) -def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> int: - """Reject the call if the breaker is open, else return the swallowed-failure count to compare against.""" +class RedisCircuitBreakerOpenError(Exception): + pass + + +def log_redis_failure( + logger: logging.Logger, level: int, message: str, exc: BaseException, with_traceback: bool = False +) -> None: + if isinstance(exc, RedisCircuitBreakerOpenError): + logger.debug("%s: %s", message, exc) + return + logger.log(level, "%s: %s", message, exc, exc_info=exc if with_traceback else None) + + +@dataclass(frozen=True, slots=True) +class _BreakerAdmission: + swallowed_before: int + generation: int + + +def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> _BreakerAdmission: + """Reject the call if the breaker is open, else record what its success may later prove.""" if breaker.is_open(): - raise Exception(f"Redis circuit breaker is open — skipping {name}") - return _swallowed_redis_failures.get() + raise RedisCircuitBreakerOpenError(f"Redis circuit breaker is open — skipping {name}") + return _BreakerAdmission(swallowed_before=_swallowed_redis_failures.get(), generation=breaker.generation) -def _exit_circuit_breaker(breaker: RedisCircuitBreaker, swallowed_before: int) -> None: - """Record success only when nothing failed while the call ran. +def _exit_circuit_breaker(breaker: RedisCircuitBreaker, admission: _BreakerAdmission) -> None: + """Record success only when nothing failed while the call ran and the breaker has not moved since. Several Redis methods catch their own connection errors and return a default, so a - method that returned is not on its own proof of a healthy Redis. + method that returned is not on its own proof of a healthy Redis. A success also vouches + only for the breaker state that admitted the call: a call admitted before the breaker + opened, or a probe admitted before a later failure reopened it, finishes knowing nothing + about whether Redis has recovered since, so only the current probe may close the breaker. """ - if _swallowed_redis_failures.get() == swallowed_before: - breaker.record_success() + if _swallowed_redis_failures.get() != admission.swallowed_before: + return + if breaker.generation != admission.generation: + return + breaker.record_success() async def _run_under_circuit_breaker( @@ -398,14 +452,14 @@ async def _run_under_circuit_breaker( Shared by the method decorator and the Lua script executor so both feed the same health signal. """ - swallowed_before: Final = _enter_circuit_breaker(breaker, name) + admission: Final = _enter_circuit_breaker(breaker, name) try: result: Final = await call() except Exception as e: if _is_redis_health_failure(e): breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) raise - _exit_circuit_breaker(breaker, swallowed_before) + _exit_circuit_breaker(breaker, admission) return result @@ -415,14 +469,14 @@ def _run_under_circuit_breaker_sync( call: Callable[[], _RedisCallResult], ) -> _RedisCallResult: """Run one blocking Redis call under a circuit breaker, feeding the same health signal as the async path.""" - swallowed_before: Final = _enter_circuit_breaker(breaker, name) + admission: Final = _enter_circuit_breaker(breaker, name) try: result: Final = call() except Exception as e: if _is_redis_health_failure(e): - breaker.record_failure() + breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) raise - _exit_circuit_breaker(breaker, swallowed_before) + _exit_circuit_breaker(breaker, admission) return result @@ -736,6 +790,43 @@ class RedisCache(BaseCache): ) raise e + @_redis_circuit_breaker_guard_sync + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call. + + A counter whose key expired while a request was still in flight would otherwise be + recreated negative by that request's decrement. Clamping inside the same call is what + keeps it safe: a separate corrective write could land after another pod's increment and + erase it. + + The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was + created rather than ``ttl`` after it was last touched. Refreshing it on every touch + would keep a count a dead worker never decremented alive for as long as the group + takes traffic. Returns the resulting count. + """ + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval + _INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl + ) + return _LUA_COUNT.validate_python(count) + + @_redis_circuit_breaker_guard_sync + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Read integer counters for ``key_list``, in order, raising when Redis cannot answer. + + ``batch_get_cache`` swallows every failure and returns an empty dict, which the caller + cannot tell apart from "every counter is unset". A caller that has to fall back to its + own numbers when Redis is unreachable needs the failure, not a dict of zeros. + """ + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys)) + + @_redis_circuit_breaker_guard + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Async twin of ``batch_get_counts``, raising on failure the same way.""" + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) + @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: start_time: Final = time.time() @@ -1241,6 +1332,14 @@ class RedisCache(BaseCache): result = result.decode() return float(result) + @_redis_circuit_breaker_guard + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees.""" + _redis_client: Final = self._async_commands() + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl) + return _LUA_COUNT.validate_python(count) + async def flush_cache_buffer(self): print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}") await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) @@ -1258,6 +1357,7 @@ class RedisCache(BaseCache): except Exception: return ast.literal_eval(decoded) + @_redis_circuit_breaker_guard_sync def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs): try: key = self.check_and_fix_namespace(key=key) @@ -1277,8 +1377,8 @@ class RedisCache(BaseCache): print_verbose(f"Got Redis Cache: key: {key}, cached_response {cached_response}") return self._get_cache_logic(cached_response=cached_response) except Exception as e: - # NON blocking - notify users Redis is throwing an exception - verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: ", e) + verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: %s", e) + _record_swallowed_redis_failure(self._circuit_breaker, e) def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]: """ @@ -1315,12 +1415,12 @@ class RedisCache(BaseCache): key_value_dict = {} _key_list: Final = [key for key in key_list if key is not None] start_time: Final = time.time() + admission: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache") try: - swallowed_before: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache") _keys: Final = [self.check_and_fix_namespace(key=cache_key or "") for cache_key in _key_list] results: Final = self._run_redis_mget_operation(keys=_keys) - _exit_circuit_breaker(self._circuit_breaker, swallowed_before) + _exit_circuit_breaker(self._circuit_breaker, admission) end_time: Final = time.time() _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index d70f947469a..5a6debc4af5 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -5,8 +5,11 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args +from openai.types.chat import ChatCompletion +from openai.types.responses import Response from openai.types.responses.custom_tool_param import CustomToolParam from openai.types.responses.response_input_param import ( FunctionCallOutput, @@ -22,7 +25,7 @@ import litellm from litellm import ModelResponse from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( - responses_reasoning_item_from_thinking_blocks, + responses_reasoning_items_from_thinking_blocks, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( @@ -33,7 +36,7 @@ from litellm.responses.sse_output_recovery import ( record_output_item_chunk, record_output_text_chunk, ) -from litellm.responses.utils import normalize_responses_api_stream_options +from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_responses_api_stream_options from litellm.types.llms.openai import ( REASONING_EFFORT, ChatCompletionAnnotation, @@ -43,6 +46,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParamFunctionChunk, Reasoning, ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, ResponsesAPIStreamEvents, ) from litellm.types.utils import GenericStreamingChunk, ModelResponseStream @@ -54,7 +58,7 @@ if TYPE_CHECKING: ) from pydantic import BaseModel - from litellm import LiteLLMLoggingObj, ModelResponse + from litellm import LiteLLMLoggingObj from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import ( ALL_RESPONSES_API_TOOL_PARAMS, @@ -69,6 +73,28 @@ if TYPE_CHECKING: from litellm.types.utils import Choices +_CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage")) +_RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *ResponsesAPIResponse.model_fields)) - frozenset( + ChatCompletion.model_fields +) + + +def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]: + return MappingProxyType( + { + key: value + for key, value in (response_fields.items() if response_fields else ()) + if value is not None and key not in _CHAT_COMPLETION_FIELDS and key not in _RESPONSES_API_ONLY_FIELDS + } + ) + + +def _upstream_response_id(response_id: str | None) -> str | None: + if response_id is None: + return None + return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id) + + class _ReasoningSummaryText(TypedDict): type: str text: str @@ -103,8 +129,8 @@ def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]: return stored raw_blocks: Final = msg.get("thinking_blocks") or () blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json - from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks) - return [] if from_thinking is None else [dict(from_thinking)] # mutable-ok: API message payload + replayed: Final = responses_reasoning_items_from_thinking_blocks(blocks) + return [dict(item) for item in replayed] # mutable-ok: API message payload def _build_reasoning_item( @@ -201,7 +227,7 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False): provider_specific_fields: Mapping[str, object] -def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict: +def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict: """Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw string payload in ``input`` rather than ``arguments``; both map to @@ -344,7 +370,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): and isinstance(tool_call.get("custom"), dict) ) - for msg in messages: + leading_system_count: Final = next( + (index for index, msg in enumerate(messages) if msg.get("role") != "system"), + len(messages), + ) + + for index, msg in enumerate(messages): role = msg.get("role") content = msg.get("content", "") tool_calls = msg.get("tool_calls") @@ -352,7 +383,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if role == "system": # Extract system message as instructions - if isinstance(content, str): + if isinstance(content, str) and index < leading_system_count: if instructions: # Concatenate multiple system prompts with a space instructions = f"{instructions} {content}" @@ -724,7 +755,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Tool calls accumulate into the single trailing tool_calls choice # like the typed branches above; a choice per call would hide every # call after choices[0] from chat clients - accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index)) + accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index)) tool_call_index += 1 elif handle_raw_dict_callback is not None: choice, index = handle_raw_dict_callback(item=raw_item, index=index) @@ -904,6 +935,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage), ) + model_response.id = _upstream_response_id(raw_response.id) or raw_response.id + for key, value in _provider_metadata(raw_response.model_extra).items(): + setattr(model_response, key, value) + # Preserve hidden params from the ResponsesAPIResponse, especially the headers # which contain important provider information like x-request-id raw_response_hidden_params: Final = getattr(raw_response, "_hidden_params", {}) @@ -1166,6 +1201,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Cast to Any to match the expected union type for tools list items tools.append(cast(Any, web_search_tool)) + def transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None": + return self._transform_response_format_to_text_format(response_format) + def _transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None": """ Transform Chat Completion response_format parameter to Responses API text.format parameter. @@ -1359,20 +1397,22 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if event_type == "response.created": # Initial response creation event verbose_logger.debug("Chat provider: response.created -> %s", parsed_chunk) + created_response: Final = parsed_chunk.get("response") return ModelResponseStream( + id=_upstream_response_id(created_response.get("id")) if created_response else None, choices=[ StreamingChoices( index=0, delta=Delta(content=""), finish_reason=None, ) - ] + ], ) elif event_type == "response.output_item.added": # New output item added output_item = parsed_chunk.get("item", {}) if output_item.get("type") in ("function_call", "custom_tool_call"): - converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0)) + converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0)) provider_specific_fields: Final = converted.get("provider_specific_fields") function_chunk: Final = ChatCompletionToolCallFunctionChunk( @@ -1447,7 +1487,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): index=0, delta=Delta( tool_calls=( - _tool_call_dict_from_output_item( + tool_call_dict_from_output_item( output_item, parsed_chunk.get("output_index", 0) ), ) @@ -1534,6 +1574,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): from litellm.responses.utils import ResponseAPILoggingUtils usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_data.get("usage")) + provider_metadata: Final = _provider_metadata(response_data) return ModelResponseStream( choices=[ StreamingChoices( @@ -1546,6 +1587,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) ], usage=usage, + provider_specific_fields=dict(provider_metadata) or None, # mutable-ok: field is typed dict ) else: pass diff --git a/litellm/constants.py b/litellm/constants.py index ce744e9c58a..6b984c2673c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16 MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) +budget_reservation_disabled_info_emitted = False DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)) DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) SQS_SEND_MESSAGE_ACTION: Final = "SendMessage" @@ -72,6 +73,9 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096)) DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3)) DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1)) DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5)) +DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = float( + os.getenv("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", "1") +) DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) @@ -139,6 +143,7 @@ DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float( os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3) ) MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)) +MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH: Final = 2048 DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS: Final = 2000 @@ -193,6 +198,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [ "x-litellm-adaptive-router-model", "x-litellm-applied-guardrails", "x-litellm-guardrail-scan-id", + "x-litellm-guardrail-scan-metadata", "x-litellm-cache-key", ] @@ -329,6 +335,7 @@ DEFAULT_SSL_CIPHERS: Final = os.getenv( ########### v2 Architecture constants for managing writing updates to the database ########### REDIS_UPDATE_BUFFER_KEY: Final = "litellm_spend_update_buffer" +REDIS_GATEWAY_REQUESTS_BUFFER_KEY: Final = "litellm_gateway_requests_buffer" REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_spend_update_buffer" REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_update_buffer" REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer" @@ -391,6 +398,18 @@ TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range( minimum=1, maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS, ) +TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( + "TOKEN_COUNTER_MAX_EXACT_CHARS", + default=4_000_000, + minimum=1, + maximum=1_000_000_000, +) +TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = get_env_int_in_range( + "TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", + default=4, + minimum=1, + maximum=256, +) MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) @@ -563,6 +582,7 @@ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float( LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100) LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000) LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0 +AWS_SIGNING_MAX_THREADS: Final = 16 DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv( "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield" ) @@ -1370,6 +1390,7 @@ bedrock_embedding_models: Final[set] = set( "cohere.embed-multilingual-v3", "cohere.embed-v4:0", "twelvelabs.marengo-embed-2-7-v1:0", + "twelvelabs.marengo-embed-3-0-v1:0", ] ) @@ -1457,7 +1478,10 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata" OLD_LITELLM_METADATA_FIELD: Final = "metadata" RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name" SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" +OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens", "max_output_tokens"}) +CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" +ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" @@ -1630,6 +1654,7 @@ CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data" MAVVRIK_FOCUS_EXPORT_JOB_NAME: Final = "mavvrik_focus_export_usage_data" CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)) SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup" +BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME: Final = "background_health_check_db_save" KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job" EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job" WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job" @@ -1658,6 +1683,7 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INT RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3"))) RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2")) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) +PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, int(os.getenv("PROXY_DB_LOOKUP_MAX_CONCURRENCY", "25"))) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500"))) @@ -1759,6 +1785,10 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) +SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 +SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30 +SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000 +SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding # staleness so a transient DB error (which surfaces as an empty list) cannot @@ -1901,6 +1931,16 @@ HTTP_FRAMING_HEADERS: Final[frozenset[str]] = frozenset( } ) +PROVIDER_REQUEST_ID_HEADERS: Final[tuple[str, ...]] = ( + "x-amzn-requestid", + "x-request-id", + "request-id", + "x-ms-request-id", + "apim-request-id", + "x-goog-request-id", + "cf-ray", +) + # Browser-facing security headers that a malicious or misconfigured upstream # provider must not be able to set on the proxy's own response. BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( @@ -1960,6 +2000,10 @@ NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset( } ) +UNKNOWN_MODEL_SPEND_LOG_MODEL: Final[str] = "unknown-model" +MAX_SPEND_LOG_MODEL_NAME_LENGTH: Final[int] = 256 +MCP_SPEND_LOG_MODEL_PREFIX: Final[str] = "MCP: " + # PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this # sentinel api_key so PTU flat cost stays distinguishable from real per-request # spend under the table's composite unique constraint. diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index cb7da32f857..814eaaf76f7 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import TranscriptionUsageObjectTransformation, ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( + BilledTokenRates, CostCalculatorUtils, _generic_cost_per_character, _get_regional_uplift_multiplier, @@ -45,6 +46,9 @@ from litellm.llms.azure.cost_calculation import ( from litellm.llms.azure_ai.cost_calculator import ( cost_per_token as azure_ai_cost_per_token, ) +from litellm.llms.azure_ai.cost_calculator import ( + is_azure_model_router as azure_ai_is_model_router_name, +) from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.llms.bedrock.cost_calculation import ( cost_per_token as bedrock_cost_per_token, @@ -81,6 +85,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 +501,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 @@ -1114,6 +1126,7 @@ def _store_cost_breakdown_in_logging_obj( service_tier: str | None = None, data_residency: str | None = None, vertex_location: str | None = None, + billed_token_rates: BilledTokenRates | None = None, ) -> None: """ Helper function to store cost breakdown in the logging object. @@ -1158,6 +1171,7 @@ def _store_cost_breakdown_in_logging_obj( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + billed_token_rates=billed_token_rates, ) except Exception as breakdown_error: @@ -1651,11 +1665,10 @@ def completion_cost( data_residency=data_residency, vertex_location=vertex_location, response=completion_response, - request_model=request_model_for_cost, ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) - if custom_llm_provider == "azure_ai": + if custom_llm_provider == "azure_ai" and not azure_ai_is_model_router_name(model): model_for_additional_costs = request_model_for_cost if completion_response is not None: hidden_params = getattr(completion_response, "_hidden_params", None) or {} @@ -1727,6 +1740,7 @@ def completion_cost( _reasoning_cost: float | None = None _cache_read_cost: float | None = None _cache_creation_cost: float | None = None + _billed_token_rates: BilledTokenRates | None = None if cost_per_token_usage_object is not None and model: _breakdown_provider: str | None = ( custom_llm_provider if isinstance(custom_llm_provider, str) else None @@ -1738,10 +1752,12 @@ def completion_cost( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + custom_cost_per_token=custom_cost_per_token, ) _reasoning_cost = _token_type_breakdown.reasoning_cost _cache_read_cost = _token_type_breakdown.cache_read_cost _cache_creation_cost = _token_type_breakdown.cache_creation_cost + _billed_token_rates = _token_type_breakdown.rates _store_cost_breakdown_in_logging_obj( litellm_logging_obj=litellm_logging_obj, prompt_tokens_cost_usd_dollar=prompt_tokens_cost_usd_dollar, @@ -1761,6 +1777,7 @@ def completion_cost( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + billed_token_rates=_billed_token_rates, ) return _final_cost @@ -2406,6 +2423,46 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): ) +_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"}) + + +class _ResponsesWsEventResponse(BaseModel): + usage: Mapping[str, object] | None = None + + +class _ResponsesWsEvent(BaseModel): + type: str = "" + response: _ResponsesWsEventResponse | None = None + + +class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): + @staticmethod + def collect_usage_from_responses_ws_results( + results: Sequence[Mapping[str, object]], + ) -> tuple[Usage, ...]: + events: Final = tuple(_ResponsesWsEvent.model_validate(result) for result in results) + return tuple( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the realtime processor uses + event.response.usage + ) + for event in events + if event.type in _RESPONSES_WS_BILLABLE_EVENT_TYPES + and event.response is not None + and event.response.usage is not None + ) + + @staticmethod + def collect_and_combine_usage_from_responses_ws_results( + results: Sequence[Mapping[str, object]], + ) -> Usage: + collected_usage_objects: Final = ResponsesWebSocketTokenUsageProcessor.collect_usage_from_responses_ws_results( + results + ) + return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects( + list(collected_usage_objects) # mutable-ok: combine_usage_objects requires a list parameter + ) + + _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 16202321709..3f22a4b2dcd 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -10,12 +10,14 @@ ## LiteLLM versions of the OpenAI Exception Types import enum +from collections.abc import Sequence from typing import Any, Final import httpx import openai from litellm.types.utils import LiteLLMCommonStrings +from litellm.types.vector_stores import VectorStoreSearchFailure class RateLimitErrorCategory(str, enum.Enum): @@ -288,6 +290,29 @@ class ImageFetchError(BadRequestError): ) +VECTOR_STORE_SEARCH_FAILED_CODE: Final = "vector_store_search_failed" + + +class VectorStoreSearchError(BadRequestError): + def __init__( + self, + failures: Sequence[VectorStoreSearchFailure], + model: str | None = None, + llm_provider: str | None = None, + ) -> None: + self.failures: Final[tuple[VectorStoreSearchFailure, ...]] = tuple(failures) + detail: Final = "; ".join(f"{failure['vector_store_id']}: {failure['error']}" for failure in self.failures) + super().__init__( + message=( + "The request could not be grounded in every configured vector store. " + f"{len(self.failures)} vector store search(es) failed: {detail}" + ), + model=model, + llm_provider=llm_provider, + body={"type": "invalid_request_error", "code": VECTOR_STORE_SEARCH_FAILED_CODE}, + ) + + class UnprocessableEntityError(openai.UnprocessableEntityError): def __init__( self, @@ -338,6 +363,7 @@ class Timeout(openai.APITimeoutError): num_retries: int | None = None, headers: dict | None = None, exception_status_code: int | None = None, + response: httpx.Response | None = None, ): request: Final = httpx.Request( method="POST", @@ -352,6 +378,8 @@ class Timeout(openai.APITimeoutError): self.max_retries = max_retries self.num_retries = num_retries self.headers = headers + if response is not None: + self.response = response # custom function to convert to str def __str__(self): diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 3503468c735..6cecc1e6157 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -10,6 +10,7 @@ from contextlib import AbstractAsyncContextManager from datetime import timedelta from functools import partial from importlib import metadata +from types import MappingProxyType from typing import Any, Final, Protocol, TypeAlias, TypeVar import httpx @@ -18,6 +19,7 @@ from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServ from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client from mcp.shared.message import SessionMessage +from mcp.shared.session import RequestResponder from typing_extensions import Unpack _TransportStreams: TypeAlias = tuple[ @@ -53,15 +55,22 @@ def missing_streamable_http_client_error() -> ImportError: ) -from mcp.types import CallToolRequestParams as MCPCallToolRequestParams -from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import ( + METHOD_NOT_FOUND, + ClientResult, GetPromptRequestParams, GetPromptResult, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, Prompt, ResourceTemplate, + ServerNotification, + ServerRequest, TextContent, ) +from mcp.types import CallToolRequestParams as MCPCallToolRequestParams +from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import Tool as MCPTool from pydantic import AnyUrl @@ -69,6 +78,7 @@ from litellm._logging import verbose_logger from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT from litellm.experimental_mcp_client.tools import list_tools_with_pagination from litellm.llms.custom_httpx.http_handler import get_ssl_configuration +from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( MCPAuth, @@ -146,8 +156,8 @@ _SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT) otherwise carries JSON-RPC error codes.""" -def _as_read_timeout(exc: BaseException) -> TimeoutError | None: - """The session read timeout elapsing, re-expressed as a ``TimeoutError``, or ``None``. +def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None: + """Normalize an MCP SDK read timeout for client and gateway diagnostics, or return ``None``. The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error @@ -442,6 +452,18 @@ class MCPClient: in_flight_error: BaseException | None = None try: read_stream, write_stream = transport[0], transport[1] + stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future() + + async def receive_message( + message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, + ) -> None: + if not isinstance(message, (ValueError, httpx.RequestError, OSError)): + return + if not stream_error.done(): + stream_error.set_result(message) + # The SDK closes pending requests when its message handler raises. + raise RuntimeError("MCP response stream failed") + # Build session kwargs with optional callbacks session_kwargs: Final[dict[str, Any]] = {} if self._sampling_callback is not None: @@ -456,6 +478,7 @@ class MCPClient: read_stream, write_stream, read_timeout_seconds=timedelta(seconds=self.timeout), + message_handler=receive_message, **session_kwargs, ) session: Final = await session_ctx.__aenter__() @@ -467,6 +490,10 @@ class MCPClient: if isinstance(ins, str) and ins.strip(): self._last_initialize_instructions = ins.strip() return await operation(session) + except McpError: + if stream_error.done(): + raise stream_error.result() + raise finally: try: await session_ctx.__aexit__(None, None, None) @@ -501,11 +528,10 @@ class MCPClient: transport_ctx, http_client = self._create_transport_context() return await self._execute_session_operation(transport_ctx, operation) except Exception as e: - read_timeout: Final = _as_read_timeout(e) + read_timeout: Final = as_mcp_read_timeout(e) if read_timeout is not None: verbose_logger.warning( - "MCP client timed out after %ss waiting for %s to answer; the server accepted the " - "request and ended its response stream without a JSON-RPC reply", + "MCP client timed out after %ss waiting for a valid MCP response from %s", self.timeout, self.server_url or "stdio", ) @@ -607,7 +633,9 @@ class MCPClient: auth=effective_auth, verify=ssl_config, follow_redirects=True, - event_hooks={"request": [guard]} if guard else {}, + event_hooks=MappingProxyType( + {"response": [capture_upstream_error_response], "request": [guard] if guard else []} + ), # mutable-ok: httpx types require lists of hooks ) return factory @@ -757,8 +785,19 @@ class MCPClient: """List available prompts from the server.""" verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_prompts_operation(session: ClientSession): - return await session.list_prompts() + async def _list_prompts_operation(session: ClientSession) -> ListPromptsResult: + capabilities: Final = session.get_server_capabilities() + if capabilities is not None and capabilities.prompts is None: + return ListPromptsResult(prompts=[]) + try: + return await session.list_prompts() + except McpError as error: + if error.error.code != METHOD_NOT_FOUND: + raise + verbose_logger.debug( + "MCP client list_prompts is unsupported by %s: %s", self.server_url or "stdio", error + ) + return ListPromptsResult(prompts=[]) try: result: Final = await self.run_with_session(_list_prompts_operation) @@ -834,8 +873,19 @@ class MCPClient: """List available resources from the server.""" verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio") - async def _list_resources_operation(session: ClientSession): - return await session.list_resources() + async def _list_resources_operation(session: ClientSession) -> ListResourcesResult: + capabilities: Final = session.get_server_capabilities() + if capabilities is not None and capabilities.resources is None: + return ListResourcesResult(resources=[]) + try: + return await session.list_resources() + except McpError as error: + if error.error.code != METHOD_NOT_FOUND: + raise + verbose_logger.debug( + "MCP client list_resources is unsupported by %s: %s", self.server_url or "stdio", error + ) + return ListResourcesResult(resources=[]) try: result: Final = await self.run_with_session(_list_resources_operation) @@ -870,8 +920,19 @@ class MCPClient: """List available resource templates from the server.""" verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio") - async def _list_resource_templates_operation(session: ClientSession): - return await session.list_resource_templates() + async def _list_resource_templates_operation(session: ClientSession) -> ListResourceTemplatesResult: + capabilities: Final = session.get_server_capabilities() + if capabilities is not None and capabilities.resources is None: + return ListResourceTemplatesResult(resourceTemplates=[]) + try: + return await session.list_resource_templates() + except McpError as error: + if error.error.code != METHOD_NOT_FOUND: + raise + verbose_logger.debug( + "MCP client list_resource_templates is unsupported by %s: %s", self.server_url or "stdio", error + ) + return ListResourceTemplatesResult(resourceTemplates=[]) try: result: Final = await self.run_with_session(_list_resource_templates_operation) diff --git a/litellm/files/main.py b/litellm/files/main.py index 19da77b7364..218518eb3cd 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -31,7 +31,7 @@ FileCreateProvider = Literal[ FileRetrieveProvider = Literal[ "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic" ] -FileDeleteProvider = Literal["openai", "azure", "gemini", "litellm_proxy", "manus", "anthropic"] +FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"] FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"] import litellm from litellm import get_secret_str diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 3519240dda9..4f9b18713d0 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -16,6 +16,8 @@ from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement @@ -23,6 +25,7 @@ from litellm.integrations.prompt_management_base import PromptManagementClient from litellm.litellm_core_utils.prompt_templates.common_utils import ( with_prompt_cache_breakpoint, ) +from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request from litellm.types.integrations.anthropic_cache_control_hook import ( GATEWAY_INJECTED_CACHE_METADATA_KEY, GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT, @@ -62,10 +65,26 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset( ) OPENAI_API_HOST: Final = "api.openai.com" OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE") +_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues +def _validated_object_mapping(value: object) -> dict[object, object] | None: + try: + return _OBJECT_MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _validated_object_list(value: object) -> list[object] | None: + try: + return _OBJECT_LIST_ADAPTER.validate_python(value) + except ValidationError: + return None + + def supports_openai_prompt_cache_breakpoint(model: str) -> bool: model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model) if model_map_flag is not None: @@ -114,6 +133,36 @@ CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_ class AnthropicCacheControlHook(CustomPromptManagement): + @staticmethod + def _request_value(request_kwargs: object, key: str) -> object: + request_mapping: Final = _validated_object_mapping(request_kwargs) + if request_mapping is None: + return None + return request_mapping.get(key) + + @staticmethod + def _request_user_agent(request_kwargs: object) -> str | None: + proxy_server_request: Final = AnthropicCacheControlHook._request_value(request_kwargs, "proxy_server_request") + proxy_server_request_mapping: Final = _validated_object_mapping(proxy_server_request) + if proxy_server_request_mapping is None: + return None + headers: Final = proxy_server_request_mapping.get("headers") + headers_mapping: Final = _validated_object_mapping(headers) + if headers_mapping is None: + return None + user_agent: Final = next( + (value for key, value in headers_mapping.items() if isinstance(key, str) and key.lower() == "user-agent"), + None, + ) + return user_agent if isinstance(user_agent, str) else None + + @staticmethod + def _request_system(request_kwargs: object) -> str | list[object] | None: + system: Final = AnthropicCacheControlHook._request_value(request_kwargs, "system") + if isinstance(system, str): + return system + return _validated_object_list(system) + def get_chat_completion_prompt( self, model: str, @@ -520,12 +569,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): points: Sequence[CacheControlInjectionPoint], messages: list[AllMessageValues], tools: list[object] | None, + cache_control: object, model: str, custom_llm_provider: str | None, api_base: object, prompt_cache_options: object, ) -> Sequence[Mapping[str, object]] | None: - if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools): + if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control): return None return AnthropicCacheControlHook._stamped_with_dialect( points, model, custom_llm_provider, api_base, prompt_cache_options @@ -561,6 +611,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): messages: list[AllMessageValues], system: str | list | None, tools: list | None, + cache_control: object = None, ) -> bool: """Whether configured injection points must yield to client-set cache_control. @@ -573,13 +624,14 @@ class AnthropicCacheControlHook(CustomPromptManagement): """ if all(point.get("_litellm_judged") for point in points): return False - return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools) + return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control) @staticmethod def _request_has_cache_control( messages: list[AllMessageValues], system: str | list | None, tools: list | None = None, + cache_control: object = None, ) -> bool: """Return True if the request already carries any client-supplied cache_control. @@ -591,6 +643,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): carry the mark either at the top level (Anthropic shape) or nested under ``function`` (OpenAI shape); the Anthropic chat transform accepts both. """ + if cache_control is not None: + return True if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0: return True if tools is not None: @@ -612,6 +666,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider: str | None, tools: list | None = None, enable_prompt_caching: bool | None = None, + cache_control: object = None, + request_kwargs: object = None, ) -> list[CacheControlInjectionPoint]: """Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on. @@ -649,7 +705,12 @@ class AnthropicCacheControlHook(CustomPromptManagement): if not supports_prompt_caching(model=model, custom_llm_provider=provider): return [] - if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools): + if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control): + return [] + + if is_claude_code_one_shot_subagent_request( + messages, system, tools, AnthropicCacheControlHook._request_user_agent(request_kwargs) + ): return [] control: Final = AnthropicCacheControlHook._default_control() @@ -665,6 +726,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): models: Iterable[str], tools: list[AllToolParamValues] | None = None, enable_prompt_caching: bool | None = None, + request_kwargs: object = None, ) -> list[AllMessageValues]: """Return the messages auto prompt caching will send, default breakpoints included. @@ -681,11 +743,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): for candidate in ( AnthropicCacheControlHook.get_default_injection_points( messages=messages, - system=None, model=model, custom_llm_provider=None, tools=tools, enable_prompt_caching=enable_prompt_caching, + system=AnthropicCacheControlHook._request_system(request_kwargs), + cache_control=AnthropicCacheControlHook._request_value(request_kwargs, "cache_control"), + request_kwargs=request_kwargs, ) for model in models ) @@ -730,6 +794,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): non_default_params["cache_control_injection_points"], messages, tools, + non_default_params.get("cache_control"), model, custom_llm_provider, api_base, @@ -747,6 +812,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider=custom_llm_provider, tools=tools, enable_prompt_caching=enable_prompt_caching, + cache_control=non_default_params.get("cache_control"), + request_kwargs=non_default_params, ) if points: non_default_params["cache_control_injection_points"] = points @@ -853,10 +920,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): enable_prompt_caching: Final = cast( # cast-ok: kwargs is untyped; key stamped as bool by the proxy bool | None, kwargs.pop("enable_prompt_caching", None) ) + cache_control: Final = kwargs.get("cache_control") configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None) ) - if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools): + if configured and AnthropicCacheControlHook._should_stand_down( + configured, typed_messages, system, tools, cache_control + ): return messages, system injection_points: list[CacheControlInjectionPoint] = configured or [] if not injection_points and model is not None: @@ -867,6 +937,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): model=model, custom_llm_provider=custom_llm_provider, enable_prompt_caching=enable_prompt_caching, + cache_control=cache_control, + request_kwargs=kwargs, ) if not injection_points: return messages, system diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index eba6c862f7a..db5f790615f 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -21,6 +21,8 @@ from types import MappingProxyType from typing import Final, TypeVar from urllib.parse import urlparse +import httpx + from litellm._logging import verbose_logger from litellm.integrations.batch_utils import ( BatchSendCancelled, @@ -418,7 +420,7 @@ class AzureSentinelLogger(CustomBatchLogger): "Content-Type": "application/json", } - async def _send_batch(batch: Sequence[_QueuedPayload]): + async def _send_batch(batch: Sequence[_QueuedPayload]) -> httpx.Response: body: Final = safe_dumps(batch) return await self.async_httpx_client.post( url=api_endpoint, diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 7a2295a35ae..85bfcc6e7ed 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -356,15 +356,49 @@ "description": "OpenTelemetry collector endpoint URL", "required": true }, + "otel_traces_endpoint": { + "type": "text", + "ui_name": "Traces Endpoint URL", + "description": "Complete trace export URL used verbatim when the collector does not serve /v1/traces (OTel v2 only)", + "required": false + }, "otel_headers": { "type": "text", "ui_name": "Headers", "description": "Headers for OTEL exporter (e.g., x-honeycomb-team=YOUR_API_KEY)", "required": false + }, + "otel_exporter_otlp_protocol": { + "type": "select", + "ui_name": "Export Protocol", + "description": "OTLP wire format for trace exports. Use http/json for collectors that cannot decode protobuf", + "options": ["http/protobuf", "http/json"], + "required": false } }, "description": "OpenTelemetry Logging Integration" }, + { + "id": "pointfive", + "displayName": "PointFive", + "logo": "pointfive.png", + "supports_key_team_logging": false, + "dynamic_params": { + "POINTFIVE_API_KEY": { + "type": "password", + "ui_name": "API Key", + "description": "PointFive API key, used to request an upload url for each batch of logs", + "required": true + }, + "POINTFIVE_API_URL": { + "type": "text", + "ui_name": "API URL", + "description": "PointFive API endpoint. Leave blank to use https://api.pointfive.co/api/v1/ingestion", + "required": false + } + }, + "description": "PointFive Logging Integration" + }, { "id": "s3", "displayName": "S3", diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2d66a280663..77bf4820a1a 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -601,6 +601,12 @@ class CustomGuardrail(CustomLogger): event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None, supported_event_hooks: list[GuardrailEventHooks], ) -> None: + allowed_hooks: Final = frozenset(supported_event_hooks) | ( + frozenset((GuardrailEventHooks.logging_only,)) + if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks + else frozenset() + ) + def _validate_event_hook_list_is_in_supported_event_hooks( event_hook: list[GuardrailEventHooks] | list[str], supported_event_hooks: list[GuardrailEventHooks], @@ -608,7 +614,7 @@ class CustomGuardrail(CustomLogger): for hook in event_hook: if isinstance(hook, str): hook = GuardrailEventHooks(hook) - if hook not in supported_event_hooks: + if hook not in allowed_hooks: raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}") if event_hook is None: @@ -629,7 +635,7 @@ class CustomGuardrail(CustomLogger): default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default] _validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks) elif isinstance(event_hook, GuardrailEventHooks): - if event_hook not in supported_event_hooks: + if event_hook not in allowed_hooks: raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}") @staticmethod @@ -773,7 +779,7 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail - def _deployment_pre_call_target(self) -> "CustomLogger": + def _deployment_hook_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self try: @@ -802,7 +808,7 @@ class CustomGuardrail(CustomLogger): # CHECK IF GUARDRAIL REJECTS THE REQUEST if call_type == CallTypes.completion or call_type == CallTypes.acompletion: - target: Final = self._deployment_pre_call_target() + target: Final = self._deployment_hook_target() if target is not self: kwargs["guardrail_to_apply"] = self result: Final = await target.async_pre_call_hook( @@ -844,18 +850,24 @@ class CustomGuardrail(CustomLogger): if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: return None - # CHECK IF GUARDRAIL REJECTS THE REQUEST - result: Final = await self.async_post_call_success_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id=request_data.get("user_api_key_user_id"), - team_id=request_data.get("user_api_key_team_id"), - end_user_id=request_data.get("user_api_key_end_user_id"), - api_key=request_data.get("user_api_key_hash"), - request_route=request_data.get("user_api_key_request_route"), - ), - data=request_data, - response=response, - ) + target: Final = self._deployment_hook_target() + try: + if target is not self: + request_data["guardrail_to_apply"] = self # rebind-ok: dispatch consumes this key + result: Final = await target.async_post_call_success_hook( + user_api_key_dict=UserAPIKeyAuth( + user_id=request_data.get("user_api_key_user_id"), + team_id=request_data.get("user_api_key_team_id"), + end_user_id=request_data.get("user_api_key_end_user_id"), + api_key=request_data.get("user_api_key_hash"), + request_route=request_data.get("user_api_key_request_route"), + ), + data=request_data, + response=response, + ) + finally: + if target is not self: + request_data.pop("guardrail_to_apply", None) if not self._is_valid_response_type(result): return None diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 8f03e08f02d..62ca6b0254e 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -3,12 +3,13 @@ import re import traceback from collections.abc import AsyncGenerator, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional from pydantic import BaseModel from litellm._logging import verbose_logger -from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, EMPTY_MAPPING from litellm.types.integrations.argilla import ArgillaItem from litellm.types.integrations.custom_logger import AgenticLoopPlan from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest @@ -885,22 +886,25 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return LITELLM_METADATA_FIELD return OLD_LITELLM_METADATA_FIELD + def redacts_messages_itself(self) -> bool: + return False + def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict: """ Redacts or excludes fields from StandardLoggingPayload before callbacks receive it. This method handles two features: - 1. turn_off_message_logging: When True, redacts messages and responses + 1. turn_off_message_logging: When True, redacts messages and responses (unless the callback + redacts them itself, see `redacts_messages_itself`) 2. standard_logging_payload_excluded_fields: Removes specified fields entirely Return a modified copy of the provided logging payload. This is useful for logging payloads that contain sensitive information. """ - from copy import copy - import litellm from litellm import Choices, Message, ModelResponse + from litellm.litellm_core_utils.classifier_logging import CLASSIFIER_AUDIT_FIELDS, without_classifier_audit turn_off_message_logging: Final[bool] = getattr(self, "turn_off_message_logging", False) excluded_fields: Final[list[str] | None] = getattr(litellm, "standard_logging_payload_excluded_fields", None) @@ -909,30 +913,25 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac if turn_off_message_logging is False and not excluded_fields: return model_call_details - # Only make a shallow copy of the top-level dict to avoid deepcopy issues - # with complex objects like AuthenticationError that may be present - model_call_details_copy: Final = copy(model_call_details) standard_logging_object: Final = model_call_details.get("standard_logging_object") if standard_logging_object is None: - return model_call_details_copy + return model_call_details.copy() # Make a copy of just the standard_logging_object to avoid modifying the original - standard_logging_object_copy: Final = copy(standard_logging_object) - - # Handle excluded fields - remove them entirely from the payload - if excluded_fields: - for field in excluded_fields: - if field in standard_logging_object_copy: - del standard_logging_object_copy[field] + standard_logging_object_copy: Final = { + key: value + for key, value in standard_logging_object.items() + if key not in (excluded_fields or ()) and not (turn_off_message_logging and key in CLASSIFIER_AUDIT_FIELDS) + } # Handle turn_off_message_logging - redact messages and responses (if not already excluded) - if turn_off_message_logging: + if turn_off_message_logging and not self.redacts_messages_itself(): redacted_str: Final = "redacted-by-litellm" - if "messages" not in (excluded_fields or []) and standard_logging_object_copy.get("messages") is not None: + if "messages" not in (excluded_fields or ()) and standard_logging_object_copy.get("messages") is not None: standard_logging_object_copy["messages"] = [Message(content=redacted_str).model_dump()] - if "response" not in (excluded_fields or []) and standard_logging_object_copy.get("response") is not None: + if "response" not in (excluded_fields or ()) and standard_logging_object_copy.get("response") is not None: response: Final = standard_logging_object_copy["response"] # Check if this is a ResponsesAPIResponse (has "output" field) if isinstance(response, dict) and "output" in response: @@ -956,8 +955,18 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac model_response_dict: Final = model_response.model_dump() standard_logging_object_copy["response"] = model_response_dict - model_call_details_copy["standard_logging_object"] = standard_logging_object_copy - return model_call_details_copy + params: Final = model_call_details.get("litellm_params") + request: Final = params.get("proxy_server_request") if isinstance(params, dict) else None + redacted_params: Final = ( + MappingProxyType({"litellm_params": {**params, "proxy_server_request": without_classifier_audit(request)}}) + if turn_off_message_logging and isinstance(params, dict) and isinstance(request, dict) + else EMPTY_MAPPING + ) + return { + **model_call_details, + **redacted_params, + "standard_logging_object": standard_logging_object_copy, + } async def get_proxy_server_request_from_cold_storage_with_object_key( self, diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 728bf41856f..c64a12c6d75 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -165,18 +165,54 @@ def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, An ) -def _redact_messages(messages: Sequence[Message]) -> tuple[Message, ...]: - """Each message's shape with its content replaced and tool payloads dropped; no message is invented.""" - return tuple( - { - "role": role if isinstance(role, str) and role in _SAFE_REDACTED_MESSAGE_ROLES else "", - "content": REDACTED_BY_LITELLM, - } - for message in messages - for role in (message.get("role", ""),) +def _safe_identifier(value: object) -> str: + return value if isinstance(value, str) else "" + + +def _redact_tool_call(tool_call: ToolCall) -> ToolCall: + return ToolCall( + name=_safe_identifier(tool_call.get("name")), + arguments=REDACTED_BY_LITELLM, + tool_id=_safe_identifier(tool_call.get("tool_id")), + type=_safe_identifier(tool_call.get("type")), ) +def _redact_tool_result(tool_result: ToolResult) -> ToolResult: + return ToolResult( + name=_safe_identifier(tool_result.get("name")), + result=REDACTED_BY_LITELLM, + tool_id=_safe_identifier(tool_result.get("tool_id")), + type=_safe_identifier(tool_result.get("type")), + ) + + +def _redact_message(message: Message) -> Message: + role: Final = message.get("role", "") + tool_calls: Final = message.get("tool_calls", ()) + tool_results: Final = message.get("tool_results", ()) + redacted: Final[Message] = { + "role": role if isinstance(role, str) and role in _SAFE_REDACTED_MESSAGE_ROLES else "", + "content": REDACTED_BY_LITELLM, + **({"tool_calls": tuple(_redact_tool_call(call) for call in tool_calls)} if tool_calls else {}), + **({"tool_results": tuple(_redact_tool_result(result) for result in tool_results)} if tool_results else {}), + } + return redacted + + +def _redact_messages(messages: Sequence[Message]) -> tuple[Message, ...]: + return tuple(_redact_message(message) for message in messages) + + +def _tool_output_tokens(messages: Sequence[Message], model: str) -> float | None: + results: Final = tuple( + result.get("result", "") for message in messages for result in message.get("tool_results", ()) + ) + if not results: + return None + return float(sum(litellm.token_counter(model=model, text=result) for result in results)) + + def _cost_dimension_tags( standard_logging_payload: StandardLoggingPayload, router_fields: Mapping[str, object] ) -> tuple[str, ...]: @@ -583,6 +619,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): standard_logging_payload=standard_logging_payload, call_type=standard_logging_payload.get("call_type"), ) + tool_output_tokens: Final = _tool_output_tokens(input_messages, standard_logging_payload.get("model") or "") input_meta: Final = InputMeta(messages=_redact_messages(input_messages) if redact_payload else input_messages) output_meta: Final = OutputMeta( messages=_redact_messages(output_messages) if redact_payload else output_messages @@ -618,7 +655,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): **({"tool_definitions": tool_definitions} if tool_definitions else {}), } - metrics: Final = self._assemble_metrics(standard_logging_payload) + metrics: Final = self._assemble_metrics(standard_logging_payload, tool_output_tokens) payload: Final[LLMObsPayload] = LLMObsPayload( parent_id=metadata_parent_id if metadata_parent_id else "undefined", @@ -676,6 +713,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): ) return error_info + def redacts_messages_itself(self) -> bool: + return True + def _payload_logging_is_off(self, kwargs: Mapping[str, Any]) -> bool: return ( bool(self.turn_off_message_logging) @@ -683,7 +723,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): or should_redact_message_logging(dict(kwargs)) ) - def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics: + def _assemble_metrics( + self, standard_logging_payload: StandardLoggingPayload, tool_output_tokens: float | None + ) -> LLMMetrics: """ Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from. @@ -721,6 +763,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): else {} ), **({"reasoning_output_tokens": reasoning_output_tokens} if reasoning_output_tokens else {}), + **({"tool_output_tokens": tool_output_tokens} if tool_output_tokens is not None else {}), } return metrics diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 9576eabaa34..b75369965de 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -2,6 +2,7 @@ # On success, logs events to Langfuse import inspect import os +import re import traceback from collections.abc import Callable, Iterable, Mapping from datetime import datetime @@ -63,6 +64,44 @@ def _object_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, dict) else None +def _widened_items(mapping: Mapping[str, object]) -> Iterable[tuple[object, object]]: + """Header pairs with the key type widened back to what a caller-supplied dict can actually hold.""" + return mapping.items() + + +def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_request: object) -> bool: + if not isinstance(trace_id, str) or not isinstance(session_id, str): + return False + request: Final = _object_mapping(proxy_server_request) + raw_headers: Final = _object_mapping(request.get("headers")) if request is not None else None + if raw_headers is None: + return False + headers: Final = MappingProxyType( + {key.lower(): value for key, value in _widened_items(raw_headers) if isinstance(key, str)} + ) + if headers.get("x-litellm-trace-id"): + return False + if headers.get("langfuse_trace_id") is not None: + return False + if trace_id != session_id and headers.get("langfuse_session_id") != session_id: + return False + if headers.get("x-litellm-session-id") == trace_id: + return True + if re.fullmatch(r"[a-zA-Z0-9_\-]{8,}", trace_id) is None: + return False + user_agent: Final = headers.get("user-agent") + codex: Final = isinstance(user_agent, str) and re.match(r"^codex[-_ /]", user_agent, re.IGNORECASE) is not None + return any( + value == trace_id + and ( + key == "x-session-id" + or re.fullmatch(r"x-.+-session-id", key) is not None + or (codex and key in ("session-id", "session_id", "thread-id", "conversation_id")) + ) + for key, value in headers.items() + ) + + class _UsageObject(Protocol): """Token-count surface the Langfuse logger reads off a response usage payload.""" @@ -609,6 +648,18 @@ class LangFuseLogger: # This allows continuing an existing trace while still returning the correct trace_id if existing_trace_id is not None: trace_id = existing_trace_id + resolved_trace_id: Final = ( + litellm_call_id or trace_id + if existing_trace_id is None + and _is_session_header_trace(trace_id, session_id, litellm_params.get("proxy_server_request")) + else trace_id + ) + if resolved_trace_id != trace_id: + verbose_logger.debug( + "Langfuse: trace_id %s came from a session header; using call id %s so each call gets its own trace", + trace_id, + resolved_trace_id, + ) requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ())) update_trace_keys: Final = ( requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else () @@ -663,7 +714,7 @@ class LangFuseLogger: trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm" else: # don't overwrite an existing trace trace_params = { - "id": trace_id, + "id": resolved_trace_id, "name": trace_name, "session_id": session_id, "input": masked_input if not mask_input else "redacted-by-litellm", @@ -845,13 +896,13 @@ class LangFuseLogger: # Verify langfuse accepted our trace_id; if it differs, log a warning but still return our intended value # to match expected test behavior if hasattr(generation_client, "trace_id") and generation_client.trace_id: - if generation_client.trace_id != trace_id: + if generation_client.trace_id != resolved_trace_id: verbose_logger.warning( "Langfuse trace_id mismatch: set %s, but langfuse returned %s. Using our intended trace_id for consistency.", - trace_id, + resolved_trace_id, generation_client.trace_id, ) - return trace_id, generation_id + return resolved_trace_id, generation_id except Exception: verbose_logger.error("Langfuse Layer Error - %s", traceback.format_exc()) return None, None diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index a2f0b7cf39c..8731e96440f 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -133,17 +133,17 @@ class MlflowLogger(CustomLogger): if final_response: end_time_ns: Final = int(end_time.timestamp() * 1e9) - self._extract_and_set_chat_attributes(span, kwargs, final_response) - self._end_span_or_trace( - span=span, - outputs=final_response, - status=SpanStatusCode.OK, - end_time_ns=end_time_ns, - ) - - # Remove the stream_id from the map - with self._lock: - self._stream_id_to_span.pop(litellm_call_id) + try: + self._extract_and_set_chat_attributes(span, kwargs, final_response) + self._end_span_or_trace( + span=span, + outputs=final_response, + status=SpanStatusCode.OK, + end_time_ns=end_time_ns, + ) + finally: + with self._lock: + self._stream_id_to_span.pop(litellm_call_id, None) def _add_chunk_events(self, span, response_obj): from mlflow.entities import SpanEvent @@ -282,15 +282,15 @@ class MlflowLogger(CustomLogger): """End an MLflow span or a trace.""" if span.parent_id is None: self._client.end_trace( - trace_id=span.request_id, + span.request_id, outputs=outputs, status=status, end_time_ns=end_time_ns, ) else: self._client.end_span( - trace_id=span.request_id, - span_id=span.span_id, + span.request_id, + span.span_id, outputs=outputs, status=status, end_time_ns=end_time_ns, diff --git a/litellm/integrations/newrelic/newrelic_metrics.py b/litellm/integrations/newrelic/newrelic_metrics.py index 25dbfc2bdb2..da952b78d3f 100644 --- a/litellm/integrations/newrelic/newrelic_metrics.py +++ b/litellm/integrations/newrelic/newrelic_metrics.py @@ -5,8 +5,9 @@ NR Reference API: https://docs.newrelic.com/docs/data-apis/ingest-apis/metric-ap `async_log_success_event` / `async_log_failure_event` queue one record per request; at flush the queue is aggregated by (team, model group, model, provider, status) -into count/summary metrics. `interval.ms` is the real window between flushes, -computed at flush time. +into count/summary metrics, plus one max/remaining budget gauge pair per team +taken from the team's latest record. `interval.ms` is the real window between +flushes, computed at flush time. Team-scoped by construction: the ingest key is injected explicitly and there is deliberately no environment-variable fallback, so a team's metrics are never sent @@ -47,11 +48,14 @@ from litellm.types.integrations.newrelic import ( NEWRELIC_METRIC_PROMPT_TOKENS, NEWRELIC_METRIC_REQUEST_DURATION_MS, NEWRELIC_METRIC_REQUESTS, + NEWRELIC_METRIC_TEAM_MAX_BUDGET, + NEWRELIC_METRIC_TEAM_REMAINING_BUDGET, NEWRELIC_METRIC_TOTAL_TOKENS, NEWRELIC_METRICS_MAX_BATCH_SIZE, NEWRELIC_METRICS_MAX_DRAIN_PASSES, NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE, NewRelicCountMetric, + NewRelicGaugeMetric, NewRelicMetric, NewRelicMetricCommon, NewRelicMetricEnvelope, @@ -98,6 +102,8 @@ def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload) completion_tokens=int(standard_logging_object.get("completion_tokens") or 0), total_tokens=int(standard_logging_object.get("total_tokens") or 0), duration_ms=float(standard_logging_object.get("response_time") or 0.0) * 1000.0, + team_max_budget=metadata.get("user_api_key_team_max_budget") if metadata else None, + team_spend=metadata.get("user_api_key_team_spend") if metadata else None, ) @@ -140,6 +146,33 @@ def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[N return (*count_metrics, summary_metric) +def _team_budget_gauges(record: NewRelicMetricRecord) -> tuple[NewRelicMetric, ...]: + team_max_budget: Final = record.team_max_budget + if team_max_budget is None: + return () + attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType + key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN] + for key, value in (("team_id", record.team_id), ("team_alias", record.team_alias)) + if value + } + remaining_budget: Final = team_max_budget - (record.team_spend or 0.0) - record.response_cost + return ( + NewRelicGaugeMetric( + name=NEWRELIC_METRIC_TEAM_MAX_BUDGET, type="gauge", value=team_max_budget, attributes=attributes + ), + NewRelicGaugeMetric( + name=NEWRELIC_METRIC_TEAM_REMAINING_BUDGET, type="gauge", value=remaining_budget, attributes=attributes + ), + ) + + +def _team_budget_metrics(records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]: + latest_by_team: Final[Mapping[str, NewRelicMetricRecord]] = MappingProxyType( + {record.team_id: record for record in records if record.team_id} + ) + return tuple(gauge for record in latest_by_team.values() for gauge in _team_budget_gauges(record)) + + def build_metric_payload( records: tuple[NewRelicMetricRecord, ...], *, @@ -158,7 +191,7 @@ def build_metric_payload( "timestamp": int(window_start * 1000), "interval.ms": interval_ms, } - return (NewRelicMetricEnvelope(common=common, metrics=metrics),) + return (NewRelicMetricEnvelope(common=common, metrics=(*metrics, *_team_budget_metrics(records))),) class NewRelicMetricsLogger(CustomBatchLogger): diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 5519896a961..630aa313dc9 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -16,9 +16,11 @@ from opentelemetry.trace import ( Span, Tracer, get_current_span, + get_tracer_provider, set_span_in_context, use_span, ) +from opentelemetry.trace import TracerProvider as ApiTracerProvider import litellm from litellm._logging import verbose_logger @@ -63,6 +65,7 @@ from litellm.integrations.otel.plumbing.metrics import ( create_genai_metrics, ) from litellm.integrations.otel.plumbing.providers import ( + attach_tenant_fan_out, build_tracer_provider, get_event_logger, get_meter, @@ -85,6 +88,7 @@ if TYPE_CHECKING: ) LITELLM_TRACER_NAME: Final = "litellm" +_published_v2_provider: ApiTracerProvider | None = None def _span_error_from_exception( @@ -180,7 +184,9 @@ class OpenTelemetryV2(CustomLogger): self.config: OpenTelemetryV2Config = config or OpenTelemetryV2Config(**kwargs) self.callback_name = callback_name self._tracer_provider: TracerProvider = ( - tracer_provider if tracer_provider is not None else build_tracer_provider(self.config) + tracer_provider + if tracer_provider is not None + else build_tracer_provider(self.config, tenant_overrides=True) ) self.tracer: Tracer = get_tracer(self._tracer_provider, LITELLM_TRACER_NAME) self._metrics_recorder = self._init_metrics(meter_provider) @@ -195,6 +201,11 @@ class OpenTelemetryV2(CustomLogger): self._open_llm_calls: OrderedDict[str, _LLMCallSpan] = OrderedDict() self._init_otel_logger_on_litellm_proxy() + @property + def tracer_provider(self) -> TracerProvider: + """The provider this logger emits through, read-only to its callers.""" + return self._tracer_provider + def _init_metrics(self, meter_provider: "MeterProvider | None") -> "GenAIMetricRecorder | None": """Create the six GenAI histograms when metrics are enabled, else ``None``. @@ -863,12 +874,33 @@ def publish_global_otel_v2_provider( ``opentelemetry.trace.set_tracer_provider``) are injected so the publish step is unit-testable without reading or mutating real global OTel state. Returns the logger whose provider was published. + + The published provider is also the one that fans spans out to key/team + destinations, because it is the only provider the whole request tree passes + through; see :func:`attach_tenant_fan_out`. It is remembered for + :func:`fan_out_provider` because neither the OTel global (``set_tracer_provider`` + keeps the first provider it was ever handed) nor + ``proxy_server.open_telemetry_logger`` (a legacy v1 logger can hold that slot) + reliably leads back to it. """ + global _published_v2_provider logger: Final = select_global_otel_v2_logger(in_memory_loggers, registered=registered) - set_global_provider(logger._tracer_provider) + attach_tenant_fan_out(logger.tracer_provider, *_v2_configs(in_memory_loggers, logger)) + set_global_provider(logger.tracer_provider) + _published_v2_provider = logger.tracer_provider # rebind-ok: startup records the one provider carrying the fan-out return logger +def _v2_configs(in_memory_loggers: Sequence[object], logger: "OpenTelemetryV2") -> tuple[OpenTelemetryV2Config, ...]: + """Every v2 logger's config, the published logger's first. + + Each preset keeps its own provider and exporters, so the accounts the operator + writes to are spread over all of them, not held by the published logger alone. + """ + others: Final = tuple(cb.config for cb in in_memory_loggers if isinstance(cb, OpenTelemetryV2) and cb is not logger) + return (logger.config, *others) + + def _registered_v2_logger() -> "OpenTelemetryV2 | None": try: from litellm.proxy import proxy_server @@ -904,6 +936,25 @@ def seed_request_identity(user_api_key_dict: object, model: str | None = None) - logger.seed_request_identity(user_api_key_dict, model=model) +def fan_out_provider() -> ApiTracerProvider: + """The provider :func:`publish_global_otel_v2_provider` gave the tenant fan-out. + + Read off the publish itself, not the OTel global and not the registered logger: + the global keeps whichever provider claimed it first (auto-instrumentation, a + legacy logger), and the registered slot can hold a v1 logger while the publish + picked a v2 one from ``_in_memory_loggers``. Either detour lands on a provider + with no fan-out and drops every destination at auth. + """ + published: Final = _published_v2_provider + if published is not None: + return published + logger: Final = _registered_v2_logger() + if logger is not None: + attach_tenant_fan_out(logger.tracer_provider, logger.config) + return logger.tracer_provider + return get_tracer_provider() + + @contextmanager def phase_span(name: str) -> "Iterator[Span | None]": logger: Final = _registered_v2_logger() diff --git a/litellm/integrations/otel/mappers/legacy.py b/litellm/integrations/otel/mappers/legacy.py index 37475acb8f7..d25c25cd127 100644 --- a/litellm/integrations/otel/mappers/legacy.py +++ b/litellm/integrations/otel/mappers/legacy.py @@ -23,6 +23,7 @@ from litellm.integrations.otel.model.payloads import ( ServiceSpanData, ToolDefinition, ) +from litellm.integrations.otel.model.semconv import Error # Attribute keys in the semconv-ai / Traceloop vocabulary. _LEGACY_SYSTEM: Final = "gen_ai.system" @@ -36,7 +37,7 @@ _LEGACY_PRESENCE_PENALTY: Final = "llm.presence_penalty" _LEGACY_STOP_SEQUENCES: Final = "llm.chat.stop_sequences" _LEGACY_SERVICE: Final = "service" _LEGACY_CALL_TYPE: Final = "call_type" -_LEGACY_ERROR: Final = "error" +_LEGACY_ERROR: Final = Error.MESSAGE_LEGACY class LegacyMapper: diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index 9e3064c2bff..bd542ddc20c 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -69,9 +69,17 @@ class ExporterSpec(BaseModel): kind: str = Field( default="console", - description="console | in_memory | otlp_http | otlp_grpc | ", + description="console | in_memory | otlp_http | http/json | otlp_grpc | ", ) endpoint: str | None = None + traces_endpoint: str | None = Field( + default=None, + description=( + "Complete OTLP/HTTP trace URL, used verbatim. Set this when the " + "collector serves traces on a path other than ``/v1/traces``; " + "``endpoint`` is a base URL the signal path is appended to." + ), + ) headers: str | None = None owner: ExporterOwner | None = Field( default=None, @@ -127,6 +135,14 @@ class OpenTelemetryV2Config(BaseSettings): default=None, validation_alias=AliasChoices("OTEL_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"), ) + traces_endpoint: str | None = Field( + default=None, + validation_alias=AliasChoices("OTEL_TRACES_ENDPOINT", "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"), + description=( + "Complete OTLP/HTTP trace URL for the single-destination shorthand, " + "used verbatim instead of ``endpoint`` + ``/v1/traces``." + ), + ) headers: str | None = Field( default=None, validation_alias=AliasChoices("OTEL_HEADERS", "OTEL_EXPORTER_OTLP_HEADERS"), @@ -250,17 +266,22 @@ class OpenTelemetryV2Config(BaseSettings): @model_validator(mode="after") def _normalize(self) -> "OpenTelemetryV2Config": # An endpoint with the default exporter kind implies OTLP/HTTP. - if self.endpoint and self.exporter == "console": + if (self.endpoint or self.traces_endpoint) and self.exporter == "console": self.exporter = "otlp_http" # When no explicit destinations are given, fold the single-destination - # shorthand into one spec so the provider always has a destination. + # shorthand into one spec so the provider always has a destination. A spec + # with no fields set is how the presets tell "nothing configured" from an + # operator who asked for the console by name. if not self.exporters: self.exporters = [ ExporterSpec( kind=self.exporter, endpoint=self.endpoint, + traces_endpoint=self.traces_endpoint, headers=self.headers, ) + if not self.model_fields_set.isdisjoint(("exporter", "endpoint", "headers")) + else ExporterSpec() ] # Ensure ``genai`` is always present and first. names = list(self.mapper_names) diff --git a/litellm/integrations/otel/model/destination.py b/litellm/integrations/otel/model/destination.py new file mode 100644 index 00000000000..299253cac77 --- /dev/null +++ b/litellm/integrations/otel/model/destination.py @@ -0,0 +1,49 @@ +"""The resolved OTLP destination a request's traces export to. + +Backend-agnostic on purpose: every OTEL backend reduces to an endpoint plus auth +headers. The per-backend field mapping lives in ``presets.destinations``. +""" + +from collections.abc import Mapping +from typing import Final +from urllib.parse import quote + +from pydantic import BaseModel, ConfigDict, Field + + +class OtelDestination(BaseModel): + model_config = ConfigDict(frozen=True) + + endpoint: str + headers: Mapping[str, str] = Field(default_factory=dict) + resource_attributes: Mapping[str, str] = Field(default_factory=dict) + callback_name: str | None = None + protocol: str | None = Field( + default=None, + description=( + "OTLP transport, defaulting to the backend's own. Not derivable from the " + "scheme: Arize's ``https://otlp.arize.com/v1`` is gRPC." + ), + ) + + def header_string(self) -> str: + """Render headers as the ``k=v,k2=v2`` form an ``ExporterSpec`` expects. + + Values are percent-encoded because ``providers.parse_headers`` decodes them + with the SDK's W3C-Baggage parser: a value carrying a ``,`` or ``=`` (a + Langfuse project name, a base64 Authorization payload ending in ``==``) + would otherwise be split into bogus pairs on the way back out. + """ + return ",".join(f"{key}={quote(value, safe='')}" for key, value in self.headers.items()) + + def cache_key(self) -> tuple[str, tuple[tuple[str, str], ...], tuple[tuple[str, str], ...], str | None]: + """Identity for processor reuse, so one destination means one exporter.""" + return ( + self.endpoint, + tuple(sorted(self.headers.items())), + tuple(sorted(self.resource_attributes.items())), + self.protocol, + ) + + +NO_DESTINATIONS: Final[tuple[OtelDestination, ...]] = () diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index af5327cbd41..d3628005bac 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -204,6 +204,9 @@ class Error: TYPE: Final = "error.type" MESSAGE: Final = "error.message" + # The same text under the bare key the semconv-ai / Traceloop vocabulary uses + # (see ``LegacyMapper``), so anything reading or redacting error text covers both. + MESSAGE_LEGACY: Final = "error" class LiteLLMError: diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index aa7cc8e2afd..21e61c71fb7 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -1,8 +1,9 @@ """Trace-context + Baggage helpers.""" +import os from collections.abc import Mapping from contextvars import ContextVar, Token -from typing import Final +from typing import TYPE_CHECKING, Final from opentelemetry import baggage from opentelemetry.context import Context, get_current @@ -21,6 +22,9 @@ from opentelemetry.trace.propagation.tracecontext import ( from litellm.integrations.otel.model.semconv import HTTP +if TYPE_CHECKING: + from litellm.integrations.otel.model.destination import OtelDestination + _PROPAGATOR: Final = TraceContextTextMapPropagator() # The request's root span — the FastAPI-owned SERVER span — captured ONCE when the @@ -304,3 +308,65 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None: return None carrier: Final = {str(key).lower(): value for key, value in headers.items()} return _PROPAGATOR.extract(carrier) + + +# The OTLP destinations this request's key or team pointed its traces at, resolved +# once during auth. A ``ContextVar`` for the same reason the root span above is one: +# it rides the request task's context into the ``asyncio.create_task`` children that +# close the LLM span, and it is visible to every ``SpanProcessor.on_end`` that fires +# on the request task. Stateful MCP handlers set and reset it per message; the +# request-task value otherwise dies with that task. +_request_destinations: Final['ContextVar[tuple["OtelDestination", ...]]'] = ContextVar( + "litellm_otel_request_destinations", default=() +) + + +def set_request_destinations(destinations: 'tuple["OtelDestination", ...]') -> "Token[tuple[OtelDestination, ...]]": + """Anchor the destinations this request exports to and return a reset token.""" + return _request_destinations.set(destinations) + + +def reset_request_destinations(token: "Token[tuple[OtelDestination, ...]]") -> None: + _request_destinations.reset(token) + + +def request_destinations() -> 'tuple["OtelDestination", ...]': + """The destinations resolved for this request, empty outside a proxy request.""" + return _request_destinations.get() + + +#: ``litellm_settings: otel_tenant_destination_mode`` and its env equivalent. +ADDITIVE_DESTINATION_MODE: Final = "additive" +OTEL_TENANT_DESTINATION_MODE_ENV: Final = "LITELLM_OTEL_TENANT_DESTINATION_MODE" + + +def tenant_destinations_are_additive() -> bool: + """Whether a tenant destination exports alongside the operator's own exporter. + + Override is the default: the tenant's traffic reaches the tenant's account and + nowhere else. Operators running one org-wide backend across every team set this + to ``additive`` so the same trace lands in both places. + """ + import litellm + + configured: Final = litellm.otel_tenant_destination_mode or os.environ.get(OTEL_TENANT_DESTINATION_MODE_ENV) + return isinstance(configured, str) and configured.strip().lower() == ADDITIVE_DESTINATION_MODE + + +def destination_backends() -> frozenset[str]: + """Backends this request resolved a tenant destination for. + + The fan-out already carries the whole trace to those destinations, so the + per-request tracer route must never send a second copy, in either mode. + """ + return frozenset(d.callback_name for d in _request_destinations.get() if d.callback_name) + + +def suppressed_backends() -> frozenset[str]: + """Backends whose operator-level exporters this request must NOT reach. + + Empty under ``additive``, where the operator keeps its copy of every span. + """ + if tenant_destinations_are_additive(): + return frozenset() + return destination_backends() diff --git a/litellm/integrations/otel/plumbing/otlp_json.py b/litellm/integrations/otel/plumbing/otlp_json.py new file mode 100644 index 00000000000..b4b659f1e01 --- /dev/null +++ b/litellm/integrations/otel/plumbing/otlp_json.py @@ -0,0 +1,70 @@ +"""OTLP/HTTP span exporter that sends the OTLP/JSON encoding instead of protobuf. + +The SDK only ships a protobuf OTLP/HTTP exporter; this reuses its transport and +retry loop and swaps the payload for OTLP/JSON (enums as integers, ids as hex). +""" + +import base64 +import json +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, TypeAlias + +from google.protobuf.json_format import MessageToDict +from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +from opentelemetry.sdk.trace import ReadableSpan + +JSON_CONTENT_TYPE: Final = "application/json" +_HEX_ID_KEYS: Final = frozenset({"traceId", "spanId", "parentSpanId"}) + +_JsonValue: TypeAlias = "Mapping[str, _JsonValue] | Sequence[_JsonValue] | str | int | float | bool | None" +_JsonObject: TypeAlias = Mapping[str, "_JsonValue"] + + +def _objects(node: _JsonObject, key: str) -> tuple[_JsonObject, ...]: + items: Final = node.get(key) + if isinstance(items, str) or not isinstance(items, Sequence): + return () + return tuple(item for item in items if isinstance(item, Mapping)) + + +def _hex_ids(node: _JsonObject) -> _JsonObject: + return MappingProxyType( + { + key: base64.b64decode(item).hex() if key in _HEX_ID_KEYS and isinstance(item, str) else item + for key, item in node.items() + } + ) + + +def _hex_span(span: _JsonObject) -> _JsonObject: + links: Final = _objects(span, "links") + if not links: + return _hex_ids(span) + return MappingProxyType({**_hex_ids(span), "links": tuple(_hex_ids(link) for link in links)}) + + +def _hex_scope_spans(scope: _JsonObject) -> _JsonObject: + return MappingProxyType({**scope, "spans": tuple(_hex_span(span) for span in _objects(scope, "spans"))}) + + +def _hex_resource_spans(resource: _JsonObject) -> _JsonObject: + scope_spans: Final = tuple(_hex_scope_spans(scope) for scope in _objects(resource, "scopeSpans")) + return MappingProxyType({**resource, "scopeSpans": scope_spans}) + + +def encode_spans_json(spans: Sequence[ReadableSpan]) -> bytes: + payload: Final[_JsonObject] = MessageToDict(encode_spans(spans), use_integers_for_enums=True) + resource_spans: Final = tuple(_hex_resource_spans(resource) for resource in _objects(payload, "resourceSpans")) + hexed: Final[_JsonObject] = MappingProxyType({**payload, "resourceSpans": resource_spans}) + return json.dumps(hexed, default=dict, separators=(",", ":")).encode() + + +class OTLPJsonSpanExporter(OTLPSpanExporter): + def __init__(self, endpoint: str | None, headers: dict[str, str]) -> None: # mutable-ok: SDK __init__ takes Dict + super().__init__(endpoint=endpoint, headers=headers) + self._session.headers["Content-Type"] = JSON_CONTENT_TYPE + + def _serialize_spans(self, spans: Sequence[ReadableSpan]) -> bytes: + return encode_spans_json(spans) diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index fb74ff85e5b..81f22c8c642 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -1,9 +1,14 @@ """Provider / exporter factory + the Baggage span processor.""" -from collections.abc import Callable, Iterable +import queue +import threading +import time +from collections import OrderedDict +from collections.abc import Callable, Iterable, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal -from opentelemetry import _logs, baggage, metrics +from opentelemetry import _logs, baggage, metrics, trace from opentelemetry._events import EventLogger from opentelemetry._logs import LoggerProvider, NoOpLoggerProvider from opentelemetry.context import Context @@ -19,7 +24,8 @@ from opentelemetry.sdk._logs.export import ( ) from opentelemetry.sdk.metrics import MeterProvider as SDKMeterProvider from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace import Event, ReadableSpan, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace import Span as SDKSpan from opentelemetry.sdk.trace.export import ( BatchSpanProcessor, ConsoleSpanExporter, @@ -29,18 +35,35 @@ from opentelemetry.sdk.trace.export import ( from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( InMemorySpanExporter, ) -from opentelemetry.trace import Span, SpanKind, Tracer +from opentelemetry.trace import Span, SpanKind, Status, Tracer from opentelemetry.util.re import parse_env_headers +from opentelemetry.util.types import Attributes, AttributeValue +from litellm._logging import verbose_logger from litellm._version import version as litellm_version from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config -from litellm.integrations.otel.model.semconv import LiteLLM +from litellm.integrations.otel.model.semconv import ( + DB, + MCP, + Error, + ExceptionEvent, + GenAI, + LiteLLM, + LiteLLMError, + Server, +) from litellm.integrations.otel.model.spans import LiteLLMSpanKind +from litellm.integrations.otel.plumbing.context import ( + request_destinations, + suppressed_backends, +) if TYPE_CHECKING: from opentelemetry.metrics import Meter from opentelemetry.sdk.metrics.export import MetricReader + from litellm.integrations.otel.model.destination import OtelDestination + _SPAN_KIND_BY_ROLE_KIND: Final[dict[LiteLLMSpanKind, SpanKind]] = { LiteLLMSpanKind.SERVER: SpanKind.SERVER, LiteLLMSpanKind.CLIENT: SpanKind.CLIENT, @@ -136,7 +159,8 @@ def parse_headers(raw: str | None) -> dict[str, str]: _IN_MEMORY_KINDS: Final = ("in_memory", "inmemory", "memory") -_OTLP_HTTP_KINDS: Final = ("otlp_http", "http", "http/protobuf", "http/json") +_OTLP_HTTP_JSON_KINDS: Final = ("http/json",) +_OTLP_HTTP_KINDS: Final = ("otlp_http", "http", "http/protobuf", *_OTLP_HTTP_JSON_KINDS) _OTLP_GRPC_KINDS: Final = ("otlp_grpc", "grpc") @@ -164,13 +188,20 @@ def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter: return factory(spec) if kind in _IN_MEMORY_KINDS: return InMemorySpanExporter() + if kind in _OTLP_HTTP_JSON_KINDS: + from litellm.integrations.otel.plumbing.otlp_json import OTLPJsonSpanExporter + + return OTLPJsonSpanExporter( + endpoint=spec.traces_endpoint or _otlp_traces_endpoint(spec.endpoint), + headers=parse_headers(spec.headers), + ) if kind in _OTLP_HTTP_KINDS: from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter as HTTPExporter, ) return HTTPExporter( - endpoint=_otlp_traces_endpoint(spec.endpoint), + endpoint=spec.traces_endpoint or _otlp_traces_endpoint(spec.endpoint), headers=parse_headers(spec.headers), ) if kind in _OTLP_GRPC_KINDS: @@ -194,6 +225,555 @@ def _processor_for(exporter: SpanExporter, use_simple: bool | None) -> SpanProce return SimpleSpanProcessor(exporter) if use_simple else BatchSpanProcessor(exporter) +#: Distinct tenant destinations whose exporters stay alive. Each holds a connection +#: pool and a batch thread, so the cache is bounded and evicts least-recently-used. +_MAX_CACHED_DESTINATION_PROCESSORS: Final = 32 + +#: Workers closing shed destination processors, bounding the threads a tenant can +#: create by cycling its destination config. +_DRAIN_WORKERS: Final = 2 + +#: Shed processors waiting to be closed before the fan-out stops building new ones. +#: Each still owns a batch thread until its close returns, and a collector that never +#: answers makes every close take the exporter's full timeout, so past this many the +#: operator's exporter keeps the span instead (see ``deliverable``). +_MAX_PENDING_DRAINS: Final = 64 + +#: How long ``shutdown`` waits for spans already being forwarded, so teardown closes +#: no processor under one. Bounded: an exporter that never returns must not hold the +#: proxy open. +_SHUTDOWN_DRAIN_SECONDS: Final = 5.0 + +#: An exporter's account: its normalized endpoint and the credentials it presents. +_SinkKey = tuple[str, tuple[tuple[str, str], ...]] + +#: Header names that spell one credential two ways. Arize's operator exporter sends +#: ``space_id`` where a tenant destination sends ``arize-space-id``. +_CREDENTIAL_ALIASES: Final = MappingProxyType({"arize_space_id": "space_id"}) + + +class _DrainPool: + """Closes shed destination processors off the span-export path. + + ``shutdown`` flushes over the network and is reached from ``on_end``, so closing + one inline would let a single unreachable tenant collector stall every other + tenant's spans behind it. A fixed set of workers rather than a thread per + processor means a tenant cycling its destination config cannot spawn threads as + fast as it can send requests; slow shutdowns queue behind each other. + + The workers are daemons and belong to the fan-out that sheds the processors, so + neither an unreachable collector nor a lazily built process-wide singleton can + hold the proxy open on the way down. + """ + + def __init__( + self, + workers: int = _DRAIN_WORKERS, + pending: "queue.Queue[SpanProcessor | None] | None" = None, + capacity: int = _MAX_PENDING_DRAINS, + ) -> None: + self._workers: Final = workers + self._capacity: Final = capacity + self._lock: Final = threading.Lock() + self._closed = False + self._backlog = 0 # guarded by ``_lock``: submitted processors whose close has not returned + self._pending: Final[queue.Queue[SpanProcessor | None]] = pending if pending is not None else queue.Queue() + self._threads: Final = tuple( + threading.Thread(target=self._drain_until_closed, daemon=True, name="litellm-otel-destination-drain") + for _ in range(workers) + ) + for worker in self._threads: + worker.start() + + def submit(self, processor: SpanProcessor) -> None: + """Queue ``processor`` for closing, or hand it off once the pool is retired. + + The check and the put share one lock. Reading a closed flag on its own leaves + room for :meth:`close` to run in between, and the processor would land behind + the sentinels every worker has already exited on. + + Past close there is no worker left to take it, and the caller is whichever + thread just ended a span, so closing it inline would park that thread on a + network flush the shutdown deadline has already stopped waiting for. The extra + thread is bounded by the same close: the fan-out stops handing processors out + at that point, so only the ones already exporting when it happened arrive here. + """ + with self._lock: + if not self._closed: + self._backlog += 1 + self._pending.put(processor) + return + threading.Thread( + target=_shutdown_quietly, + args=(processor,), + daemon=True, + name="litellm-otel-destination-drain-straggler", + ).start() + + def saturated(self) -> bool: + """Whether enough closes are outstanding that building another processor must wait. + + The workers close in order and each close blocks for as long as its exporter + does, so a collector that stopped answering would otherwise turn every new + destination into one more batch thread parked behind them, for as long as the + tenants keep rotating. Holding the count here rather than reading the queue + keeps the two processors a worker is mid-close on in the total. + """ + with self._lock: + return self._backlog >= self._capacity + + def close(self, timeout: float | None = None) -> None: + """Retire the workers once they have closed everything already queued. + + A proxy that rebuilds its telemetry builds another fan-out, so workers that + outlive the one that started them are two more threads per reload, forever. + + ``timeout`` bounds how long the caller waits for that draining to finish. The + workers are daemons, so whatever is still flushing when it expires is dropped + by the interpreter rather than holding it open. + """ + with self._lock: + if self._closed: + return + self._closed = True + for _ in range(self._workers): + self._pending.put(None) + if timeout is None: + return + deadline: Final = time.monotonic() + timeout + for worker in self._threads: + worker.join(timeout=max(0.0, deadline - time.monotonic())) + + def _drain_until_closed(self) -> None: + while True: + processor: SpanProcessor | None = self._pending.get() # rebind-ok: loop variable + if processor is None: + return + _shutdown_quietly(processor) + with self._lock: + self._backlog -= 1 + + +_NO_ATTRIBUTES: Final[Mapping[str, AttributeValue]] = MappingProxyType({}) +_DB_SYSTEM_KEYS: Final = frozenset({DB.SYSTEM_NAME, DB.SYSTEM_LEGACY}) +# Keys on a database span that describe the proxy's own datastore: its host, its +# port, and its schema. +_DATASTORE_ENDPOINT_KEYS: Final = frozenset({Server.ADDRESS, Server.PORT, DB.NAMESPACE}) +# A span carrying one of these describes the tenant's own call (the model call, the +# MCP call, the guardrail), so its error text is theirs to see. Every other span is +# the proxy's own work, whose error text names the operator's infrastructure. +_TENANT_OWNED_KEYS: Final = frozenset({GenAI.OPERATION_NAME, MCP.METHOD_NAME, LiteLLM.GUARDRAIL_NAME}) +_PROXY_ERROR_TEXT_KEYS: Final = frozenset({Error.MESSAGE, Error.MESSAGE_LEGACY}) +# A guardrail that never answered carries the exception it raised as its response, +# which names the operator's guardrail endpoint. The second spelling is the legacy +# status the request-level logger still maps. +_GUARDRAIL_UNREACHABLE_STATUSES: Final = frozenset({"guardrail_failed_to_respond", "failure"}) +# Attribute prefixes the FastAPI instrumentor uses for headers the operator opted to +# capture (``OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_*``). The request +# side carries the caller's bearer token verbatim. +_CAPTURED_HEADER_PREFIXES: Final = ("http.request.header.", "http.response.header.") +# The instrumentor stamps the request URL on the server span with its query string, +# under the old convention and the new one, and litellm accepts a virtual key as a +# ``?key=`` query parameter. +_URL_KEYS: Final = frozenset({"http.url", "http.target", "url.full"}) +_URL_QUERY_KEY: Final = "url.query" + + +class _TenantSpanView(ReadableSpan): + """A ``ReadableSpan`` view for one destination, leaving the operator's own span alone.""" + + def __init__( + self, + inner: ReadableSpan, + resource: Resource, + attributes: Attributes, + events: Sequence[Event], + status: Status, + ) -> None: + super().__init__( + name=inner.name, + context=inner.context, + parent=inner.parent, + resource=resource, + attributes=attributes, + events=events, + links=inner.links, + kind=inner.kind, + status=status, + start_time=inner.start_time, + end_time=inner.end_time, + instrumentation_scope=inner.instrumentation_scope, + ) + + +def _is_database_span(attributes: Mapping[str, AttributeValue]) -> bool: + return any(key in attributes for key in _DB_SYSTEM_KEYS) + + +def _is_tenant_owned_span(attributes: Mapping[str, AttributeValue]) -> bool: + return any(key in attributes for key in _TENANT_OWNED_KEYS) + + +def _guardrail_unreachable(attributes: Mapping[str, AttributeValue]) -> bool: + return attributes.get(LiteLLM.GUARDRAIL_STATUS) in _GUARDRAIL_UNREACHABLE_STATUSES + + +def _tenant_visible(key: str, database: bool, owned: bool, unreachable_guardrail: bool) -> bool: + if key.startswith(_CAPTURED_HEADER_PREFIXES) or key in (LiteLLMError.STACK_TRACE, _URL_QUERY_KEY): + return False + if database and key in _DATASTORE_ENDPOINT_KEYS: + return False + if unreachable_guardrail and key == LiteLLM.GUARDRAIL_RESPONSE: + return False + return owned or key not in _PROXY_ERROR_TEXT_KEYS + + +def _without_query(key: str, value: AttributeValue) -> AttributeValue: + if key not in _URL_KEYS or not isinstance(value, str): + return value + return value.partition("?")[0] + + +def _same_attributes(kept: Mapping[str, AttributeValue], attributes: Mapping[str, AttributeValue]) -> bool: + return len(kept) == len(attributes) and all(kept[key] is value for key, value in attributes.items()) + + +def _without_stack_trace(event: Event) -> Event: + attributes: Final = event.attributes or _NO_ATTRIBUTES + if ExceptionEvent.STACKTRACE not in attributes: + return event + return Event( + name=event.name, + attributes=MappingProxyType( + {key: value for key, value in attributes.items() if key != ExceptionEvent.STACKTRACE} + ), + timestamp=event.timestamp, + ) + + +def _for_destination(span: ReadableSpan, destination: "OtelDestination") -> ReadableSpan: + """The view of ``span`` a tenant destination receives. + + A span the tenant's own call produced keeps its error text. Every other span is + the proxy's own work (the request root, auth, the database), and its error text, + its events and its status description come off, since a Prisma failure there + spells out the operator's Postgres endpoint. A database span loses that endpoint + too, and a guardrail that failed to respond loses its response text, which is the + exception it raised and names the operator's guardrail endpoint. Stack traces walk + the operator's install and come off every span, as do the headers the operator + captures on the server span, whose request side holds the caller's bearer token, + and the query string of the request URL, which can hold the same key. The span + itself stays, so the tenant still gets the whole trace tree. + """ + extra: Final = destination.resource_attributes + attributes: Final = span.attributes or _NO_ATTRIBUTES + database: Final = _is_database_span(attributes) + owned: Final = _is_tenant_owned_span(attributes) + unreachable: Final = _guardrail_unreachable(attributes) + kept: Final = MappingProxyType( + { + key: _without_query(key, value) + for key, value in attributes.items() + if _tenant_visible(key, database, owned, unreachable) + } + ) + recorded: Final = span.events + events: Final = tuple(_without_stack_trace(event) for event in recorded) if owned else () + unchanged: Final = owned and _same_attributes(kept, attributes) and all(a is b for a, b in zip(events, recorded)) + if not extra and unchanged: + return span + resource: Final = span.resource.merge(Resource(extra)) if extra else span.resource + status: Final = span.status if owned else Status(span.status.status_code) + return _TenantSpanView(span, resource, kept, events, status) + + +class TenantFanOutSpanProcessor(SpanProcessor): + """Export every finished span to each destination this request resolved. + + Destinations ride a request-scoped ``ContextVar`` set during auth, so concurrent + requests stay isolated. The forwarded view keeps the original trace and parent + ids, so the tenant gets the same tree the operator would have received. + + Exactly one provider carries this processor, the one published as the OTel global + (see :func:`attach_tenant_fan_out`). That provider is the only one every span + passes through: the FastAPI server span, the auth span and the post-call database + spans are emitted on the global, while a second v2 logger's provider sees only + that logger's own gen-AI span. Attaching the fan-out per logger would hand a + tenant a one-span trace whenever its backend is not the global one, and two + copies of the model call whenever it is. + """ + + def __init__( + self, + processor_factory: 'Callable[["OtelDestination"], SpanProcessor | None] | None' = None, + shutdown_drain_seconds: float = _SHUTDOWN_DRAIN_SECONDS, + operator_sinks: frozenset[_SinkKey] = frozenset(), + pending_drains: int = _MAX_PENDING_DRAINS, + drain_pool: _DrainPool | None = None, + ) -> None: + self._operator_sinks: Final = operator_sinks + self._drain_seconds: Final = shutdown_drain_seconds + self._lock: Final = threading.Condition() + self._closed = False # guarded by ``_lock``: an unlocked read races the teardown it gates + self._build: Final = processor_factory if processor_factory is not None else _destination_processor + self._processors: OrderedDict[object, SpanProcessor] = OrderedDict() # mutable-ok: bounded LRU + self._retired: OrderedDict[int, SpanProcessor] = OrderedDict() # mutable-ok: drains as exports finish + self._exporting: dict[int, int] = {} # mutable-ok: per-processor in-flight export count + self._drain: Final = drain_pool if drain_pool is not None else _DrainPool(capacity=pending_drains) + + def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None: + return None + + def on_end(self, span: ReadableSpan) -> None: + suppressed: Final = suppressed_backends() + for destination in request_destinations(): + if self._operator_already_writes(destination, suppressed): + continue + processor = self._acquire(destination) # rebind-ok: loop variable; pyright forbids Final in a loop + if processor is None: + continue + try: + processor.on_end(_for_destination(span, destination)) + except Exception as exc: # noqa: BLE001 # one destination's failure must not cost the others their span + verbose_logger.debug("OTel V2 fan-out: forwarding to %s failed: %s", destination.endpoint, exc) + finally: + self._release(processor) + + def _operator_already_writes(self, destination: "OtelDestination", suppressed: frozenset[str]) -> bool: + """Whether the operator's own exporter is sending this span to the same account. + + Only reachable under ``additive``, where nothing is suppressed: a team that + names the operator's own project would otherwise have every span written + there twice, once by the operator's exporter and once by the fan-out. + """ + return ( + destination.callback_name not in suppressed + and _sink_key(destination.endpoint, destination.headers) in self._operator_sinks + ) + + def shutdown(self) -> None: + """Close every destination processor, once the spans in flight have landed. + + ``on_end`` runs on whichever thread ends a span and can reach this fan-out + while the SDK is tearing the provider down, so closing blind would drop a + trace mid-forward and would hand the next caller a fresh exporter nothing + will ever close. Refusing new work and then waiting out the in-flight ones + keeps both from happening. A straggler past the bound is retired instead of + closed: the thread still exporting it closes it through the drain as soon as + its export returns, so no span is dropped mid-forward. + + Every close then goes to the drain rather than running here. Closing a + destination processor flushes it over the network and the SDK joins its own + worker with no timeout of its own, so one tenant collector that answers but + never finishes a response would otherwise hold process teardown open for as + long as it likes. The drain's workers are daemons, and the whole teardown + shares one deadline. + """ + deadline: Final = time.monotonic() + self._drain_seconds + with self._lock: + self._closed = True + self._lock.wait_for(lambda: not self._exporting, timeout=self._drain_seconds) + live: Final = tuple((id(p), p) for p in (*self._processors.values(), *self._retired.values())) + closing: Final = tuple(p for ident, p in live if ident not in self._exporting) + self._processors.clear() + self._retired = OrderedDict( # mutable-ok: the same bounded map, keeping only what is still exporting + (ident, p) for ident, p in live if ident in self._exporting + ) + for processor in closing: + self._drain.submit(processor) + self._drain.close(timeout=max(0.0, deadline - time.monotonic())) + + def force_flush(self, timeout_millis: int = 30000) -> bool: + results: Final = tuple(self._flush_one(processor, timeout_millis) for processor in self._snapshot()) + return all(results) + + def _snapshot(self) -> tuple[SpanProcessor, ...]: + with self._lock: + return (*self._processors.values(), *self._retired.values()) + + @staticmethod + def _flush_one(processor: SpanProcessor, timeout_millis: int) -> bool: + try: + return processor.force_flush(timeout_millis) + except Exception: # noqa: BLE001 # one exporter's flush failure must not fail the whole flush + return False + + def deliverable(self, destinations: Iterable["OtelDestination"]) -> tuple["OtelDestination", ...]: + """The subset of ``destinations`` this fan-out can actually export to. + + A destination whose exporter will not build (a protocol whose package is not + installed, a malformed endpoint) has to be dropped before the request anchors + it, not when its first span ends. By then the operator's own exporter has been + told to hold that backend's spans back for this request, so dropping there + loses the span outright instead of leaving it where it would have gone with no + override at all. + """ + return tuple(destination for destination in destinations if self._buildable(destination)) + + def _buildable(self, destination: "OtelDestination") -> bool: + """Whether a processor for ``destination`` exists or can be built right now.""" + with self._lock: + if self._closed: + return False + built: Final = self._cached_or_built_locked(destination, anchored=False) + drained: Final = self._drainable_locked() + for shed in drained: + self._drain.submit(shed) + return built is not None + + def _acquire(self, destination: "OtelDestination") -> SpanProcessor | None: + """The processor for ``destination``, marked busy until ``_release``. + + The build happens under the same lock that reads the cache, so a cold cache + met by a burst of concurrent requests yields one exporter rather than one per + thread with all but the winner shed. Building an exporter opens no connection, + so the cost of holding the lock is a constructor, once per destination. + """ + with self._lock: + if self._closed: + return None + processor: Final = self._cached_or_built_locked(destination, anchored=True) + if processor is None: + return None + self._exporting[id(processor)] = self._exporting.get(id(processor), 0) + 1 + drained: Final = self._drainable_locked() + for shed in drained: + self._drain.submit(shed) + return processor + + def _cached_or_built_locked(self, destination: "OtelDestination", *, anchored: bool) -> SpanProcessor | None: + """The cached processor for ``destination``, or a new one if the drain can take it. + + Every build past the cache cap sheds one processor into the drain, so while the + shed ones are stuck closing against a collector that stopped answering, a + destination that is not yet anchored is refused rather than parked behind them: + ``deliverable`` then leaves its spans with the operator's exporter until the + drain catches up. One the request already anchored is rebuilt regardless. The + operator's exporter has stood down for it, so refusing here would drop the span, + and other tenants' auths can evict it in the meantime, with that eviction being + what tips the drain over. Eviction holds while the drain is saturated, so such a + rebuild costs the cache one entry rather than shedding another processor, and + the total stays at one per destination in flight. + """ + key: Final = destination.cache_key() + if (cached := self._processors.get(key)) is not None: + self._processors.move_to_end(key) + self._retire_overflow_locked() + return cached + if not anchored and self._drain.saturated(): + verbose_logger.debug("OTel V2 fan-out: drain saturated, not building for %s", destination.endpoint) + return None + return self._build_locked(destination, key) + + def _build_locked(self, destination: "OtelDestination", key: object) -> SpanProcessor | None: + built: Final = self._build(destination) + if built is None: + return None + self._processors[key] = built + self._retire_overflow_locked() + return built + + def _release(self, processor: SpanProcessor) -> None: + with self._lock: + remaining: Final = self._exporting.get(id(processor), 1) - 1 + if remaining > 0: + self._exporting[id(processor)] = remaining + else: + self._exporting.pop(id(processor), None) + if not self._exporting: + self._lock.notify_all() + drained: Final = self._drainable_locked() + for retired in drained: + self._drain.submit(retired) + + def _retire_overflow_locked(self) -> None: + """Move the LRU processor out of the cache once it is past the cap, drain permitting. + + Eviction is what feeds the drain, and a destination a request already anchored + is rebuilt on its next span, which would shed another one. While the shed ones + are stuck closing against a collector that stopped answering, evicting would + churn the cache at one more processor, and one more batch thread, per span. + Holding above the cap instead keeps the total at one processor per destination + in flight, since ``deliverable`` anchors no new destination while the drain is + saturated. Once it has room again, every hit and build trims one entry. + """ + if len(self._processors) <= _MAX_CACHED_DESTINATION_PROCESSORS or self._drain.saturated(): + return + _, evicted = self._processors.popitem(last=False) + self._retired[id(evicted)] = evicted + + def _drainable_locked(self) -> tuple[SpanProcessor, ...]: + """Retired processors no thread is exporting through, removed from the list. + + ``on_end`` holds a processor across an export, so closing an evicted one there + drops the span it is holding. A retiree is out of the cache and can never be + handed out again, so once its export count reaches zero it stays there. + """ + idle: Final = tuple(key for key in self._retired if self._exporting.get(key, 0) == 0) + return tuple(self._retired.pop(key) for key in idle) + + +def _destination_processor(destination: "OtelDestination") -> SpanProcessor | None: + """A batching OTLP processor aimed at ``destination``, or ``None`` if unbuildable. + + A protocol that resolves to a headerless exporter is unbuildable too: the + console fallback would swallow the tenant's credentials and print its spans to + the proxy's stdout while the operator's exporter stands down for them. + """ + kind: Final = destination.protocol or "otlp_http" + if exporter_transport(kind) == "headerless": + verbose_logger.debug("OTel V2 fan-out: no OTLP transport for protocol %r at %s", kind, destination.endpoint) + return None + try: + spec: Final = ExporterSpec( + kind=kind, + endpoint=destination.endpoint, + headers=destination.header_string(), + owner=None, + ) + return _processor_for(_exporter_from_spec(spec), use_simple=False) + except Exception as exc: # noqa: BLE001 # a malformed destination must not break the request or the other destinations + verbose_logger.debug("OTel V2 fan-out: no processor for %s: %s", destination.endpoint, exc) + return None + + +def _shutdown_quietly(processor: SpanProcessor) -> None: + try: + processor.shutdown() + except Exception as exc: # noqa: BLE001 # defensive: shedding a spare processor must not raise + verbose_logger.debug("OTel V2 fan-out: discarding processor failed: %s", exc) + + +class _OverriddenBackendFilter(SpanProcessor): + """Hold a span back from ``owner``'s operator-level exporter when the request + pointed ``owner`` at a tenant's own account. + + Wrapping is the only place this works: ``SynchronousMultiSpanProcessor.on_end`` + ignores return values, so a sibling processor can never veto the export. + + Under ``additive`` mode nothing is suppressed, so the wrapper passes every span + straight through and the operator keeps its copy. + """ + + def __init__(self, inner: SpanProcessor, owner: str) -> None: + self._inner: Final = inner + self._owner: Final = owner + + def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None: + self._inner.on_start(span, parent_context) + + def on_end(self, span: ReadableSpan) -> None: + if self._owner in suppressed_backends(): + return + self._inner.on_end(span) + + def shutdown(self) -> None: + self._inner.shutdown() + + def force_flush(self, timeout_millis: int = 30000) -> bool: + return self._inner.force_flush(timeout_millis) + + def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter: """Build a single exporter from the top-level config fields. @@ -201,7 +781,14 @@ def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter: ``exporter`` / ``endpoint`` / ``headers`` fields. To configure multiple exporters, populate ``config.exporters`` directly. """ - return _exporter_from_spec(ExporterSpec(kind=config.exporter, endpoint=config.endpoint, headers=config.headers)) + return _exporter_from_spec( + ExporterSpec( + kind=config.exporter, + endpoint=config.endpoint, + traces_endpoint=config.traces_endpoint, + headers=config.headers, + ) + ) def _otlp_metrics_endpoint(endpoint: str | None) -> str | None: @@ -437,6 +1024,7 @@ def build_tracer_provider( exporter: SpanExporter | None = None, baggage_processor: SpanProcessor | None = None, use_simple_processor: bool | None = None, + tenant_overrides: bool = False, ) -> TracerProvider: """Build the shared :class:`TracerProvider`. @@ -445,6 +1033,13 @@ def build_tracer_provider( ``config.exporters`` entry — this is what fans spans out to multiple backends. ``exporter`` and ``use_simple_processor`` are explicit overrides: pass a single exporter to attach exactly that one (used by tests). + + ``tenant_overrides`` wraps each owned exporter so a request that pointed that + backend at a key's or team's own account skips it. Every v2 logger's provider + wants it, since any of them may own the overridden backend; delivering to the + tenant is a separate job, done once by :func:`attach_tenant_fan_out`. The + per-tenant providers this same function builds must leave it off, or they would + filter out the very spans they exist to carry. """ provider: Final = TracerProvider(resource=build_resource(config)) if baggage_processor is None: @@ -461,15 +1056,107 @@ def build_tracer_provider( if spec.requires_headers and not spec.headers: continue exp = _exporter_from_spec(spec) + processor = _processor_for( + exp, + (spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor), + ) + owner = spec.owner.value if spec.owner is not None else None provider.add_span_processor( - _processor_for( - exp, - (spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor), - ) + _OverriddenBackendFilter(processor, owner) if tenant_overrides and owner is not None else processor ) return provider +_FAN_OUT_ATTACH_LOCK: Final = threading.Lock() + + +def attach_tenant_fan_out(provider: TracerProvider, *configs: OpenTelemetryV2Config) -> None: + """Give ``provider`` the fan-out that delivers spans to key/team destinations. + + Called on the one provider published as the OTel global, and idempotent so a + second publish (a test, a re-initialized proxy) cannot double-export. Concurrent + first calls (requests racing to anchor before any publish) serialize on one lock + so exactly one fan-out lands. ``configs`` name the operator's own exporters, one + config per v2 logger since each keeps its own provider and still writes its + account, so an additive destination pointing at any of them is delivered once + rather than twice. + """ + with _FAN_OUT_ATTACH_LOCK: + if any(isinstance(processor, TenantFanOutSpanProcessor) for processor in _attached_processors(provider)): + return + provider.add_span_processor(TenantFanOutSpanProcessor(operator_sinks=operator_sink_keys(*configs))) + + +def deliverable_destinations( + destinations: Iterable["OtelDestination"], + provider: trace.TracerProvider | None = None, +) -> tuple["OtelDestination", ...]: + """The destinations a request can anchor, given what is published to carry them. + + Anchoring a destination is what tells the operator's own exporter to stand down + for that backend, so one nothing can deliver has to be dropped here: with no + fan-out attached, or with an exporter that will not build, the request keeps + exactly the routing it would have had without any override. + """ + fan_out: Final = next( + ( + processor + for processor in _attached_processors(provider if provider is not None else trace.get_tracer_provider()) + if isinstance(processor, TenantFanOutSpanProcessor) + ), + None, + ) + return fan_out.deliverable(destinations) if fan_out is not None else () + + +def operator_sink_keys(*configs: OpenTelemetryV2Config) -> frozenset[_SinkKey]: + """The accounts the operator's own exporters write to, in destination terms. + + Every v2 logger's config counts, since each logger exports through its own + provider. An exporter with no endpoint of its own resolves one from the + environment at export time, so it has no comparable identity and is left out, + and so is one that never reaches the wire: a console kind ignores the endpoint, + and a header-gated spec with no credentials is skipped when the provider is built. + """ + return frozenset( + key + for config in configs + for spec in config.exporters + if _exports_to_the_wire(spec) and (key := _sink_key(spec.endpoint, parse_headers(spec.headers))) is not None + ) + + +def _exports_to_the_wire(spec: ExporterSpec) -> bool: + """Whether ``build_tracer_provider`` gives ``spec`` an exporter that sends OTLP.""" + return exporter_transport(spec.kind) != "headerless" and not (spec.requires_headers and not spec.headers) + + +def _sink_key(endpoint: str | None, headers: Mapping[str, str]) -> "_SinkKey | None": + """The account an exporter writes to, or ``None`` when it has no fixed one. + + Normalized on the three counts that make one account look like two: the operator's + spec carries the signal path a tenant destination leaves for the exporter to + append, header names survive one round trip lowercased and the other not, and one + credential answers to more than one name (see :data:`_CREDENTIAL_ALIASES`). + """ + normalized: Final = _otlp_traces_endpoint(endpoint) + if normalized is None: + return None + return (normalized, tuple(sorted((_credential_name(name), value) for name, value in headers.items()))) + + +def _credential_name(header: str) -> str: + """The credential a header carries, under whichever name the backend spells it.""" + normalized: Final = header.strip().lower().replace("-", "_") + return _CREDENTIAL_ALIASES.get(normalized, normalized) + + +def _attached_processors(provider: trace.TracerProvider) -> "tuple[SpanProcessor, ...]": + """The processors already on ``provider``, or empty when the SDK hides them.""" + multi: Final = getattr(provider, "_active_span_processor", None) + return tuple(getattr(multi, "_span_processors", ())) + + def get_tracer(provider: TracerProvider, name: str = "litellm") -> Tracer: # Stamp the instrumentation scope with the LiteLLM package version so every # emitted span carries a deterministic ``scope.version`` (the standard OTel diff --git a/litellm/integrations/otel/plumbing/routing.py b/litellm/integrations/otel/plumbing/routing.py index 227e18f3663..f78d18d943c 100644 --- a/litellm/integrations/otel/plumbing/routing.py +++ b/litellm/integrations/otel/plumbing/routing.py @@ -25,6 +25,7 @@ from opentelemetry.trace import Tracer from litellm._logging import verbose_logger from litellm.constants import OTEL_SERVICE_NAME_METADATA_KEYS from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config +from litellm.integrations.otel.plumbing.context import destination_backends from litellm.integrations.otel.plumbing.providers import ( build_tracer_provider, exporter_transport, @@ -231,10 +232,21 @@ class TenantTracerCache: concurrent overflow eviction can't shut it down between selection and the caller's span start. The caller must ``release`` it exactly once. """ + # A backend with a destination is delivered by the fan-out processor, which + # carries the whole trace and already carries this tenant's credentials and + # service name. Routing here too would detach this span onto a second provider, + # so the tenant would get the request tree plus a stray one-span trace. + if self._callback_name is not None and self._callback_name in destination_backends(): + return TenantRoute(tracer=default, detached=False) credential_headers: Final = self._credential_headers(dynamic_params) project_headers: Final = self._project_headers(auth_metadata) service_name: Final = tenant_service_name(auth_metadata) - if not credential_headers and not project_headers and service_name is None: + tenant_account: Final = bool(credential_headers) or bool(project_headers) + # A service name on its own only relabels the operator's own backend, so moving + # the span to a second provider for it while some other backend has a + # destination would drop the model call out of the trace the fan-out delivers. + # The destination stamps the same service name itself. + if not tenant_account and (service_name is None or destination_backends()): return TenantRoute(tracer=default, detached=False) # A fixed per-integration region endpoint (New Relic us/eu), never a # caller-supplied host; ``None`` keeps the preset's own endpoint. @@ -255,7 +267,7 @@ class TenantTracerCache: _shutdown_provider(evicted) return TenantRoute( tracer=get_tracer(provider, self._tracer_name), - detached=bool(project_headers) or bool(credential_headers), + detached=tenant_account, provider=provider, ) diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index f45b1cd3cff..965213f2ee4 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -39,6 +39,7 @@ class _AgentOpsSettings(BaseSettings): def agentops_preset( *, config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, ) -> OpenTelemetryV2Config: """Build the AgentOps config without any network I/O. diff --git a/litellm/integrations/otel/presets/arize.py b/litellm/integrations/otel/presets/arize.py index ee0de675657..d7ce87f5552 100644 --- a/litellm/integrations/otel/presets/arize.py +++ b/litellm/integrations/otel/presets/arize.py @@ -26,10 +26,12 @@ class _ArizeSettings(BaseSettings): def arize_preset( *, config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, ) -> OpenTelemetryV2Config: + base: Final = config_overrides or OpenTelemetryV2Config() + mappers: Final = ensure_mappers(base.mapper_names, "openinference") arize_cfg: Final = _V1ArizeLogger.get_arize_config() headers: Final = _arize_headers(arize_cfg) - base: Final = config_overrides or OpenTelemetryV2Config() return base.model_copy( update={ "exporters": [ @@ -41,7 +43,7 @@ def arize_preset( owner=ExporterOwner.ARIZE_AX, ), ], - "mapper_names": ensure_mappers(base.mapper_names, "openinference"), + "mapper_names": mappers, "resource_attributes": { **base.resource_attributes, **({"model_id": arize_cfg.project_name} if arize_cfg.project_name else {}), diff --git a/litellm/integrations/otel/presets/base.py b/litellm/integrations/otel/presets/base.py index 3b9991f86a4..3a768a08a4f 100644 --- a/litellm/integrations/otel/presets/base.py +++ b/litellm/integrations/otel/presets/base.py @@ -18,6 +18,18 @@ class Preset(Protocol): ``config_overrides`` lets one preset layer onto another's config (or onto test-supplied defaults); the factory calls presets with no arguments. + + ``allow_missing_credentials`` lets a credential-mandatory backend (langfuse and + weave) degrade to an exporter-less, mapper-only config instead of raising when the + operator set no env credentials of their own. That is a real + deployment: every team brings its own account and the operator keeps none, and + without it the whole V2 path silently falls back to the legacy integration, so + no team destination is ever reached. Credential-optional backends ignore it. """ - def __call__(self, *, config_overrides: OpenTelemetryV2Config | None = None) -> OpenTelemetryV2Config: ... + def __call__( + self, + *, + config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, + ) -> OpenTelemetryV2Config: ... diff --git a/litellm/integrations/otel/presets/destinations.py b/litellm/integrations/otel/presets/destinations.py new file mode 100644 index 00000000000..2bf9bfa5261 --- /dev/null +++ b/litellm/integrations/otel/presets/destinations.py @@ -0,0 +1,152 @@ +"""Map a key's or team's callback vars to the OTLP destination its traces export to. + +Header building is delegated to each preset's existing ``*_dynamic_headers`` builder, +so a destination authenticates exactly the way the per-request tracer route already +did; only the endpoint and transport need a per-backend rule. +""" + +import os +from collections.abc import Callable, Mapping +from functools import lru_cache +from types import MappingProxyType +from typing import Final + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.otel.model.destination import OtelDestination +from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host +from litellm.types.utils import StandardCallbackDynamicParams + +#: An endpoint plus the OTLP transport to reach it with, or ``None`` when the backend +#: names no destination. The transport is ``None`` where the backend has only one. +_Destination = tuple[str, str | None] + + +@lru_cache(maxsize=128) +def _warn_host_not_allowlisted(host: str) -> None: + """Cached so one misconfigured team logs once rather than once per request.""" + verbose_logger.warning( + "OTel V2: not exporting to key/team Langfuse host '%s'. Add it to " + "litellm_settings.provider_url_destination_allowed_hosts to permit it", + host, + ) + + +def _langfuse_destination(params: StandardCallbackDynamicParams) -> "_Destination | None": + """The tenant's own Langfuse host, else the operator's, else Langfuse US cloud. + + A host the tenant named has to be allowlisted by the operator, the same way a + URL-valued ``model`` is: anyone who can mint a key can write it, and it becomes an + endpoint the proxy posts the request's whole trace to, carrying the tenant's own + credentials. The operator's own ``LANGFUSE_HOST`` is not checked, since an internal + collector there is a deployment choice. + """ + from litellm.integrations.langfuse.langfuse_otel import ( + LANGFUSE_CLOUD_US_ENDPOINT, + LangfuseOtelLogger, + ) + + tenant_host: Final = params.get("langfuse_host") or None + host: Final = tenant_host or LangfuseOtelLogger._get_langfuse_otel_host() # pyright: ignore[reportPrivateUsage] # reuse the backend's own env host resolver rather than duplicating it + if not host: + return (LANGFUSE_CLOUD_US_ENDPOINT, None) + normalized: Final = host if host.startswith("http") else f"https://{host}" + endpoint: Final = f"{normalized.rstrip('/')}/api/public/otel" + if tenant_host is None: + return (endpoint, None) + if not is_url_destination_allowed_by_host(endpoint, litellm.provider_url_destination_allowed_hosts): + _warn_host_not_allowlisted(host) + return None + return (endpoint, None) + + +def _arize_destination(params: StandardCallbackDynamicParams) -> "_Destination | None": + from litellm.integrations.arize.arize import ArizeLogger + + config: Final = ArizeLogger.get_arize_config() + return (config.endpoint, config.protocol) + + +def _weave_destination(params: StandardCallbackDynamicParams) -> "_Destination | None": + from litellm.integrations.weave.weave_otel import weave_otel_endpoint + + return (weave_otel_endpoint(os.environ.get("WANDB_HOST")), None) + + +def _newrelic_destination(params: StandardCallbackDynamicParams) -> "_Destination | None": + from litellm.integrations.otel.presets.newrelic import newrelic_dynamic_endpoint + + endpoint: Final = newrelic_dynamic_endpoint(params) + return (endpoint, None) if endpoint else None + + +#: Callback name -> destination resolver. A backend is destination-capable exactly +#: when it appears here AND in ``DYNAMIC_HEADERS_BY_CALLBACK``: without a header +#: builder the destination would carry no tenant credentials, and the exporter +#: would post the tenant's traffic to the operator's account. +_DESTINATION_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicParams], "_Destination | None"]]] = ( + MappingProxyType( + { + "langfuse_otel": _langfuse_destination, + "arize": _arize_destination, + "weave_otel": _weave_destination, + "newrelic": _newrelic_destination, + } + ) +) + +#: Headers a destination must carry to authenticate. Several dynamic-header builders +#: gate each credential independently, so a half-configured backend yields a non-empty +#: but unusable header set; accepting it would suppress the operator's own exporter and +#: send the request's whole trace where it cannot be stored. +_REQUIRED_HEADERS_BY_CALLBACK: Final[Mapping[str, frozenset[str]]] = MappingProxyType( + { + "langfuse_otel": frozenset({"Authorization"}), + "arize": frozenset({"arize-space-id", "api_key"}), + "weave_otel": frozenset({"Authorization", "project_id"}), + "newrelic": frozenset({"api-key"}), + } +) + +_NO_ATTRS: Final[Mapping[str, str]] = MappingProxyType({}) + + +def destination_capable_backends() -> frozenset[str]: + """Backends a key or team can point at its own account.""" + from litellm.integrations.otel.presets import DYNAMIC_HEADERS_BY_CALLBACK + + return frozenset(_DESTINATION_BY_CALLBACK) & frozenset(DYNAMIC_HEADERS_BY_CALLBACK) + + +def destination_for( + callback_name: str, + params: StandardCallbackDynamicParams, + service_name: str | None = None, +) -> OtelDestination | None: + """The destination ``params`` names for ``callback_name``, or ``None``. + + ``None`` means the caller configured nothing usable for this backend, so the + request keeps the operator's global exporters. ``service_name`` is the key's or + team's ``otel_service_name``, which the per-request tracer route applies when the + backend is not overridden and the destination has to apply once it is. + """ + from litellm.integrations.otel.presets import DYNAMIC_HEADERS_BY_CALLBACK + + header_builder: Final = DYNAMIC_HEADERS_BY_CALLBACK.get(callback_name) + destination_builder: Final = _DESTINATION_BY_CALLBACK.get(callback_name) + if header_builder is None or destination_builder is None: + return None + headers: Final = header_builder(params) + if not headers or not _REQUIRED_HEADERS_BY_CALLBACK[callback_name] <= frozenset(headers): + return None + resolved: Final = destination_builder(params) + if resolved is None: + return None + endpoint, protocol = resolved + return OtelDestination( + endpoint=endpoint, + headers=MappingProxyType(dict(headers)), # mutable-ok: MappingProxyType needs a concrete mapping to wrap + resource_attributes=MappingProxyType({"service.name": service_name}) if service_name else _NO_ATTRS, + callback_name=callback_name, + protocol=protocol, + ) diff --git a/litellm/integrations/otel/presets/langfuse.py b/litellm/integrations/otel/presets/langfuse.py index c2f64422eff..9149e0c0d94 100644 --- a/litellm/integrations/otel/presets/langfuse.py +++ b/litellm/integrations/otel/presets/langfuse.py @@ -10,17 +10,32 @@ from litellm.integrations.otel.model.config import ( ExporterSpec, OpenTelemetryV2Config, ) -from litellm.integrations.otel.presets.utils import ensure_mappers +from litellm.integrations.otel.presets.utils import ( + credential_gated_exporters, + ensure_mappers, +) from litellm.types.utils import StandardCallbackDynamicParams def langfuse_preset( *, config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, ) -> OpenTelemetryV2Config: - cfg: Final = _V1Langfuse.get_langfuse_otel_config() - kind: Final = cfg.exporter if isinstance(cfg.exporter, str) else "otlp_http" base: Final = config_overrides or OpenTelemetryV2Config() + mappers: Final = ensure_mappers(base.mapper_names, "langfuse") + try: + cfg: Final = _V1Langfuse.get_langfuse_otel_config() + except Exception: + if not allow_missing_credentials: + raise + return base.model_copy( + update={ # mutable-ok: pydantic model_copy takes a plain update mapping + "exporters": credential_gated_exporters(base.exporters, ExporterOwner.LANGFUSE_OTEL), + "mapper_names": mappers, + } + ) + kind: Final = cfg.exporter if isinstance(cfg.exporter, str) else "otlp_http" return base.model_copy( update={ "exporters": [ @@ -32,7 +47,7 @@ def langfuse_preset( owner=ExporterOwner.LANGFUSE_OTEL, ), ], - "mapper_names": ensure_mappers(base.mapper_names, "langfuse"), + "mapper_names": mappers, } ) diff --git a/litellm/integrations/otel/presets/langtrace.py b/litellm/integrations/otel/presets/langtrace.py index c88e4715ab0..2312575f04a 100644 --- a/litellm/integrations/otel/presets/langtrace.py +++ b/litellm/integrations/otel/presets/langtrace.py @@ -9,6 +9,7 @@ from litellm.integrations.otel.presets.utils import ensure_mappers def langtrace_preset( *, config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, ) -> OpenTelemetryV2Config: """Compose the Langtrace mapper on top of the customer's OTLP destination. diff --git a/litellm/integrations/otel/presets/levo.py b/litellm/integrations/otel/presets/levo.py index 41b3758cf3e..c1580cf7a5b 100644 --- a/litellm/integrations/otel/presets/levo.py +++ b/litellm/integrations/otel/presets/levo.py @@ -13,6 +13,7 @@ from litellm.integrations.otel.model.config import ( def levo_preset( *, config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, ) -> OpenTelemetryV2Config: cfg: Final = _V1Levo.get_levo_config() base: Final = config_overrides or OpenTelemetryV2Config() diff --git a/litellm/integrations/otel/presets/newrelic.py b/litellm/integrations/otel/presets/newrelic.py index 4660a707355..771b3f643c8 100644 --- a/litellm/integrations/otel/presets/newrelic.py +++ b/litellm/integrations/otel/presets/newrelic.py @@ -44,6 +44,7 @@ class _NewRelicSettings(BaseSettings): def newrelic_preset( *, config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, ) -> OpenTelemetryV2Config: settings: Final = _NewRelicSettings() base: Final = config_overrides or OpenTelemetryV2Config() diff --git a/litellm/integrations/otel/presets/phoenix.py b/litellm/integrations/otel/presets/phoenix.py index eef407b6c1b..f4f34ee7525 100644 --- a/litellm/integrations/otel/presets/phoenix.py +++ b/litellm/integrations/otel/presets/phoenix.py @@ -60,6 +60,7 @@ def phoenix_project_headers(auth_metadata: Mapping[str, str] | None) -> Mapping[ def phoenix_preset( *, config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, ) -> OpenTelemetryV2Config: cfg: Final = _V1Phoenix.get_arize_phoenix_config() headers: Final = cfg.otlp_auth_headers if hasattr(cfg, "otlp_auth_headers") else None diff --git a/litellm/integrations/otel/presets/utils.py b/litellm/integrations/otel/presets/utils.py index 328569d3daf..1270c41e77b 100644 --- a/litellm/integrations/otel/presets/utils.py +++ b/litellm/integrations/otel/presets/utils.py @@ -3,6 +3,8 @@ from collections.abc import Iterable from typing import Final +from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec + def ensure_mappers(mapper_names: Iterable[str], *names: str) -> list[str]: """Return ``mapper_names`` with each of ``names`` appended if not already present. @@ -15,3 +17,32 @@ def ensure_mappers(mapper_names: Iterable[str], *names: str) -> list[str]: if name not in result: result.append(name) return result + + +def credential_gated_exporters( + exporters: "Iterable[ExporterSpec]", owner: "ExporterOwner" +) -> "tuple[ExporterSpec, ...]": + """``exporters`` with the operator's destination replaced by a header-gated one. + + Used when a credential-mandatory backend is asked to build without the operator's + own credentials, so only key/team destinations receive spans. Two things have to + happen for that to mean "export nowhere": the placeholder console spec that + ``OpenTelemetryV2Config`` folds in for an empty exporter list is dropped, or every + span would be printed to stdout, and the gated spec keeps the owner so the + override filter still recognises which backend this provider speaks for. + """ + return ( + *(spec for spec in exporters if not is_unconfigured_placeholder(spec)), + ExporterSpec(owner=owner, requires_headers=True), + ) + + +def is_unconfigured_placeholder(spec: "ExporterSpec") -> bool: + """Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured. + + No field set is what says the operator asked for nothing: an exporter they did + configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default, + and so does the gated spec this module appends, which would otherwise eat itself + when one preset layers onto another. + """ + return not spec.model_fields_set diff --git a/litellm/integrations/otel/presets/weave.py b/litellm/integrations/otel/presets/weave.py index 51d0ad01093..644cd39ad36 100644 --- a/litellm/integrations/otel/presets/weave.py +++ b/litellm/integrations/otel/presets/weave.py @@ -7,7 +7,10 @@ from litellm.integrations.otel.model.config import ( ExporterSpec, OpenTelemetryV2Config, ) -from litellm.integrations.otel.presets.utils import ensure_mappers +from litellm.integrations.otel.presets.utils import ( + credential_gated_exporters, + ensure_mappers, +) from litellm.integrations.weave.weave_otel import ( _get_weave_authorization_header, get_weave_otel_config, @@ -18,9 +21,21 @@ from litellm.types.utils import StandardCallbackDynamicParams def weave_preset( *, config_overrides: OpenTelemetryV2Config | None = None, + allow_missing_credentials: bool = False, ) -> OpenTelemetryV2Config: - weave_cfg: Final = get_weave_otel_config() base: Final = config_overrides or OpenTelemetryV2Config() + mappers: Final = ensure_mappers(base.mapper_names, "openinference", "weave") + try: + weave_cfg: Final = get_weave_otel_config() + except Exception: + if not allow_missing_credentials: + raise + return base.model_copy( + update={ # mutable-ok: pydantic model_copy takes a plain update mapping + "exporters": credential_gated_exporters(base.exporters, ExporterOwner.WEAVE_OTEL), + "mapper_names": mappers, + } + ) return base.model_copy( update={ "exporters": [ @@ -33,7 +48,7 @@ def weave_preset( ), ], # Weave consumes OpenInference + a small Weave-specific overlay. - "mapper_names": ensure_mappers(base.mapper_names, "openinference", "weave"), + "mapper_names": mappers, } ) diff --git a/litellm/integrations/pointfive/__init__.py b/litellm/integrations/pointfive/__init__.py new file mode 100644 index 00000000000..1f3ca3c65c7 --- /dev/null +++ b/litellm/integrations/pointfive/__init__.py @@ -0,0 +1,5 @@ +"""PointFive logging integration for LiteLLM.""" + +from litellm.integrations.pointfive.logger import PointFiveLogger + +__all__ = ("PointFiveLogger",) diff --git a/litellm/integrations/pointfive/logger.py b/litellm/integrations/pointfive/logger.py new file mode 100644 index 00000000000..c352dac11e7 --- /dev/null +++ b/litellm/integrations/pointfive/logger.py @@ -0,0 +1,304 @@ +""" +PointFive logging integration. + +Buffers ``StandardLoggingPayload`` records and ships each flush as one gzipped +newline-delimited JSON object, rather than one object per request. Uploads go through a +presigned URL issued by the PointFive API, so the proxy needs no cloud credentials and +runs unchanged wherever it is hosted. +""" + +import asyncio +from collections.abc import Mapping +from datetime import datetime +from typing import Final + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.pointfive.payload import chunk_lines, encode_lines, serialize_records +from litellm.integrations.pointfive.upload_client import PointFiveUploadClient, PointFiveUploadError +from litellm.litellm_core_utils.redact_messages import ( + redacted_standard_logging_payload, + should_redact_message_logging, +) +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client, httpxSpecialProvider +from litellm.secret_managers.main import get_secret_str +from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus +from litellm.types.integrations.pointfive import DEFAULT_API_URL, PointFiveInitParams, PointFiveUploadFailure + +_ENV_REFERENCE_PREFIX: Final = "os.environ/" + + +def _resolved_secret(value: str | None) -> str | None: + """ + Resolve a config value that may name a secret, in any shape the secret manager accepts. + + A reference that resolves to nothing stays unresolved rather than falling back to its own + text, so an unset ``os.environ/NAME`` reports a missing key instead of being sent as one. + """ + if value is None: + return None + resolved: Final = get_secret_str(value) + if resolved: + return resolved + return None if value.startswith(_ENV_REFERENCE_PREFIX) else value + + +def _configured_params() -> PointFiveInitParams: + """Read ``litellm.pointfive_params``, validating a raw config dict on the way through.""" + configured: Final = litellm.pointfive_params + if isinstance(configured, PointFiveInitParams): + return configured + if isinstance(configured, Mapping): + return PointFiveInitParams.model_validate(configured) + return PointFiveInitParams() + + +def _resolved_api_key(params: PointFiveInitParams) -> str | None: + """Prefer the configured key, falling back to the environment the proxy UI writes.""" + return _resolved_secret(params.api_key) or get_secret_str("POINTFIVE_API_KEY") + + +def _resolved_api_url(params: PointFiveInitParams) -> str: + """Prefer the configured url, then the environment, then the public endpoint.""" + return _resolved_secret(params.api_url) or get_secret_str("POINTFIVE_API_URL") or DEFAULT_API_URL + + +def _upload_client_for(params: PointFiveInitParams) -> PointFiveUploadClient: + """ + Build an upload client for the key and url configured right now. + + Resolved per call rather than kept: the proxy ui writes new values into the + environment of a running proxy, and reading them once would need a restart to take + effect. ``get_async_httpx_client`` is cached, so this reuses the same connections. + """ + api_key: Final = _resolved_api_key(params) + if not api_key: + raise ValueError( + "pointfive logging requires an api key. Set POINTFIVE_API_KEY, or " + "litellm_settings.pointfive_params.api_key in config.yaml" + ) + return PointFiveUploadClient( + api_key=api_key, + api_url=_resolved_api_url(params), + http_client=get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback), + max_retries=params.max_upload_retries, + ) + + +class PointFiveLogger(CustomBatchLogger): + """Batching callback that ships LiteLLM request logs to PointFive.""" + + preserve_events_added_during_flush = True + + def __init__( + self, + params: PointFiveInitParams | None = None, + upload_client: PointFiveUploadClient | None = None, + start_periodic_flush: bool = True, + ) -> None: + resolved: Final = params if params is not None else _configured_params() + self.max_batch_bytes: Final = resolved.max_batch_bytes + self.params: Final = resolved + self.given_upload_client: Final = upload_client + if upload_client is None: + _upload_client_for(resolved) # refuse to start without a key, rather than at the first flush + super().__init__( + flush_lock=asyncio.Lock(), + batch_size=resolved.batch_size, + flush_interval=resolved.flush_interval, + turn_off_message_logging=bool(resolved.turn_off_message_logging), + ) + self._flushing: bool = False + self._batch_flush_task: asyncio.Task[None] | None = None + self._periodic_flush_task: asyncio.Task[None] | None = ( + self._start_periodic_flush_task() if start_periodic_flush else None + ) + + @property + def upload_client(self) -> PointFiveUploadClient: + """The client for the currently configured key and url, so a ui edit needs no restart.""" + if self.given_upload_client is not None: + return self.given_upload_client + return _upload_client_for(self.params) + + def _start_periodic_flush_task(self) -> asyncio.Task[None] | None: + """Start the periodic flush only once an event loop is actually running.""" + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return None + return loop.create_task(self.periodic_flush()) + + def _start_batch_flush_task(self) -> None: + """ + Upload a full batch in the background, so no request waits on PointFive. + + Awaiting it here put the upload, its retries and their backoff on the caller's + path, and a hung api held a response open for as long as the attempts took. + """ + if self._batch_flush_task is not None and not self._batch_flush_task.done(): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + self._batch_flush_task = loop.create_task(self.flush_queue(skip_if_flushing=True)) + + def _flush_task_is_alive(self) -> bool: + """A task whose loop has been closed never runs again, yet never reports itself done.""" + task: Final = self._periodic_flush_task + return task is not None and not task.done() and not task.get_loop().is_closed() + + async def periodic_flush(self) -> None: + """ + Report in straight away, then flush on the interval as usual. + + The inherited loop sleeps first, so a proxy that has just loaded the callback says + nothing for a whole interval, five minutes by default. PointFive shows the integration + as still waiting for its first call for all that time, which reads as a broken setup + rather than an idle one. An empty queue makes this first cycle a ping, so a proxy with + no traffic yet announces itself without uploading an object that holds no records. + """ + await self.flush_queue(skip_if_flushing=True) + await super().periodic_flush() + + async def async_log_success_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + await self._enqueue(kwargs) + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + await self._enqueue(kwargs) + + async def _enqueue(self, kwargs: Mapping[str, object]) -> None: + """Buffer one record, flushing early once the batch threshold is reached.""" + try: + if not self._flush_task_is_alive(): + self._periodic_flush_task = self._start_periodic_flush_task() + + record: Final = self._record_for(kwargs) + if record is None: + verbose_logger.debug("pointfive: event carried no standard_logging_object, skipping") + return + + self.log_queue.append(record) + self._drop_overflow() + if len(self.log_queue) >= self.batch_size: + self._start_batch_flush_task() + except Exception: # noqa: BLE001 # logging must never break the request path + verbose_logger.exception("pointfive: failed to queue an event") + + def _record_for(self, kwargs: Mapping[str, object]) -> Mapping[str, object] | None: + """ + The record to buffer, redacted the way the framework would have redacted it. + + A success reaches a callback already redacted, an async failure does not, so both + the excluded-field list and this callback's own setting are applied here, then the + global, per-request and header settings that only the framework's predicate knows. + """ + details: Final = self.redact_standard_logging_payload_from_model_call_details( + dict(kwargs) # mutable-ok: both framework helpers take the call details as a dict + ) + payload: Final = details.get("standard_logging_object") + if not isinstance(payload, dict): + return None + if should_redact_message_logging(details): + return redacted_standard_logging_payload(payload) + return payload + + def _drop_overflow(self) -> None: + """ + Hold the queue to its cap as records arrive, not only after a flush has failed. + + Never while a flush is running: it holds a snapshot taken by length, and trimming + the front underneath it would make the post-flush drain remove records that arrived + during the upload and were never sent. The next arrival after the flush trims. + """ + if self._flushing: + return + overflow: Final = len(self.log_queue) - self.max_queue_size + if overflow <= 0: + return + del self.log_queue[:overflow] + verbose_logger.warning("pointfive: queue over %s records, dropped %s oldest", self.max_queue_size, overflow) + + async def flush_queue(self, skip_if_flushing: bool = False) -> None: + """ + Flush as usual, or report liveness when there is nothing to send. + + ``CustomBatchLogger`` skips an empty queue entirely, so without this an idle proxy + would look identical to a dead one. + + ``skip_if_flushing`` is what a full batch, and the loop's opening cycle, pass. Uploading one takes seconds, and + every event arriving meanwhile crosses the threshold too, so each would queue on the + flush lock and then ship the handful of records left behind it. That turns one burst + into a stream of tiny objects, which is what batching exists to avoid. The running + flush already carries what is queued, and the interval catches whatever it missed. + """ + if not self.log_queue: + await self._ping() + return + if skip_if_flushing and self._flushing: + return + + self._flushing = True + try: + await super().flush_queue() + finally: + self._flushing = False + + async def async_health_check(self) -> IntegrationHealthCheckStatus: + """Answer the proxy ui test button by asking the api whether it accepts this key.""" + try: + failure: Final = await self.upload_client.ping() + except ValueError as missing_key: + return IntegrationHealthCheckStatus(status="unhealthy", error_message=str(missing_key)) + if failure is not None: + return IntegrationHealthCheckStatus(status="unhealthy", error_message=failure.detail) + return IntegrationHealthCheckStatus(status="healthy", error_message=None) + + async def _ping(self) -> None: + """Report liveness, never failing the flush over it.""" + try: + failure: Final = await self.upload_client.ping() + except ValueError as missing_key: + verbose_logger.warning("pointfive: liveness ping skipped, %s", missing_key) + return + if failure is not None: + verbose_logger.warning("pointfive: liveness ping failed, %s", failure.detail) + + async def async_send_batch(self) -> None: + """ + Upload everything queued, split into objects of at most ``max_batch_bytes``. + + A retryable failure propagates so ``CustomBatchLogger`` keeps the rest of the batch + for the next flush; the records already shipped or already refused leave the queue + first, so a retry re-sends at most the object that failed. A rejection the server + will refuse again drops that object, since holding it would block every record + queued behind it. + """ + pending: Final = tuple(self.log_queue) + if not pending: + return + + client: Final = self.upload_client + chunks: Final = chunk_lines(serialize_records(pending), self.max_batch_bytes) + for index, chunk in enumerate(chunks): + outcome = await client.upload(await encode_lines(chunk)) + if not isinstance(outcome, PointFiveUploadFailure): + continue + if outcome.retryable: + del self.log_queue[: sum(len(shipped) for shipped in chunks[:index])] + raise PointFiveUploadError(outcome.detail) + verbose_logger.error("pointfive: dropping %s records, %s", len(chunk), outcome.detail) diff --git a/litellm/integrations/pointfive/payload.py b/litellm/integrations/pointfive/payload.py new file mode 100644 index 00000000000..e3362eda4ee --- /dev/null +++ b/litellm/integrations/pointfive/payload.py @@ -0,0 +1,53 @@ +"""Turns buffered log records into the gzipped NDJSON objects that get uploaded.""" + +import gzip +from collections.abc import Iterator, Mapping, Sequence +from itertools import accumulate, groupby, islice +from typing import Final + +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + +_NEWLINE_BYTES: Final = 1 + + +def serialize_records(records: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + """Serialize each record to one JSON line.""" + return tuple(safe_dumps(record) for record in records) + + +def _encoded_size(line: str) -> int: + return len(line.encode("utf-8")) + _NEWLINE_BYTES + + +def _object_indices(sizes: Sequence[int], max_bytes: int) -> Iterator[int]: + """Number each line with the object it belongs to, opening a new one on overflow.""" + + def advance(state: tuple[int, int], size: int) -> tuple[int, int]: + index, used = state + return (index + 1, size) if used and used + size > max_bytes else (index, used + size) + + return (index for index, _ in islice(accumulate(sizes, advance, initial=(0, 0)), 1, None)) + + +def chunk_lines(lines: Sequence[str], max_bytes: int) -> tuple[tuple[str, ...], ...]: + """ + Group serialized lines into objects of at most ``max_bytes`` uncompressed. + + A line above the bound on its own still becomes its own object. A record cannot be + split, and holding it back would stall every record queued behind it. + """ + sizes: Final = tuple(_encoded_size(line) for line in lines) + numbered: Final = zip(_object_indices(sizes, max_bytes), lines, strict=True) + return tuple(tuple(line for _, line in group) for _, group in groupby(numbered, lambda pair: pair[0])) + + +async def encode_lines(lines: Sequence[str]) -> bytes: + """ + Join lines as NDJSON and gzip them off the event loop. + + An object can be several megabytes, and compressing that inline would block the + proxy for as long as it takes. + """ + compress: Final = asyncify(gzip.compress) + return await compress("\n".join(lines).encode("utf-8")) diff --git a/litellm/integrations/pointfive/upload_client.py b/litellm/integrations/pointfive/upload_client.py new file mode 100644 index 00000000000..56ba6689017 --- /dev/null +++ b/litellm/integrations/pointfive/upload_client.py @@ -0,0 +1,194 @@ +""" +Uploads one batch to PointFive through a presigned URL. + +The proxy holds no cloud credentials. For every batch it asks the PointFive API for a +single-use presigned URL and PUTs the bytes there, so the same plugin runs unchanged on +AWS, GCP, Azure or on-prem. The server picks the object key, so the proxy never chooses +where its data lands. +""" + +import asyncio +from collections.abc import Awaitable, Callable +from types import MappingProxyType +from typing import Final + +import httpx +from pydantic import BaseModel, Field, ValidationError + +import litellm +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import SSRFError, validate_url +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.integrations.pointfive import ( + RETRYABLE_UPLOAD_STATUS_CODES, + PointFiveUploadFailure, + PointFiveUploadTarget, +) + +UPLOAD_KIND: Final = "LITELLM" +UPLOAD_URL_PATH: Final = "/upload-url" +PING_PATH: Final = "/ping" +PUT_HEADERS: Final = MappingProxyType({"Content-Type": "application/x-ndjson", "Content-Encoding": "gzip"}) + + +class _PresignRequest(BaseModel): + kind: str = UPLOAD_KIND + byte_count: int = Field(serialization_alias="byteCount") + + +class _PingRequest(BaseModel): + kind: str = UPLOAD_KIND + + +class _TargetPayload(BaseModel): + upload_url: str = Field(alias="uploadUrl") + object_key: str = Field(alias="objectKey") + + +class _ErrorPayload(BaseModel): + error: str = "" + + +class PointFiveUploadError(Exception): + """A batch could not be uploaded and the failure is worth retrying.""" + + +def _failure_for(response: httpx.Response, what: str) -> PointFiveUploadFailure: + detail: Final = f"{what} returned {response.status_code}" + reason: Final = _refusal_reason(response.text) + return PointFiveUploadFailure( + f"{detail}, {reason}" if reason else detail, + retryable=response.status_code in RETRYABLE_UPLOAD_STATUS_CODES, + ) + + +def _refusal_reason(body: str) -> str: + try: + return _ErrorPayload.model_validate_json(body).error + except ValidationError: + return "" + + +def _parse_target(body: str) -> PointFiveUploadTarget | PointFiveUploadFailure: + try: + target: Final = _TargetPayload.model_validate_json(body) + except ValidationError: + return PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False) + return PointFiveUploadTarget(upload_url=target.upload_url, object_key=target.object_key) + + +class PointFiveUploadClient: + """Presigns and uploads one batch at a time.""" + + def __init__( + self, + api_key: str, + api_url: str, + http_client: AsyncHTTPHandler, + max_retries: int, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + validate_upload_url: Callable[[str], tuple[str, str]] = validate_url, + ) -> None: + self.api_key: Final = api_key + self.api_url: Final = api_url.rstrip("/") + self.http_client: Final = http_client + self.max_retries: Final = max_retries + self.sleep: Final = sleep + self.validate_upload_url: Final = validate_upload_url + + async def upload(self, body: bytes) -> str | PointFiveUploadFailure: + """ + Upload one gzipped batch, returning the object key it landed at. + + Every attempt presigns again, so a retry never reuses a URL that has expired or + has already been consumed. + """ + for attempt in range(self.max_retries): + match await self._upload_once(body): + case PointFiveUploadFailure(retryable=True) as failure: + if attempt + 1 >= self.max_retries: + return PointFiveUploadFailure( + f"{failure.detail}, gave up after {self.max_retries} attempts", retryable=True + ) + await self.sleep(float(1 << attempt)) + case outcome: + return outcome + return PointFiveUploadFailure("max_upload_retries must be at least 1", retryable=False) + + async def _upload_once(self, body: bytes) -> str | PointFiveUploadFailure: + target: Final = await self._presign(len(body)) + if isinstance(target, PointFiveUploadFailure): + return target + + rejection: Final = await self._put(target, body) + if rejection is not None: + return rejection + + verbose_logger.debug("pointfive: uploaded %s gzipped bytes to %s", len(body), target.object_key) + return target.object_key + + async def ping(self) -> PointFiveUploadFailure | None: + """Report that the proxy is alive when it has nothing to upload.""" + body: Final = await self._post(PING_PATH, _PingRequest()) + if isinstance(body, PointFiveUploadFailure): + return body + return None + + async def _presign(self, byte_count: int) -> PointFiveUploadTarget | PointFiveUploadFailure: + """Ask the PointFive API for a presigned URL sized to this batch.""" + body: Final = await self._post(UPLOAD_URL_PATH, _PresignRequest(byte_count=byte_count)) + if isinstance(body, PointFiveUploadFailure): + return body + return _parse_target(body) + + async def _post(self, path: str, request: BaseModel) -> str | PointFiveUploadFailure: + """POST one JSON request to the PointFive ingestion API and return its raw body.""" + try: + response: Final = await self.http_client.post( + self.api_url + path, + json=request.model_dump(by_alias=True), + headers={ # mutable-ok: AsyncHTTPHandler.post types headers as dict + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + ) + except httpx.HTTPStatusError as e: + return _failure_for(e.response, "pointfive api") + except Exception as e: # noqa: BLE001 # a transport fault is worth another attempt + return PointFiveUploadFailure(f"pointfive api unreachable: {type(e).__name__}", retryable=True) + return response.text + + async def _put(self, target: PointFiveUploadTarget, body: bytes) -> PointFiveUploadFailure | None: + """ + PUT the batch to the presigned URL, which carries its own authorization. + + The server chose that URL, so it is treated like any other externally supplied + destination: the host is checked against blocked networks before connecting, and + a redirect is refused rather than followed. A presigned URL never legitimately + redirects, and following one would let a compromised endpoint point the proxy at + an internal service. + """ + destination: Final = self._destination(target.upload_url) + if isinstance(destination, PointFiveUploadFailure): + return destination + url, host = destination + headers: Final = dict(PUT_HEADERS, Host=host) if host else dict(PUT_HEADERS) # mutable-ok: put wants dict + try: + await self.http_client.put(url, data=body, headers=headers, follow_redirects=False) + except httpx.HTTPStatusError as e: + if e.response.is_redirect: + return PointFiveUploadFailure( + f"presigned upload redirected with {e.response.status_code}, refusing to follow", retryable=False + ) + return _failure_for(e.response, "presigned upload") + except Exception as e: # noqa: BLE001 # a transport fault is worth another attempt + return PointFiveUploadFailure(f"presigned upload unreachable: {type(e).__name__}", retryable=True) + return None + + def _destination(self, upload_url: str) -> tuple[str, str | None] | PointFiveUploadFailure: + if not getattr(litellm, "user_url_validation", True): + return upload_url, None + try: + return self.validate_upload_url(upload_url) + except SSRFError as e: + return PointFiveUploadFailure(f"presigned upload url refused: {e}", retryable=False) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 6766d246894..540ce6738fc 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -246,6 +246,7 @@ class PrometheusLogger(CustomLogger): # logger so toggling these flags only takes effect after a # restart, keeping init-time and runtime label sets in sync. self._cached_metric_labels: dict[str, list[str]] = {} + self._emit_input_sequence_length_label = litellm.prometheus_emit_input_sequence_length_label is True _custom_buckets: Final = litellm.prometheus_latency_buckets self.latency_buckets = tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS @@ -1522,6 +1523,11 @@ class PrometheusLogger(CustomLogger): # 2. Pyright does not allow us to run isinstance(standard_logging_payload, StandardLoggingPayload) <- this would be ideal enum_values=enum_values, label_context=label_context, + input_sequence_length=( + self._get_input_sequence_length(standard_logging_payload, kwargs, response_obj) + if self._emit_input_sequence_length_label + else None + ), ) # set x-ratelimit headers @@ -2192,6 +2198,36 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set(remaining_tokens) + @staticmethod + def _get_input_sequence_length( + standard_logging_payload: StandardLoggingPayload, + kwargs: Mapping[str, object], + response_obj: object, + ) -> str: + prompt_tokens: Final = standard_logging_payload.get("prompt_tokens") + if prompt_tokens: + return get_input_sequence_length_bucket(prompt_tokens) + combined_usage: Final = kwargs.get("combined_usage_object") + if ( + combined_usage is not None + and getattr(kwargs.get("_litellm_upstream_reported_usage"), "total_tokens", None) is not None + ): + return get_input_sequence_length_bucket(None) + reported_usage: Final = ( + response_obj.get("usage") if isinstance(response_obj, dict) else getattr(response_obj, "usage", None) + ) + if reported_usage is None and combined_usage is None: + return get_input_sequence_length_bucket(None) + usage_metadata: Final = standard_logging_payload["metadata"].get("usage_object") + if isinstance(usage_metadata, Mapping): + return get_input_sequence_length_bucket(usage_metadata.get("prompt_tokens")) + if combined_usage is None and isinstance(response_obj, dict): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + normalized_usage: Final[Mapping[str, object]] = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj) + return get_input_sequence_length_bucket(normalized_usage.get("prompt_tokens")) + return get_input_sequence_length_bucket(prompt_tokens) + def _set_latency_metrics( self, kwargs: dict, @@ -2202,7 +2238,16 @@ class PrometheusLogger(CustomLogger): user_api_team_alias: str | None, enum_values: UserAPIKeyLabelValues, label_context: PrometheusLabelFactoryContext | None = None, + input_sequence_length: str | None = None, ): + latency_enum_values: Final = ( + replace(enum_values, input_sequence_length=input_sequence_length) + if input_sequence_length is not None + else enum_values + ) + latency_label_context: Final = ( + PrometheusLabelFactoryContext(latency_enum_values) if input_sequence_length is not None else label_context + ) # latency metrics end_time: Final[datetime] = kwargs.get("end_time") or datetime.now() start_time: Final[datetime | None] = kwargs.get("start_time") @@ -2220,8 +2265,8 @@ class PrometheusLogger(CustomLogger): supported_enum_labels=self.get_labels_for_metric( metric_name="litellm_llm_api_time_to_first_token_metric" ), - enum_values=enum_values, - label_context=label_context, + enum_values=latency_enum_values, + label_context=latency_label_context, ) self.litellm_llm_api_time_to_first_token_metric.labels(**_ttft_labels).observe(time_to_first_token_seconds) self._track_end_user_metric_series( @@ -2241,8 +2286,8 @@ class PrometheusLogger(CustomLogger): if api_call_total_time_seconds is not None: _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_llm_api_latency_metric"), - enum_values=enum_values, - label_context=label_context, + enum_values=latency_enum_values, + label_context=latency_label_context, ) self.litellm_llm_api_latency_metric.labels(**_labels).observe(api_call_total_time_seconds) self._track_end_user_metric_series( @@ -2272,8 +2317,8 @@ class PrometheusLogger(CustomLogger): ) _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_request_total_latency_metric"), - enum_values=enum_values, - label_context=label_context, + enum_values=latency_enum_values, + label_context=latency_label_context, ) self.litellm_request_total_latency_metric.labels(**_labels).observe(_observed_total_time_seconds) self._track_end_user_metric_series( diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 712ce41d09e..972ac79e306 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -10,9 +10,11 @@ import asyncio import time from collections.abc import Mapping from datetime import datetime -from typing import Final, cast +from typing import TYPE_CHECKING, Final, cast from urllib.parse import quote +import httpx + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS @@ -24,7 +26,7 @@ from litellm.integrations.s3 import ( from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -35,6 +37,9 @@ from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload from .custom_batch_logger import CustomBatchLogger +if TYPE_CHECKING: + from botocore.credentials import Credentials + class S3Logger(CustomBatchLogger, BaseAWSLLM): def __init__( @@ -232,6 +237,26 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): f"{get_aws_dns_suffix(self.s3_region_name)}/{encoded_key}" ) + def _sign_put( + self, credentials: "Credentials", url: str, json_string: str, headers: Mapping[str, str] + ) -> dict[str, str]: # mutable-ok: [LIT001] AsyncHTTPHandler.put/HTTPHandler.put only accept dict headers + """ + ``RefreshableCredentials`` (IMDS roles) may refresh between the access key, secret and token + reads SigV4 performs, producing a mixed-generation signature that S3 rejects with 403. + Freezing first makes the three values one atomic snapshot. + """ + from botocore.auth import S3SigV4Auth + from botocore.awsrequest import AWSRequest + from botocore.credentials import RefreshableCredentials + + frozen: Final = ( + credentials.get_frozen_credentials() if isinstance(credentials, RefreshableCredentials) else credentials + ) + aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=dict(headers)) + aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) + S3SigV4Auth(frozen, "s3", aws_region_name).add_auth(aws_request) + return dict(aws_request.headers.items()) + def _sse_headers(self) -> Mapping[str, str]: candidates: Final = { "x-amz-server-side-encryption": self.s3_server_side_encryption, @@ -317,26 +342,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): try: import base64 import hashlib - - from botocore.auth import S3SigV4Auth - from botocore.awsrequest import AWSRequest except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") try: from litellm.litellm_core_utils.asyncify import asyncify asyncified_get_credentials: Final = asyncify(self.get_credentials) - credentials: Final = await asyncified_get_credentials( - aws_access_key_id=self.s3_aws_access_key_id, - aws_secret_access_key=self.s3_aws_secret_access_key, - aws_session_token=self.s3_aws_session_token, - aws_region_name=self.s3_region_name, - aws_session_name=self.s3_aws_session_name, - aws_profile_name=self.s3_aws_profile_name, - aws_role_name=self.s3_aws_role_name, - aws_web_identity_token=self.s3_aws_web_identity_token, - aws_sts_endpoint=self.s3_aws_sts_endpoint, - ) verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key) verbose_logger.debug("s3_v2 logger - s3_verify setting: %s", self.s3_verify) @@ -363,19 +374,28 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): **self._sse_headers(), } - # Sign the request - aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) - aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) + async def signed_put() -> httpx.Response: + credentials: Final = await asyncified_get_credentials( + aws_access_key_id=self.s3_aws_access_key_id, + aws_secret_access_key=self.s3_aws_secret_access_key, + aws_session_token=self.s3_aws_session_token, + aws_region_name=self.s3_region_name, + aws_session_name=self.s3_aws_session_name, + aws_profile_name=self.s3_aws_profile_name, + aws_role_name=self.s3_aws_role_name, + aws_web_identity_token=self.s3_aws_web_identity_token, + aws_sts_endpoint=self.s3_aws_sts_endpoint, + ) + signed_headers: Final = await run_aws_signing(self._sign_put, credentials, url, json_string, headers) + try: + return await self.async_httpx_client.put(url, data=json_string, headers=signed_headers) + except httpx.HTTPStatusError as error: + return error.response - # Prepare the signed headers - signed_headers: Final = dict(aws_request.headers.items()) - - # Make the request with retry for transient S3 errors (500/503) max_retries: Final = 3 for attempt in range(max_retries): - response = await self.async_httpx_client.put(url, data=json_string, headers=signed_headers) - if response.status_code in (500, 503) and attempt < max_retries - 1: + response = await signed_put() + if response.status_code in (403, 500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s", @@ -479,20 +499,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): try: import base64 import hashlib - - from botocore.auth import S3SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") try: verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key) - credentials: Final[Credentials] = self.get_credentials( - aws_access_key_id=self.s3_aws_access_key_id, - aws_secret_access_key=self.s3_aws_secret_access_key, - aws_session_token=self.s3_aws_session_token, - aws_region_name=self.s3_region_name, - ) url: Final = self._build_object_url(batch_logging_element.s3_object_key) @@ -516,22 +526,24 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): **self._sse_headers(), } - # Sign the request - aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) - aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) - - # Prepare the signed headers - signed_headers: Final = dict(aws_request.headers.items()) - httpx_client: Final = _get_httpx_client( params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None) ) - # Make the request with retry for transient S3 errors (500/503) + + def signed_put() -> httpx.Response: + credentials: Final = self.get_credentials( + aws_access_key_id=self.s3_aws_access_key_id, + aws_secret_access_key=self.s3_aws_secret_access_key, + aws_session_token=self.s3_aws_session_token, + aws_region_name=self.s3_region_name, + ) + signed_headers: Final = self._sign_put(credentials, url, json_string, headers) + return httpx_client.put(url, data=json_string, headers=signed_headers) + max_retries: Final = 3 for attempt in range(max_retries): - response = httpx_client.put(url, data=json_string, headers=signed_headers) - if response.status_code in (500, 503) and attempt < max_retries - 1: + response = signed_put() + if response.status_code in (403, 500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s", @@ -597,7 +609,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="GET", url=url, headers=headers) - S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request) + await run_aws_signing(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 2b4c8c9928d..d787375ca3c 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -22,7 +22,7 @@ from litellm.constants import ( SQS_SEND_MESSAGE_ACTION, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): data=prepped.body, headers=prepped.headers, ) - SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth(aws_request) + await run_aws_signing(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request) signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 12ff38ce4ba..b2243060c6c 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -5,20 +5,29 @@ This hook is called before making an LLM request when a vector store is configur It searches the vector store for relevant context and appends it to the messages. """ -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_args + +from pydantic import TypeAdapter, ValidationError +from typing_extensions import assert_never import litellm import litellm.vector_stores from litellm._logging import verbose_logger +from litellm.exceptions import VectorStoreSearchError from litellm.integrations.custom_logger import CustomLogger -from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionUserMessage, + ResponsesAPIResponse, +) from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.utils import CallTypes, StandardCallbackDynamicParams from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, - VectorStoreResultContent, + VectorStoreSearchFailure, + VectorStoreSearchFailureMode, VectorStoreSearchResponse, VectorStoreSearchResult, ) @@ -30,6 +39,10 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +SEARCH_FAILURES_FIELD: Final = "vector_store_search_failures" +_DEFAULT_FAILURE_MODE: Final[VectorStoreSearchFailureMode] = "annotate" +_FAILURE_MODE_ADAPTER: Final = TypeAdapter(VectorStoreSearchFailureMode) + class ProxyRuntime(Protocol): def llm_router(self) -> "Router | None": ... @@ -54,11 +67,31 @@ class ProxyServerRuntime: return prisma_client +@dataclass(frozen=True, slots=True) +class SearchSucceeded: + response: VectorStoreSearchResponse + + +@dataclass(frozen=True, slots=True) +class SearchFailed: + failure: VectorStoreSearchFailure + + +SearchOutcome = SearchSucceeded | SearchFailed + + +@dataclass(frozen=True, slots=True) +class VectorStoreAugmentation: + messages: tuple[AllMessageValues, ...] + search_results: tuple[VectorStoreSearchResponse, ...] + failures: tuple[VectorStoreSearchFailure, ...] + + class VectorStorePreCallHook(CustomLogger): CONTENT_PREFIX_STRING = "Context:\n\n" """ Custom logger that handles vector store searches before LLM calls. - + When a vector store is configured, this hook: 1. Extracts the query from the last user message 2. Calls litellm.vector_stores.search() to get relevant context @@ -101,100 +134,153 @@ class VectorStorePreCallHook(CustomLogger): Returns: Tuple of (model, modified_messages, non_default_params) """ + requested_vector_store_ids: Final = _requested_vector_store_ids(non_default_params) try: - # Check if vector store is configured - if litellm.vector_store_registry is None: - return model, messages, non_default_params - - prisma_client: Final = self.proxy_runtime.prisma_client() - llm_router: Final = self.proxy_runtime.llm_router() - - # Use database fallback to ensure synchronization across instances - vector_stores_to_run: list[ - LiteLLM_ManagedVectorStore - ] = await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback( + augmentation: VectorStoreAugmentation | None = await self._augment_messages( + messages=messages, non_default_params=non_default_params, tools=tools, - prisma_client=prisma_client, + litellm_logging_obj=litellm_logging_obj, ) - - if not vector_stores_to_run: - return model, messages, non_default_params - - # Extract the query from the last user message - query: Final = self._extract_query_from_messages(messages) - - if not query: - verbose_logger.debug("No query found in messages for vector store search") - return model, messages, non_default_params - - modified_messages: list[AllMessageValues] = messages.copy() - all_search_results: Final[list[VectorStoreSearchResponse]] = [] - - for vector_store_to_run in vector_stores_to_run: - # Get vector store id from the vector store config - vector_store_id = vector_store_to_run.get("vector_store_id", "") - custom_llm_provider = vector_store_to_run.get("custom_llm_provider") - litellm_params_for_vector_store = vector_store_to_run.get("litellm_params", {}) or {} - request_litellm_params = litellm_logging_obj.model_call_details.get("litellm_params", {}) - request_metadata = ( - request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {} - ) - if llm_router is not None: - search_function = cast( # cast-ok: normalize router search callable - Callable[..., Awaitable[VectorStoreSearchResponse]], - llm_router.avector_store_search, - ) - else: - search_function = cast( # cast-ok: normalize SDK search callable - Callable[..., Awaitable[VectorStoreSearchResponse]], - litellm.vector_stores.asearch, - ) - try: - search_response = await search_function( - **{ - "vector_store_id": vector_store_id, - "query": query, - "custom_llm_provider": custom_llm_provider, - "metadata": request_metadata, - **litellm_params_for_vector_store, - }, - ) - except Exception as search_error: - verbose_logger.warning( - "Vector store search failed for vector_store_id=%s, continuing without its context: %s", - vector_store_id, - search_error, - ) - continue - - verbose_logger.debug("search_response: %s", search_response) - - # Store search results for later use in citations - all_search_results.append(search_response) - - # Process search results and append as context - modified_messages = self._append_search_results_to_messages( - messages=modified_messages, search_response=search_response - ) - - # Get the number of results for logging - num_results = 0 - num_results = len(search_response.get("data", []) or []) - verbose_logger.debug("Vector store search completed. Added context from %s results", num_results) - - # Store search results as-is (already in OpenAI-compatible format) - if litellm_logging_obj and all_search_results: - litellm_logging_obj.model_call_details["search_results"] = all_search_results - - return model, modified_messages, non_default_params - except Exception as e: - verbose_logger.exception("Error in VectorStorePreCallHook: %s", e) - # Return original parameters on error + verbose_logger.exception( + "Error in VectorStorePreCallHook for vector_store_ids=%s: %s", + requested_vector_store_ids, + e, + ) return model, messages, non_default_params - def _extract_query_from_messages(self, messages: list[AllMessageValues]) -> str | None: + if augmentation is None: + return model, messages, non_default_params + + for detail, value in ( + ("search_results", list(augmentation.search_results)), + (SEARCH_FAILURES_FIELD, augmentation.failures), + ): + if value: + litellm_logging_obj.model_call_details[detail] = value + + if augmentation.failures: + failure_mode: Final = _configured_failure_mode() + match failure_mode: + case "error": + raise VectorStoreSearchError(failures=augmentation.failures, model=model) + case "annotate": + pass + case _: + assert_never(failure_mode) + + return model, list(augmentation.messages), non_default_params + + async def _augment_messages( + self, + messages: Sequence[AllMessageValues], + non_default_params: dict, + tools: list[dict] | None, + litellm_logging_obj: LiteLLMLoggingObj, + ) -> VectorStoreAugmentation | None: + if litellm.vector_store_registry is None: + return None + + prisma_client: Final = self.proxy_runtime.prisma_client() + llm_router: Final = self.proxy_runtime.llm_router() + + # Use database fallback to ensure synchronization across instances + vector_stores_to_run: Final[ + Sequence[LiteLLM_ManagedVectorStore] + ] = await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params=non_default_params, + tools=tools, + prisma_client=prisma_client, + ) + + if not vector_stores_to_run: + return None + + query: Final = self._extract_query_from_messages(messages) + + if not query: + verbose_logger.debug("No query found in messages for vector store search") + return None + + request_litellm_params: Final = litellm_logging_obj.model_call_details.get("litellm_params", {}) + request_metadata: Final = ( + request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {} + ) + search_function: Final = ( + cast( # cast-ok: normalize router search callable + Callable[..., Awaitable[VectorStoreSearchResponse]], + llm_router.avector_store_search, + ) + if llm_router is not None + else cast( # cast-ok: normalize SDK search callable + Callable[..., Awaitable[VectorStoreSearchResponse]], + litellm.vector_stores.asearch, + ) + ) + + outcomes: Final = tuple( + [ + await self._search_one( + vector_store=vector_store_to_run, + query=query, + request_metadata=request_metadata, + search_function=search_function, + ) + for vector_store_to_run in vector_stores_to_run + ] + ) + search_results: Final = tuple(outcome.response for outcome in outcomes if isinstance(outcome, SearchSucceeded)) + failures: Final = tuple(outcome.failure for outcome in outcomes if isinstance(outcome, SearchFailed)) + + return VectorStoreAugmentation( + messages=self._messages_with_context(messages=messages, search_results=search_results), + search_results=search_results, + failures=failures, + ) + + async def _search_one( + self, + vector_store: LiteLLM_ManagedVectorStore, + query: str, + request_metadata: Mapping[str, object], + search_function: Callable[..., Awaitable[VectorStoreSearchResponse]], + ) -> SearchOutcome: + vector_store_id: Final = vector_store.get("vector_store_id", "") + custom_llm_provider: Final = vector_store.get("custom_llm_provider") + litellm_params_for_vector_store: Final = vector_store.get("litellm_params", {}) or {} + try: + search_response: Final = await search_function( + **{ + "vector_store_id": vector_store_id, + "query": query, + "custom_llm_provider": custom_llm_provider, + "metadata": request_metadata, + **litellm_params_for_vector_store, + }, + ) + except Exception as search_error: + verbose_logger.warning( + "Vector store search failed for vector_store_id=%s, continuing without its context: %s", + vector_store_id, + search_error, + ) + return SearchFailed( + failure=VectorStoreSearchFailure( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider, + error=str(search_error), + ) + ) + + verbose_logger.debug( + "Vector store search completed for vector_store_id=%s. Added context from %s results", + vector_store_id, + len(search_response.get("data") or ()), + ) + return SearchSucceeded(response=search_response) + + def _extract_query_from_messages(self, messages: Sequence[AllMessageValues]) -> str | None: """ Extract the query from the last user message. @@ -223,48 +309,40 @@ class VectorStorePreCallHook(CustomLogger): return None - def _append_search_results_to_messages( + def _messages_with_context( self, - messages: list[AllMessageValues], - search_response: VectorStoreSearchResponse, - ) -> list[AllMessageValues]: - """ - Append search results as context to the messages. + messages: Sequence[AllMessageValues], + search_results: Sequence[VectorStoreSearchResponse], + ) -> tuple[AllMessageValues, ...]: + context_messages: Final = tuple( + context_message + for search_response in search_results + if (context_message := self._context_message(search_response)) is not None + ) + if not context_messages: + return tuple(messages) + return (*messages[:-1], *context_messages, *messages[-1:]) - Args: - messages: Original list of messages - search_response: Response from vector store search - - Returns: - Modified list of messages with context appended - """ - search_response_data: Final[list[VectorStoreSearchResult] | None] = search_response.get("data") + def _context_message(self, search_response: VectorStoreSearchResponse) -> AllMessageValues | None: + """Build the context message for one vector store's results, or None when it returned nothing usable.""" + search_response_data: Final[Sequence[VectorStoreSearchResult] | None] = search_response.get("data") if not search_response_data: - return messages + return None - context_content = self.CONTENT_PREFIX_STRING + context_texts: Final = tuple( + content_text + for result in search_response_data + for content_item in (result.get("content") or ()) + if (content_text := content_item.get("text")) + ) + if not context_texts: + return None - for result in search_response_data: - result_content: list[VectorStoreResultContent] | None = result.get("content") - if result_content: - for content_item in result_content: - content_text: str | None = content_item.get("text") - if content_text: - context_content += content_text + "\n\n" - - # Only add context if we found any content - if context_content != "Context:\n\n": - # Create a copy of messages to avoid modifying the original - modified_messages: Final = messages.copy() - # Add context as a new message before the last user message - context_message: Final[ChatCompletionUserMessage] = { - "role": "user", - "content": context_content, - } - modified_messages.insert(-1, cast(AllMessageValues, context_message)) - return modified_messages - - return messages + context_message: Final[ChatCompletionUserMessage] = { + "role": "user", + "content": self.CONTENT_PREFIX_STRING + "".join(f"{text}\n\n" for text in context_texts), + } + return cast(AllMessageValues, context_message) async def async_post_call_success_deployment_hook( self, @@ -287,34 +365,34 @@ class VectorStorePreCallHook(CustomLogger): verbose_logger.debug("No litellm_logging_obj in request_data") return None - verbose_logger.debug("model_call_details keys: %s", list(litellm_logging_obj.model_call_details.keys())) - # Get search results from model_call_details (already in OpenAI format) - search_results: Final[list[VectorStoreSearchResponse] | None] = litellm_logging_obj.model_call_details.get( - "search_results" + search_results: Final[Sequence[VectorStoreSearchResponse] | None] = ( + litellm_logging_obj.model_call_details.get("search_results") + ) + search_failures: Final[Sequence[VectorStoreSearchFailure] | None] = ( + litellm_logging_obj.model_call_details.get(SEARCH_FAILURES_FIELD) ) - verbose_logger.debug("Search results found: %s", search_results is not None) - - if not search_results: - verbose_logger.debug("No search results found") + if not search_results and not search_failures: + verbose_logger.debug("No search results or search failures found") return None + if isinstance(response, ResponsesAPIResponse): + if search_failures: + setattr(response, SEARCH_FAILURES_FIELD, list(search_failures)) + return response + # Add search results to response object if hasattr(response, "choices") and response.choices: for choice in response.choices: if hasattr(choice, "message") and choice.message: - # Get existing provider_specific_fields or create new dict provider_fields = getattr(choice.message, "provider_specific_fields", None) or {} - - # Add search results (already in OpenAI-compatible format) - provider_fields["search_results"] = search_results - - # Set the provider_specific_fields + if search_results: + provider_fields["search_results"] = search_results + if search_failures: + provider_fields[SEARCH_FAILURES_FIELD] = search_failures setattr(choice.message, "provider_specific_fields", provider_fields) - verbose_logger.debug("Added %s search results to response", len(search_results)) - # Return modified response return response @@ -339,29 +417,24 @@ class VectorStorePreCallHook(CustomLogger): verbose_logger.debug("VectorStorePreCallHook.async_post_call_streaming_deployment_hook called") # Get search results from model_call_details (already in OpenAI format) - search_results: Final[list[VectorStoreSearchResponse] | None] = request_data.get("search_results") + search_results: Final[Sequence[VectorStoreSearchResponse] | None] = request_data.get("search_results") + search_failures: Final[Sequence[VectorStoreSearchFailure] | None] = request_data.get(SEARCH_FAILURES_FIELD) - verbose_logger.debug("Search results found for streaming chunk: %s", search_results is not None) - - if not search_results: - verbose_logger.debug("No search results found for streaming chunk") + if not search_results and not search_failures: + verbose_logger.debug("No search results or search failures found for streaming chunk") return response_chunk # Add search results to streaming chunk if hasattr(response_chunk, "choices") and response_chunk.choices: for choice in response_chunk.choices: if hasattr(choice, "delta") and choice.delta: - # Get existing provider_specific_fields or create new dict provider_fields = getattr(choice.delta, "provider_specific_fields", None) or {} - - # Add search results (already in OpenAI-compatible format) - provider_fields["search_results"] = search_results - - # Set the provider_specific_fields + if search_results: + provider_fields["search_results"] = search_results + if search_failures: + provider_fields[SEARCH_FAILURES_FIELD] = search_failures choice.delta.provider_specific_fields = provider_fields - verbose_logger.debug("Added %s search results to streaming chunk", len(search_results)) - # Return modified chunk return response_chunk @@ -369,3 +442,23 @@ class VectorStorePreCallHook(CustomLogger): verbose_logger.exception("Error adding search results to streaming chunk: %s", e) # Don't fail the request if search results fail to be added return response_chunk + + +def _requested_vector_store_ids(non_default_params: Mapping[str, object]) -> tuple[str, ...]: + requested: Final = non_default_params.get("vector_store_ids") + if not isinstance(requested, (list, tuple)): + return () + return tuple(str(vector_store_id) for vector_store_id in requested) + + +def _configured_failure_mode() -> VectorStoreSearchFailureMode: + try: + return _FAILURE_MODE_ADAPTER.validate_python(litellm.vector_store_search_failure_mode) + except ValidationError: + verbose_logger.warning( + "Unsupported vector_store_search_failure_mode=%r, falling back to %r. Supported modes: %s", + litellm.vector_store_search_failure_mode, + _DEFAULT_FAILURE_MODE, + ", ".join(get_args(VectorStoreSearchFailureMode)), + ) + return _DEFAULT_FAILURE_MODE diff --git a/litellm/integrations/weave/weave_otel.py b/litellm/integrations/weave/weave_otel.py index f2cc64a9ba2..50289263f38 100644 --- a/litellm/integrations/weave/weave_otel.py +++ b/litellm/integrations/weave/weave_otel.py @@ -117,6 +117,14 @@ def _get_weave_authorization_header(api_key: str) -> str: return f"Basic {auth_header}" +def weave_otel_endpoint(host: str | None) -> str: + """The OTLP traces endpoint for a self-managed ``host``, else Weave cloud.""" + if not host: + return WEAVE_BASE_URL + WEAVE_OTEL_ENDPOINT + normalized: Final = host if host.startswith("http") else f"https://{host}" + return normalized.rstrip("/") + WEAVE_OTEL_ENDPOINT + + def get_weave_otel_config() -> WeaveOtelConfig: """ Retrieves the Weave OpenTelemetry configuration based on environment variables. @@ -134,7 +142,6 @@ def get_weave_otel_config() -> WeaveOtelConfig: """ api_key: Final = os.getenv("WANDB_API_KEY") project_id: Final = os.getenv("WANDB_PROJECT_ID") - host = os.getenv("WANDB_HOST") if not api_key: raise ValueError("WANDB_API_KEY must be set for Weave OpenTelemetry integration.") @@ -144,15 +151,8 @@ def get_weave_otel_config() -> WeaveOtelConfig: "WANDB_PROJECT_ID must be set for Weave OpenTelemetry integration. Format: /" ) - if host: - if not host.startswith("http"): - host = "https://" + host - # Self-managed instances use a different path - endpoint = host.rstrip("/") + WEAVE_OTEL_ENDPOINT - verbose_logger.debug("Using Weave OTEL endpoint from host: %s", endpoint) - else: - endpoint = WEAVE_BASE_URL + WEAVE_OTEL_ENDPOINT - verbose_logger.debug("Using Weave cloud endpoint: %s", endpoint) + endpoint: Final = weave_otel_endpoint(os.getenv("WANDB_HOST")) + verbose_logger.debug("Using Weave OTEL endpoint: %s", endpoint) # Weave uses Basic auth with format: api: auth_header: Final = _get_weave_authorization_header(api_key=api_key) diff --git a/litellm/litellm_core_utils/api_route_to_call_types.py b/litellm/litellm_core_utils/api_route_to_call_types.py index e3562095d7f..428d7563d4a 100644 --- a/litellm/litellm_core_utils/api_route_to_call_types.py +++ b/litellm/litellm_core_utils/api_route_to_call_types.py @@ -14,21 +14,48 @@ from typing import Final from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes +def _segment_matches(route_segment: str, pattern_segment: str) -> bool: + """ + Match one concrete path segment against one pattern segment. + A bare placeholder ({param}) matches any segment; a placeholder with a + literal suffix ({model}:generateContent) requires the segment to end with + that suffix and have a non-empty value before it. + """ + if not pattern_segment.startswith("{"): + return route_segment == pattern_segment + placeholder_end: Final = pattern_segment.find("}") + if placeholder_end == -1: + return route_segment == pattern_segment + literal_suffix: Final = pattern_segment[placeholder_end + 1 :] + if not literal_suffix: + return True + return route_segment.endswith(literal_suffix) and len(route_segment) > len(literal_suffix) + + +def _pattern_tail_spans_segments(pattern_tail: str) -> bool: + """ + Whether the pattern's last segment is a suffixed placeholder + ({model}:generateContent) that may absorb extra route segments, mirroring + FastAPI's {model_name:path} converter for slash-containing model names. + """ + return pattern_tail.startswith("{") and "}" in pattern_tail and not pattern_tail.endswith("}") + + def _route_matches_pattern(route: str, pattern: str) -> bool: """ Return True if the concrete route matches the pattern. - Pattern segments like {param} match any single path segment. + Pattern segments like {param} match any single path segment, and a + suffixed placeholder in the last segment may span multiple segments. """ route_parts: Final = route.strip("/").split("/") pattern_parts: Final = pattern.strip("/").split("/") - if len(route_parts) != len(pattern_parts): + if len(route_parts) < len(pattern_parts): return False - for r, p in zip(route_parts, pattern_parts): - if p.startswith("{") and p.endswith("}"): - continue - if r != p: - return False - return True + if len(route_parts) > len(pattern_parts) and not _pattern_tail_spans_segments(pattern_parts[-1]): + return False + head_count: Final = len(pattern_parts) - 1 + merged_parts: Final = (*route_parts[:head_count], "/".join(route_parts[head_count:])) + return all(_segment_matches(r, p) for r, p in zip(merged_parts, pattern_parts)) def get_call_types_for_route(route: str) -> Sequence[CallTypes] | None: diff --git a/litellm/litellm_core_utils/classifier_logging.py b/litellm/litellm_core_utils/classifier_logging.py new file mode 100644 index 00000000000..fdc1cac26c0 --- /dev/null +++ b/litellm/litellm_core_utils/classifier_logging.py @@ -0,0 +1,67 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN, ClassifierAudit + +CLASSIFIER_AUDIT_FIELDS: Final = ("classifier_input", "originating_request_masked") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def classifier_input_snapshot(value: object, *, openai_sdk: bool = False) -> Mapping[str, JsonValue] | None: + try: + if openai_sdk and isinstance(value, Mapping): + body: Final = MappingProxyType( + {key: item for key, item in value.items() if key not in ("extra_headers", "extra_query", "extra_body")} + ) + extra_body: Final = value.get("extra_body") + return _JSON_OBJECT.validate_python( + MappingProxyType({**body, **extra_body}) if isinstance(extra_body, Mapping) else body + ) + return ( + _JSON_OBJECT.validate_json(value) + if isinstance(value, (str, bytes)) + else _JSON_OBJECT.validate_python(value) + ) + except ValidationError: + return None + + +def is_classifier_call(call_type: str, params: Mapping[str, object]) -> bool: + return call_type in ("completion", "acompletion", "responses", "aresponses") and any( + isinstance(metadata := params.get(key), Mapping) + and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == AUTOROUTER_CLASSIFIER_CALL_ORIGIN + for key in ("metadata", "litellm_metadata") + ) + + +def masked_originating_request(request_kwargs: Mapping[str, object] | None) -> Mapping[str, JsonValue] | None: + request: Final = request_kwargs.get("proxy_server_request") if request_kwargs is not None else None + body: Final = request.get("body") if isinstance(request, Mapping) else None + if not isinstance(body, Mapping): + return None + serializable: Final = classifier_input_snapshot(safe_dumps(body)) + return classifier_input_snapshot(redact_credentials_in_payload(serializable)) if serializable is not None else None + + +def classifier_audit_fields(payload: Mapping[str, object]) -> ClassifierAudit: + classifier_input: Final = classifier_input_snapshot(payload.get("classifier_input")) + originating_request: Final = classifier_input_snapshot(payload.get("originating_request_masked")) + if classifier_input is None: + return ( + ClassifierAudit(originating_request_masked=originating_request) + if originating_request is not None + else ClassifierAudit() + ) + if originating_request is None: + return ClassifierAudit(classifier_input=classifier_input) + return ClassifierAudit(classifier_input=classifier_input, originating_request_masked=originating_request) + + +def without_classifier_audit(payload: Mapping[str, object]) -> dict[str, object]: + return {key: value for key, value in payload.items() if key not in CLASSIFIER_AUDIT_FIELDS} diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 2cdcfe4879c..aa7d6ca1699 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -1,10 +1,13 @@ # What is this? ## Helper utilities import copy +import logging +import re from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal import httpx +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason @@ -19,6 +22,13 @@ else: Span = Any +_CODEX_CLIENT_PREFIX_RE: Final = re.compile(r"^codex[-_ /]", re.IGNORECASE) + + +def is_codex_user_agent(user_agent: str) -> bool: + return bool(_CODEX_CLIENT_PREFIX_RE.match(user_agent)) + + def safe_divide_seconds(seconds: float, denominator: float, default: float | None = None) -> float | None: """ Safely divide seconds by denominator, handling zero division. @@ -37,6 +47,41 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non return float(seconds / denominator) +_DROP_PARAMS_BOOL: Final = TypeAdapter(bool) + + +def normalize_drop_params(value: object) -> bool | None: + if value is None or isinstance(value, bool): + return value + try: + return _DROP_PARAMS_BOOL.validate_python(value.strip() if isinstance(value, str) else value) + except ValidationError: + return None + + +def drop_params_flag(value: object, source: str, logger: logging.Logger) -> bool: + normalized: Final = normalize_drop_params(value) + if normalized is None and value is not None: + logger.warning("%s=%r is not a flag value, treating it as off", source, value) + return bool(normalized) + + +DROP_PARAMS_ENV_VAR: Final = "LITELLM_DROP_PARAMS" + + +def drop_params_env_flag(environ: Mapping[str, str], logger: logging.Logger) -> bool: + configured: Final = environ.get(DROP_PARAMS_ENV_VAR, "").strip() + if configured == "": + return False + normalized: Final = normalize_drop_params(configured) + if normalized is None: + logger.warning( + "%s=%r is not a flag value, treating it as on. Set it to true or false", DROP_PARAMS_ENV_VAR, configured + ) + return True + return normalized + + def safe_divide( numerator: float, denominator: float, diff --git a/litellm/litellm_core_utils/credential_accessor.py b/litellm/litellm_core_utils/credential_accessor.py index 7071750b970..7b4e8240b69 100644 --- a/litellm/litellm_core_utils/credential_accessor.py +++ b/litellm/litellm_core_utils/credential_accessor.py @@ -7,16 +7,19 @@ from litellm.types.utils import CredentialItem class CredentialAccessor: + @staticmethod + def find_credential(credential_name: str) -> CredentialItem | None: + return next( + (credential for credential in litellm.credential_list if credential.credential_name == credential_name), + None, + ) + @staticmethod def get_credential_values(credential_name: str) -> dict: """Safe accessor for credentials.""" - if not litellm.credential_list: - return {} - for credential in litellm.credential_list: - if credential.credential_name == credential_name: - return credential.credential_values.copy() - return {} + credential: Final = CredentialAccessor.find_credential(credential_name) + return {} if credential is None else credential.credential_values.copy() @staticmethod def upsert_credentials(credentials: list[CredentialItem]): diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 6449aa4d46e..6294f3bc577 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -43,6 +43,7 @@ from litellm.integrations.newrelic import NewRelicLogger from litellm.integrations.openmeter import OpenMeterLogger from litellm.integrations.opentelemetry import OpenTelemetry from litellm.integrations.opik.opik import OpikLogger +from litellm.integrations.pointfive import PointFiveLogger from litellm.integrations.posthog import PostHogLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.s3_v2 import S3Logger @@ -95,6 +96,7 @@ class CustomLoggerRegistry: "agentops": AgentOps, "deepeval": DeepEvalLogger, "s3_v2": S3Logger, + "pointfive": PointFiveLogger, "aws_sqs": SQSLogger, "dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler, "dynamic_rate_limiter_v3": _PROXY_DynamicRateLimitHandlerV3, diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 8f8c955d971..82708d412c9 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -860,6 +860,7 @@ def _map_bedrock_exception( message=mantle_context_window_message, model=model, llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), ) if ( "too many tokens" in error_str @@ -873,6 +874,7 @@ def _map_bedrock_exception( message=f"BedrockException: Context Window Error - {error_str}", model=model, llm_provider="bedrock", + response=getattr(original_exception, "response", None), ) elif "Conversation blocks and tool result blocks cannot be provided in the same turn." in error_str: raise BadRequestError( @@ -924,12 +926,14 @@ def _map_bedrock_exception( message=f"BedrockException: Timeout Error - {error_str}", model=model, llm_provider="bedrock", + response=getattr(original_exception, "response", None), ) elif "Could not process image" in error_str: raise litellm.InternalServerError( message=f"BedrockException - {error_str}", model=model, llm_provider="bedrock", + response=getattr(original_exception, "response", None), ) elif hasattr(original_exception, "status_code"): if original_exception.status_code == 500: @@ -937,10 +941,7 @@ def _map_bedrock_exception( message=f"BedrockException - {original_exception.message}", llm_provider="bedrock", model=model, - response=httpx.Response( - status_code=500, - request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), - ), + response=getattr(original_exception, "response", None), ) elif original_exception.status_code == 401: raise AuthenticationError( @@ -969,6 +970,7 @@ def _map_bedrock_exception( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), ) elif original_exception.status_code == 422: raise BadRequestError( @@ -1001,6 +1003,7 @@ def _map_bedrock_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, exception_status_code=original_exception.status_code, + response=getattr(original_exception, "response", None), ) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 389e6f7f501..edd2e88f95c 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -2,6 +2,7 @@ from collections.abc import Mapping, MutableMapping from types import MappingProxyType from typing import Final +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.llms.openai.data_residency import infer_openai_data_residency AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( @@ -16,18 +17,16 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( "aws_web_identity_token", "aws_sts_endpoint", "aws_external_id", + "aws_session_tags", "aws_bedrock_runtime_endpoint", "aws_bedrock_project_id", } ) -# The per-deployment Rust opt-in. -RUST_KWARG_KEY: Final = "rust" - # Keys `completion()` forwards from its own kwargs into `get_litellm_params`, # which are otherwise invisible to it because that call site passes explicit # named arguments rather than `**kwargs`. -FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY}) +FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS # Pre-define optional kwargs keys as frozenset for O(1) lookups # These are extracted from kwargs only if present, avoiding unnecessary .get() calls @@ -58,10 +57,6 @@ OPTIONAL_KWARGS_KEYS: Final = ( "itpm", "otpm", "use_xai_oauth", - # The per-deployment Rust opt-in. `all_litellm_params` keeps it out - # of the provider body; this keeps it *in* litellm_params, which is - # where the chat completions handlers read it from. - RUST_KWARG_KEY, } ) | AWS_CREDENTIAL_KWARGS_KEYS @@ -120,7 +115,7 @@ def get_litellm_params( custom_prompt_dict: dict | None = None, litellm_metadata: dict | None = None, disable_add_transform_inline_image_block: bool | None = None, - drop_params: bool | None = None, + drop_params: bool | str | None = None, prompt_id: str | None = None, prompt_variables: dict | None = None, async_call: bool | None = None, @@ -182,7 +177,7 @@ def get_litellm_params( "custom_prompt_dict": custom_prompt_dict, "litellm_metadata": litellm_metadata, "disable_add_transform_inline_image_block": disable_add_transform_inline_image_block, - "drop_params": drop_params, + "drop_params": normalize_drop_params(drop_params), "prompt_id": prompt_id, "prompt_variables": prompt_variables, "async_call": async_call, diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 9cba5db8ab7..91a22144805 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -2,6 +2,7 @@ Pulls the cost + context window + provider route for known models from https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json This can be disabled by setting the LITELLM_LOCAL_MODEL_COST_MAP environment variable to True. +The ``lite`` and ``litellm-proxy`` CLI entry points also use the bundled map without fetching. ``` export LITELLM_LOCAL_MODEL_COST_MAP=True @@ -9,17 +10,22 @@ export LITELLM_LOCAL_MODEL_COST_MAP=True """ import asyncio +import hashlib import json import os import random +import sys +import threading import time from collections.abc import Awaitable, Callable -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timezone from importlib.resources import files +from pathlib import Path from typing import Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict from litellm import verbose_logger from litellm.constants import ( @@ -31,6 +37,12 @@ from litellm.litellm_core_utils.fallback_generalizations import ( ) FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" +_CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) + + +def _is_cli_process() -> bool: + return Path(sys.argv[0]).stem in _CLI_ENTRYPOINT_NAMES + # Reserved top-level keys that are not model entries. They must be excluded # from the model-count integrity check so a real upstream shrink can't be masked. @@ -42,6 +54,10 @@ def _count_model_entries(model_cost: dict) -> int: return sum(1 for key in model_cost if key not in RESERVED_TOP_LEVEL_KEYS) +def git_blob_id(body: bytes) -> str: + return hashlib.sha1(b"blob %d\0" % len(body) + body, usedforsecurity=False).hexdigest() + + class GetModelCostMap: """ Handles fetching, validating, and loading the model cost map. @@ -53,13 +69,24 @@ class GetModelCostMap: _backup_model_count: int = -1 # -1 = not yet loaded + @staticmethod + def read_local_model_cost_map_bytes() -> bytes: + return files("litellm").joinpath("model_prices_and_context_window_backup.json").read_bytes() + + @staticmethod + def read_local_model_cost_map_text() -> str: + return GetModelCostMap.read_local_model_cost_map_bytes().decode("utf-8") + + @staticmethod + def load_local_model_cost_map_with_revision() -> "ModelCostMapReloaded": + body: Final = GetModelCostMap.read_local_model_cost_map_bytes() + content: Final = json.loads(body) + return ModelCostMapReloaded(model_cost_map=content, revision=git_blob_id(body)) + @staticmethod def load_local_model_cost_map() -> dict: """Load the local backup model cost map bundled with the package.""" - content: Final = json.loads( - files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text(encoding="utf-8") - ) - return content + return GetModelCostMap.load_local_model_cost_map_with_revision().model_cost_map @classmethod def _get_backup_model_count(cls) -> int: @@ -159,11 +186,18 @@ class GetModelCostMap: RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504}) MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3 MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS: Final = 30.0 +_litellm_import_complete = threading.Event() + + +def mark_litellm_import_complete() -> None: + _litellm_import_complete.set() @dataclass(frozen=True, slots=True) class ModelCostMapReloaded: model_cost_map: dict # mutable-ok: adopted as litellm.model_cost, whose consumer contract is a plain mutable dict + revision: str | None = None + etag: str | None = None @dataclass(frozen=True, slots=True) @@ -252,7 +286,9 @@ def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemp return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}") if not isinstance(parsed, dict): return ModelCostMapReloadUnavailable(reason=f"expected a JSON object from {url}, got {type(parsed).__name__}") - return ModelCostMapReloaded(model_cost_map=parsed) + return ModelCostMapReloaded( + model_cost_map=parsed, revision=git_blob_id(response.content), etag=response.headers.get("etag") + ) def _next_retry_wait( @@ -293,12 +329,13 @@ async def _fetch_remote_model_cost_map_with_retry( def _fetch_remote_model_cost_map_with_retry_sync( url: str, timeout: int, - max_attempts: int, + attempts: range, sleep: Callable[[float], None], rng: random.Random, client: _SyncGetClient, ) -> ModelCostMapReloadResult: - for attempt in range(1, max_attempts + 1): + max_attempts: Final = attempts.stop - 1 + for attempt in attempts: outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout) if not isinstance(outcome, _FetchAttemptRetryable): return outcome @@ -326,13 +363,12 @@ async def refetch_model_cost_map( map they already have. """ if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) _cost_map_source_info.source = "local" _cost_map_source_info.url = None _cost_map_source_info.is_env_forced = True _cost_map_source_info.fallback_reason = None - return ModelCostMapReloaded( - model_cost_map=_finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()) - ) + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()) result: Final = await _fetch_remote_model_cost_map_with_retry( url=url, @@ -353,11 +389,12 @@ async def refetch_model_cost_map( backup_model_count=GetModelCostMap._get_backup_model_count(), ): return ModelCostMapReloadUnavailable(reason=f"model cost map from {url} failed integrity validation") + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) _cost_map_source_info.source = "remote" _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False _cost_map_source_info.fallback_reason = None - return ModelCostMapReloaded(model_cost_map=_finalize_model_cost_map(result.model_cost_map)) + return _finalize_loaded_model_cost_map(result) class ModelCostMapSourceInfo: @@ -368,13 +405,35 @@ class ModelCostMapSourceInfo: is_env_forced: bool = False fallback_reason: str | None = None loaded_at: "datetime | None" = None + source_revision: str | None = None + etag: str | None = None # Module-level singleton tracking the source of the current cost map _cost_map_source_info: Final = ModelCostMapSourceInfo() -def get_model_cost_map_source_info() -> dict: +class CostMapProvenance(TypedDict): + source_revision: ReadOnly[str | None] + etag: ReadOnly[str | None] + + +class CostMapSourceInfo(CostMapProvenance): + source: ReadOnly[str] + url: ReadOnly[str | None] + is_env_forced: ReadOnly[bool] + fallback_reason: ReadOnly[str | None] + loaded_at: ReadOnly[str | None] + + +def get_model_cost_map_provenance() -> CostMapProvenance: + return { + "source_revision": _cost_map_source_info.source_revision, + "etag": _cost_map_source_info.etag, + } + + +def get_model_cost_map_source_info() -> CostMapSourceInfo: """ Return metadata about where the current model cost map was loaded from. @@ -383,12 +442,19 @@ def get_model_cost_map_source_info() -> dict: - url: the remote URL attempted (or None for local-only) - is_env_forced: True if LITELLM_LOCAL_MODEL_COST_MAP=True forced local usage - fallback_reason: human-readable reason if remote failed and local was used + - loaded_at: ISO 8601 time this process last loaded the map + - source_revision: git blob id of the loaded file's bytes + - etag: the ETag of the remote fetch (None for the bundled backup) """ + loaded_at: Final = _cost_map_source_info.loaded_at return { "source": _cost_map_source_info.source, "url": _cost_map_source_info.url, "is_env_forced": _cost_map_source_info.is_env_forced, "fallback_reason": _cost_map_source_info.fallback_reason, + "loaded_at": loaded_at.isoformat() if loaded_at is not None else None, + "source_revision": _cost_map_source_info.source_revision, + "etag": _cost_map_source_info.etag, } @@ -464,6 +530,74 @@ def _finalize_model_cost_map(model_cost: dict) -> dict: return _expand_model_aliases(model_cost) +def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded: + _cost_map_source_info.source_revision = loaded.revision + _cost_map_source_info.etag = loaded.etag + return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map)) + + +def adopt_model_cost_map( + new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract +) -> int: + import litellm + from litellm import utils + + litellm.model_cost = new_model_cost_map + utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation + litellm.add_known_models(model_cost_map=new_model_cost_map) + fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0 + utils.reapply_runtime_model_cost_registrations() + return fetched_model_count + + +def _retry_remote_fetch_in_background( + url: str, + timeout: int, + max_attempts: int, + sleep: Callable[[float], None], + rng: random.Random, + client: _SyncGetClient, + first_outcome: _FetchAttemptRetryable, +) -> None: + try: + first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng) + if isinstance(first_wait, ModelCostMapReloadUnavailable): + return + sleep(first_wait) + result: Final = _fetch_remote_model_cost_map_with_retry_sync( + url=url, + timeout=timeout, + attempts=range(2, max_attempts + 1), + sleep=sleep, + rng=rng, + client=client, + ) + if isinstance(result, ModelCostMapReloadUnavailable): + verbose_logger.warning( + "LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup", + url, + max_attempts, + ) + return + _litellm_import_complete.wait() + if not GetModelCostMap.validate_model_cost_map( + fetched_map=result.model_cost_map, + backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", + url, + ) + return + finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map + _cost_map_source_info.source = "remote" + _cost_map_source_info.fallback_reason = None + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) + adopt_model_cost_map(finalized) + except Exception as e: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays + verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) + + def get_model_cost_map( url: str, timeout: int = 5, @@ -475,10 +609,12 @@ def get_model_cost_map( """ Public entry point — returns the model cost map dict. - 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only. + 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set or this is a ``lite`` / + ``litellm-proxy`` CLI process, uses the local backup only. 2. Otherwise fetches from ``url``, retrying transient HTTP errors - (429/5xx/transport) with Retry-After-aware backoff, validates - integrity, and falls back to the local backup on any failure. + (429/5xx/transport) with Retry-After-aware backoff in a background + thread, validates integrity, and falls back to the local backup on any + failure. Only the backup model count is cached (a single int) for validation. The full backup dict is only parsed when it must be *returned* as a @@ -487,34 +623,44 @@ def get_model_cost_map( _cost_map_source_info.loaded_at = datetime.now(timezone.utc) # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. - if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": + if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): _cost_map_source_info.source = "local" _cost_map_source_info.url = None _cost_map_source_info.is_env_forced = True _cost_map_source_info.fallback_reason = None - return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()) + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False - result: Final = _fetch_remote_model_cost_map_with_retry_sync( - url=url, - timeout=timeout, - max_attempts=max_attempts, - sleep=sleep, - rng=rng if rng is not None else random.Random(), - client=client if client is not None else httpx, - ) - if isinstance(result, ModelCostMapReloadUnavailable): + fetch_client: Final = client if client is not None else httpx + fetch_rng: Final = rng if rng is not None else random.Random() + outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout) + if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1: + threading.Thread( + target=_retry_remote_fetch_in_background, + kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping + "url": url, + "timeout": timeout, + "max_attempts": max_attempts, + "sleep": sleep, + "rng": fetch_rng, + "client": fetch_client, + "first_outcome": outcome, + }, + name="litellm-model-cost-map-retry", + daemon=True, + ).start() + if not isinstance(outcome, ModelCostMapReloaded): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, - result.reason, + outcome.reason, ) _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}" - return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()) - content: Final = result.model_cost_map + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {outcome.reason}" + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map + content: Final = outcome.model_cost_map # Validate using cached count (cheap int comparison, no file I/O) if not GetModelCostMap.validate_model_cost_map( @@ -527,8 +673,8 @@ def get_model_cost_map( ) _cost_map_source_info.source = "local" _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" - return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()) + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map _cost_map_source_info.source = "remote" _cost_map_source_info.fallback_reason = None - return _finalize_model_cost_map(content) + return _finalize_loaded_model_cost_map(outcome).model_cost_map diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 989777e9412..498d662a906 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -14,10 +14,10 @@ from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime as dt_object from functools import lru_cache from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast from httpx import Response -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue import litellm from litellm import ( @@ -36,17 +36,20 @@ 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 ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, + EMPTY_MAPPING, + PROVIDER_REQUEST_ID_HEADERS, SENTRY_DENYLIST, SENTRY_PII_DENYLIST, ) from litellm.cost_calculator import ( RealtimeAPITokenUsageProcessor, + ResponsesWebSocketTokenUsageProcessor, _select_model_name_for_cost_calc, ) from litellm.exceptions import ( @@ -62,6 +65,11 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger +from litellm.litellm_core_utils.classifier_logging import ( + classifier_audit_fields, + classifier_input_snapshot, + is_classifier_call, +) from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.internal_call_metadata import ( @@ -87,6 +95,7 @@ from litellm.litellm_core_utils.redact_messages import ( redact_message_input_output_from_custom_logger, redact_message_input_output_from_logging, redact_streaming_responses_for_custom_logger, + should_redact_message_logging, ) from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.llms.base_llm.search.transformation import SearchResponse @@ -118,7 +127,6 @@ from litellm.types.utils import ( CachingDetails, CallTypes, CostBreakdown, - CostResponseTypes, CustomPricingLiteLLMParams, DynamicPromptManagementParamLiteral, EmbeddingResponse, @@ -181,6 +189,7 @@ from ..integrations.lunary import LunaryLogger from ..integrations.newrelic import NewRelicLogger from ..integrations.openmeter import OpenMeterLogger from ..integrations.opik.opik import OpikLogger +from ..integrations.pointfive import PointFiveLogger from ..integrations.posthog import PostHogLogger from ..integrations.prompt_layer import PromptLayerLogger from ..integrations.s3 import S3Logger @@ -200,7 +209,9 @@ if TYPE_CHECKING: from mcp.types import EmbeddedResource, ImageContent, TextContent from litellm.integrations.otel.logger import OpenTelemetryV2 - from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config + from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates + from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -255,11 +266,38 @@ _in_memory_loggers: Final[list[CustomLogger]] = [] _STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggingMetadata.__annotations__.keys()) + +def _get_provider_request_id(original_exception: Exception) -> str | None: + try: + error_response: Final = getattr(original_exception, "response", None) + header_sources: Final = ( + _get_response_headers(original_exception), + getattr(error_response, "headers", None), + getattr(original_exception, "litellm_response_headers", None), + ) + return next( + ( + str(value) + for expected_header_name in PROVIDER_REQUEST_ID_HEADERS + for headers in header_sources + if isinstance(headers, Mapping) + for header_name, value in headers.items() + if isinstance(header_name, str) and header_name.lower() == expected_header_name and value + ), + None, + ) + except Exception: + return None + + ### GLOBAL VARIABLES ### # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS +_UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( + ("user_api_key_auth", "user_api_key_budget_reservation") +) sentry_sdk_instance = None capture_exception = None @@ -422,6 +460,13 @@ def _provider_response_id(source: object) -> str | None: return candidate if isinstance(candidate, str) and candidate else None +def mask_api_base_credentials(api_base: str) -> str: + if "key=" not in api_base: + return api_base + key_end: Final = api_base.find("key=") + 4 + return api_base[:key_end] + "*" * 5 + api_base[-4:] + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -438,6 +483,7 @@ class Logging(LiteLLMLoggingBaseClass): stream_options = None litellm_request_debug: bool = False streamed_anthropic_message_id: str | None = None + classifier_input: Mapping[str, JsonValue] | None = None def __init__( self, @@ -556,6 +602,7 @@ class Logging(LiteLLMLoggingBaseClass): # Initialize cost breakdown field self.cost_breakdown: CostBreakdown | None = None + self.billed_token_rates: BilledTokenRates | None = None # Init Caching related details self.caching_details: CachingDetails | None = None @@ -1164,14 +1211,7 @@ class Logging(LiteLLMLoggingBaseClass): return data def _get_masked_api_base(self, api_base: str) -> str: - if "key=" in api_base: - # Find the position of "key=" in the string - key_index: Final = api_base.find("key=") + 4 - # Mask the last 5 characters after "key=" - masked_api_base = api_base[:key_index] + "*" * 5 + api_base[-4:] - else: - masked_api_base = api_base - return str(masked_api_base) + return str(mask_api_base_credentials(api_base)) def _pre_call(self, input, api_key, model=None, additional_args={}): """ @@ -1182,6 +1222,14 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["api_key"] = api_key self.model_call_details["additional_args"] = additional_args self.model_call_details["log_event_type"] = "pre_api_call" + if is_classifier_call(self.call_type, self.model_call_details.get("litellm_params") or EMPTY_MAPPING): + self.classifier_input = ( + None + if should_redact_message_logging(self.model_call_details) + else classifier_input_snapshot( + additional_args.get("complete_input_dict"), openai_sdk=additional_args.get("openai_sdk") is True + ) + ) if model: # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model self.model_call_details["litellm_params"]["api_base"] = self._get_masked_api_base( @@ -1560,6 +1608,7 @@ class Logging(LiteLLMLoggingBaseClass): service_tier: str | None = None, data_residency: str | None = None, vertex_location: str | None = None, + billed_token_rates: "BilledTokenRates | None" = None, ) -> None: """ Helper method to store cost breakdown in the logging object. @@ -1579,8 +1628,10 @@ class Logging(LiteLLMLoggingBaseClass): service_tier: Tier the costs above were priced on, already resolved data_residency: Region uplift the costs above were priced on, already resolved vertex_location: Vertex AI location the costs above were priced on, already resolved + billed_token_rates: Per-token rates the costs above were billed at, already resolved """ + self.billed_token_rates = billed_token_rates self.cost_breakdown = CostBreakdown( input_cost=input_cost, output_cost=output_cost, @@ -2003,6 +2054,17 @@ class Logging(LiteLLMLoggingBaseClass): results=result, ) + elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped + combined_ws_usage: Final = ( + ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( + results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream + ) + ) + logging_result = LiteLLMRealtimeStreamLoggingObject( + usage=combined_ws_usage, + results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream + ) + elif ( self.call_type == CallTypes.llm_passthrough_route.value or self.call_type == CallTypes.allm_passthrough_route.value @@ -2334,52 +2396,17 @@ class Logging(LiteLLMLoggingBaseClass): for scope in [key for key in spans_logged if isinstance(key, tuple) and key[-1:] == ("success",)]: del spans_logged[scope] - def _flush_passthrough_collected_chunks_helper( - self, - raw_bytes: list[bytes], - provider_config: "BasePassthroughConfig", - ) -> Optional["CostResponseTypes"]: - all_chunks: Final = provider_config._convert_raw_bytes_to_str_lines(raw_bytes) - complete_streaming_response: Final = provider_config.handle_logging_collected_chunks( - all_chunks=all_chunks, - litellm_logging_obj=self, - model=self.model, - custom_llm_provider=self.model_call_details.get("custom_llm_provider", ""), - endpoint=self.model_call_details.get("endpoint", ""), - ) - return complete_streaming_response - - def flush_passthrough_collected_chunks( - self, - raw_bytes: list[bytes], - provider_config: "BasePassthroughConfig", - ): + def flush_passthrough_collected_chunks(self, collector: "PassthroughStreamCollector"): """ - Flush collected chunks from the logging object - This is used to log the collected chunks once streaming is done on passthrough endpoints - - 1. Decode the raw bytes to string lines - 2. Get the complete streaming response from the provider config - 3. Log the complete streaming response (trigger success handler) - This is used for passthrough endpoints + Log the response a passthrough stream collector assembled once streaming is done (trigger success handler) """ - complete_streaming_response: Final = self._flush_passthrough_collected_chunks_helper( - raw_bytes=raw_bytes, - provider_config=provider_config, - ) + complete_streaming_response: Final = collector.build_logged_response(litellm_logging_obj=self) if complete_streaming_response is not None: self.success_handler(result=complete_streaming_response) - async def async_flush_passthrough_collected_chunks( - self, - raw_bytes: list[bytes], - provider_config: "BasePassthroughConfig", - ): - complete_streaming_response: Final = self._flush_passthrough_collected_chunks_helper( - raw_bytes=raw_bytes, - provider_config=provider_config, - ) + async def async_flush_passthrough_collected_chunks(self, collector: "PassthroughStreamCollector"): + complete_streaming_response: Final = collector.build_logged_response(litellm_logging_obj=self) if complete_streaming_response is not None: await self.async_success_handler(result=complete_streaming_response) @@ -2899,13 +2926,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,15 +2933,26 @@ 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 result._hidden_params["batch_successful_requests"] = batch_successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same result._hidden_params pattern as response_cost/batch_models above result._hidden_params["batch_failed_requests"] = batch_failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_usage + batch_prompt_cost: Final = kwargs.get("batch_prompt_cost", None) + batch_completion_cost: Final = kwargs.get("batch_completion_cost", None) + if ( + isinstance(batch_prompt_cost, float) + and isinstance(batch_completion_cost, float) + and isinstance(batch_cost, float) + ): + self.set_cost_breakdown( + input_cost=batch_prompt_cost, + output_cost=batch_completion_cost, + total_cost=batch_cost, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) elif should_compute_batch_data: batch_result: Final = await _handle_completed_batch( @@ -2937,6 +2968,12 @@ class Logging(LiteLLMLoggingBaseClass): result._hidden_params["batch_successful_requests"] = batch_result.successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_result.usage + self.set_cost_breakdown( + input_cost=batch_result.prompt_cost, + output_cost=batch_result.completion_cost, + total_cost=batch_result.cost, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) self.truncated_messages_for_logging = await truncate_base64_in_messages_async( StandardLoggingPayloadSetup.append_system_prompt_messages( @@ -3918,11 +3955,12 @@ class Logging(LiteLLMLoggingBaseClass): LiteLLMResponsesTransformationHandler, ) + served_id: Final = _provider_response_id(result) try: - return LiteLLMResponsesTransformationHandler().transform_response( + translated: Final = LiteLLMResponsesTransformationHandler().transform_response( model=self.model, raw_response=result, - model_response=litellm.ModelResponse(id=_provider_response_id(result)), + model_response=litellm.ModelResponse(id=served_id), logging_obj=self, request_data={}, messages=[], @@ -3930,6 +3968,8 @@ class Logging(LiteLLMLoggingBaseClass): litellm_params={}, encoding=litellm.encoding, ) + translated.id = served_id or translated.id + return translated except Exception as e: verbose_logger.debug( "Responses API -> ModelResponse translation failed for " @@ -3937,7 +3977,7 @@ class Logging(LiteLLMLoggingBaseClass): "usage-only ModelResponse to keep the spend_logs row.", str(e), ) - model_response: Final = litellm.ModelResponse(id=_provider_response_id(result)) + model_response: Final = litellm.ModelResponse(id=served_id) model_response.model = self.model usage: Final = getattr(result, "usage", None) if usage is not None and ResponseAPILoggingUtils._is_response_api_usage(usage): @@ -4321,6 +4361,14 @@ def _init_custom_logger_compatible_class( _s3_v2_logger: Final = S3V2Logger() _in_memory_loggers.append(_s3_v2_logger) return _s3_v2_logger + elif logging_integration == "pointfive": + for callback in _in_memory_loggers: + if isinstance(callback, PointFiveLogger): + return callback + + _pointfive_logger: Final = PointFiveLogger() + _in_memory_loggers.append(_pointfive_logger) + return _pointfive_logger elif logging_integration == "aws_sqs": for callback in _in_memory_loggers: if isinstance(callback, SQSLogger): @@ -4800,31 +4848,83 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom Returns ``None`` when V2 is off OR when there's no preset registered for ``callback_name`` — callers should then fall through to the legacy path. + + A preset that needs operator credentials it cannot find is allowed to build + only when this request has a key/team destination for that backend and another + V2 logger is already registered to carry the fan-out. The resulting logger keeps + only its credential-gated exporter, while the registered logger owns operator + delivery. Without that carrier, a preset that raises or that ends up with nothing + but its gated exporter and the default console placeholder returns ``None``, so the + caller falls through to the legacy path exactly as before V2 landed. """ from litellm.integrations.otel.model.config import is_otel_v2_enabled if not is_otel_v2_enabled(): return None from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger + from litellm.integrations.otel.plumbing.context import destination_backends from litellm.integrations.otel.presets import PRESET_BY_CALLBACK preset_fn: Final = PRESET_BY_CALLBACK.get(callback_name) if preset_fn is None: return None + serves_a_destination: Final = callback_name in destination_backends() + has_v2_logger: Final = any(isinstance(callback, OpenTelemetryV2) for callback in _in_memory_loggers) + carried: Final = serves_a_destination and has_v2_logger for callback in _in_memory_loggers: - if isinstance(callback, OpenTelemetryV2) and getattr(callback, "callback_name", None) == callback_name: + if ( + isinstance(callback, OpenTelemetryV2) + and getattr(callback, "callback_name", None) == callback_name + and (serves_a_destination or not _exports_nowhere(callback.config)) + ): return callback try: - config: Final = preset_fn() + built: Final = preset_fn(allow_missing_credentials=carried) except Exception: # If env vars are missing or the preset raises, defer to the legacy path # so customers get the same error story they had before V2 landed. return None + gated: Final = _is_credential_gated(built) + if gated and not carried and not _has_operator_exporter(built): + return None + config: Final = _only_the_gated_exporter(built) if gated and carried else built + if _exports_nowhere(config): + verbose_logger.warning( + "OTel V2: no operator credentials for '%s'; only key/team destinations will receive its traces", + callback_name, + ) v2_logger: Final = build_otel_v2_logger(config=config, callback_name=callback_name) _in_memory_loggers.append(v2_logger) return v2_logger +def _exports_nowhere(config: "OpenTelemetryV2Config") -> bool: + """Whether every exporter in ``config`` is waiting on credentials it never got.""" + return all(_is_gated(spec) for spec in config.exporters) + + +def _is_credential_gated(config: "OpenTelemetryV2Config") -> bool: + """Whether the preset built without the operator's own credentials for its backend.""" + return any(_is_gated(spec) for spec in config.exporters) + + +def _has_operator_exporter(config: "OpenTelemetryV2Config") -> bool: + """Whether the operator configured somewhere real to export, beyond the default console placeholder.""" + from litellm.integrations.otel.presets.utils import is_unconfigured_placeholder + + return any(not _is_gated(spec) and not is_unconfigured_placeholder(spec) for spec in config.exporters) + + +def _only_the_gated_exporter(config: "OpenTelemetryV2Config") -> "OpenTelemetryV2Config": + return config.model_copy( + update={"exporters": [spec for spec in config.exporters if _is_gated(spec)]} # mutable-ok: model_copy update + ) + + +def _is_gated(spec: "ExporterSpec") -> bool: + return spec.requires_headers and not spec.headers + + def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list[CustomLogger]) -> None: """ Auto-initialize ArizePhoenixLogger when Phoenix env vars are detected. @@ -4957,6 +5057,10 @@ def get_custom_logger_compatible_class( for callback in _in_memory_loggers: if isinstance(callback, S3V2Logger): return callback + elif logging_integration == "pointfive": + for callback in _in_memory_loggers: + if isinstance(callback, PointFiveLogger): + return callback elif logging_integration == "aws_sqs": for callback in _in_memory_loggers: if isinstance(callback, SQSLogger): @@ -5250,23 +5354,23 @@ class StandardLoggingPayloadSetup: Returns: dict: Merged metadata with user API key fields taking precedence """ - merged_metadata: Final[dict] = {} - - # Start with metadata (user API key fields) - but skip non-serializable objects - if litellm_params.get("metadata") and isinstance(litellm_params.get("metadata"), dict): - for key, value in litellm_params["metadata"].items(): - # Skip non-serializable objects like UserAPIKeyAuth - if key in {"user_api_key_auth", "user_api_key_budget_reservation"}: - continue - merged_metadata[key] = value - - # Then merge litellm_metadata (model-related fields) - this will NOT overwrite existing keys - if litellm_params.get("litellm_metadata") and isinstance(litellm_params.get("litellm_metadata"), dict): - for key, value in litellm_params["litellm_metadata"].items(): - if key not in merged_metadata: # Don't overwrite existing keys from metadata - merged_metadata[key] = value - - return merged_metadata + metadata: Final = litellm_params.get("metadata") + litellm_metadata: Final = litellm_params.get("litellm_metadata") + user_metadata: Final = MappingProxyType( + { + key: value + for key, value in (metadata.copy().items() if isinstance(metadata, dict) else ()) + if key not in _UNSERIALIZABLE_METADATA_KEYS + } + ) + model_metadata: Final = MappingProxyType( + { + key: value + for key, value in (litellm_metadata.copy().items() if isinstance(litellm_metadata, dict) else ()) + if key not in user_metadata + } + ) + return {**user_metadata, **model_metadata} # mutable-ok: function contract returns a plain dict @staticmethod def get_standard_logging_metadata( @@ -5524,7 +5628,7 @@ class StandardLoggingPayloadSetup: additional_logging_headers[key] = additiona_headers[_key] # Preserve all remaining headers verbatim (e.g. llm_provider-x-request-id) - for k, v in additiona_headers.items(): + for k, v in additiona_headers.copy().items(): if k.lower() not in typed_keys: additional_logging_headers[k] = v @@ -5670,6 +5774,7 @@ class StandardLoggingPayloadSetup: rate_limit_category: Final = validate_rate_limit_category(getattr(original_exception, "category", None)) rate_limit_type: Final = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None)) budget_error: Final = original_exception if isinstance(original_exception, BudgetExceededError) else None + provider_request_id: Final = _get_provider_request_id(original_exception) if original_exception else None return StandardLoggingPayloadErrorInformation( error_code=error_status, @@ -5677,6 +5782,7 @@ class StandardLoggingPayloadSetup: llm_provider=_llm_provider_in_exception, traceback=_redact_string(traceback_info), error_message=_redact_string(error_message), + error_provider_request_id=provider_request_id, error_rate_limit_category=rate_limit_category, error_rate_limit_type=rate_limit_type, error_budget_entity_type=budget_error.entity_type if budget_error else None, @@ -6171,6 +6277,18 @@ def get_standard_logging_object_payload( ) payload: Final[StandardLoggingPayload] = StandardLoggingPayload( + **( + classifier_audit_fields( + MappingProxyType( + { + "classifier_input": logging_obj.classifier_input, + "originating_request_masked": proxy_server_request.get("originating_request_masked"), + } + ) + ) + if is_classifier_call(call_type or "", litellm_params) and not should_redact_message_logging(kwargs) + else EMPTY_MAPPING + ), id=str(id), litellm_call_id=kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), trace_id=StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( @@ -6352,7 +6470,7 @@ def _get_traceback_str_for_error(error_str: str) -> str: from decimal import Decimal # used for unit testing -from typing import Any, Optional, Union +from typing import Any, Union def create_dummy_standard_logging_payload() -> StandardLoggingPayload: diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 8e24302b440..e5977ca4156 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -10,6 +10,7 @@ from typing import Any, Final, Literal, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import litellm +from litellm._internal_context import current_billing_time from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import ( select_tier_for_input, @@ -19,6 +20,7 @@ from litellm.types.utils import ( CacheCreationTokenDetails, CallTypes, CompletionTokensDetailsWrapper, + CostPerToken, DataResidency, ImageResponse, ModelInfo, @@ -26,6 +28,7 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, ServiceTier, Usage, + text_tokens_without_nested_reasoning, ) from litellm.utils import get_model_info @@ -304,7 +307,7 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_ than being localised, so callers must pass datetime.now(timezone.utc), never datetime.now(), or every window shifts by the host's offset. """ - reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + reference: Final = current_time if current_time is not None else current_billing_time() now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time() windows: Final = (off_peak_hours_utc,) if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc for window in windows: @@ -391,7 +394,7 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = rules: the flat hours_utc windows, which apply every day, or any entry in windows, whose hours apply only on its weekdays. """ - reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + reference: Final = current_time if current_time is not None else current_billing_time() reference_utc: Final = ( reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference.replace(tzinfo=timezone.utc) ) @@ -513,6 +516,7 @@ def _get_token_base_cost( current_time: datetime | None = None, *, threshold_is_inclusive: bool = False, + missing_cache_read_uses_input: bool = False, ) -> tuple[float, float, float, float, float]: """ Return prompt cost, completion cost, and cache costs for a given model and usage. @@ -523,6 +527,9 @@ def _get_token_base_cost( `threshold_is_inclusive` switches that comparison to >=, for providers such as xAI that bill the higher tier once the prompt reaches the threshold. + `missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved + input rate instead of 0.0; an explicit 0.0 rate stays a real price either way. + Returns: Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost) """ @@ -550,29 +557,16 @@ def _get_token_base_cost( float, _get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"), ) - cache_read_cost = cast(float, _get_cost_per_unit(model_info, cache_read_cost_key)) + cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None) ## CHECK IF ABOVE THRESHOLD # Optimization: collect threshold keys first to avoid sorting all model_info keys. - # Most models don't have threshold pricing, so we can return early. # Exclude service_tier-specific variants (e.g. input_cost_per_token_above_200k_tokens_priority) # so that the threshold detection loop only processes standard keys. The # service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key. threshold_keys: Final = [ k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES) ] - if not threshold_keys: - return _apply_off_peak_to_base_costs( - model_info, - current_time, - ( - prompt_base_cost, - completion_base_cost, - cache_creation_cost, - cache_creation_cost_above_1hr, - cache_read_cost, - ), - ) # Only sort the threshold keys (typically 1-2 keys instead of 66+) threshold: float | None = None @@ -661,10 +655,7 @@ def _get_token_base_cost( ), ) - cache_read_cost = cast( - float, - _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost), - ) + cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost) break except (IndexError, ValueError): @@ -672,6 +663,17 @@ def _get_token_base_cost( except Exception: continue + if cache_read_cost is None: + cache_read_cost = ( + _off_peak_rate( + _open_off_peak_block(model_info, current_time) or MappingProxyType({}), + "input_cost_per_token", + prompt_base_cost, + ) + if missing_cache_read_uses_input + else 0.0 + ) + return _apply_off_peak_to_base_costs( model_info, current_time, @@ -780,6 +782,7 @@ class PromptTokensDetailsResult(TypedDict): image_count: int video_length_seconds: float audio_length_seconds: float + query_count: int def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: @@ -828,6 +831,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or 0.0 ) + query_count: Final = _coerce_token_count(getattr(usage.prompt_tokens_details, "query_count", 0)) return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, @@ -841,6 +845,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: image_count=image_count, video_length_seconds=float(video_length_seconds), audio_length_seconds=float(audio_length_seconds), + query_count=query_count, ) @@ -860,7 +865,7 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu ) or 0 ) - text_tokens: Final = ( + reported_text_tokens: Final = ( cast( int | None, getattr(usage.completion_tokens_details, "text_tokens", None), @@ -882,6 +887,12 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu or 0 ) video_tokens: Final = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0)) + text_tokens: Final = text_tokens_without_nested_reasoning( + completion_tokens=usage.completion_tokens, + text_tokens=reported_text_tokens, + reasoning_tokens=reasoning_tokens, + other_modality_tokens=audio_tokens + image_tokens + video_tokens, + ) return CompletionTokensDetailsResult( audio_tokens=audio_tokens, @@ -972,6 +983,11 @@ def _calculate_input_cost( prompt_tokens_details["audio_length_seconds"], ) + if prompt_tokens_details["query_count"]: + prompt_cost += calculate_cost_component( + model_info, "input_cost_per_query", prompt_tokens_details["query_count"] + ) + return prompt_cost @@ -1143,6 +1159,7 @@ def generic_cost_per_token( image_count=0, video_length_seconds=0.0, audio_length_seconds=0.0, + query_count=0, ) if usage.prompt_tokens_details: prompt_tokens_details = parse_prompt_tokens_details(usage) @@ -1180,7 +1197,7 @@ def generic_cost_per_token( usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0 ) - billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) + billing_time: Final = current_time if current_time is not None else current_billing_time() ( prompt_base_cost, completion_base_cost, @@ -1294,42 +1311,90 @@ def _coerce_token_count(value: object) -> int: return value if isinstance(value, int) and value > 0 else 0 +@dataclass(frozen=True, slots=True) +class BilledTokenRates: + """Per-token rates one request's usage bills at, after token tiers, off-peak windows and the + regional multipliers the totals apply, so each cost line equals its token count times its rate.""" + + input_cost_per_token: float + output_cost_per_token: float + cache_read_input_token_cost: float + cache_creation_input_token_cost: float + cache_creation_input_token_cost_above_1hr: float + output_cost_per_reasoning_token: float + + def scaled(self, multiplier: float) -> "BilledTokenRates": + if multiplier == 1.0: + return self + return BilledTokenRates( + input_cost_per_token=self.input_cost_per_token * multiplier, + output_cost_per_token=self.output_cost_per_token * multiplier, + cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier, + cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier, + cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier, + output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier, + ) + + @dataclass(frozen=True, slots=True) class TokenTypeCostBreakdown: reasoning_cost: float cache_read_cost: float cache_creation_cost: float + rates: BilledTokenRates | None = None + """Rates these lines were billed at, so a caller reporting both cannot resolve them a second, + differently-argued way. None when the model's pricing could not be resolved.""" -def get_token_type_cost_breakdown( - model: str, - custom_llm_provider: str | None, +def _reasoning_token_count(usage: Usage) -> int: + parsed: Final = ( + parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 + ) + return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) + + +def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]: + """(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details + first, then the private top-level counters the Usage constructor mirrors cache tokens onto for + providers/callers that bypass the details.""" + parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None + parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0 + parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0 + return ( + parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)), + parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)), + parsed["cache_creation_token_details"] if parsed is not None else None, + ) + + +def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRates: + """Flat custom pricing has no tiers, uplifts or reasoning rate: cache tokens bill at the configured + cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does.""" + input_rate: Final = custom_cost_per_token["input_cost_per_token"] + output_rate: Final = custom_cost_per_token["output_cost_per_token"] + cache_creation_rate: Final = custom_cost_per_token.get("cache_creation_input_token_cost", input_rate) + return BilledTokenRates( + input_cost_per_token=input_rate, + output_cost_per_token=output_rate, + cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate), + cache_creation_input_token_cost=cache_creation_rate, + cache_creation_input_token_cost_above_1hr=cache_creation_rate, + output_cost_per_reasoning_token=output_rate, + ) + + +def _cost_map_billed_rates( + model_info: ModelInfo, usage: Usage, - service_tier: str | None = None, - data_residency: str | None = None, - vertex_location: str | None = None, - current_time: datetime | None = None, -) -> TokenTypeCostBreakdown: - """ - Provider-agnostic cost of reasoning and cache tokens, derived from the usage - object and model pricing alone. - - This works for every provider, including Perplexity/Cerebras/Dashscope whose - cost calculators bypass ``generic_cost_per_token``, because cache tokens always - land on ``prompt_tokens_details`` (via the Usage constructor and provider - transformations) and reasoning tokens on ``completion_tokens_details``. It reuses - the same rate-resolution primitives as the total-cost path so the breakdown can - never drift from the totals. Returns zeros (never raises) when the model or its - pricing cannot be resolved. - """ - try: - model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - return TokenTypeCostBreakdown(0.0, 0.0, 0.0) - - billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) + custom_llm_provider: str | None, + service_tier: str | None, + data_residency: str | None, + vertex_location: str | None, + current_time: datetime | None, +) -> BilledTokenRates: + billing_time: Final = current_time if current_time is not None else current_billing_time() ( - _prompt_base_cost, + prompt_base_cost, completion_base_cost, cache_creation_cost_rate, cache_creation_cost_above_1hr_rate, @@ -1341,13 +1406,6 @@ def get_token_type_cost_breakdown( current_time=billing_time, threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) - - reasoning_tokens = ( - parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 - ) - if not reasoning_tokens: - reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) - reasoning_rate: Final = _resolve_billed_reasoning_rate( model_info=model_info, usage=usage, @@ -1355,60 +1413,157 @@ def get_token_type_cost_breakdown( completion_base_cost=completion_base_cost, current_time=billing_time, ) - reasoning_cost = float(reasoning_tokens) * reasoning_rate + multiplier: Final = ( + _get_regional_uplift_multiplier(model_info, data_residency) + * get_vertex_regional_endpoint_uplift(model_info, vertex_location) + * get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) + ) + return BilledTokenRates( + input_cost_per_token=prompt_base_cost, + output_cost_per_token=completion_base_cost, + cache_read_input_token_cost=cache_read_cost_rate, + cache_creation_input_token_cost=cache_creation_cost_rate, + cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate, + output_cost_per_reasoning_token=reasoning_rate, + ).scaled(multiplier) - cache_read_tokens = 0 - cache_creation_tokens = 0 - cache_creation_token_details: CacheCreationTokenDetails | None = None - if usage.prompt_tokens_details is not None: - prompt_tokens_details: Final = parse_prompt_tokens_details(usage) - cache_read_tokens = prompt_tokens_details["cache_hit_tokens"] - cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"] - cache_creation_token_details = prompt_tokens_details["cache_creation_token_details"] - # Fall back to the private top-level counters the Usage constructor mirrors cache - # tokens onto, so providers/callers that bypass prompt_tokens_details are covered. - if not cache_read_tokens: - cache_read_tokens = _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)) - if not cache_creation_tokens: - cache_creation_tokens = _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)) - cache_read_cost = float(cache_read_tokens) * cache_read_cost_rate - cache_creation_cost = calculate_cache_writing_cost( - cache_creation_tokens=cache_creation_tokens, - cache_creation_token_details=cache_creation_token_details, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr_rate, - cache_creation_cost=cache_creation_cost_rate, +def get_billed_token_rates( + model: str, + custom_llm_provider: str | None, + usage: Usage, + service_tier: str | None = None, + data_residency: str | None = None, + vertex_location: str | None = None, + current_time: datetime | None = None, + custom_cost_per_token: CostPerToken | None = None, +) -> BilledTokenRates | None: + """Rates the cost calculator bills ``usage`` at, resolved exactly as the totals and the token-type + breakdown resolve them. None when the model's pricing cannot be resolved.""" + if custom_cost_per_token is not None: + return _custom_pricing_rates(custom_cost_per_token) + try: + model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates + return None + return _cost_map_billed_rates( + model_info=model_info, + usage=usage, + custom_llm_provider=custom_llm_provider, + service_tier=service_tier, + data_residency=data_residency, + vertex_location=vertex_location, + current_time=current_time, ) - # Apply the same flat regional-processing uplift the totals get, so per-type - # costs stay reconciled with input_cost/output_cost for regionalized OpenAI hosts. - uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency) - if uplift != 1.0: - reasoning_cost *= uplift - cache_read_cost *= uplift - cache_creation_cost *= uplift - vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) - if vertex_uplift != 1.0: - reasoning_cost *= vertex_uplift - cache_read_cost *= vertex_uplift - cache_creation_cost *= vertex_uplift +def get_token_type_cost_breakdown( + model: str, + custom_llm_provider: str | None, + usage: Usage, + service_tier: str | None = None, + data_residency: str | None = None, + vertex_location: str | None = None, + current_time: datetime | None = None, + custom_cost_per_token: CostPerToken | None = None, +) -> TokenTypeCostBreakdown: + """ + Provider-agnostic cost of reasoning and cache tokens, derived from the usage + object and model pricing alone. - # Mirror the provider-specific geo uplift (e.g. Anthropic us: 1.1) the totals - # apply, so cache and reasoning line items stay reconciled with them. - geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) - if geo_multiplier != 1.0: - reasoning_cost *= geo_multiplier - cache_read_cost *= geo_multiplier - cache_creation_cost *= geo_multiplier + This works for every provider, including Perplexity/Cerebras/Dashscope whose + cost calculators bypass ``generic_cost_per_token``, because cache tokens always + land on ``prompt_tokens_details`` (via the Usage constructor and provider + transformations) and reasoning tokens on ``completion_tokens_details``. It reuses + the same rate resolution as the total-cost path (``get_billed_token_rates``) so the + breakdown can never drift from the totals. A deployment billed by + ``custom_cost_per_token`` is priced from those flat rates instead of the cost map and, + like its totals, bills cache writes flat rather than by their 5m/1h split. + Returns zeros (never raises) when the model or its pricing cannot be resolved. + """ + rates: Final = get_billed_token_rates( + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + service_tier=service_tier, + data_residency=data_residency, + vertex_location=vertex_location, + current_time=current_time, + custom_cost_per_token=custom_cost_per_token, + ) + if rates is None: + return TokenTypeCostBreakdown(0.0, 0.0, 0.0) + cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage) + cache_creation_cost: Final = ( + float(cache_creation_tokens) * rates.cache_creation_input_token_cost + if custom_cost_per_token is not None + else calculate_cache_writing_cost( + cache_creation_tokens=cache_creation_tokens, + cache_creation_token_details=cache_creation_token_details, + cache_creation_cost_above_1hr=rates.cache_creation_input_token_cost_above_1hr, + cache_creation_cost=rates.cache_creation_input_token_cost, + ) + ) return TokenTypeCostBreakdown( - reasoning_cost=reasoning_cost, - cache_read_cost=cache_read_cost, + reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token, + cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost, cache_creation_cost=cache_creation_cost, + rates=rates, ) +def calculate_prompt_caching_savings( + model_info: ModelInfo, + usage: Usage, + custom_llm_provider: str | None, + service_tier: str | None = None, + data_residency: str | None = None, + vertex_location: str | None = None, + billed_at: datetime | None = None, +) -> float: + """Read discount minus write premium, using the biller's rate and TTL resolution. + + Missing reads and unpublished (missing/zero) writes claim no saving or premium; + explicit zero reads remain free. An unpublished 1h price uses the ordinary write rate. + ``billed_at`` is the request's completion time, so off-peak windows resolve as the + biller saw them rather than at the later spend write. + """ + prompt_base_cost, _, cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost = _get_token_base_cost( + model_info=model_info, + usage=usage, + service_tier=service_tier, + current_time=billed_at, + threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), + missing_cache_read_uses_input=True, + ) + write_rate: Final = cache_creation_cost or prompt_base_cost + write_rate_1h: Final = cache_creation_cost_above_1hr or write_rate + prompt_tokens_details: Final = parse_prompt_tokens_details(usage) + cache_read_tokens: Final = max(prompt_tokens_details["cache_hit_tokens"], 0) + cache_creation_tokens: Final = max(prompt_tokens_details["cache_creation_tokens"], 0) + details: Final = prompt_tokens_details["cache_creation_token_details"] + cache_creation_details: Final = ( + CacheCreationTokenDetails( + ephemeral_5m_input_tokens=max(details.ephemeral_5m_input_tokens or 0, 0), + ephemeral_1h_input_tokens=max(details.ephemeral_1h_input_tokens or 0, 0), + ) + if details is not None + else None + ) + read_discount: Final = cache_read_tokens * max(prompt_base_cost - cache_read_cost, 0.0) + write_premium: Final = calculate_cache_writing_cost( + cache_creation_tokens=cache_creation_tokens, + cache_creation_token_details=cache_creation_details, + cache_creation_cost_above_1hr=write_rate_1h - prompt_base_cost, + cache_creation_cost=write_rate - prompt_base_cost, + ) + uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift( + model_info, vertex_location + ) + return (read_discount - write_premium) * uplift + + def calculate_image_response_cost_from_usage( model: str, image_response: ImageResponse, diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index a6f10e1ede3..87524d86c61 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -3,7 +3,7 @@ import json import re import time import traceback -from collections.abc import Iterable, Sequence +from collections.abc import Mapping, Sequence from typing import Final, Literal, cast import litellm @@ -151,6 +151,16 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None: del choice.enhancements +def _invalid_choices_message(response_object: Mapping[str, object]) -> str: + raw_keys: Final = list(response_object.keys()) + if "choices" not in response_object: + return f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {raw_keys}" + return ( + f"LiteLLM: provider returned 'choices' that is not a list ({type(response_object['choices']).__name__}). " + f"Raw keys: {raw_keys}" + ) + + async def convert_to_streaming_response_async( response_object: dict | None = None, ): @@ -179,14 +189,12 @@ async def convert_to_streaming_response_async( choice_list: Final[list[StreamingChoices]] = [] - if not response_object.get("choices"): + if not isinstance(response_object.get("choices"), list): from litellm.exceptions import APIError raise APIError( status_code=500, - message=( - f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}" - ), + message=_invalid_choices_message(response_object), llm_provider="", model="", ) @@ -287,14 +295,12 @@ def convert_to_streaming_response( model_response_object: Final = ModelResponseStream() choice_list: Final[list[StreamingChoices]] = [] - if not response_object.get("choices"): + if not isinstance(response_object.get("choices"), list): from litellm.exceptions import APIError raise APIError( status_code=500, - message=( - f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}" - ), + message=_invalid_choices_message(response_object), llm_provider="", model="", ) @@ -623,15 +629,12 @@ def convert_to_model_response_object( return convert_to_streaming_response(response_object=response_object) choice_list: Final[list[Choices]] = [] - if not response_object.get("choices") or not isinstance(response_object["choices"], Iterable): + if not isinstance(response_object.get("choices"), list): from litellm.exceptions import APIError raise APIError( status_code=500, - message=( - "LiteLLM: provider returned a response with no 'choices'. " - f"Raw keys: {list(response_object.keys())}" - ), + message=_invalid_choices_message(response_object), llm_provider="", model="", ) diff --git a/litellm/litellm_core_utils/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py index f1ae6492e4e..d04abcb6e7b 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py @@ -1,7 +1,8 @@ +from collections.abc import Mapping from typing import Final -def get_response_headers(_response_headers: dict | None = None) -> dict: +def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict: """ Sets the Appropriate OpenAI headers for the response and forward all headers as llm_provider-{header} @@ -31,7 +32,7 @@ def get_response_headers(_response_headers: dict | None = None) -> dict: return {**llm_provider_headers, **openai_headers} -def _get_llm_provider_headers(response_headers: dict) -> dict: +def _get_llm_provider_headers(response_headers: Mapping[str, str]) -> dict: """ Adds a llm_provider-{header} to all headers that are not already prefixed with llm_provider diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 9b612993a69..31523c9309d 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -441,7 +441,15 @@ class LoggingCallbackManager: return result + def get_callback_objects(self) -> tuple[tuple[str, CustomLogger | Callable], ...]: + return tuple( + (self._get_callback_string(callback), callback) + for callback in self._get_all_callbacks() + if not isinstance(callback, str) + ) + def _get_callback_string(self, callback: CustomLogger | Callable | str) -> str: + from litellm.integrations.opentelemetry import OpenTelemetry from litellm.litellm_core_utils.custom_logger_registry import ( CustomLoggerRegistry, ) @@ -449,6 +457,8 @@ class LoggingCallbackManager: """Convert a callback to its string representation""" if isinstance(callback, str): return callback + elif isinstance(callback, OpenTelemetry) and callback.callback_name is not None: + return callback.callback_name elif isinstance(callback, CustomLogger): # Try to get the string representation from the registry callback_str: Final = CustomLoggerRegistry.get_callback_str_from_class_type(type(callback)) diff --git a/litellm/litellm_core_utils/private_json.py b/litellm/litellm_core_utils/private_json.py index 30f64c8fc27..4cd4a9b4f82 100644 --- a/litellm/litellm_core_utils/private_json.py +++ b/litellm/litellm_core_utils/private_json.py @@ -36,6 +36,21 @@ def stage_private_json(path: str, data: Mapping[str, object]) -> str: return tmp_path +def stage_private_bytes(path: str, data: bytes) -> str: + parent: Final = Path(path).parent + parent.mkdir(parents=True, exist_ok=True) + fd, tmp_path = tempfile.mkstemp(dir=str(parent), prefix=".tmp-") + try: + with os.fdopen(fd, "wb") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + except BaseException: + Path(tmp_path).unlink(missing_ok=True) + raise + return tmp_path + + def commit_staged_json(staged: str, path: str) -> None: """Move a staged file into place, replacing whatever is there in one step""" try: @@ -68,3 +83,8 @@ def discard_staged_json(staged: str) -> None: def write_private_json(path: str, data: Mapping[str, object]) -> None: """Atomically write JSON to path with owner-only permissions (0600)""" commit_staged_json(stage_private_json(path, data), path) + + +def write_private_bytes(path: str, data: bytes) -> None: + """Atomically write bytes to path with owner-only permissions (0600); a reader holding the old file keeps it whole""" + commit_staged_json(stage_private_bytes(path, data), path) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 00b80839dde..2485896184e 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,8 +6,8 @@ import io import json import mimetypes import re -from collections.abc import Iterable, Mapping, Sequence -from itertools import groupby +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from itertools import groupby, islice from os import PathLike from pathlib import Path from types import MappingProxyType @@ -1320,17 +1320,128 @@ def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mappin return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo -def tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]: +_SUBSCHEMA_KEYWORDS: Final = frozenset( + { + "additionalItems", + "additionalProperties", + "contains", + "else", + "if", + "items", + "not", + "propertyNames", + "then", + "unevaluatedItems", + "unevaluatedProperties", + } +) +_SUBSCHEMA_LIST_KEYWORDS: Final = frozenset({"allOf", "anyOf", "items", "oneOf", "prefixItems"}) +_SUBSCHEMA_MAP_KEYWORDS: Final = frozenset( + {"$defs", "definitions", "dependentSchemas", "patternProperties", "properties"} +) + +_MAX_SCHEMA_NESTING: Final = 1024 + + +def drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]: + """Drop every regex in a schema position that Python's ``re`` cannot compile. + + OpenAI validates tool ``parameters`` against the 2020-12 metaschema with + ``jsonschema``'s format checker, which hands each ``pattern`` value and each + ``patternProperties`` key to ``re.compile``, so a regex written for an + ECMA-262 engine (Unicode property escapes such as ``\\p{Cc}``, as in Claude + Code's ``Artifact`` tool) is refused with "'...' is not a 'regex'" by every + model family on both the chat and Responses wires. Only schema positions are + walked (properties, items, combinators, ``$defs`` and the other applicators), + so a ``pattern`` key inside ``default``, ``examples``, ``const`` or vendor + extensions is data and stays. Outside strict mode the keyword is only a + hint, so dropping it costs the model a constraint and the caller nothing. + Compilable regexes and everything else pass through, the input is never + mutated, and the same object comes back when nothing was dropped. The walk + is level-order rather than recursive, rebuilt deepest level first, and stops + at more schema levels than a JSON parser admits, so a cyclic schema built in + code cannot spin it. + """ + rebuilt: dict[int, Mapping[str, object]] = {} # mutable-ok: per-call memo of rewritten nodes, deepest level first + for level in reversed(tuple(islice(_schema_levels(schema), _MAX_SCHEMA_NESTING))): + rebuilt.update( + (id(node), rewritten) + for node in level + if (rewritten := _node_without_non_python_regex(node, rebuilt)) is not node + ) + return rebuilt.get(id(schema), schema) + + +def _schema_levels(schema: Mapping[str, object]) -> Iterator[tuple[Mapping[str, object], ...]]: + frontier: tuple[Mapping[str, object], ...] = (schema,) # rebind-ok: level-order cursor, one level a round + while frontier: + yield frontier + frontier = tuple(child for node in frontier for child in _subschemas(node)) + + +def _subschemas(node: Mapping[str, object]) -> Iterator[Mapping[str, object]]: + for key, value in node.items(): + if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict): + yield from (sub for sub in value.values() if isinstance(sub, dict)) + elif key in _SUBSCHEMA_LIST_KEYWORDS and isinstance(value, list): + yield from (sub for sub in value if isinstance(sub, dict)) + elif key in _SUBSCHEMA_KEYWORDS and isinstance(value, dict): + yield value + + +def _node_without_non_python_regex( + node: Mapping[str, object], rebuilt: Mapping[int, Mapping[str, object]] +) -> Mapping[str, object]: + kept: Final = { # mutable-ok: tool parameters are JSON dicts + key: _keyword_value_rebuilt(key, value, rebuilt) + for key, value in node.items() + if key != "pattern" or not isinstance(value, str) or _is_python_regex(value) + } + return node if len(kept) == len(node) and all(kept[key] is node[key] for key in kept) else kept + + +def _keyword_value_rebuilt(key: str, value: object, rebuilt: Mapping[int, Mapping[str, object]]) -> object: + if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict): + kept: Final = { # mutable-ok: tool parameters are JSON dicts + name: rebuilt.get(id(sub), sub) + for name, sub in value.items() + if key != "patternProperties" or not isinstance(name, str) or _is_python_regex(name) + } + return value if len(kept) == len(value) and all(kept[name] is value[name] for name in kept) else kept + if key in _SUBSCHEMA_LIST_KEYWORDS and isinstance(value, list): + items: Final = [rebuilt.get(id(sub), sub) for sub in value] # mutable-ok: tool parameters are JSON lists + return value if all(new is old for new, old in zip(items, value, strict=True)) else items + if key in _SUBSCHEMA_KEYWORDS and isinstance(value, dict): + return rebuilt.get(id(value), value) + return value + + +def _is_python_regex(pattern: str) -> bool: + try: + re.compile(pattern) + except (re.error, RecursionError): + return False + return True + + +def flatten_combinators_and_drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]: + return flatten_top_level_schema_combinators(drop_non_python_regex_patterns(schema)) + + +def tool_with_sanitized_parameters( + tool: Mapping[str, object], + sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], +) -> Mapping[str, object]: function: Final = tool.get("function") if not isinstance(function, dict): return tool parameters: Final = function.get("parameters") if not isinstance(parameters, dict): return tool - flattened: Final = flatten_top_level_schema_combinators(parameters) - if flattened is parameters: + sanitized: Final = sanitize(parameters) + if sanitized is parameters: return tool - return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts + return {**tool, "function": {**function, "parameters": sanitized}} # mutable-ok: request tools are JSON dicts def _get_image_mime_type_from_url(url: str) -> str | None: @@ -1823,14 +1934,11 @@ def _extract_reasoning_content(message: dict) -> tuple[str | None, str | None]: return None, message_content -def _readable_thinking_text( - block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, -) -> str: +def _readable_thinking_text(block: Mapping[str, object]) -> str: """The text a chat model can read back, empty for redacted blocks and malformed ones.""" if block.get("type") != "thinking": return "" - thinking: Final = cast(ChatCompletionThinkingBlock, block).get("thinking") # cast-ok: narrowed by the type tag - return str(thinking or "") + return str(block.get("thinking") or "") def reasoning_content_from_thinking_blocks( @@ -1843,24 +1951,125 @@ def reasoning_content_from_thinking_blocks( return "\n".join(text for block in thinking_blocks if (text := _readable_thinking_text(block))) -def responses_reasoning_item_from_thinking_blocks( - thinking_blocks: Iterable[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], -) -> ChatCompletionReasoningItem | None: - """Build a Responses API `reasoning` input item from Anthropic thinking blocks. +ENCRYPTED_REASONING_SIGNATURE_PREFIX: Final = "litellm_encrypted_reasoning:" - The item carries no `id`: the Responses API rejects an empty one and 404s on any id it - did not mint itself, while an item without an id is always accepted. + +def encrypted_reasoning_signature(encrypted_content: str) -> str: + """The opaque value a Responses API reasoning item's `encrypted_content` travels in. + + Anthropic clients echo a thinking block's `signature` and a redacted block's `data` + back verbatim, so either field can carry the encrypted reasoning across turns; the + prefix tells the two apart from a signature Anthropic minted. """ + return f"{ENCRYPTED_REASONING_SIGNATURE_PREFIX}{encrypted_content}" + + +def _carries_encrypted_reasoning(signature: object) -> bool: + return isinstance(signature, str) and signature.startswith(ENCRYPTED_REASONING_SIGNATURE_PREFIX) + + +def encrypted_content_from_signature(signature: object) -> str | None: + if not isinstance(signature, str) or not _carries_encrypted_reasoning(signature): + return None + return signature.removeprefix(ENCRYPTED_REASONING_SIGNATURE_PREFIX) or None + + +def _encrypted_reasoning_field(block: Mapping[str, object]) -> object: + match block.get("type"): + case "thinking": + return block.get("signature") + case "redacted_thinking": + return block.get("data") + case _: + return None + + +def encrypted_content_of_block(block: Mapping[str, object]) -> str | None: + return encrypted_content_from_signature(_encrypted_reasoning_field(block)) + + +def is_encrypted_reasoning_block(block: object) -> bool: + """A thinking or redacted_thinking block carrying Responses API encrypted reasoning. + + Only the Responses API that minted the content can read it back, so an Anthropic + backend has to drop such a block rather than fail signature verification on it. + """ + if not isinstance(block, Mapping): + return False + mapping: Final = cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance + return _carries_encrypted_reasoning(_encrypted_reasoning_field(mapping)) + + +def strip_encrypted_reasoning_from_messages(messages: object) -> None: + """Drop the bridge-tagged reasoning blocks a routed deployment cannot decrypt from + Anthropic-shaped history. + + The whole block goes, the way #40280 drops undecryptable Responses ``input`` items: a + provider that did not mint the block rejects it signed (a foreign signature) and unsigned + (a missing signature) alike, so keeping its text as an unsigned thinking block only moves + the 400 from the router to the provider. + + Mutates the content lists in place: the router's fallback snapshot shares these + message objects, so a rebound list would replay the stripped blocks on the fallback hop. + """ + if not isinstance(messages, list): + return + for content in _anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json + _strip_encrypted_reasoning_from_blocks(content) + + +def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: + return ( + cast(list[object], content) # cast-ok: narrowed by isinstance + for message in messages + if isinstance(message, Mapping) + for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance + if isinstance(content, list) + ) + + +def _strip_encrypted_reasoning_from_blocks(content: object) -> None: + blocks: Final = cast(list[object], content) # cast-ok: narrowed by the caller's isinstance + kept: Final = tuple(block for block in blocks if not is_encrypted_reasoning_block(block)) + blocks[:] = kept # rebind-ok: shared with fallback snapshot + + +def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: + index, block = indexed_block + return f"encrypted:{index}" if is_encrypted_reasoning_block(block) else "summary" + + +def _reasoning_item_from_block_group(group: tuple[Mapping[str, object], ...]) -> ChatCompletionReasoningItem | None: summary: Final[list[ChatCompletionReasoningSummaryTextBlock]] = [ # mutable-ok: API message payload ChatCompletionReasoningSummaryTextBlock(type="summary_text", text=text) - for block in thinking_blocks + for block in group if (text := _readable_thinking_text(block)) ] + encrypted_content: Final = encrypted_content_of_block(group[0]) + if encrypted_content is not None: + return ChatCompletionReasoningItem(type="reasoning", summary=summary, encrypted_content=encrypted_content) if not summary: return None return ChatCompletionReasoningItem(type="reasoning", summary=summary) +def responses_reasoning_items_from_thinking_blocks( + thinking_blocks: Iterable[Mapping[str, object]], +) -> tuple[ChatCompletionReasoningItem, ...]: + """Build Responses API `reasoning` input items from Anthropic thinking blocks. + + A block carrying encrypted reasoning replays the item it came from byte for byte; + a run of plain thinking blocks collapses into one summary-only item. No item carries + an `id`: the Responses API 404s on any id it did not mint itself and rejects an empty + one, while an item without an id is always accepted. + """ + return tuple( + item + for _, group in groupby(enumerate(thinking_blocks), key=_reasoning_replay_group_key) + if (item := _reasoning_item_from_block_group(tuple(block for _, block in group))) is not None + ) + + def _parse_content_for_reasoning( message_text: str | None, ) -> tuple[str | None, str | None]: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 56c1d605700..ece619e3883 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -46,6 +46,7 @@ from litellm.types.utils import GenericImageParsingChunk from .common_utils import ( convert_content_list_to_str, infer_content_type_from_url_and_content, + is_encrypted_reasoning_block, is_non_content_values_set, parse_tool_call_arguments, ) @@ -2299,13 +2300,16 @@ def sanitize_messages_for_tool_calling( def _is_unsignable_thinking_block(block: object) -> bool: - """A `thinking` block that Anthropic cannot accept on input. + """A thinking block that Anthropic cannot accept on input. Anthropic verifies the thinking signature cryptographically, so a block whose signature is null, empty, or missing (e.g. from an open-source reasoning model) - is rejected with a 400 and must be dropped rather than blanked or repaired. - `redacted_thinking` blocks carry no signature and are always kept. + is rejected with a 400 and must be dropped rather than blanked or repaired, and + so is a block whose signature or data carries another provider's encrypted + reasoning. A `redacted_thinking` block Anthropic minted is always kept. """ + if is_encrypted_reasoning_block(block): + return True if not isinstance(block, dict) or block.get("type") != "thinking": return False signature: Final = block.get("signature") diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 81132fa89a5..c9933422cc3 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -2,7 +2,11 @@ Helper functions to handle images passed in messages """ +import asyncio import base64 +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType from typing import Final from httpx import Response @@ -11,9 +15,11 @@ import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB -from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get +from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 +MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20 in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) @@ -72,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str: return result +def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": + verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict) + return litellm.ImageFetchError( + "Error: Unable to fetch image from URL. The proxy could not resolve this host or its URL policy rejected it; " + f"an admin can check the proxy log and `user_url_allowed_hosts` in general_settings. url={url}" + ) + + async def async_convert_url_to_base64(url: str) -> str: if url.startswith("data:") and ";base64," in url: return url @@ -93,6 +107,8 @@ async def async_convert_url_to_base64(url: str) -> str: return _process_image_response(response, url) except litellm.ImageFetchError: raise + except SSRFError as e: + raise _rejected_image_fetch(url, e) from e except Exception: pass raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") @@ -119,8 +135,197 @@ def convert_url_to_base64(url: str) -> str: return _process_image_response(response, url) except litellm.ImageFetchError: raise + except SSRFError as e: + raise _rejected_image_fetch(url, e) from e except Exception as e: verbose_logger.exception(e) raise litellm.ImageFetchError( f"Error: Unable to fetch image from URL after 3 attempts. url={url}", ) + + +_REMOTE_URL_PREFIXES: Final = ("http://", "https://") + + +@dataclass(frozen=True, slots=True) +class _RemoteImage: + part: Mapping[str, object] + image_url: Mapping[str, object] | None + url: str + + +@dataclass(frozen=True, slots=True) +class _RemoteFile: + part: Mapping[str, object] + file: Mapping[str, object] + url: str + + +def _as_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None # pyright: ignore[reportUnknownVariableType] # fields are parsed one by one + + +def _remote_url(candidate: object) -> str | None: + return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None + + +_ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"}) + + +@dataclass(frozen=True, slots=True) +class _RemoteSource: + part: Mapping[str, object] + source: Mapping[str, object] + url: str + + +@dataclass(frozen=True, slots=True) +class RemoteMedia: + url: str + fields: Mapping[str, object] + part_type: str + + +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def inline_every_remote_url(_media: RemoteMedia) -> bool: + return True + + +def inline_remote_image_urls(media: RemoteMedia) -> bool: + return media.part_type == "image_url" + + +def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None: + if fields.get("type") != "image_url": + return None + image_url: Final = fields.get("image_url") + image_url_fields: Final = _as_mapping(image_url) + url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url) + return _RemoteImage(fields, image_url_fields, url) if url is not None else None + + +def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None: + file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None + url: Final = _remote_url(file.get("file_id")) if file is not None else None + return _RemoteFile(fields, file, url) if file is not None and url is not None else None + + +def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None: + source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None + url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None + return _RemoteSource(fields, source, url) if source is not None and url is not None else None + + +def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None: + fields: Final = _as_mapping(part) + if fields is None: + return None + return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields) + + +def _remote_media(remote: _RemoteImage | _RemoteFile | _RemoteSource) -> RemoteMedia: + match remote: + case _RemoteImage(_, image_url, url): + return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS, "image_url") + case _RemoteFile(_, file, url): + return RemoteMedia(url, file, "file") + case _RemoteSource(part, source, url): + return RemoteMedia(url, source, str(part.get("type"))) + + +_PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"}) + + +def _inferred_format(file: Mapping[str, object], url: str) -> Mapping[str, str]: + return _PDF_FORMAT if "format" not in file and url.lower().endswith(".pdf") else MappingProxyType({}) + + +def _inlined_image_url(image_url: Mapping[str, object] | None, data_url: str) -> Mapping[str, object] | str: + return {**image_url, "url": data_url} if image_url is not None else data_url # mutable-ok: json-serialized part + + +def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mapping[str, object]: + kept: Final = {k: v for k, v in file.items() if k != "file_id"} # mutable-ok: json-serialized message part + return {**kept, **_inferred_format(file, url), "file_data": data_url} # mutable-ok: json-serialized part + + +def _base64_source(url: str, data_url: str) -> Mapping[str, str]: + fetched_media_type, data = data_url.removeprefix("data:").split(";base64,", 1) + media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched_media_type + return {"type": "base64", "media_type": media_type, "data": data} # mutable-ok: json-serialized message part + + +def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -> Mapping[str, object]: + match remote: + case _RemoteImage(part, image_url, _): + return {**part, "image_url": _inlined_image_url(image_url, data_url)} # mutable-ok: json-serialized part + case _RemoteFile(part, file, url): + return {**part, "file": _inlined_file(file, url, data_url)} # mutable-ok: json-serialized message part + case _RemoteSource(part, _, url): + return {**part, "source": _base64_source(url, data_url)} # mutable-ok: json-serialized message part + + +def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]: + content: Final = message.get("content") + return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one + + +def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]) -> object: + remote: Final = _parse_remote_part(part) + if remote is None or not should_inline(_remote_media(remote)): + return part + data_url: Final = data_urls.get(remote.url) + return _inline(remote, data_url) if data_url is not None else part + + +def _inline_message( + message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool] +) -> AllMessageValues: + parts: Final = _content_parts(message) + if not parts: + return message + inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks + _inline_part(part, data_urls, should_inline) for part in parts + ] + inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message + return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined + + +async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str: + async with in_flight: + return await async_convert_url_to_base64(url) + + +async def _fetch_data_urls(remote_urls: tuple[str, ...]) -> tuple[str, ...]: + in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) + fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight)) for url in remote_urls) + try: + return tuple(await asyncio.gather(*fetches)) + except BaseException: + for fetch in fetches: + fetch.cancel() + await asyncio.gather(*fetches, return_exceptions=True) + raise + + +async def async_inline_remote_media( + messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues] + should_inline: Callable[[RemoteMedia], bool] = inline_every_remote_url, +) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues] + remote_urls: Final = tuple( + dict.fromkeys( + remote.url + for message in messages + for part in _content_parts(message) + if (remote := _parse_remote_part(part)) is not None and should_inline(_remote_media(remote)) + ) + ) + if not remote_urls: + return messages + data_urls: Final = await _fetch_data_urls(remote_urls) + inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True))) + return [ # mutable-ok: transform_request takes a list + _inline_message(message, inlined, should_inline) for message in messages + ] diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 9402d465712..9d22a5ddef5 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm.constants import REDACTED_BY_LITELLM from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.classifier_logging import without_classifier_audit from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) @@ -162,12 +163,19 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): output_item["arguments"] = redacted_str -def _redact_standard_logging_object(model_call_details: dict): - """Redact messages and response inside standard_logging_object if present.""" - standard_logging_object: Final = model_call_details.get("standard_logging_object") - if standard_logging_object is None: - return +def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]: + """ + Return a copy of a ``StandardLoggingPayload`` with its messages and response redacted. + The success path redacts through ``perform_redaction`` before a callback ever sees the + payload, but the failure path does not, so a callback that batches both has to redact + the ones it is handed. + """ + return _redact_standard_logging_object(payload) + + +def _redact_standard_logging_object(payload: Mapping[str, object]) -> dict[str, object]: + standard_logging_object: Final = copy.deepcopy(without_classifier_audit(payload)) redacted_str: Final = REDACTED_BY_LITELLM if standard_logging_object.get("messages") is not None: @@ -190,6 +198,7 @@ def _redact_standard_logging_object(model_call_details: dict): else: # For other formats (empty dict, None, etc.), use simple text format standard_logging_object["response"] = {"text": redacted_str} + return standard_logging_object def _redact_tool_calls_dict(message: Mapping[str, object]) -> None: @@ -241,10 +250,16 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons copy via redact_streaming_responses_for_custom_logger instead. """ # Redact model_call_details + params: Final = model_call_details.get("litellm_params") + request: Final = params.get("proxy_server_request") if isinstance(params, dict) else None + if isinstance(params, dict) and isinstance(request, Mapping): + model_call_details["litellm_params"] = {**params, "proxy_server_request": without_classifier_audit(request)} model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}] model_call_details["prompt"] = "" model_call_details["input"] = "" - _redact_standard_logging_object(model_call_details) + standard_logging_object: Final = model_call_details.get("standard_logging_object") + if isinstance(standard_logging_object, Mapping): + model_call_details["standard_logging_object"] = _redact_standard_logging_object(standard_logging_object) redact_vertex_ai_metadata_from_litellm_params(model_call_details) # Redact streaming response diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 22dd4170963..b4c1beea33e 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -15,6 +15,7 @@ _DEFAULT_SENSITIVE_PATTERNS: Final = frozenset( "token", "auth", "authorization", + "cookie", "credential", # Plural form: Vertex uses ``vertex_credentials``; segment-exact # matching otherwise misses it because "credential" != "credentials". diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0e01577b20e..90698296142 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,6 +1,6 @@ import base64 import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast @@ -24,6 +24,7 @@ from litellm.types.utils import ( Choices, CompletionTokensDetails, CompletionTokensDetailsWrapper, + Delta, Function, FunctionCall, ModelResponse, @@ -209,7 +210,7 @@ def apply_grounding_request_counts( class ChunkProcessor: - def __init__(self, chunks: list, messages: list | None = None): + def __init__(self, chunks: list, messages: Sequence | None = None): self.chunks = self._sort_chunks(chunks) self.messages = messages self.first_chunk = chunks[0] @@ -326,6 +327,18 @@ class ChunkProcessor: return chunk_id return "" + @staticmethod + def _get_role_from_chunks(chunks: Sequence["_BaseChunk"]) -> str: + return ChunkProcessor._role_of_choice(next((c["choices"][0] for c in chunks if c.get("choices")), None)) + + @staticmethod + def _role_of_choice(choice: object) -> str: + match choice: + case StreamingChoices(delta=Delta(role=str() as role)) | {"delta": {"role": str() as role}} if role: + return role + case _: + return "assistant" + @staticmethod def _get_model_from_chunks(chunks: Sequence["_BaseChunk"], first_chunk_model: str) -> str: """ @@ -353,8 +366,7 @@ class ChunkProcessor: model: Final = ChunkProcessor._get_model_from_chunks(chunks, first_chunk_model) system_fingerprint: Final = chunk.get("system_fingerprint", None) - first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), chunk) - role: Final = first_chunk_with_choices["choices"][0]["delta"]["role"] + role: Final = ChunkProcessor._get_role_from_chunks(chunks) finish_reason = "stop" for chunk in chunks: if "choices" in chunk and len(chunk["choices"]) > 0: @@ -992,8 +1004,9 @@ class ChunkProcessor: chunks: Sequence["_UsageBearingChunk | ModelResponse"], model: str, completion_output: str, - messages: list | None = None, + messages: Sequence | None = None, reasoning_tokens: int | None = None, + count_prompt_tokens: Callable[[], int] | None = None, ) -> Usage: """ Calculate usage for the given chunks. @@ -1018,7 +1031,9 @@ class ChunkProcessor: cost: Final[float | None] = calculated_usage_per_chunk["cost"] try: - returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages) + returned_usage.prompt_tokens = prompt_tokens or ( + count_prompt_tokens() if count_prompt_tokens else token_counter(model=model, messages=messages) + ) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") returned_usage.prompt_tokens = 0 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6ae17bac6ff..6a5a8832cc6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -236,9 +236,11 @@ class CustomStreamWrapper: stream_options=None, make_call: Callable | None = None, _response_headers: dict | httpx.Headers | None = None, + count_prompt_tokens: Callable[[], int] | None = None, ): self.model = model self.make_call = make_call + self.count_prompt_tokens = count_prompt_tokens self.custom_llm_provider = custom_llm_provider self.logging_obj: LiteLLMLoggingObject = logging_obj self.completion_stream = completion_stream @@ -1473,17 +1475,14 @@ class CustomStreamWrapper: self.received_finish_reason = response_obj["finish_reason"] elif self.custom_llm_provider == "cached_response": cached_chunk: Final = cast(ModelResponseStream, chunk) - chunk_finish_reason: Final = cached_chunk.choices[0].finish_reason + cached_choice: Final = cached_chunk.choices[0] if cached_chunk.choices else None + chunk_finish_reason: Final = cached_choice.finish_reason if cached_choice is not None else None response_obj = { - "text": cached_chunk.choices[0].delta.content, + "text": cached_choice.delta.content if cached_choice is not None else None, "is_finished": chunk_finish_reason is not None, "finish_reason": chunk_finish_reason, "original_chunk": cached_chunk, - "tool_calls": ( - cached_chunk.choices[0].delta.tool_calls - if hasattr(cached_chunk.choices[0].delta, "tool_calls") - else None - ), + "tool_calls": (getattr(cached_choice.delta, "tool_calls", None) if cached_choice is not None else None), } completion_obj["content"] = response_obj["text"] @@ -1644,7 +1643,7 @@ class CustomStreamWrapper: except Exception: model_response.choices[0].delta = Delta() else: - if self.stream_options is not None and self.stream_options["include_usage"] is True: + if self.send_stream_usage is True: model_response.choices = [] return model_response self._record_usage_only_chunk(model_response=model_response) @@ -1999,6 +1998,7 @@ class CustomStreamWrapper: chunks=self.chunks, messages=self.messages, logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, ) except Exception as e: # stream_chunk_builder can re-raise (as APIError) on large agentic @@ -2251,6 +2251,7 @@ class CustomStreamWrapper: chunks=self.chunks, messages=self.messages, logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, ) except Exception as e: # see sync __next__: a raise from stream_chunk_builder inside this @@ -2374,6 +2375,7 @@ class CustomStreamWrapper: chunks=self.chunks, messages=self.messages if isinstance(self.messages, list) else None, logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, ) if partial_response is None: return diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 3732ffd734c..1de2533d514 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -3,11 +3,15 @@ import base64 import io import struct -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Final, Literal, cast +import anyio +import anyio.lowlevel import httpx import tiktoken +from tokenizers import Tokenizer +from typing_extensions import ParamSpec, TypeVar import litellm from litellm import verbose_logger @@ -21,7 +25,10 @@ from litellm.constants import ( MAX_TILE_HEIGHT, MAX_TILE_WIDTH, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, + TOKEN_COUNTER_MAX_CONCURRENT_COUNTS, + TOKEN_COUNTER_MAX_EXACT_CHARS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.default_encoding import encoding as default_encoding from litellm.litellm_core_utils.url_utils import safe_get from litellm.llms.custom_httpx.http_handler import _get_httpx_client @@ -172,6 +179,13 @@ def calculate_tiles_needed( return total_tiles +def high_detail_image_token_upper_bound(base_tokens: int = 85) -> int: + largest_tile_count: Final = calculate_tiles_needed( + MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES + ) + return base_tokens + (base_tokens * 2) * largest_tile_count + + def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]: return struct.unpack(fmt, buffer) @@ -317,6 +331,32 @@ TokenCounterFunction = Callable[[str], int] Type for a function that counts tokens in a string. """ +EXTRAPOLATION_SAMPLES: Final = 16 +T_ParamSpec: Final = ParamSpec("T_ParamSpec") +T_Retval = TypeVar("T_Retval") +_COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter") + + +def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter: + existing: Final = _COUNT_OFFLOAD_LIMITER.get(None) + if existing is not None: + return existing + created: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + _COUNT_OFFLOAD_LIMITER.set(created) + return created + + +def offload_token_count( + function: Callable[T_ParamSpec, T_Retval], +) -> Callable[T_ParamSpec, Awaitable[T_Retval]]: + async def offloaded( + *args: T_ParamSpec.args, + **kwargs: T_ParamSpec.kwargs, # kwargs-ok: ParamSpec keeps the wrapped function's own keyword contract + ) -> T_Retval: + return await asyncify(function, limiter=_count_offload_limiter_for_this_loop())(*args, **kwargs) + + return offloaded + def _get_tiktoken_count_function( encode_length: Callable[[str], int], @@ -538,9 +578,40 @@ def _count_extra( return num_tokens +def _get_extrapolating_count_function( + count_exactly: TokenCounterFunction, + max_exact_chars: int = TOKEN_COUNTER_MAX_EXACT_CHARS, +) -> TokenCounterFunction: + def count_tokens(text: str) -> int: + if len(text) <= max_exact_chars: + return count_exactly(text) + samples: Final = _evenly_spaced_samples(text, max_exact_chars) + sampled_chars: Final = sum(len(sample) for sample in samples) + return round(sum(count_exactly(sample) for sample in samples) * len(text) / sampled_chars) + + return count_tokens + + +def _evenly_spaced_samples(text: str, total_chars: int) -> tuple[str, ...]: + sample_count: Final = min(EXTRAPOLATION_SAMPLES, total_chars) + sample_chars: Final = total_chars // sample_count + last_start: Final = len(text) - sample_chars + return tuple( + text[start : start + sample_chars] + for start in (last_start * index // max(sample_count - 1, 1) for index in range(sample_count)) + ) + + def _get_count_function( model: str | None, custom_tokenizer: dict | SelectTokenizerResponse | None = None, +) -> TokenCounterFunction: + return _get_extrapolating_count_function(_get_exact_count_function(model, custom_tokenizer)) + + +def _get_exact_count_function( + model: str | None, + custom_tokenizer: dict | SelectTokenizerResponse | None = None, ) -> TokenCounterFunction: """ Get the function to count tokens based on the model and custom tokenizer.""" @@ -549,10 +620,10 @@ def _get_count_function( if model is not None or custom_tokenizer is not None: tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model) if tokenizer_json["type"] == "huggingface_tokenizer": + tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"] def count_tokens(text: str) -> int: - enc: Final = tokenizer_json["tokenizer"].encode(text) - return len(enc.ids) + return len(tokenizer.encode_batch_fast([text])[0]) return count_tokens elif tokenizer_json["type"] == "openai_tokenizer": diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 1e43117933d..fa070a648f5 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -19,6 +19,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): check but still resolve DNS and still rewrite HTTP to the resolved IP. """ +import asyncio import socket from ipaddress import ip_address, ip_network from typing import Any, Final, Protocol @@ -471,7 +472,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response kwargs.pop("follow_redirects", None) headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} for _ in range(_MAX_REDIRECTS): - validated_url, original_host = validate_url(url) + validated_url, original_host = await asyncio.to_thread(validate_url, url) response = await fetcher.get( validated_url, headers={**headers_view["headers"], "Host": original_host}, diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index c23797f72af..9d50345d70d 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,9 +13,11 @@ Pattern Overview: """ import json -from collections.abc import Mapping, Sequence +from collections.abc import Mapping, MutableSequence, Sequence from copy import deepcopy from dataclasses import dataclass +from itertools import chain, repeat +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable from typing_extensions import ReadOnly, TypedDict, assert_never @@ -29,6 +31,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, + StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( anthropic_tool_name, @@ -39,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, + stream_item_field, stream_item_fingerprint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( @@ -151,6 +155,46 @@ class ExtractedInput: EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) +@dataclass(frozen=True, slots=True) +class _ToolCallShape: + name: str | None + arguments: str + + +@dataclass(frozen=True, slots=True) +class _SSEFieldRewrite: + """One field of one nested section of a buffered SSE event, rewritten.""" + + section: str + field: str + value: object + + +class _SSEEventRewriter(Protocol): + def __call__(self, event: Mapping[str, object]) -> _SSEFieldRewrite | None: ... + + +def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewriter) -> Mapping[str, object]: + rewrite: Final = rewrite_event(event) + section: Final = None if rewrite is None else event.get(rewrite.section) + if rewrite is None or not isinstance(section, Mapping): + return event + return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} # mutable-ok: json.dumps needs a dict + + +def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]: + """The guardrail-visible shape of each tool call, whether the guardrail handed + back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts.""" + functions: Final = tuple(stream_item_field(tool_call, "function") for tool_call in tool_calls) + return tuple( + _ToolCallShape( + name=name if isinstance(name := stream_item_field(function, "name"), str) else None, + arguments=arguments if isinstance(arguments := stream_item_field(function, "arguments"), str) else "", + ) + for function in functions + ) + + class _AnthropicSSEDelta(TypedDict, total=False): type: ReadOnly[str] text: ReadOnly[str] @@ -168,10 +212,18 @@ class AnthropicMessagesHandler(BaseTranslation): them through guardrail rewrites; downstream provider handling is out of scope. """ + delivers_ended_stream_rewrites = True + assembles_streamed_response = True + def __init__(self): super().__init__() self.adapter = LiteLLMAnthropicMessagesAdapter() + def post_call_hook_response(self, response: object) -> object: + if not isinstance(response, ModelResponse): + return response + return self.adapter.translate_openai_response_to_anthropic(response) + @staticmethod def _build_streaming_usage_response( responses_so_far: Sequence[object], @@ -1014,11 +1066,17 @@ class AnthropicMessagesHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, + stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> Sequence[object]: """ Process output streaming response by applying guardrails to text content. Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far. + With ``deliver_ended_stream_rewrites``, an ended stream whose guardrail rewrote the text gets the rewrite + written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked); + a rewrite on a stream that never reported a ``stop_reason`` has no write-back and is reported as + undeliverable, so the pipeline executor discards it and releases the original chunks. """ from litellm.integrations.custom_guardrail import ModifyResponseException @@ -1040,6 +1098,7 @@ class AnthropicMessagesHandler(BaseTranslation): first_choice.message.tool_calls, ) string_so_far = first_choice.message.content + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ()) guardrail_inputs: Final = GenericGuardrailAPIInputs() if string_so_far: guardrail_inputs["texts"] = [string_so_far] @@ -1065,6 +1124,28 @@ class AnthropicMessagesHandler(BaseTranslation): responses_so_far, request_data ) raise + guardrailed_texts: Final = _guardrailed_inputs.get("texts") + if ( + deliver_ended_stream_rewrites + and isinstance(string_so_far, str) + and string_so_far + and guardrailed_texts + and guardrailed_texts[0] != string_so_far + ): + self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0]) + if deliver_ended_stream_rewrites: + returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls") + self._write_ended_stream_tool_call_rewrites( + responses_so_far, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=_tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) + and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else tool_calls_list or () + ), + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) else: verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") return responses_so_far @@ -1087,6 +1168,11 @@ class AnthropicMessagesHandler(BaseTranslation): if e.original_response is None: e.original_response = self._build_streaming_usage_response(responses_so_far, request_data) raise + unended_texts: Final = _guardrailed_inputs.get("texts") + if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") return responses_so_far def _prepare_request_data( @@ -1180,6 +1266,139 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs + @staticmethod + def _write_ended_stream_text_rewrite( + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place + rewritten_text: str, + ) -> None: + """Deliver an ended-stream guardrail text rewrite by rewriting the + buffered chunks in place: the first ``text_delta`` carries the full + rewritten text and every later one is blanked, leaving the surrounding + message and content-block framing untouched.""" + replacements: Final = chain((rewritten_text,), repeat("")) + + def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + delta: Final = event.get("delta") + if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping): + return None + if delta.get("type") != "text_delta": + return None + return _SSEFieldRewrite("delta", "text", next(replacements)) + + AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta) + + @classmethod + def _write_ended_stream_tool_call_rewrites( + cls, + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place + *, + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str, + ) -> None: + """Deliver ended-stream guardrail tool-call rewrites by rewriting the + buffered chunks in place: the rebuilt response lists tool calls in the + order of the stream's ``tool_use`` blocks, so the nth rewritten call lands + on the nth block, its first ``input_json_delta`` carrying the full rewritten + arguments, every later one blanked, and ``content_block_start`` carrying the + rewritten name. Blocks that do not line up with the rebuilt tool calls make + the rewrite undeliverable, so the pipeline executor discards it and releases + the original chunks.""" + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + block_indices: Final = tuple( + index + for item in responses_so_far + for event in cls._iter_sse_events(item) + if event.get("type") == "content_block_start" + and isinstance(block := event.get("content_block"), Mapping) + and block.get("type") == "tool_use" + and isinstance(index := event.get("index"), int) + ) + if len(block_indices) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + rewrites_by_block: Final = MappingProxyType( + { + index: after + for index, before, after in zip(block_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if after != before + } + ) + argument_replacements: Final = MappingProxyType( + {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()} + ) + + def rewrite_tool_use(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + index: Final = event.get("index") + if not isinstance(index, int) or index not in rewrites_by_block: + return None + match event.get("type"): + case "content_block_start": + name: Final = rewrites_by_block[index].name + if name is None: + return None + return _SSEFieldRewrite("content_block", "name", name) + case "content_block_delta": + delta: Final = event.get("delta") + if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta": + return None + return _SSEFieldRewrite("delta", "partial_json", next(argument_replacements[index])) + case _: + return None + + cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use) + + @staticmethod + def _rewrite_ended_stream_events( + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place + rewrite_event: _SSEEventRewriter, + ) -> None: + """Replace every buffered event ``rewrite_event`` returns a rewrite for, in + both chunk formats this stream carries (parsed event dicts and raw SSE + bytes), leaving every other event and the framing untouched.""" + rewritten_items: Final = tuple( + AnthropicMessagesHandler._rewrite_buffered_item(item, rewrite_event) for item in responses_so_far + ) + responses_so_far[:] = rewritten_items # rebind-ok: delivers the rewrites into the caller's buffer + + @staticmethod + def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object: + if isinstance(item, dict): + return _rewritten_event(_as_str_mapping(item), rewrite_event) + if isinstance(item, (bytes, bytearray)): + return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event) + return item + + @staticmethod + def _rewrite_sse_events(sse_bytes: bytes, rewrite_event: _SSEEventRewriter) -> bytes: + """Rewrite the data lines of one SSE chunk that ``rewrite_event`` rewrites, + leaving all other events and framing byte-identical.""" + try: + decoded: Final = sse_bytes.decode("utf-8") + except UnicodeDecodeError: + return sse_bytes + return "\n\n".join( + "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, rewrite_event) for line in block.split("\n")) + for block in decoded.split("\n\n") + ).encode("utf-8") + + @staticmethod + def _rewrite_sse_line(line: str, rewrite_event: _SSEEventRewriter) -> str: + if not line.startswith("data:"): + return line + try: + data: Final[str | int | float | bool | None | Sequence[object] | Mapping[str, object]] = json.loads( + line[len("data:") :].strip() + ) + except json.JSONDecodeError: + return line + if not isinstance(data, dict): + return line + rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event) + return line if rewritten is data else "data: " + json.dumps(rewritten) + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: stream_ended: Final = self._check_streaming_has_ended(responses_so_far) return StreamingScanKey( diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index c82be07a5c5..82189461403 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -368,27 +368,92 @@ class AnthropicChatCompletion(BaseLLM): if config is None: raise ValueError(f"Provider config not found for model: {model} and provider: {custom_llm_provider}") - def build_request() -> tuple[dict, dict]: # mutable-ok: rewritten in place downstream - """Translate the request the Python way, returning `(headers, data)`. + transform_params: Final = {**optional_params, "is_vertex_request": is_vertex_request} + + def finish_request(request_data: dict) -> tuple[dict, dict]: # mutable-ok: rewritten in place downstream + """Filter beta headers and emit pre_call, returning `(headers, data)`. The pair stays mutable because the streaming path rewrites it in - place (`data["stream"] = True`) before sending. - - Shared by the normal path and by the Rust path's fallback, which - builds it only when the Rust call did not serve the request. + place (`data["stream"] = True`) before sending. A Rust attempt that + declined already emitted pre_call for this request, so skip it there. """ - request_data: Final = config.transform_request( - model=model, - messages=messages, - optional_params={**optional_params, "is_vertex_request": is_vertex_request}, - litellm_params=litellm_params, - headers=headers, - ) - return update_request_with_filtered_beta( + request_headers, data = update_request_with_filtered_beta( headers=headers, request_data=request_data, provider=custom_llm_provider, ) + if not serves_via_rust: + logging_obj.pre_call( + input=messages, + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": request_headers, + }, + ) + print_verbose(f"_is_function_call: {_is_function_call}") + return request_headers, data + + async def acompletion_dispatch() -> "ModelResponse | CustomStreamWrapper": + """Translate then send, so the provider config can inline remote media off the event loop.""" + request_headers, data = finish_request( + await config.async_transform_request( + model=model, + messages=messages, + optional_params=transform_params, + litellm_params=litellm_params, + headers=headers, + ) + ) + if ( + stream is True + ): # if function call - fake the streaming (need complete blocks for output parsing in openai format) + print_verbose("makes async anthropic streaming POST request") + data["stream"] = stream + return await self.acompletion_stream_function( + model=model, + messages=messages, + data=data, + api_base=api_base, + custom_prompt_dict=custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + encoding=encoding, + api_key=api_key, + logging_obj=logging_obj, + optional_params=optional_params, + stream=stream, + _is_function_call=_is_function_call, + json_mode=json_mode, + litellm_params=litellm_params, + logger_fn=logger_fn, + headers=request_headers, + timeout=timeout, + client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), + ) + return await self.acompletion_function( + model=model, + messages=messages, + data=data, + api_base=api_base, + custom_prompt_dict=custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + encoding=encoding, + api_key=api_key, + provider_config=config, + logging_obj=logging_obj, + optional_params=optional_params, + stream=stream, + _is_function_call=_is_function_call, + litellm_params=litellm_params, + logger_fn=logger_fn, + headers=request_headers, + client=client, + json_mode=json_mode, + timeout=timeout, + ) # The Rust core owns the whole call for the subset it accepts, so ask # before transforming: whichever path runs emits pre_call exactly once. @@ -424,35 +489,6 @@ class AnthropicChatCompletion(BaseLLM): additional_args=rust_logging_args, ) if acompletion is True: - - async def python_fallback() -> "ModelResponse | CustomStreamWrapper": - # pre_call already fired for this request above. The Rust - # path only declines before the provider is called, so this - # is the same attempt continuing, not a second one. - fallback_headers, fallback_data = build_request() - return await self.acompletion_function( - model=model, - messages=messages, - data=fallback_data, - api_base=api_base, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - encoding=encoding, - api_key=api_key, - provider_config=config, - logging_obj=logging_obj, - optional_params=optional_params, - stream=stream, - _is_function_call=_is_function_call, - litellm_params=litellm_params, - logger_fn=logger_fn, - headers=fallback_headers, - client=client, - json_mode=json_mode, - timeout=timeout, - ) - return rust_chat_completions_bridge.achat_completions_or_fallback( model=model, messages=messages, @@ -464,7 +500,7 @@ class AnthropicChatCompletion(BaseLLM): extra_headers=headers, timeout=timeout, on_response=log_rust_post_call, - python_fallback=python_fallback, + python_fallback=acompletion_dispatch, ) rust_response: Final = rust_chat_completions_bridge.chat_completions( model=model, @@ -481,74 +517,18 @@ class AnthropicChatCompletion(BaseLLM): if rust_response is not None: return rust_response - headers, data = build_request() - - ## LOGGING - # Reaching here with `serves_via_rust` set means the Rust attempt - # declined at call time, before the provider was called, and already - # logged this request. That is the same attempt continuing. - if not serves_via_rust: - logging_obj.pre_call( - input=messages, - api_key=api_key, - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": headers, - }, - ) - print_verbose(f"_is_function_call: {_is_function_call}") if acompletion is True: - if ( - stream is True - ): # if function call - fake the streaming (need complete blocks for output parsing in openai format) - print_verbose("makes async anthropic streaming POST request") - data["stream"] = stream - return self.acompletion_stream_function( - model=model, - messages=messages, - data=data, - api_base=api_base, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - encoding=encoding, - api_key=api_key, - logging_obj=logging_obj, - optional_params=optional_params, - stream=stream, - _is_function_call=_is_function_call, - json_mode=json_mode, - litellm_params=litellm_params, - logger_fn=logger_fn, - headers=headers, - timeout=timeout, - client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), - ) - else: - return self.acompletion_function( - model=model, - messages=messages, - data=data, - api_base=api_base, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - encoding=encoding, - api_key=api_key, - provider_config=config, - logging_obj=logging_obj, - optional_params=optional_params, - stream=stream, - _is_function_call=_is_function_call, - litellm_params=litellm_params, - logger_fn=logger_fn, - headers=headers, - client=client, - json_mode=json_mode, - timeout=timeout, - ) + return acompletion_dispatch() else: + headers, data = finish_request( + config.transform_request( + model=model, + messages=messages, + optional_params=transform_params, + litellm_params=litellm_params, + headers=headers, + ) + ) ## COMPLETION CALL if ( stream is True @@ -744,10 +724,11 @@ class ModelResponseIterator: content_block: Final = ContentBlockDelta(**chunk) thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] = [] - self.content_blocks.append(content_block) if "text" in content_block["delta"]: text = content_block["delta"]["text"] - elif "partial_json" in content_block["delta"]: + return text, tool_use, thinking_blocks, provider_specific_fields, reasoning_content + self.content_blocks.append(content_block) + if "partial_json" in content_block["delta"]: # Only emit tool calls if we're in a tool_use or server_tool_use block # web_search_tool_result blocks also have input_json_delta but should not be treated as tool calls # See: https://github.com/BerriAI/litellm/issues/17254 diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 5f7ac73c919..5463f1862ad 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -26,6 +26,11 @@ from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.prompt_templates.common_utils import ( sanitize_input_schema_for_anthropic, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import ( + RemoteMedia, + async_inline_remote_media, + inline_remote_image_urls, +) from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.anthropic import ( @@ -1840,6 +1845,25 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): break return headers + def inlines_remote_media(self, media: RemoteMedia) -> bool: + return inline_remote_image_urls(media) and media.url.startswith("http://") + + async def async_transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, object], # mutable-ok: BaseConfig signature + litellm_params: dict[str, object], # mutable-ok: BaseConfig signature + headers: dict[str, object], # mutable-ok: BaseConfig signature + ) -> dict[str, object]: # mutable-ok: BaseConfig signature + return self.transform_request( + model=model, + messages=await async_inline_remote_media(messages, should_inline=self.inlines_remote_media), + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + def transform_request( self, model: str, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 6079b709bcc..87c4ec8938e 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -21,6 +21,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, + is_encrypted_reasoning_block, ) from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, @@ -67,6 +68,101 @@ _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$") _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$") _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$") _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)") +_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" +_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) +_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) + + +_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") + + +def is_claude_code_user_agent(user_agent: str) -> bool: + """Claude Code sends its API calls through the Anthropic SDK as `claude-cli/` and its own + fetches, such as gateway model discovery, as `claude-code/`""" + return user_agent.startswith(_CLAUDE_CODE_USER_AGENT_PREFIXES) + + +def _validated_claude_code_mapping(value: object) -> dict[object, object] | None: + try: + return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _validated_claude_code_list(value: object) -> list[object] | None: + try: + return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None: + stripped: Final = text.strip() + if "\n" in stripped or "\r" in stripped or not stripped.startswith(_CLAUDE_CODE_BILLING_HEADER_PREFIX): + return None + fields: Final = tuple( + field + for raw_field in stripped.removeprefix(_CLAUDE_CODE_BILLING_HEADER_PREFIX).split(";") + if (field := raw_field.strip()) + ) + if not fields or any("=" not in field for field in fields): + return None + parsed_fields: Final = tuple( + (parts[0].strip(), parts[1].strip()) for field in fields for parts in (field.split("=", 1),) + ) + if any(not key or not value for key, value in parsed_fields): + return None + return parsed_fields + + +def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None: + if isinstance(system, str): + return (system,) + blocks: Final = _validated_claude_code_list(system) + if blocks is None: + return None + block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks) + if any(block is None for block in block_mappings): + return None + text_values: Final = tuple( + block.get("text") for block in block_mappings if block is not None and block.get("type") == "text" + ) + if len(text_values) != len(blocks) or any(not isinstance(text, str) for text in text_values): + return None + meaningful_text: Final = tuple(text for text in text_values if isinstance(text, str) and text.strip()) + return meaningful_text or None + + +def _is_claude_code_subagent_billing_system(system: object) -> bool: + billing_texts: Final = _claude_code_billing_texts(system) + if billing_texts is None: + return False + billing_fields: Final = tuple( + fields for text in billing_texts if (fields := _claude_code_billing_fields(text)) is not None + ) + if len(billing_fields) != len(billing_texts): + return False + subagent_values: Final = tuple( + value for fields in billing_fields for key, value in fields if key == "cc_is_subagent" + ) + return subagent_values == ("true",) + + +def is_claude_code_one_shot_subagent_request( + messages: list[AllMessageValues], + system: object, + tools: object, + user_agent: str | None, +) -> bool: + only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None + return ( + user_agent is not None + and is_claude_code_user_agent(user_agent) + and not tools + and only_message is not None + and only_message.get("role") == "user" + and _is_claude_code_subagent_billing_system(system) + ) def _strip_bedrock_id_suffixes(model: str) -> str: @@ -1111,6 +1207,32 @@ def strip_thinking_blocks_from_anthropic_messages(messages: list[Any]) -> list[A return out +def _without_encrypted_reasoning_blocks(message: dict) -> dict | None: # mutable-ok: Anthropic message payload shape + if not isinstance(message, Mapping): + return message + content: Final = message.get("content") + if not isinstance(content, list): + return message + kept: Final = [b for b in content if not is_encrypted_reasoning_block(b)] # mutable-ok: API message payload + if len(kept) == len(content): + return message + if not kept: + return None + return {**message, "content": kept} # mutable-ok: API message payload + + +def strip_encrypted_reasoning_blocks_from_anthropic_messages( + messages: Sequence[dict], # mutable-ok: Anthropic message payload shape +) -> list[dict]: # mutable-ok: AnthropicMessagesRequest.messages is typed list[dict] + """ + Drop thinking / redacted_thinking blocks that carry another provider's encrypted + reasoning (a turn the Responses API bridge served) before the request reaches + Anthropic, which cannot verify them. Anthropic's own signed blocks are kept. + """ + stripped: Final = (_without_encrypted_reasoning_blocks(m) for m in messages) + return [m for m in stripped if m is not None] # mutable-ok: API message payload + + def strip_thinking_blocks_from_anthropic_messages_request_dict( data: dict[str, Any], ) -> None: @@ -1411,6 +1533,25 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format +def _without_provider_specific_fields(block: object) -> object: + if not isinstance(block, dict) or "provider_specific_fields" not in block: + return block + return {k: v for k, v in block.items() if k != "provider_specific_fields"} # mutable-ok: JSON wire format + + +def _strip_provider_specific_fields_in_message(message: object) -> object: + if not isinstance(message, dict) or not isinstance(message.get("content"), list): + return message + content: Final = [_without_provider_specific_fields(b) for b in message["content"]] # mutable-ok: JSON wire format + return {**message, "content": content} # mutable-ok: JSON wire format + + +def strip_provider_specific_fields_from_anthropic_messages( + messages: Sequence[object], +) -> Sequence[object]: + return [_strip_provider_specific_fields_in_message(m) for m in messages] # mutable-ok: JSON wire format + + def _normalized_cache_control(cache_control: object) -> dict[str, str] | None: # mutable-ok: JSON wire format if not isinstance(cache_control, Mapping): return None @@ -1520,11 +1661,16 @@ def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: def _anthropic_model_entry( - model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str] + model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str], listed_ids: Mapping[str, str] ) -> Mapping[str, object]: + listed_id: Final = listed_ids.get(model["id"]) + source: Final[Mapping[str, object]] = ( + MappingProxyType({"source_model": model["id"]}) if listed_id is not None else MappingProxyType({}) + ) return { # mutable-ok: JSON response body, serialized by the route and never mutated "type": "model", - "id": model["id"], + "id": listed_id or model["id"], + **source, "display_name": display_names.get(model["id"], model["id"]), "created_at": created_at, "max_input_tokens": model.get("max_input_tokens"), @@ -1535,6 +1681,7 @@ def _anthropic_model_entry( def create_anthropic_model_list_response( models: Sequence[ModelInfoResponse], display_names: Mapping[str, str] = MappingProxyType({}), + listed_ids: Mapping[str, str] = MappingProxyType({}), ) -> Mapping[str, object]: """Build the Anthropic-native /v1/models envelope. @@ -1544,17 +1691,19 @@ def create_anthropic_model_list_response( over from the OpenAI-shaped listing, named as the Messages API names them, and are always present because the vendor shape declares them nullable, not optional. display_names maps a listed model id to a configured human-readable name; ids - without an entry fall back to the id itself, matching the vendor behavior + without an entry fall back to the id itself, matching the vendor behavior. + listed_ids maps a model id to the id the caller should see it under (the Claude + Code view); ids without an entry are listed as they are """ created_at: Final = ( datetime.fromtimestamp(DEFAULT_MODEL_CREATED_AT_TIME, tz=timezone.utc).isoformat().replace("+00:00", "Z") ) data: Final = [ # mutable-ok: JSON response body, serialized by the route and never mutated - _anthropic_model_entry(model, created_at, display_names) for model in models + _anthropic_model_entry(model, created_at, display_names, listed_ids) for model in models ] return { # mutable-ok: JSON response body, serialized by the route and never mutated "data": data, "has_more": False, - "first_id": models[0]["id"] if models else None, - "last_id": models[-1]["id"] if models else None, + "first_id": data[0]["id"] if data else None, + "last_id": data[-1]["id"] if data else None, } diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 32b8cbe6343..9158ff4569f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -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, diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 594fac512e6..9b3a57cc422 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -113,10 +113,15 @@ from litellm.litellm_core_utils.reasoning_effort_utils import ( from litellm.llms.anthropic.common_utils import ( is_empty_unsigned_thinking_block, normalize_anthropic_tool_use_id, + strip_encrypted_reasoning_blocks_from_anthropic_messages, ) 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, @@ -413,7 +418,8 @@ class LiteLLMAnthropicMessagesAdapter: model: str | None = None, ) -> list: new_messages: Final[list[AllMessageValues]] = [] - for m in messages: + replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) + for m in replayable_messages: user_message: ChatCompletionUserMessage | None = None tool_message_list: list[ChatCompletionToolMessage] = [] new_user_content_list: list[ChatCompletionTextObject | ChatCompletionImageObject] = [] @@ -1314,6 +1320,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: @@ -1346,7 +1354,7 @@ class LiteLLMAnthropicMessagesAdapter: # Add provider_specific_fields if signature is present if provider_specific_fields: tool_use_block.provider_specific_fields = provider_specific_fields - new_content.append(tool_use_block.model_dump()) + new_content.append(tool_use_block.model_dump(exclude_none=True)) return new_content @@ -1472,13 +1480,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( - openai_finish_reason=response.choices[0].finish_reason + openai_finish_reason: Final = response.choices[0].finish_reason if response.choices else "stop" + translated_finish_reason: Final = self._translate_openai_finish_reason_to_anthropic( + openai_finish_reason=openai_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") @@ -1501,6 +1519,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 +1560,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 +1634,10 @@ class LiteLLMAnthropicMessagesAdapter: elif reasoning_content: return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content) else: - return "text_delta", ContentTextBlockDelta(type="text_delta", text=text) + refusal_text: Final = "".join( + refusal for choice in choices if (refusal := openai_chat_refusal_text(choice.delta)) is not None + ) + return "text_delta", ContentTextBlockDelta(type="text_delta", text=text + refusal_text) def translate_streaming_openai_response_to_anthropic( self, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index b82903d6f87..9d1e921cce4 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -18,6 +18,7 @@ from litellm.llms.anthropic.common_utils import ( flatten_unencrypted_web_search_results_in_anthropic_messages, sanitize_tool_use_ids_in_anthropic_messages, strip_empty_content_blocks_from_anthropic_messages, + strip_provider_specific_fields_from_anthropic_messages, ) from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, @@ -650,7 +651,7 @@ def anthropic_messages_handler( return base_llm_http_handler.anthropic_messages_handler( model=model, - messages=messages, + messages=strip_provider_specific_fields_from_anthropic_messages(messages), anthropic_messages_provider_config=anthropic_messages_provider_config, anthropic_messages_optional_request_params=dict(anthropic_messages_optional_request_params), _is_async=is_async, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py index c1f10c245f8..d9cc65e730f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py @@ -82,7 +82,7 @@ async def anthropic_messages_with_mcp( LiteLLM_Proxy_MCP_Handler, ) - mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + mcp_references, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_references: return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 8267da157ad..9f9346fad4d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -25,6 +25,7 @@ from ...common_utils import ( AnthropicModelInfo, optionally_handle_anthropic_oauth, strip_advisor_blocks_from_messages, + strip_encrypted_reasoning_blocks_from_anthropic_messages, ) DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -613,7 +614,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): messages = strip_advisor_blocks_from_messages(messages) anthropic_messages_request: Final[AnthropicMessagesRequest] = AnthropicMessagesRequest( - messages=messages, + messages=strip_encrypted_reasoning_blocks_from_anthropic_messages(messages), max_tokens=max_tokens, model=model, **anthropic_messages_optional_request_params, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 242300c7b6d..7545dff1408 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -1,8 +1,11 @@ -from collections.abc import Mapping +from collections.abc import Iterable, Mapping, Sequence from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints -from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams +from litellm.types.llms.anthropic import ( + AnthropicMessagesRequestOptionalParams, + AnthropicStopDetails, +) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -25,6 +28,69 @@ def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | return stop_details if isinstance(stop_details, dict) else None +def refusal_stop_details(explanation: str | None) -> AnthropicStopDetails: + """The ``stop_details`` object accompanying a translated ``stop_reason: "refusal"``.""" + return AnthropicStopDetails(type="refusal", category=None, explanation=explanation) + + +def _mapping_field(container: object, key: str) -> object | None: + """One key of a raw provider payload, or None when the payload is not a mapping.""" + if not isinstance(container, Mapping): + return None + return cast(Mapping[str, object], container).get(key) # cast-ok: raw payload, callers re-check every value + + +def _mapping_str_field(container: object, key: str) -> str | None: + value: Final = _mapping_field(container, key) + return value if isinstance(value, str) and value else None + + +def openai_chat_refusal_text(message_or_delta: object) -> str | None: + """ + Refusal text carried by an OpenAI Chat Completions message or streaming delta, + read from ``refusal`` or from the ``provider_specific_fields`` LiteLLM parks it + in, or None when the turn is not a refusal. + """ + refusal: Final = getattr(message_or_delta, "refusal", None) + if isinstance(refusal, str) and refusal: + return refusal + return _mapping_str_field(getattr(message_or_delta, "provider_specific_fields", None), "refusal") + + +def _responses_message_refusal_text(item: object) -> str | None: + from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal + + if isinstance(item, ResponseOutputMessage): + return next( + (part.refusal for part in item.content if isinstance(part, ResponseOutputRefusal) and part.refusal), + None, + ) + raw_parts: Final = _mapping_field(item, "content") + if _mapping_str_field(item, "type") != "message" or not isinstance(raw_parts, Sequence): + return None + return next( + ( + refusal + for part in cast(Sequence[object], raw_parts) # cast-ok: members re-validated below + if _mapping_str_field(part, "type") == "refusal" + and isinstance(refusal := _mapping_str_field(part, "refusal"), str) + ), + None, + ) + + +def responses_output_refusal_text(output: Iterable[object]) -> str | None: + """ + Refusal text carried by an OpenAI Responses ``output`` list, in typed + (``ResponseOutputRefusal``) or raw-dictionary shape, or None when none of the + output messages refused. + """ + return next( + (text for item in output if (text := _responses_message_refusal_text(item)) is not None), + None, + ) + + def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "ContentPolicyViolationError": """The exception a safeguard-refused Anthropic response converts into so the content-policy fallback chain can re-dispatch it.""" diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index 0445c23ed8c..7731c883d9f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -19,6 +19,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.utils import ProviderConfigManager from ..utils import litellm_logging_obj_from_kwargs, local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper @@ -34,6 +35,15 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, return extra_kwargs or {} +def _provider_returns_encrypted_reasoning(model: str, custom_llm_provider: object) -> bool: + provider: Final = ( + custom_llm_provider if isinstance(custom_llm_provider, str) else litellm.get_llm_provider(model=model)[1] + ) + provider_model: Final = local_model_name(model, provider) + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config(provider, provider_model) + return responses_config is not None and "include" in responses_config.get_supported_openai_params(provider_model) + + def _build_responses_kwargs( *, max_tokens: int, @@ -85,8 +95,13 @@ def _build_responses_kwargs( request_data["output_format"] = output_format anthropic_request: Final = AnthropicMessagesRequest(**request_data) - responses_kwargs: Final = _ADAPTER.translate_request(anthropic_request) forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) + responses_kwargs: Final = _ADAPTER.translate_request( + anthropic_request, + include_encrypted_reasoning=_provider_returns_encrypted_reasoning( + model, forwarded_kwargs.get("custom_llm_provider") + ), + ) # Normalize reasoning effort based on model capabilities # (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported) @@ -111,7 +126,7 @@ def _build_responses_kwargs( responses_kwargs["stream"] = True # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.) - excluded: Final = {"anthropic_messages"} + excluded: Final = frozenset(("anthropic_messages",)) for key, value in forwarded_kwargs.items(): if key == "litellm_logging_obj" and value is not None: from litellm.litellm_core_utils.litellm_logging import ( @@ -132,6 +147,14 @@ def _build_responses_kwargs( if explicit_prompt_cache_key is not None: responses_kwargs["prompt_cache_key"] = explicit_prompt_cache_key + deployment_include: Final = forwarded_kwargs.get("include") + bridge_include: Final = responses_kwargs.get("include") + if isinstance(deployment_include, list) and isinstance(bridge_include, list): + responses_kwargs["include"] = [ + *bridge_include, + *(item for item in deployment_include if item not in bridge_include), + ] + return responses_kwargs diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index b0aba386753..f753e87fee3 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -1,16 +1,27 @@ # 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.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, +) +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 +from .transformation import ( + REASONING_SUMMARY_PART_SEPARATOR, + LiteLLMAnthropicToResponsesAPIAdapter, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject @@ -24,9 +35,10 @@ class AnthropicResponsesStreamWrapper: response.created -> message_start response.output_item.added -> content_block_start (if message/function_call) response.output_text.delta -> content_block_delta (text_delta) + response.reasoning_summary_part.added -> content_block_delta (thinking_delta separator) response.reasoning_summary_text.delta -> content_block_delta (thinking_delta) response.function_call_arguments.delta -> content_block_delta (input_json_delta) - response.output_item.done -> content_block_stop + response.output_item.done -> content_block_delta (signature_delta) + content_block_stop response.completed -> message_delta + message_stop """ @@ -49,6 +61,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 { @@ -87,6 +101,38 @@ class AnthropicResponsesStreamWrapper: ) return block_idx + @staticmethod + def _field(source: object, name: str) -> object: + return source.get(name) if isinstance(source, dict) else getattr(source, name, None) + + def _close_reasoning_item(self, item: object, item_id: str | None) -> None: + block_idx: Final = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + encrypted_content: Final = self._field(item, "encrypted_content") + signature: Final = ( + encrypted_reasoning_signature(encrypted_content) + if isinstance(encrypted_content, str) and encrypted_content + else None + ) + if block_idx < 0 and signature is None: + return + if block_idx < 0: + redacted_idx: Final = self._open_block( + item_id, + {"type": "redacted_thinking", "data": signature}, # mutable-ok: API message payload + ) + stop: Final = {"type": "content_block_stop", "index": redacted_idx} # mutable-ok: API message payload + self._chunk_queue.append(stop) + return + if signature is not None: + self._chunk_queue.append( + { # mutable-ok: API message payload + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "signature_delta", "signature": signature}, # mutable-ok: API message payload + } + ) + self._chunk_queue.append({"type": "content_block_stop", "index": block_idx}) # mutable-ok: API message payload + def _process_event(self, event: object) -> None: """Convert one Responses API event into zero or more Anthropic chunks queued for emission.""" event_type = getattr(event, "type", None) @@ -131,6 +177,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) @@ -150,6 +214,26 @@ class AnthropicResponsesStreamWrapper: ) return + if event_type == "response.reasoning_summary_part.added": + part_item_id: Final = self._field(event, "item_id") + summary_index: Final = self._field(event, "summary_index") + part_block_idx: Final = ( + self._item_id_to_block_index.get(part_item_id, -1) if isinstance(part_item_id, str) else -1 + ) + if part_block_idx < 0 or not isinstance(summary_index, int) or summary_index == 0: + return + self._chunk_queue.append( + { # mutable-ok: API message payload + "type": "content_block_delta", + "index": part_block_idx, + "delta": { # mutable-ok: API message payload + "type": "thinking_delta", + "thinking": REASONING_SUMMARY_PART_SEPARATOR, + }, + } + ) + return + # ---- reasoning summary text delta ---- if event_type == "response.reasoning_summary_text.delta": item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) @@ -195,6 +279,9 @@ class AnthropicResponsesStreamWrapper: item_id = ( getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None ) + if self._field(item, "type") == "reasoning": + self._close_reasoning_item(item, item_id) + return block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index if block_idx < 0: return @@ -215,34 +302,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 +366,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: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 0eb0e38a46e..1fdb0318bab 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -13,12 +13,17 @@ from typing import Any, Final, cast from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, - responses_reasoning_item_from_thinking_blocks, + encrypted_reasoning_signature, + responses_reasoning_items_from_thinking_blocks, with_prompt_cache_breakpoint, ) 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, @@ -29,6 +34,7 @@ from litellm.types.llms.anthropic import ( AnthropicFinishReason, AnthropicMessagesRequest, AnthropicMessagesToolChoice, + AnthropicResponseContentBlockRedactedThinking, AnthropicResponseContentBlockText, AnthropicResponseContentBlockThinking, AnthropicResponseContentBlockToolUse, @@ -39,11 +45,13 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicUsage, ) from litellm.types.llms.openai import ( - ChatCompletionThinkingBlock, ResponseAPIUsage, ResponsesAPIResponse, ) +REASONING_SUMMARY_PART_SEPARATOR: Final = "\n\n" +RESPONSES_INCLUDE_ENCRYPTED_REASONING: Final = "reasoning.encrypted_content" + class LiteLLMAnthropicToResponsesAPIAdapter: """ @@ -159,49 +167,55 @@ class LiteLLMAnthropicToResponsesAPIAdapter: return str(getattr(part, "text", None) or "") @classmethod - def _thinking_blocks_from_reasoning_item( + def _thinking_block_from_reasoning_item( cls, summary: Iterable[object], - ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload - """Anthropic thinking blocks for one Responses reasoning item. + encrypted_content: object, + ) -> dict[str, Any] | None: # mutable-ok: API message payload + """The one Anthropic block for a Responses reasoning item. - The signature stays empty: only Anthropic can sign a thinking block, and a stand-in - value would be replayed as a real one and rejected by every backend that verifies it. + The item's encrypted reasoning rides the block's opaque field (`signature`, or + `data` when there is no summary text) so the client echoes it back and the next + turn replays the very item OpenAI produced; without it the signature stays empty, + since only Anthropic can sign a thinking block. """ - return tuple( - AnthropicResponseContentBlockThinking( - type="thinking", - thinking=text, - signature=None, - ).model_dump() - for part in summary - if (text := cls._summary_part_text(part)) + text: Final = REASONING_SUMMARY_PART_SEPARATOR.join( + part_text for part in summary if (part_text := cls._summary_part_text(part)) ) + if not isinstance(encrypted_content, str) or not encrypted_content: + if not text: + return None + return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=None).model_dump() + signature: Final = encrypted_reasoning_signature(encrypted_content) + if not text: + return AnthropicResponseContentBlockRedactedThinking(type="redacted_thinking", data=signature).model_dump() + return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=signature).model_dump() @staticmethod def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: """Group a run of consecutive thinking blocks together; keep every other block alone.""" index, block = indexed_block - return "thinking" if block.get("type") == "thinking" else f"block:{index}" + return "thinking" if block.get("type") in ("thinking", "redacted_thinking") else f"block:{index}" @classmethod - def _assistant_group_to_input_item( + def _assistant_group_to_input_items( cls, group: tuple[Mapping[str, object], ...] - ) -> dict[str, Any] | None: # mutable-ok: API message payload + ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload first: Final = group[0] btype: Final = first.get("type") - if btype == "thinking": - blocks: Final = cast(tuple[ChatCompletionThinkingBlock, ...], group) # cast-ok: untrusted client payload - reasoning_item: Final = responses_reasoning_item_from_thinking_blocks(blocks) - return None if reasoning_item is None else dict(reasoning_item) # mutable-ok: API message payload + if btype in ("thinking", "redacted_thinking"): + replayed: Final = responses_reasoning_items_from_thinking_blocks(group) + return tuple(dict(item) for item in replayed) # mutable-ok: API message payload if btype == "tool_use": - return { # mutable-ok: API message payload - "type": "function_call", - "call_id": first.get("id", ""), - "name": first.get("name", ""), - "arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload - } - return None + return ( + { # mutable-ok: API message payload + "type": "function_call", + "call_id": first.get("id", ""), + "name": first.get("name", ""), + "arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload + }, + ) + return () def translate_messages_to_responses_input( self, @@ -358,7 +372,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: input_items.extend( item for _, group in groupby(enumerate(blocks), key=self._assistant_block_group_key) - if (item := self._assistant_group_to_input_item(tuple(block for _, block in group))) is not None + for item in self._assistant_group_to_input_items(tuple(block for _, block in group)) ) asst_parts: list[dict[str, Any]] = [ # mutable-ok: API message payload {"type": "output_text", "text": block.get("text", "")} # mutable-ok: API message payload @@ -491,10 +505,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def translate_request( self, anthropic_request: AnthropicMessagesRequest, + include_encrypted_reasoning: bool = True, ) -> dict[str, Any]: """ Translate a full Anthropic /v1/messages request dict to litellm.responses() / litellm.aresponses() kwargs. + + ``include_encrypted_reasoning`` asks the provider for ``reasoning.encrypted_content`` + on every call, so a reasoning model's items can be replayed intact next turn even + when the client sent no ``thinking`` block; pass False for a provider whose + Responses API rejects ``include``. """ model: Final[str] = anthropic_request["model"] messages_list: Final = cast( @@ -524,6 +544,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter: "model": model, "input": input_items, } + if include_encrypted_reasoning: + responses_kwargs["include"] = [RESPONSES_INCLUDE_ENCRYPTED_REASONING] # mutable-ok: API request payload if system and not developer_parts: if isinstance(system, str): @@ -624,17 +646,29 @@ 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): - content.extend(self._thinking_blocks_from_reasoning_item(item.summary)) + reasoning_block = self._thinking_block_from_reasoning_item(item.summary, item.encrypted_content) + if reasoning_block is not None: + content.append(reasoning_block) 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: @@ -647,24 +681,35 @@ class LiteLLMAnthropicToResponsesAPIAdapter: id=item.call_id or item.id or "", name=item.name, input=input_data, - ).model_dump() + ).model_dump(exclude_none=True) ) stop_reason = "tool_use" 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( - cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json - ) + reasoning_block = self._thinking_block_from_reasoning_item( + cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json + item.get("encrypted_content"), ) + if reasoning_block is not None: + content.append(reasoning_block) elif item_type == "function_call": try: input_data = json.loads(item.get("arguments", "{}")) @@ -676,13 +721,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter: id=item.get("call_id") or item.get("id", ""), name=item.get("name", ""), input=input_data, - ).model_dump() + ).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 +740,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), ) diff --git a/litellm/llms/anthropic/experimental_pass_through/utils.py b/litellm/llms/anthropic/experimental_pass_through/utils.py index 55fe9c47faf..335a0e5641d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/utils.py @@ -3,6 +3,8 @@ from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Final +from pydantic import BaseModel, ConfigDict, ValidationError + import litellm from litellm.types.utils import ModelInfo @@ -21,10 +23,27 @@ _EFFORT_DEGRADATION_CHAIN: Final[Mapping[str, tuple[str, ...]]] = MappingProxyTy _THINKING_OFF: Final = "none" +class _ClaudeCodeUserId(BaseModel): + """The JSON Claude Code packs into ``metadata.user_id``; only ``session_id`` is per conversation.""" + + model_config = ConfigDict(frozen=True) + + session_id: str + + def prompt_cache_key_from_user_id(user_id: object) -> str | None: - if user_id is None: + """The per-session key Claude Code carries inside ``metadata.user_id``, or nothing. + + Anthropic defines ``user_id`` as an opaque end-user id, so a plain string names a person, not + a conversation. Keying the provider cache on it pins every parallel session and subagent of that + person to one slot, which caches worse than the provider's own prompt-prefix hashing does. + """ + if not isinstance(user_id, str): + return None + try: + return _ClaudeCodeUserId.model_validate_json(user_id).session_id[:OPENAI_MAX_PROMPT_CACHE_KEY_LENGTH] or None + except ValidationError: return None - return str(user_id)[:OPENAI_MAX_PROMPT_CACHE_KEY_LENGTH] or None def litellm_logging_obj_from_kwargs(kwargs: Mapping[str, object]) -> "LiteLLMLoggingObject | None": diff --git a/litellm/llms/azure/audio_transcriptions.py b/litellm/llms/azure/audio_transcriptions.py index 4a5ed2ccb0c..564ec94ba6b 100644 --- a/litellm/llms/azure/audio_transcriptions.py +++ b/litellm/llms/azure/audio_transcriptions.py @@ -37,6 +37,7 @@ class AzureAudioTranscription(AzureChatCompletion): azure_ad_token: str | None = None, atranscription: bool = False, litellm_params: dict | None = None, + custom_llm_provider: str = "azure", ) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]: data: Final = {"model": model, "file": audio_file, **optional_params} @@ -53,6 +54,7 @@ class AzureAudioTranscription(AzureChatCompletion): logging_obj=logging_obj, model=model, litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, ) azure_client: Final = self.get_azure_openai_client( @@ -99,7 +101,7 @@ class AzureAudioTranscription(AzureChatCompletion): additional_args={"complete_input_dict": data}, original_response=stringified_response, ) - hidden_params: Final = {"model": model, "custom_llm_provider": "azure"} + hidden_params: Final = {"model": model, "custom_llm_provider": custom_llm_provider} final_response: Final[TranscriptionResponse] = convert_to_model_response_object( response_object=stringified_response, model_response_object=model_response, @@ -122,6 +124,7 @@ class AzureAudioTranscription(AzureChatCompletion): client=None, max_retries=None, litellm_params: dict | None = None, + custom_llm_provider: str = "azure", ) -> TranscriptionResponse: response = None try: @@ -178,7 +181,7 @@ class AzureAudioTranscription(AzureChatCompletion): }, original_response=stringified_response, ) - hidden_params: Final = {"model": model, "custom_llm_provider": "azure"} + hidden_params: Final = {"model": model, "custom_llm_provider": custom_llm_provider} response = convert_to_model_response_object( _response_headers=headers, response_object=stringified_response, diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 46a9dd1a531..587165e6991 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -323,6 +323,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): "api_version": api_version, "api_base": api_base, "complete_input_dict": data, + "openai_sdk": True, }, ) if not isinstance(max_retries, int): @@ -429,6 +430,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): "api_base": api_base, "acompletion": True, "complete_input_dict": data, + "openai_sdk": True, }, ) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 880a51eb584..ed16d7f3de0 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -7,8 +7,9 @@ from httpx._models import Headers, Response import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_tool_reference_parts_from_tool_messages, + flatten_combinators_and_drop_non_python_regex_patterns, hoist_images_from_tool_messages, - tool_with_flattened_parameters, + tool_with_sanitized_parameters, ) from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, @@ -39,14 +40,17 @@ else: _NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) -def flattened_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]: +def sanitized_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]: tools: Final = optional_params.get("tools") if not isinstance(tools, list): return _NO_TOOLS_UPDATE - flattened: Final = [ # mutable-ok: request tools are a JSON list - tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + sanitized: Final = [ # mutable-ok: request tools are a JSON list + tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) + if isinstance(tool, dict) + else tool + for tool in tools ] - return MappingProxyType({"tools": flattened}) + return MappingProxyType({"tools": sanitized}) class AzureOpenAIConfig(BaseConfig): @@ -278,7 +282,7 @@ class AzureOpenAIConfig(BaseConfig): "model": model, "messages": azure_messages, **optional_params, - **flattened_tools_update(optional_params), + **sanitized_tools_update(optional_params), } def transform_response( diff --git a/litellm/llms/azure/chat/o_series_transformation.py b/litellm/llms/azure/chat/o_series_transformation.py index 246bf69cb5f..09d8075e857 100644 --- a/litellm/llms/azure/chat/o_series_transformation.py +++ b/litellm/llms/azure/chat/o_series_transformation.py @@ -20,7 +20,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.utils import get_model_info, supports_reasoning from ...openai.chat.o_series_transformation import OpenAIOSeriesConfig -from .gpt_transformation import flattened_tools_update +from .gpt_transformation import sanitized_tools_update class AzureOpenAIO1Config(OpenAIOSeriesConfig): @@ -111,6 +111,6 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name flattened_params: Final = { # mutable-ok: transform_request's contract takes a plain JSON params dict **optional_params, - **flattened_tools_update(optional_params), + **sanitized_tools_update(optional_params), } return super().transform_request(model, messages, flattened_params, litellm_params, headers) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 6cb7d09cec4..e6b3eb1f2bb 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -14,6 +14,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.constants import DEFAULT_MAX_RETRIES from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import BaseOpenAILLM from litellm.secret_managers.get_azure_ad_token_provider import ( @@ -582,7 +583,8 @@ class BaseAzureLLM(BaseOpenAILLM): if scope is None: scope = "https://cognitiveservices.azure.com/.default" - max_retries: Final = litellm_params.get("max_retries") + configured_max_retries: Final = litellm_params.get("max_retries") + max_retries: Final = DEFAULT_MAX_RETRIES if configured_max_retries is None else configured_max_retries timeout: Final = litellm_params.get("timeout") if not api_key and azure_ad_token_provider is None and tenant_id and client_id and client_secret: verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth") @@ -642,8 +644,7 @@ class BaseAzureLLM(BaseOpenAILLM): else: azure_client_params["http_client"] = self._get_sync_http_client() - if max_retries is not None: - azure_client_params["max_retries"] = max_retries + azure_client_params["max_retries"] = max_retries if timeout is not None: azure_client_params["timeout"] = timeout diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 898852e645f..1b5a4083ebe 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,24 +1,102 @@ +import re +from collections.abc import Callable, Collection, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional import httpx from httpx import Response +from pydantic import BaseModel, ValidationError from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig, + RelayShape, + logged_relay_shape, + replace_path_segment, + strip_leading_model_segment, +) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse, ResponsesTerminalEvent from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse if TYPE_CHECKING: from httpx import URL - from litellm.types.utils import CostResponseTypes + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse + + +class RelayedChatRequest(BaseModel): + messages: Sequence[Mapping[str, object]] | None = None + + +class RelayedCallDetails(BaseModel): + request_data: RelayedChatRequest | None = None + + +def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None: + try: + details: Final = RelayedCallDetails.model_validate(litellm_logging_obj.model_call_details) + except ValidationError: + return None + return details.request_data.messages if details.request_data else None + + +RESPONSES_RELAY_SHAPE: Final = RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate) + +OPENAI_RELAY_SHAPES: Final = ( + RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), + RESPONSES_RELAY_SHAPE, + RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), +) + + +def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None: + """A streaming logging object assembles the logged response from the terminal event, not from its body.""" + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks=all_chunks) + if terminal_event is None: + return None + logging_obj.call_type = ( + RESPONSES_RELAY_SHAPE.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return terminal_event + + +AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(? str | None: + parts: Final = endpoint.split("/") + if len(parts) < 2: + return None + return next((part for part in parts if part in router_models), None) + + +def foreign_azure_deployment( + endpoint: str, model_group: str, served_models: Callable[[], Collection[str]] +) -> str | None: + match: Final = AZURE_DEPLOYMENT_SEGMENT.search(endpoint) + if match is None: + return None + deployment: Final = match.group(1) + if deployment == model_group: + return None + served: Final = frozenset(name.casefold() for name in served_models()) + return None if deployment.casefold() in served else deployment + + +def without_api_version(api_base: str) -> str: + url: Final = httpx.URL(api_base) + kept_params: Final = tuple((key, value) for key, value in url.params.multi_items() if key != "api-version") + return str(url.copy_with(params=httpx.QueryParams(kept_params))) class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: - return "stream" in request_data + return bool(request_data.get("stream")) def get_complete_url( self, @@ -36,14 +114,17 @@ class AzurePassthroughConfig(BasePassthroughConfig): litellm_metadata: Final = litellm_params.get("litellm_metadata") or {} model_group: Final = litellm_metadata.get("model_group") - if model_group and model_group in endpoint: - endpoint = endpoint.replace(model_group, model) + routed_endpoint: Final = replace_path_segment(endpoint, model_group, model) if model_group else endpoint + native_endpoint: Final = strip_leading_model_segment(routed_endpoint, (model,)) + caller_api_version: Final = request_query_params.get("api-version") if request_query_params else None + relay_base: Final = without_api_version(base_target_url) if caller_api_version else base_target_url complete_url: Final = BaseAzureLLM._get_base_azure_url( - api_base=base_target_url, - litellm_params=litellm_params, - route=endpoint, - default_api_version=litellm_params.get("api_version"), + api_base=relay_base, + litellm_params=MappingProxyType( + {**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")} + ), + route=native_endpoint, ) return ( httpx.URL(complete_url), @@ -92,13 +173,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): request_data: dict, logging_obj: Logging, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> Optional["LoggedRelayResponse"]: from litellm import encoding from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.utils import ModelResponse if "chat/completions" not in endpoint: - return None + return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint) openai_chat_config: Final = OpenAIGPTConfig() @@ -116,3 +197,27 @@ class AzurePassthroughConfig(BasePassthroughConfig): ) return litellm_model_response + + def handle_logging_collected_chunks( + self, + all_chunks: Sequence[str], + litellm_logging_obj: Logging, + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> Optional["LoggedRelayResponse"]: + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, + ) + + if f"/{endpoint.strip('/')}".endswith(RESPONSES_RELAY_SHAPE.path_suffix): + return logged_responses_stream(all_chunks, litellm_logging_obj) + if "chat/completions" not in endpoint: + return None + + return OpenAIPassthroughLoggingHandler()._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only OpenAI SSE-to-ModelResponse assembler; reimplementing it would fork the parser + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + messages=_relayed_messages(litellm_logging_obj), + ) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index f2d405e9a17..00e1c1e25ba 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -2,7 +2,6 @@ import copy import enum import re from typing import TYPE_CHECKING, Final, cast -from urllib.parse import urlparse import httpx from httpx import Response @@ -15,8 +14,12 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base +from litellm.llms.azure_ai.common_utils import ( + api_key_header_for_base, + 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 +45,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/`` is the entry to read whenever the map + carries it. Most gpt-5-family names have no ``azure_ai/`` row, though, and prefixing + those anyway costs them every flag: ``get_llm_provider`` re-resolves an ``azure_ai/`` + name to the azure provider when a global AZURE_AI_API_BASE points at an + openai.azure.com host, ``azure/`` is not a key either, so the lookup lands + nowhere and every effort answer degrades to False. A missing key defers to the base + resolver instead. + """ + prefixed: Final = model if model.startswith("azure_ai/") else f"azure_ai/{model}" + return prefixed if prefixed in litellm.model_cost else super()._model_map_lookup_name(model) + + +azureAIGPT5Config: Final = AzureAIGPT5Config() + + 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 +89,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. @@ -99,11 +148,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ Returns True if the request should use `api-key` header for authentication. """ - parsed_url: Final = urlparse(api_base) - host: Final = parsed_url.hostname - if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): - return True - return False + return api_key_header_for_base(api_base) == "api-key" def get_complete_url( self, diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index aa34bab5b2e..0665b3f64c5 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -19,6 +19,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool: return "/openai/deployments" not in parsed.path +def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader: + host: Final = urlparse(api_base).hostname if api_base else None + if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): + return "api-key" + return "Authorization" + + def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: """ Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index 95f536296a2..5934525eca3 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -11,7 +11,7 @@ from litellm.types.utils import Usage from litellm.utils import get_model_info -def _is_azure_model_router(model: str) -> bool: +def is_azure_model_router(model: str) -> bool: """ Check if the model is Azure AI Foundry Model Router. @@ -31,6 +31,18 @@ def _is_azure_model_router(model: str) -> bool: return "model-router" in model_lower or "model_router" in model_lower or model_lower == "azure-model-router" +ROUTER_FEE_ENTRY_NAMES: Final = frozenset({"model-router", "model_router"}) + + +def is_router_fee_entry(model: str) -> bool: + return model.lower().removeprefix("azure_ai/") in ROUTER_FEE_ENTRY_NAMES + + +def _router_fee_entry_name(model: str) -> str: + entry_name: Final = model.lower().removeprefix("azure_ai/") + return entry_name if entry_name in ROUTER_FEE_ENTRY_NAMES else "model_router" + + def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> float: """ Calculate the flat cost for Azure AI Foundry Model Router. @@ -42,20 +54,39 @@ def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> fl Returns: float: The flat cost in USD, or 0.0 if not applicable """ - if not _is_azure_model_router(model): + if not is_azure_model_router(model): return 0.0 - - # Get the model router pricing from model_prices_and_context_window.json - # Use "model_router" as the key (without actual model name suffix) - model_info: Final = get_model_info(model="model_router", custom_llm_provider="azure_ai") + model_info: Final = get_model_info(model=_router_fee_entry_name(model), custom_llm_provider="azure_ai") router_flat_cost_per_token: Final = model_info.get("input_cost_per_token", 0) - if router_flat_cost_per_token and router_flat_cost_per_token > 0: return prompt_tokens * router_flat_cost_per_token - return 0.0 +def _response_model_cost(model: str, usage: Usage, service_tier: str | None) -> tuple[float, float]: + try: + return generic_cost_per_token( + model=model, usage=usage, custom_llm_provider="azure_ai", service_tier=service_tier + ) + except Exception as e: + if not is_azure_model_router(model): + raise + verbose_logger.debug( + "Azure AI Model Router: model '%s' not in cost map, only the routing fee applies. Error: %s", model, e + ) + return 0.0, 0.0 + + +def _router_fee_name(model: str, request_model: str | None) -> str | None: + if is_router_fee_entry(model): + return None + if is_azure_model_router(model): + return model + if request_model is not None and is_azure_model_router(request_model): + return request_model + return None + + def cost_per_token( model: str, usage: Usage, @@ -64,68 +95,31 @@ def cost_per_token( service_tier: str | None = None, ) -> tuple[float, float]: """ - Calculate the cost per token for Azure AI models. + Price the response model's own tokens for Azure AI, plus the Model Router fee exactly once when either the + priced name or request_model is a Model Router name. - For Azure AI Foundry Model Router: - - Adds a flat cost of $0.14 per million input tokens (from model_prices_and_context_window.json) - - Plus the cost of the actual model used (handled by generic_cost_per_token) + A response priced as the router entry itself already carries the fee, so nothing is added on top of it. A + router deployment name that is missing from the cost map prices at the fee alone. + + completion_cost passes only the priced name: when that name is a routed model it adds the fee itself through + AzureModelRouterConfig.calculate_additional_costs as the "Azure Model Router Flat Cost" line of the cost + breakdown, and when the name is router-shaped the fee is already in the prompt cost returned here. Args: model: str, the model name without provider prefix (from response) usage: LiteLLM Usage block response_time_ms: Optional response time in milliseconds - request_model: Optional[str], the original request model name (to detect router usage) + request_model: Optional[str], the original request model name; a Model Router name adds the routing fee + service_tier: Optional service tier the request was priced on Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd Raises: - ValueError: If the model is not found in the cost map and cost cannot be calculated - (except for Model Router models where we return just the routing flat cost) + ValueError: If a model that is not a Model Router name is missing from the cost map """ - prompt_cost = 0.0 - completion_cost = 0.0 - - # Determine if this was a model router request - # Check both the response model and the request model - is_router_request: Final = _is_azure_model_router(model) or ( - request_model is not None and _is_azure_model_router(request_model) - ) - - # Calculate base cost using generic cost calculator - # This may raise an exception if the model is not in the cost map - try: - prompt_cost, completion_cost = generic_cost_per_token( - model=model, - usage=usage, - custom_llm_provider="azure_ai", - service_tier=service_tier, - ) - except Exception as e: - # For Model Router, the model name (e.g., "azure-model-router") may not be in the cost map - # because it's a routing service, not an actual model. In this case, we continue - # to calculate just the routing flat cost. - if not _is_azure_model_router(model): - # Re-raise for non-router models - they should have pricing defined - raise - verbose_logger.debug( - "Azure AI Model Router: model '%s' not in cost map, calculating routing flat cost only. Error: %s", model, e - ) - - # Add flat cost for Azure Model Router - # The flat cost is defined in model_prices_and_context_window.json for azure_ai/model_router - if is_router_request: - # Use the request model for flat cost calculation if available, otherwise use response model - router_model_for_calc: Final = request_model if request_model else model - router_flat_cost: Final = calculate_azure_model_router_flat_cost(router_model_for_calc, usage.prompt_tokens) - - if router_flat_cost > 0: - verbose_logger.debug( - f"Azure AI Model Router flat cost: ${router_flat_cost:.6f} " - f"({usage.prompt_tokens} tokens × ${router_flat_cost / usage.prompt_tokens:.9f}/token)" - ) - - # Add flat cost to prompt cost - prompt_cost += router_flat_cost - - return prompt_cost, completion_cost + prompt_cost, completion_cost = _response_model_cost(model=model, usage=usage, service_tier=service_tier) + fee_name: Final = _router_fee_name(model=model, request_model=request_model) + if fee_name is None: + return prompt_cost, completion_cost + return prompt_cost + calculate_azure_model_router_flat_cost(fee_name, usage.prompt_tokens), completion_cost diff --git a/litellm/llms/azure_ai/image_edit/__init__.py b/litellm/llms/azure_ai/image_edit/__init__.py index 51a23859058..fda9335a5d6 100644 --- a/litellm/llms/azure_ai/image_edit/__init__.py +++ b/litellm/llms/azure_ai/image_edit/__init__.py @@ -23,7 +23,7 @@ def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig: """ Get the appropriate image edit config for an Azure AI model. - - MAI models use /mai/v1/images/edits with multipart form data and size + - MAI models use /mai/v1/images/edits with multipart form data - FLUX 2 models use JSON with base64 image - FLUX 1 models use multipart/form-data """ diff --git a/litellm/llms/azure_ai/image_edit/mai_transformation.py b/litellm/llms/azure_ai/image_edit/mai_transformation.py index e639c20292b..55b179e9591 100644 --- a/litellm/llms/azure_ai/image_edit/mai_transformation.py +++ b/litellm/llms/azure_ai/image_edit/mai_transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final import httpx from httpx._types import RequestFiles @@ -13,7 +13,6 @@ from litellm.llms.azure_ai.image_generation.mai_transformation import ( from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str -from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.openai import FileTypes from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageResponse @@ -26,65 +25,8 @@ if TYPE_CHECKING: class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): """Azure AI Foundry MAI image editing (e.g. MAI-Image-2.5).""" - DEFAULT_SIZE = "1024x1024" - def get_supported_openai_params(self, model: str) -> list: - return ["prompt", "image", "model", "n", "size"] - - def map_openai_params( - self, - image_edit_optional_params: ImageEditOptionalRequestParams, - model: str, - drop_params: bool, - ) -> dict: - optional_params: Final[dict[str, Any]] = {} - supported_params: Final = self.get_supported_openai_params(model) - - for key, value in dict(image_edit_optional_params).items(): - if value is None or key in optional_params: - continue - - if key in supported_params: - if key == "size" and value: - size_param = cast(str, value) - self._validate_size_param(size_param) - optional_params[key] = size_param - else: - optional_params[key] = value - elif not drop_params: - raise ValueError( - f"Parameter {key} is not supported for model {model}. " - f"Supported parameters are {supported_params}. " - f"Set drop_params=True to drop unsupported parameters." - ) - - if "size" not in optional_params: - optional_params["size"] = self.DEFAULT_SIZE - - return optional_params - - def _validate_size_param(self, size: str) -> None: - known_sizes: Final = { - "1024x1024", - "1792x1024", - "1024x1792", - "512x512", - "256x256", - } - - if size in known_sizes: - return - - if "x" in size: - try: - tuple(map(int, size.lower().split("x", 1))) - return - except ValueError: - raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") - - raise ValueError( - f"Unsupported size value: '{size}'. Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." - ) + return ["prompt", "image", "model", "n"] def validate_environment( self, diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 64f81956ad7..67b1a8bcab3 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.exceptions import UnsupportedParamsError from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -21,6 +22,10 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_WIDTH = 1024 DEFAULT_HEIGHT = 1024 + MAX_IMAGES_PER_REQUEST: Final = 1 + MIN_DIMENSION_PX: Final = 768 + MAX_TOTAL_PX: Final = 1_056_768 + @staticmethod def get_mai_image_generation_url( api_base: str | None, @@ -145,16 +150,27 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if k in supported_params: if k == "size" and v: - self._map_size_param(v, optional_params) + self._map_size_param(v, optional_params, model) + elif k == "n" and v is not None and self._image_count(v, model) != self.MAX_IMAGES_PER_REQUEST: + if not drop_params: + raise self._unsupported( + model, + f"n={v} is not supported for model {model}. The Azure AI MAI image " + f"endpoint returns exactly {self.MAX_IMAGES_PER_REQUEST} image per " + "request and ignores any count, so a larger value would silently " + "return fewer images than requested. Send one request per image, or " + "set drop_params=True to drop n.", + ) else: optional_params[k] = v elif k in ("width", "height"): optional_params[k] = v elif not drop_params: - raise ValueError( + raise self._unsupported( + model, f"Parameter {k} is not supported for model {model}. " f"Supported parameters are {supported_params} and width/height. " - f"Set drop_params=True to drop unsupported parameters." + f"Set drop_params=True to drop unsupported parameters.", ) if "width" not in optional_params: @@ -165,7 +181,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): optional_params.pop("size", None) return optional_params - def _map_size_param(self, size: str, optional_params: dict) -> None: + @staticmethod + def _unsupported(model: str, message: str) -> UnsupportedParamsError: + return UnsupportedParamsError(message=message, llm_provider="azure_ai", model=model) + + def _image_count(self, n: object, model: str) -> int: + if isinstance(n, int): + return n + try: + return int(str(n)) + except ValueError: + raise self._unsupported(model, f"n={n!r} is not a whole number of images for model {model}.") + + def _map_size_param(self, size: str, optional_params: dict, model: str) -> None: size_mapping: Final = { "1024x1024": (1024, 1024), "1792x1024": (1792, 1024), @@ -176,19 +204,36 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if size in size_mapping: width, height = size_mapping[size] - optional_params["width"] = width - optional_params["height"] = height elif "x" in size: try: width, height = map(int, size.lower().split("x")) - optional_params["width"] = width - optional_params["height"] = height except ValueError: - raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") + raise self._unsupported( + model, f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." + ) else: - raise ValueError( + raise self._unsupported( + model, f"Unsupported size value: '{size}'. " - f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." + f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string.", + ) + + self._validate_dimensions(model=model, size=size, width=width, height=height) + optional_params["width"] = width + optional_params["height"] = height + + def _validate_dimensions(self, model: str, size: str, width: int, height: int) -> None: + if width < self.MIN_DIMENSION_PX or height < self.MIN_DIMENSION_PX: + raise self._unsupported( + model, + f"Unsupported size value: '{size}'. Azure AI MAI image models require width and " + f"height of at least {self.MIN_DIMENSION_PX} pixels.", + ) + if width * height > self.MAX_TOTAL_PX: + raise self._unsupported( + model, + f"Unsupported size value: '{size}'. Azure AI MAI image models accept at most " + f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height}).", ) def transform_image_generation_response( diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py new file mode 100644 index 00000000000..f2be1d95593 --- /dev/null +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + api_key_header_for_base, + get_azure_ai_auth_headers, +) +from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config +from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig, + RelayShape, + logged_relay_shape, + strip_leading_model_segment, +) +from litellm.types.llms.openai import AllMessageValues +from litellm.types.rerank import RerankResponse +from litellm.types.utils import CallTypes, ImageResponse, StandardPassThroughResponseObject + +if TYPE_CHECKING: + from httpx import URL, Response + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse + + +EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({}) + + +class PassthroughMetadata(BaseModel): + model_config = ConfigDict(extra="ignore") + + model_group: str = "" + + +def model_group_from(litellm_params: Mapping[str, object]) -> str: + try: + return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group + except ValidationError: + return "" + + +def api_version_from(litellm_params: Mapping[str, object]) -> str | None: + try: + return TypeAdapter(str | None).validate_python(litellm_params.get("api_version")) + except ValidationError: + return None + + +def foundry_root(api_base: str) -> str: + url: Final = httpx.URL(api_base) + segments: Final = tuple(segment for segment in url.path.split("/") if segment) + root_segments: Final = segments[: segments.index("models")] if "models" in segments else segments + return str(url.copy_with(path="/" + "/".join(root_segments), query=None)).rstrip("/") + + +def is_repeated_native_prefix(native_segments: tuple[str, ...], overlap: int) -> bool: + return overlap == len(native_segments) or native_segments[0] == "openai" + + +def without_repeated_native_prefix(root: str, native_endpoint: str) -> str: + url: Final = httpx.URL(root) + root_segments: Final = tuple(segment for segment in url.path.split("/") if segment) + native_segments: Final = tuple(segment.casefold() for segment in native_endpoint.split("/") if segment) + overlap: Final = next( + ( + length + for length in range(min(len(root_segments), len(native_segments)), 0, -1) + if tuple(segment.casefold() for segment in root_segments[-length:]) == native_segments[:length] + and is_repeated_native_prefix(native_segments, length) + ), + 0, + ) + kept_segments: Final = root_segments[: len(root_segments) - overlap] + return str(url.copy_with(path="/" + "/".join(kept_segments), query=None)).rstrip("/") + + +def relay_query_params( + request_query_params: Mapping[str, object] | None, + deployment_api_version: str | None, + api_base: str, +) -> Mapping[str, object] | None: + if request_query_params and "api-version" in request_query_params: + return request_query_params + api_version: Final = deployment_api_version or httpx.URL(api_base).params.get("api-version") + if api_version is None: + return request_query_params + return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version}) + + +def relayed_body(httpx_response: Response) -> str | dict: + try: + body: Final[object] = httpx_response.json() + except ValueError: + return httpx_response.text + return body if isinstance(body, dict) else httpx_response.text + + +FOUNDRY_RELAY_SHAPES: Final = ( + RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate), + RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate), +) + + +class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): + def __init__(self, ocr_config_for: Callable[[str], BaseOCRConfig | None] = get_azure_ai_ocr_config) -> None: + super().__init__() + self.ocr_config_for: Final = ocr_config_for + + def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: + return bool(request_data.get("stream")) + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + endpoint: str, + request_query_params: Mapping[str, object] | None, + litellm_params: Mapping[str, object], + ) -> tuple[URL, str]: + base_target_url: Final = self.get_api_base(api_base) + if base_target_url is None: + raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") + + native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) + root: Final = without_repeated_native_prefix(foundry_root(base_target_url), native_endpoint) + query_params: Final = relay_query_params( + request_query_params, api_version_from(litellm_params), base_target_url + ) + return (self.format_url(native_endpoint, root, query_params), root) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: base class contract returns dict for httpx + auth_headers: Final = get_azure_ai_auth_headers( + api_key=api_key, + litellm_params=litellm_params, + api_key_header=api_key_header_for_base(api_base), + ) + return {**headers, **auth_headers} # mutable-ok: base class contract returns dict for httpx + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: Response, + request_data: Mapping[str, object], + logging_obj: Logging, + endpoint: str, + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig + + chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict + model=model, + custom_llm_provider=custom_llm_provider, + httpx_response=httpx_response, + request_data=dict(request_data), # mutable-ok: AzurePassthroughConfig wants a dict + logging_obj=logging_obj, + endpoint=endpoint, + ) + if chat_result is not None: + return chat_result + ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint) + if ocr_result is not None: + return ocr_result + foundry_result: Final = logged_relay_shape(FOUNDRY_RELAY_SHAPES, httpx_response, logging_obj, endpoint) + if foundry_result is not None: + return foundry_result + return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) + + def logged_ocr_response( + self, model: str, httpx_response: Response, logging_obj: Logging, endpoint: str + ) -> OCRResponse | None: + ocr_config: Final = self.ocr_config_for(model) + if ocr_config is None or httpx_response.status_code != 200: + return None + relayed_url: Final = httpx_response.request.url + relayed_origin: Final = str(relayed_url.copy_with(path="/", query=None, fragment=None)).rstrip("/") + ocr_url: Final = httpx.URL( + ocr_config.get_complete_url( + api_base=relayed_origin, + model=model, + optional_params={}, # mutable-ok: BaseOCRConfig wants a dict + ) + ) + known_prefixes: Final = (model, model_group_from(logging_obj.litellm_params)) + native_endpoint: Final = strip_leading_model_segment(endpoint, known_prefixes) + if f"/{native_endpoint.strip('/')}" != ocr_url.path: + return None + try: + ocr_response: Final = ocr_config.transform_ocr_response( + model=model, raw_response=httpx_response, logging_obj=logging_obj + ) + except (ValueError, AttributeError) as error: + verbose_logger.warning("azure_ai passthrough: OCR body from %s is not costable: %s", ocr_url, error) + return None + logging_obj.call_type = CallTypes.aocr.value # rebind-ok: routes cost calculation to the per-page OCR path + return ocr_response + + def handle_logging_collected_chunks( + self, + all_chunks: Sequence[str], + litellm_logging_obj: Logging, + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> LoggedRelayResponse | None: + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig + + return AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + custom_llm_provider=custom_llm_provider, + endpoint=endpoint, + ) diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index bbe1cc85df1..7bfc87a30d6 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -411,6 +411,10 @@ class BaseConfig(ABC): def has_custom_stream_wrapper(self) -> bool: return False + @property + def uses_async_transform_request(self) -> bool: + return False + @property def supports_stream_param_in_request_body(self) -> bool: """ diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 96152141a7c..f1143425ced 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from collections.abc import Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional if TYPE_CHECKING: from fastapi import HTTPException @@ -52,6 +52,29 @@ class StreamingScanKey: class BaseTranslation(ABC): + delivers_ended_stream_rewrites: ClassVar[bool] = False + """Whether ``process_output_streaming_response`` accepts + ``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered) + stream, writes guardrail text and tool-call rewrites back across + ``responses_so_far`` so a buffered pipeline can release rewritten chunks, + raising ``UndeliverableStreamRewrite`` for a shape it cannot place. Rewrites + on every other translation are undeliverable: the pipeline executor + discards them and releases the original chunks.""" + + assembles_streamed_response: ClassVar[bool] = False + """Whether ``process_output_streaming_response`` stores the assembled response of an + ended stream under ``request_data["response"]`` before scanning it, the way the chat, + Responses, and Messages translations do. A streaming pipeline runs a guardrail that only + has the legacy post-call hook against that response, so on a translation without it such + a guardrail keeps running on its own.""" + + def post_call_hook_response(self, response: object) -> object: + """The ``response`` this endpoint's non-streaming post-call hooks receive, derived from + the object the translation stores under ``request_data["response"]`` while scanning an + ended stream. Chat and Responses scan that shape already; a translation that scans a + different one (Messages scans an OpenAI-shaped ModelResponse) overrides this.""" + return response + @staticmethod def transform_user_api_key_dict_to_metadata( user_api_key_dict: Any | None, @@ -157,6 +180,7 @@ class BaseTranslation(ABC): user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> Any: """ Process output streaming response with guardrails. @@ -164,6 +188,11 @@ class BaseTranslation(ABC): Optional to override in subclasses. ``stream_transform_sink`` is the out-parameter used by handlers that support streaming text transformations (see ``StreamTransformSink``); base handlers ignore it. + ``deliver_ended_stream_rewrites`` is passed True only when the caller + holds the whole buffered stream and the subclass declares + ``delivers_ended_stream_rewrites``: the handler then writes + guardrail text and tool-call rewrites back across ``responses_so_far`` + instead of discarding them. """ return responses_so_far diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index 9a25d3294e0..4faf0aaaf30 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -1,5 +1,6 @@ import types from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -102,6 +103,24 @@ class BaseImageEditConfig(ABC): ) -> tuple[dict, RequestFiles]: pass + async def async_transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | None, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict, RequestFiles]: + return self.transform_image_edit_request( + model=model, + prompt=prompt, + image=image, + image_edit_optional_request_params=dict(image_edit_optional_request_params), + litellm_params=litellm_params, + headers=dict(headers), + ) + def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict: """ Last pass on the request dict after ``transform_image_edit_request``, using the diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index adbf2e126fb..ec938889b88 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,5 +1,14 @@ +from __future__ import annotations + +import re from abc import abstractmethod -from typing import TYPE_CHECKING, Final, Optional, Union +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm.types.utils import CallTypes from ..base_utils import BaseLLMModelInfo @@ -7,9 +16,100 @@ if TYPE_CHECKING: from httpx import URL, Headers, Response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.utils import CostResponseTypes + from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesTerminalEvent + from litellm.types.rerank import RerankResponse + from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException + from ..ocr.transformation import OCRResponse + + LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse | ResponsesTerminalEvent + + +RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) + + +def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: + path: Final = endpoint.lstrip("/") + for model_name in model_names: + if not model_name: + continue + if path == model_name: + return "" + if path.startswith(f"{model_name}/"): + return path[len(model_name) + 1 :] + return path + + +def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: + bounded_segment: Final = re.compile(rf"(? Mapping[str, object] | None: + if httpx_response.status_code != 200: + return None + try: + return RELAYED_JSON_OBJECT.validate_python(httpx_response.json()) + except (ValueError, ValidationError): + return None + + +@dataclass(frozen=True, slots=True) +class RelayShape: + path_suffix: str + call_type: CallTypes + parse: Callable[[Mapping[str, object]], LoggedRelayResponse] + + +def logged_relay_shape( + shapes: Sequence[RelayShape], httpx_response: Response, logging_obj: LiteLLMLoggingObj, endpoint: str +) -> LoggedRelayResponse | None: + relayed_path: Final = f"/{endpoint.strip('/')}" + shape: Final = next((candidate for candidate in shapes if relayed_path.endswith(candidate.path_suffix)), None) + body: Final = relayed_json_object(httpx_response) if shape else None + if shape is None or body is None: + return None + try: + parsed: Final = shape.parse(body) + except ValidationError: + return None + logging_obj.call_type = ( + shape.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return parsed + + +class PassthroughStreamCollector(Protocol): + """Consumes relayed stream bytes as they arrive and builds the response logged for spend tracking.""" + + def add(self, chunk: bytes) -> None: ... + + def build_logged_response(self, litellm_logging_obj: LiteLLMLoggingObj) -> LoggedRelayResponse | None: ... + + +class RawBytesStreamCollector: + def __init__( + self, provider_config: BasePassthroughConfig, model: str, custom_llm_provider: str, endpoint: str + ) -> None: + self._provider_config = provider_config + self._model = model + self._custom_llm_provider = custom_llm_provider + self._endpoint = endpoint + self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks + + def add(self, chunk: bytes) -> None: + self._raw_bytes.append(chunk) + + def build_logged_response(self, litellm_logging_obj: LiteLLMLoggingObj) -> LoggedRelayResponse | None: + all_chunks: Final = self._provider_config._convert_raw_bytes_to_str_lines(self._raw_bytes) + return self._provider_config.handle_logging_collected_chunks( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=self._model, + custom_llm_provider=self._custom_llm_provider, + endpoint=self._endpoint, + ) class BasePassthroughConfig(BaseLLMModelInfo): @@ -23,8 +123,8 @@ class BasePassthroughConfig(BaseLLMModelInfo): self, endpoint: str, base_target_url: str, - request_query_params: dict | None, - ) -> "URL": + request_query_params: Mapping[str, object] | None, + ) -> URL: """ Helper function to add query params to the url Args: @@ -58,7 +158,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): endpoint: str, request_query_params: dict | None, litellm_params: dict, - ) -> tuple["URL", str]: + ) -> tuple[URL, str]: """ Get the complete url for the request Returns: @@ -88,9 +188,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): """ return headers, None - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, "Headers"] - ) -> "BaseLLMException": + def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: from litellm.llms.base_llm.chat.transformation import BaseLLMException return BaseLLMException(status_code=status_code, message=error_message, headers=headers) @@ -99,23 +197,30 @@ class BasePassthroughConfig(BaseLLMModelInfo): self, model: str, custom_llm_provider: str, - httpx_response: "Response", + httpx_response: Response, request_data: dict, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: pass def handle_logging_collected_chunks( self, all_chunks: list[str], - litellm_logging_obj: "LiteLLMLoggingObj", + litellm_logging_obj: LiteLLMLoggingObj, model: str, custom_llm_provider: str, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> LoggedRelayResponse | None: return None + def create_stream_collector( + self, model: str, custom_llm_provider: str, endpoint: str + ) -> PassthroughStreamCollector: + return RawBytesStreamCollector( + provider_config=self, model=model, custom_llm_provider=custom_llm_provider, endpoint=endpoint + ) + def _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]: """ Converts a list of raw bytes into a list of string lines, similar to aiter_lines() diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 1365941fe2a..14f00aaaa21 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -62,6 +62,9 @@ class BaseResponsesAPIConfig(ABC): """ return False + def supports_encrypted_agent_messages(self) -> bool: + return False + def sign_request( self, headers: dict, diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index 7668c6132d6..9eca3e69909 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -264,6 +264,13 @@ class BaseSearchConfig: """ raise NotImplementedError("transform_search_response must be implemented by provider") + def get_http_error_class(self, error: httpx.HTTPStatusError) -> Exception: + return self.get_error_class( + error_message=error.response.text, + status_code=error.response.status_code, + headers=dict(error.response.headers), # mutable-ok: provider error factories require dict headers + ) + def get_error_class( self, error_message: str, diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index c8d2b7fe522..07b60cb4b72 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -121,6 +121,9 @@ class RouterVectorStoreEmbeddingExecutor: class BaseVectorStoreConfig: + def validate_create_vector_store(self) -> None: + return None + def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]: return [] diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index a197702921f..f52c1cec6a8 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1,21 +1,27 @@ +import asyncio import base64 +import contextvars import hashlib import json import os import re import urllib.parse -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from datetime import datetime +from functools import partial from threading import Lock -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload import httpx -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + AWS_SIGNING_MAX_THREADS, BEDROCK_EMBEDDING_PROVIDERS_LITERAL, BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES, BEDROCK_IAM_CACHE_MAX_ENTRIES, @@ -26,6 +32,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.aws_partition import contains_bedrock_arn, get_aws_dns_suffix from litellm.litellm_core_utils.dd_tracing import tracer from litellm.secret_managers.main import get_secret, get_secret_str +from litellm.types.llms.bedrock import AwsSessionTag if TYPE_CHECKING: from botocore.awsrequest import AWSPreparedRequest @@ -47,6 +54,47 @@ _STS_REGION_FROM_ENDPOINT_PATTERN: Final = re.compile( SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz-security-token", "date"}) +_AWS_SESSION_TAGS_ADAPTER: Final[TypeAdapter[tuple[AwsSessionTag, ...]]] = TypeAdapter(tuple[AwsSessionTag, ...]) + + +def _canonical_aws_session_tags(raw_tags: object) -> tuple[AwsSessionTag, ...] | None: + if raw_tags is None: + return None + try: + validated: Final = _AWS_SESSION_TAGS_ADAPTER.validate_python(raw_tags) + except ValidationError as e: + raise ValueError( + "Invalid 'aws_session_tags' value. Expected a list of {'Key': , 'Value': } dicts, " + f"e.g. [{{'Key': 'team', 'Value': 'genai'}}]. Got: {raw_tags!r}" + ) from e + return tuple(sorted(validated, key=lambda tag: tag["Key"])) + + +class _AssumeRoleParams(TypedDict): + RoleArn: ReadOnly[str] + RoleSessionName: ReadOnly[str] + ExternalId: ReadOnly[NotRequired[str]] + Tags: ReadOnly[NotRequired[tuple[AwsSessionTag, ...]]] + + +def _assume_role_params( + aws_role_name: str, + aws_session_name: str, + aws_external_id: str | None, + aws_session_tags: Sequence[AwsSessionTag] | None, +) -> _AssumeRoleParams: + match (aws_external_id, tuple(aws_session_tags or ())): + case (None, ()): + return _AssumeRoleParams(RoleArn=aws_role_name, RoleSessionName=aws_session_name) + case (None, tags): + return _AssumeRoleParams(RoleArn=aws_role_name, RoleSessionName=aws_session_name, Tags=tags) + case (external_id, ()): + return _AssumeRoleParams(RoleArn=aws_role_name, RoleSessionName=aws_session_name, ExternalId=external_id) + case (external_id, tags): + return _AssumeRoleParams( + RoleArn=aws_role_name, RoleSessionName=aws_session_name, ExternalId=external_id, Tags=tags + ) + class BedrockRequestTarget(BaseModel): aws_region_name: str @@ -80,7 +128,11 @@ class AwsAuthError(Exception): super().__init__(self.message) # Call the base class constructor with the parameters it needs -class BaseAWSLLM: +class SignsRequestsWithAWS: + pass + + +class BaseAWSLLM(SignsRequestsWithAWS): # Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request). # Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static # access-key + secret + region use ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); ambient @@ -120,6 +172,7 @@ class BaseAWSLLM: "aws_sts_endpoint", "aws_bedrock_runtime_endpoint", "aws_external_id", + "aws_session_tags", ] def _get_ssl_verify(self, ssl_verify: bool | str | None = None): @@ -137,7 +190,7 @@ class BaseAWSLLM: return get_ssl_verify(ssl_verify=ssl_verify) - def get_cache_key(self, credential_args: Mapping[str, str | bool | None]) -> str: + def get_cache_key(self, credential_args: Mapping[str, str | bool | tuple[AwsSessionTag, ...] | None]) -> str: """ Generate a unique cache key based on the credential arguments. """ @@ -147,7 +200,7 @@ class BaseAWSLLM: def _get_or_set_cached_credentials( self, - credential_args: Mapping[str, str | bool | None], + credential_args: Mapping[str, str | bool | tuple[AwsSessionTag, ...] | None], credential_fetcher: Callable[[], tuple[Credentials, int | None]], ) -> Any: """ @@ -222,6 +275,7 @@ class BaseAWSLLM: aws_web_identity_token: str | None = None, aws_sts_endpoint: str | None = None, aws_external_id: str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, ssl_verify: bool | str | None = None, ): """ @@ -258,6 +312,7 @@ class BaseAWSLLM: (aws_external_id, "AWS_EXTERNAL_ID"), ) ) + session_tags: Final = _canonical_aws_session_tags(aws_session_tags) verbose_logger.debug( "in get credentials\n" @@ -270,7 +325,8 @@ class BaseAWSLLM: "aws_role_name=%s\n" "aws_web_identity_token=[set=%s]\n" "aws_sts_endpoint=%s\n" - "aws_external_id=%s", + "aws_external_id=%s\n" + "aws_session_tags=%s", aws_access_key_id is not None, aws_secret_access_key is not None, aws_session_token is not None, @@ -281,6 +337,7 @@ class BaseAWSLLM: aws_web_identity_token is not None, aws_sts_endpoint, aws_external_id, + session_tags, ) args: Final = { @@ -294,6 +351,7 @@ class BaseAWSLLM: "aws_web_identity_token": aws_web_identity_token, "aws_sts_endpoint": aws_sts_endpoint, "aws_external_id": aws_external_id, + "aws_session_tags": session_tags, "ssl_verify": ssl_verify, } @@ -336,6 +394,7 @@ class BaseAWSLLM: aws_region_name=aws_region_name, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=session_tags, ssl_verify=ssl_verify, ), ) @@ -980,6 +1039,7 @@ class BaseAWSLLM: aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, aws_region_name: str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, ) -> dict: """Handle cross-account role assumption for IRSA.""" import boto3 @@ -1032,16 +1092,9 @@ class BaseAWSLLM: # Now assume the target role verbose_logger.debug("Attempting to assume target role: %s with session: %s", aws_role_name, aws_session_name) - assume_role_params: Final = { - "RoleArn": aws_role_name, - "RoleSessionName": aws_session_name, - } - - # Add ExternalId parameter if provided - if aws_external_id is not None: - assume_role_params["ExternalId"] = aws_external_id - - return sts_client_with_creds.assume_role(**assume_role_params) + return sts_client_with_creds.assume_role( + **_assume_role_params(aws_role_name, aws_session_name, aws_external_id, aws_session_tags) + ) def _handle_irsa_same_account( self, @@ -1051,6 +1104,7 @@ class BaseAWSLLM: aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, aws_region_name: str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, ) -> dict: """Handle same-account role assumption for IRSA.""" import boto3 @@ -1074,16 +1128,9 @@ class BaseAWSLLM: # Assume the role verbose_logger.debug("Attempting to assume role: %s with session: %s", aws_role_name, aws_session_name) - assume_role_params: Final = { - "RoleArn": aws_role_name, - "RoleSessionName": aws_session_name, - } - - # Add ExternalId parameter if provided - if aws_external_id is not None: - assume_role_params["ExternalId"] = aws_external_id - - return sts_client.assume_role(**assume_role_params) + return sts_client.assume_role( + **_assume_role_params(aws_role_name, aws_session_name, aws_external_id, aws_session_tags) + ) def _extract_credentials_and_ttl(self, sts_response: dict) -> tuple[Credentials, int | None]: """Extract credentials and TTL from STS response. @@ -1118,6 +1165,7 @@ class BaseAWSLLM: aws_region_name: str | None, aws_sts_endpoint: str | None, aws_external_id: str | None, + aws_session_tags: tuple[AwsSessionTag, ...] | None, ssl_verify: bool | str | None, ) -> tuple[Credentials, int | None]: """ @@ -1144,6 +1192,7 @@ class BaseAWSLLM: aws_region_name=aws_region_name, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ssl_verify=ssl_verify, ) @@ -1159,6 +1208,7 @@ class BaseAWSLLM: aws_sts_endpoint: str | None = None, aws_external_id: str | None = None, ssl_verify: bool | str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, ) -> tuple[Credentials, int | None]: """ Authenticate with AWS Role @@ -1189,6 +1239,7 @@ class BaseAWSLLM: aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, aws_region_name=aws_region_name, + aws_session_tags=aws_session_tags, ) else: sts_response = self._handle_irsa_same_account( @@ -1198,6 +1249,7 @@ class BaseAWSLLM: aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, aws_region_name=aws_region_name, + aws_session_tags=aws_session_tags, ) return self._extract_credentials_and_ttl(sts_response) @@ -1234,14 +1286,9 @@ class BaseAWSLLM: **sts_client_kwargs, ) - assume_role_params: Final = { - "RoleArn": aws_role_name, - "RoleSessionName": aws_session_name, - } - - # Add ExternalId parameter if provided - if aws_external_id is not None: - assume_role_params["ExternalId"] = aws_external_id + assume_role_params: Final = _assume_role_params( + aws_role_name, aws_session_name, aws_external_id, aws_session_tags + ) try: sts_response = sts_client.assume_role(**assume_role_params) @@ -1460,6 +1507,7 @@ class BaseAWSLLM: "aws_bedrock_runtime_endpoint", None ) # https://bedrock-runtime.{region_name}.amazonaws.com aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) if bearer_token is not None: return BearerRequestTarget( @@ -1478,6 +1526,7 @@ class BaseAWSLLM: aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) return Boto3CredentialsInfo( credentials=credentials, @@ -1623,6 +1672,7 @@ class BaseAWSLLM: aws_web_identity_token: Final = optional_params.get("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.get("aws_sts_endpoint", None) aws_external_id: Final = optional_params.get("aws_external_id", None) + aws_session_tags: Final = optional_params.get("aws_session_tags", None) aws_region_name: Final = self._get_aws_region_name(optional_params=optional_params, model=model) credentials: Final[Credentials] = self.get_credentials( @@ -1636,6 +1686,7 @@ class BaseAWSLLM: aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) sigv4: Final = SigV4Auth(credentials, service_name, aws_region_name) @@ -1668,3 +1719,52 @@ class BaseAWSLLM: request_headers_dict["Authorization"] = incoming_authorization return request_headers_dict, request.body + + +def sign_aws_json_post( + get_credentials: Callable[[], Credentials], + service_name: str, + aws_region_name: str | None, + url: str, + body: str, + headers: Mapping[str, str], +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.") + + aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers) + SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) + return aws_request.prepare() + + +_SignParams = ParamSpec("_SignParams") +_SignedRequest = TypeVar("_SignedRequest") + +AWS_SIGNING_EXECUTOR: Final = ThreadPoolExecutor(max_workers=AWS_SIGNING_MAX_THREADS, thread_name_prefix="aws-signing") + + +async def run_aws_signing( + sign: Callable[_SignParams, _SignedRequest], + /, + *args: _SignParams.args, + **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped signing signature +) -> _SignedRequest: + context: Final = contextvars.copy_context() + return await asyncio.get_running_loop().run_in_executor( + AWS_SIGNING_EXECUTOR, partial(context.run, sign, *args, **kwargs) + ) + + +async def sign_request_off_loop_if_aws( + provider_config: object, + sign_request: Callable[_SignParams, _SignedRequest], + /, + *args: _SignParams.args, + **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature +) -> _SignedRequest: + if isinstance(provider_config, SignsRequestsWithAWS): + return await run_aws_signing(sign_request, *args, **kwargs) + return sign_request(*args, **kwargs) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 4b500897642..b408d2f620c 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast @@ -6,6 +6,7 @@ from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.types.llms.bedrock import AwsSessionTag from litellm.types.utils import LiteLLMBatch if TYPE_CHECKING: @@ -116,6 +117,7 @@ class BedrockBatchesHandler: aws_web_identity_token: str | None = None, aws_sts_endpoint: str | None = None, aws_external_id: str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, **kwargs: object, # kwargs-ok: litellm.cancel_batch forwards arbitrary user kwargs verbatim ) -> "LiteLLMBatch": try: @@ -139,6 +141,7 @@ class BedrockBatchesHandler: aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) client: Final = boto3.client( @@ -163,6 +166,7 @@ class BedrockBatchesHandler: aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) try: @@ -283,7 +287,7 @@ class BedrockBatchesHandler: ``aws_session_token``, ``aws_profile_name``, ``aws_role_name``, ``aws_session_name``, ``aws_web_identity_token``, ``aws_sts_endpoint``, - ``aws_external_id``). Unknown keys are ignored. + ``aws_external_id``, ``aws_session_tags``). Unknown keys are ignored. Returns: ``LiteLLMBatch`` shaped like an OpenAI Batch resource. @@ -317,6 +321,7 @@ class BedrockBatchesHandler: aws_web_identity_token=kwargs.get("aws_web_identity_token"), aws_sts_endpoint=kwargs.get("aws_sts_endpoint"), aws_external_id=kwargs.get("aws_external_id"), + aws_session_tags=kwargs.get("aws_session_tags"), ) client: Final = boto3.client( diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 690040dd93b..6aa17372258 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -667,7 +667,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=str(response.read())) + raise BedrockError( + status_code=response.status_code, + message=str(response.read()), + headers=response.headers, + response=response, + ) # LOGGING logging_obj.post_call( @@ -690,6 +695,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): raise BedrockError( status_code=response.status_code, message=f"AgentCore: Failed to read/parse JSON response body: {e}", + headers=response.headers, ) parsed: Final = self._parse_json_response(response_json) @@ -880,7 +886,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=str(await response.aread())) + raise BedrockError( + status_code=response.status_code, + message=str(await response.aread()), + headers=response.headers, + response=response, + ) # LOGGING logging_obj.post_call( @@ -903,6 +914,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): raise BedrockError( status_code=response.status_code, message=f"AgentCore: Failed to read/parse JSON response body: {e}", + headers=response.headers, ) parsed: Final = self._parse_json_response(response_json) @@ -1031,6 +1043,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error processing response: {e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) def validate_environment( @@ -1046,7 +1059,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return headers def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) def should_fake_stream( self, diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index a75124325ae..d397420cb17 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -21,8 +21,8 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper -from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token -from ..common_utils import BedrockError, _get_all_bedrock_regions +from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing +from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call @@ -66,7 +66,12 @@ def make_sync_call( ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=str(response.read())) + raise BedrockError( + status_code=response.status_code, + message=str(response.read()), + headers=response.headers, + response=response, + ) if fake_stream: model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( @@ -131,7 +136,8 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = self.get_request_headers( + prepped: Final = await run_aws_signing( + self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", extra_headers=headers, @@ -201,7 +207,8 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = self.get_request_headers( + prepped: Final = await run_aws_signing( + self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", extra_headers=headers, @@ -247,7 +254,12 @@ class BedrockConverseLLM(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -345,6 +357,7 @@ class BedrockConverseLLM(BaseAWSLLM): aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) optional_params.pop("aws_region_name", None) litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls @@ -363,6 +376,7 @@ class BedrockConverseLLM(BaseAWSLLM): aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) ) @@ -594,7 +608,12 @@ class BedrockConverseLLM(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index e097805f54a..fa18361e44c 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -4,6 +4,7 @@ Translating between OpenAI's `/chat/completion` format and Amazon's `/converse` import copy import json +import re import time import types from collections.abc import Mapping @@ -293,6 +294,10 @@ class AmazonConverseConfig(BaseConfig): llm_provider="bedrock", ) + @staticmethod + def _is_openai_gpt_reasoning_model(model: str) -> bool: + return re.search(r"openai\.gpt-\d", model) is not None + def _is_nova_2_model(self, model: str) -> bool: """ Check if the model is a Nova 2 model that supports reasoningConfig. @@ -422,15 +427,15 @@ class AmazonConverseConfig(BaseConfig): """ Handle the reasoning_effort parameter based on the model type. - - GPT-OSS models: passed through unchanged via additionalModelRequestFields. - - OpenAI GPT-5.x models: mapped to ``reasoning.effort`` via additionalModelRequestFields. + - GPT-OSS and DeepSeek V3 models: passed through unchanged via additionalModelRequestFields. + - OpenAI GPT-5.x and GPT-6 models: mapped to ``reasoning.effort`` via additionalModelRequestFields. - Nova 2 models: transformed to reasoningConfig. - Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on adaptive Claude 4.6 / 4.7). """ - if "gpt-oss" in model: + if "gpt-oss" in model or "deepseek" in model: optional_params["reasoning_effort"] = reasoning_effort - elif "openai.gpt-5" in model: + elif self._is_openai_gpt_reasoning_model(model): reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort} optional_params["reasoning"] = reasoning elif self._is_nova_2_model(model): @@ -509,6 +514,36 @@ class AmazonConverseConfig(BaseConfig): ) thinking["budget_tokens"] = BEDROCK_MIN_THINKING_BUDGET_TOKENS + def _is_deepseek_model(self, model: str, base_model: str) -> bool: + return "deepseek" in model or "deepseek" in base_model + + def _is_deepseek_r1_model(self, model: str, base_model: str) -> bool: + return "deepseek.r1" in model or "deepseek.r1" in base_model + + def _model_accepts_anthropic_thinking_param(self, model: str, base_model: str) -> bool: + """Whether the model accepts the Anthropic-shaped ``thinking`` request field. + + Only Claude reasoning models accept it. DeepSeek advertises ``supports_reasoning`` but reasons + natively: R1 returns a 400 when the field is sent and V3 silently ignores it. + """ + if self._is_deepseek_model(model=model, base_model=base_model): + return False + return ( + "claude-3-7" in model + or "claude-sonnet-4" in model + or "claude-opus-4" in model + or supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider) + or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) + ) + + def _model_rejects_reasoning_effort_param(self, model: str, base_model: str) -> bool: + """Whether the model returns a 400 for every ``reasoning_effort`` shape on Converse. + + DeepSeek R1 always reasons and rejects any reasoning request field. DeepSeek V3 accepts a raw + ``reasoning_effort`` like gpt-oss does, and every other model maps it to a shape it accepts. + """ + return self._is_deepseek_r1_model(model=model, base_model=base_model) + def get_supported_openai_params(self, model: str) -> list[str]: from litellm.utils import supports_function_calling @@ -564,23 +599,20 @@ class AmazonConverseConfig(BaseConfig): # only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html supported_params.append("tool_choice") - if "gpt-oss" in model or "openai.gpt-5" in model or "openai.gpt-5" in base_model: + if ( + "gpt-oss" in model + or self._is_openai_gpt_reasoning_model(model) + or self._is_openai_gpt_reasoning_model(base_model) + ): supported_params.append("reasoning_effort") + elif self._is_deepseek_model(model=model, base_model=base_model): + if not self._is_deepseek_r1_model(model=model, base_model=base_model): + supported_params.append("reasoning_effort") elif self._is_nova_2_model(model): # Nova 2 models support reasoning_effort (transformed to reasoningConfig) # These models use a different reasoning structure than Anthropic's thinking parameter supported_params.append("reasoning_effort") - elif ( - "claude-3-7" in model - or "claude-sonnet-4" in model - or "claude-opus-4" in model - or "deepseek.r1" in model - or supports_reasoning( - model=model, - custom_llm_provider=self.custom_llm_provider, - ) - or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) - ): + elif self._model_accepts_anthropic_thinking_param(model=model, base_model=base_model): supported_params.append("thinking") supported_params.append("reasoning_effort") supported_params.append("output_config") @@ -872,6 +904,11 @@ class AmazonConverseConfig(BaseConfig): drop_params: bool, ) -> dict: is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) + base_model: Final = BedrockModelInfo.get_base_model(model) + drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model) + drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param( + model=model, base_model=base_model + ) for param, value in non_default_params.items(): if param == "response_format" and isinstance(value, dict): @@ -920,7 +957,12 @@ class AmazonConverseConfig(BaseConfig): optional_params["_parallel_tool_use_config"] = { "tool_choice": {"type": "auto", "disable_parallel_tool_use": not value} } - if param == "thinking" and "openai.gpt-5" not in model: + if param == "thinking" and drop_thinking_param: + verbose_logger.debug( + "Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively.", + model, + ) + elif param == "thinking" and not self._is_openai_gpt_reasoning_model(model): if ( isinstance(value, dict) and value.get("type") == "adaptive" @@ -946,6 +988,11 @@ class AmazonConverseConfig(BaseConfig): AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( model=model, optional_params=optional_params, custom_llm_provider="bedrock" ) + elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_effort_param: + verbose_logger.debug( + "Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it always reasons and rejects it.", + model, + ) elif param == "reasoning_effort" and isinstance(value, str): self._handle_reasoning_effort_parameter( model=model, reasoning_effort=value, optional_params=optional_params @@ -1805,6 +1852,7 @@ class AmazonConverseConfig(BaseConfig): data=request_data, messages=messages, encoding=encoding, + json_mode=json_mode, ) def _transform_reasoning_content(self, reasoning_content_blocks: list[BedrockConverseReasoningContentBlock]) -> str: @@ -2237,6 +2285,7 @@ class AmazonConverseConfig(BaseConfig): data: dict | str, messages: list, encoding, + json_mode: bool | None = None, ) -> ModelResponse: ## LOGGING if logging_obj is not None: @@ -2247,7 +2296,9 @@ class AmazonConverseConfig(BaseConfig): additional_args={"complete_input_dict": data}, ) - json_mode: Final[bool | None] = optional_params.get("json_mode", None) + resolved_json_mode: Final[bool | None] = ( + json_mode if json_mode is not None else optional_params.get("json_mode", None) + ) ## RESPONSE OBJECT try: completion_response: Final = ConverseResponseBlock(**response.json()) @@ -2255,6 +2306,7 @@ class AmazonConverseConfig(BaseConfig): raise BedrockError( message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", status_code=422, + headers=response.headers, ) """ @@ -2338,7 +2390,7 @@ class AmazonConverseConfig(BaseConfig): chat_completion_message["thinking_blocks"] = self._transform_thinking_blocks(reasoningContentBlocks) chat_completion_message["content"] = content_str filtered_tools: Final = self._filter_json_mode_tools( - json_mode=json_mode, + json_mode=resolved_json_mode, tools=tools, chat_completion_message=chat_completion_message, ) @@ -2362,7 +2414,7 @@ class AmazonConverseConfig(BaseConfig): # When json_mode filtered out all synthetic tool calls the response # is plain content, not a pending tool invocation. Fix finish_reason # so callers (e.g. OpenAI SDK) don't misinterpret it. - if json_mode and not filtered_tools and tools: + if resolved_json_mode and not filtered_tools and tools: initial_finish_reason = "stop" ( diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index e30ec731d8c..d489e47c3b5 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -470,6 +470,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error processing response: {e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) def validate_environment( @@ -485,7 +486,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): return headers def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) def should_fake_stream( self, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index c39c88240c5..5c489ecb360 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -42,6 +42,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk from ..common_utils import ( BedrockError, build_bedrock_stream_error, + error_response_text, get_bedrock_response_stream_shape, get_bedrock_tool_name, ) @@ -184,7 +185,12 @@ async def make_call( ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=response.text) + raise BedrockError( + status_code=response.status_code, + message=error_response_text(response), + headers=response.headers, + response=response, + ) if fake_stream: model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( @@ -228,9 +234,16 @@ async def make_call( ) return completion_stream, response.headers + except BedrockError: + raise except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") except Exception as e: @@ -270,7 +283,12 @@ def make_sync_call( ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=response.text) + raise BedrockError( + status_code=response.status_code, + message=error_response_text(response), + headers=response.headers, + response=response, + ) if fake_stream: model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( @@ -314,9 +332,16 @@ def make_sync_call( ) return completion_stream, response.headers + except BedrockError: + raise except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") except Exception as e: @@ -465,10 +490,10 @@ class AWSEventStreamDecoder: reasoning_content: str | None = None thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None = None - self.content_blocks.append(delta_obj) if "text" in delta_obj: text = delta_obj["text"] elif "toolUse" in delta_obj: + self.content_blocks.append(delta_obj) # When json_mode is True and this is the internal json_tool_call, # convert tool input to text content instead of tool call arguments if self.json_mode is True and self._current_tool_name == RESPONSE_FORMAT_TOOL_NAME: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index 04c6ec86a13..5d39b68d9d5 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -247,4 +247,4 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError: """Return the appropriate error class for Bedrock.""" - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index 1671585be2d..4bf1a1cba73 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -182,4 +182,4 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError: """Return the appropriate error class for Bedrock.""" - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index cd8066cda4d..d12c8aee48c 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -212,6 +212,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): raise BedrockError( message=f"Error parsing response: {raw_response.text}, error: {e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) verbose_logger.debug( @@ -241,6 +242,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): raise BedrockError( message=f"Error setting response content: {e}. Response: {completion_response}", status_code=raw_response.status_code, + headers=raw_response.headers, ) # Calculate usage from headers diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 67720451c00..38f280eef03 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -8,7 +8,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( - async_convert_url_to_base64, + async_inline_remote_media, convert_url_to_base64, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -172,6 +172,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): return _anthropic_request + @property + def uses_async_transform_request(self) -> bool: + return True + async def async_transform_request( self, model: str, @@ -180,26 +184,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): litellm_params: dict, headers: dict, ) -> dict: - _anthropic_request: Final = self._build_bedrock_anthropic_request_base( + return self.transform_request( model=model, - messages=messages, + messages=await async_inline_remote_media(messages), optional_params=optional_params, litellm_params=litellm_params, headers=headers, ) - await self._async_convert_document_url_sources_to_base64(_anthropic_request) - beta_list: Final = self._compute_bedrock_invoke_beta_headers( - model=model, - messages=messages, - optional_params=optional_params, - headers=headers, - ) - if beta_list: - _anthropic_request["anthropic_beta"] = beta_list - - return _anthropic_request - def _build_bedrock_anthropic_request_base( self, model: str, @@ -321,45 +313,6 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): "data": image_chunk["data"], } - async def _async_convert_document_url_sources_to_base64(self, anthropic_request: dict) -> None: - """ - Async version of document URL conversion for async completion paths. - """ - messages: Final = anthropic_request.get("messages") - if not isinstance(messages, list): - return - - for message in messages: - if not isinstance(message, dict): - continue - content = message.get("content") - if not isinstance(content, list): - continue - - for block in content: - if not isinstance(block, dict) or block.get("type") != "document": - continue - source = block.get("source") - if not isinstance(source, dict) or source.get("type") != "url": - continue - source_url = source.get("url") - if not isinstance(source_url, str): - continue - - inferred_format: str | None = None - if source_url.lower().endswith(".pdf"): - inferred_format = "application/pdf" - base64_url = await async_convert_url_to_base64(url=source_url) - image_chunk = convert_to_anthropic_image_obj( - openai_image_url=base64_url, - format=inferred_format, - ) - block["source"] = { - "type": "base64", - "media_type": image_chunk["media_type"], - "data": image_chunk["data"], - } - def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict: """ Convert tool search entries to the format supported by the Bedrock Invoke API. diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 37121d2ece7..90a2692f68a 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -295,7 +295,11 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): try: completion_response: Final = raw_response.json() except Exception: - raise BedrockError(message=raw_response.text, status_code=raw_response.status_code) + raise BedrockError( + message=raw_response.text, + status_code=raw_response.status_code, + headers=raw_response.headers, + ) verbose_logger.debug( "bedrock invoke response % s", json.dumps(completion_response, indent=4, default=str), @@ -336,6 +340,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): optional_params=optional_params, litellm_params=litellm_params, encoding=encoding, + json_mode=json_mode, ) elif provider == "twelvelabs": return litellm.AmazonTwelveLabsPegasusConfig().transform_response( @@ -363,6 +368,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error processing={raw_response.text}, Received error={e}", status_code=422, + headers=raw_response.headers, ) try: @@ -384,6 +390,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error parsing received text={outputText}.\nError-{e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) ## CALCULATING USAGE - bedrock returns usage in the headers @@ -431,7 +438,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) @track_llm_api_timing() async def get_async_custom_stream_wrapper( diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index 31fc079c0c9..7e2037c33f1 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -10,6 +10,7 @@ at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth. from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Any, Final +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -110,21 +111,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): litellm_params: dict, headers: dict, ) -> dict: - model_id: Final = model.replace("mantle/", "", 1) - - request: Final = self._build_bedrock_anthropic_request_base( - model=model_id, - messages=messages, + return self.transform_request( + model=model, + messages=await async_inline_remote_media(messages), optional_params=optional_params, litellm_params=litellm_params, headers=headers, ) - await self._async_convert_document_url_sources_to_base64(request) - return self._restore_mantle_body_fields( - request=request, - model_id=model_id, - optional_params=optional_params, - ) @staticmethod def _restore_mantle_body_fields(request: dict, model_id: str, optional_params: dict) -> dict: diff --git a/litellm/llms/bedrock/claude_platform/common_utils.py b/litellm/llms/bedrock/claude_platform/common_utils.py index 311f3a56b84..fb7f2185ec5 100644 --- a/litellm/llms/bedrock/claude_platform/common_utils.py +++ b/litellm/llms/bedrock/claude_platform/common_utils.py @@ -1,7 +1,10 @@ from typing import Final +import httpx + import litellm from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError from litellm.secret_managers.main import get_secret_str CLAUDE_PLATFORM_SERVICE_NAME: Final = "aws-external-anthropic" @@ -15,6 +18,14 @@ def strip_claude_platform_route(model: str) -> str: class BedrockClaudePlatformMixin(BaseAWSLLM): + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + @staticmethod def _get_workspace_id(optional_params: dict, litellm_params: dict) -> str | None: workspace_id = ( diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index fe675a30a00..cb2c70e74c8 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -33,8 +33,53 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues +_ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs" + + +def error_response_text(response: httpx.Response) -> str: + try: + return response.text + except httpx.ResponseNotRead: + return response.reason_phrase + + +def _synthesize_error_response( + *, status_code: int, headers: dict[str, object] | httpx.Headers, request: httpx.Request | None +) -> tuple[httpx.Request, httpx.Response]: + error_request: Final = request or httpx.Request(method="POST", url=_ERROR_REQUEST_URL) + safe_headers: Final = ( + headers + if isinstance(headers, httpx.Headers) + else tuple((key, value) for key, value in headers.items() if isinstance(value, (str, bytes))) + ) + return error_request, httpx.Response(status_code=status_code, headers=safe_headers, request=error_request) + + class BedrockError(BaseLLMException): - pass + def __init__( + self, + status_code: int, + message: str, + headers: dict[str, object] | httpx.Headers | None = None, + request: httpx.Request | None = None, + response: httpx.Response | None = None, + body: dict[str, object] | None = None, + status_code_is_synthesized: bool = False, + ) -> None: + error_request, error_response = ( + _synthesize_error_response(status_code=status_code, headers=headers, request=request) + if response is None and headers + else (request, response) + ) + super().__init__( + status_code=status_code, + message=message, + headers=headers, + request=error_request, + response=error_response, + body=body, + status_code_is_synthesized=status_code_is_synthesized, + ) _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = ( @@ -48,6 +93,7 @@ _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = ( "aws_web_identity_token", "aws_sts_endpoint", "aws_external_id", + "aws_session_tags", ) @@ -1618,6 +1664,7 @@ class CommonBatchFilesUtils: aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), aws_external_id=optional_params.get("aws_external_id"), + aws_session_tags=optional_params.get("aws_session_tags"), ) # Prepare the request data diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 2383350b3a3..d7fb510f057 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -10,9 +10,10 @@ import httpx import litellm from litellm._logging import verbose_logger +from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client class BedrockCountTokensHandler(BedrockCountTokensConfig): @@ -27,6 +28,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): request_data: dict[str, Any], litellm_params: dict[str, Any], resolved_model: str, + client: AsyncHTTPHandler | None = None, ) -> dict[str, Any]: """ Handle a CountTokens request using existing LiteLLM patterns. @@ -75,7 +77,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): # Extract api_key for bearer token auth if provided api_key: Final = litellm_params.get("api_key", None) headers: Final = {"Content-Type": "application/json"} - signed_headers, signed_body = self._sign_request( + signed_headers, signed_body = await run_aws_signing( + self._sign_request, service_name="bedrock", headers=headers, optional_params=litellm_params, @@ -85,7 +88,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): api_key=api_key, ) - async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) + async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) response: Final = await async_client.post( endpoint_url, @@ -102,6 +105,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise BedrockError( status_code=response.status_code, message=error_text, + headers=response.headers, + response=response, ) bedrock_response: Final = response.json() @@ -124,6 +129,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise BedrockError( status_code=e.response.status_code, message=e.response.text, + headers=e.response.headers, + response=e.response, ) except Exception as e: verbose_logger.error("Error in CountTokens handler: %s", e) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 5fb86d476f4..7efdfd3cebb 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -5,7 +5,7 @@ Handles embedding calls to Bedrock's `/invoke` endpoint import copy import json import urllib.parse -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Final, get_args, overload import httpx @@ -26,7 +26,7 @@ from litellm.types.llms.bedrock import ( ) from litellm.types.utils import EmbeddingResponse, LlmProviders -from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token +from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing from ..common_utils import BedrockError from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config @@ -35,12 +35,26 @@ from .amazon_titan_multimodal_transformation import ( ) from .amazon_titan_v2_transformation import AmazonTitanV2Config from .cohere_transformation import BedrockCohereEmbeddingConfig -from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig +from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig, drop_params_enabled if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +def _sign_get_request( + credentials: Credentials, url: str, headers: Mapping[str, str], aws_region_name: str +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + + request: Final = AWSRequest(method="GET", url=url, data=None, headers=headers) + SigV4Auth(credentials, "bedrock", aws_region_name).add_auth(request) + return request.prepare() + + class BedrockEmbedding(BaseAWSLLM): @overload def _load_credentials( @@ -73,6 +87,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) ### SET REGION NAME ### if aws_region_name is None: @@ -103,6 +118,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) ) return credentials, aws_region_name @@ -132,7 +148,12 @@ class BedrockEmbedding(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -161,7 +182,12 @@ class BedrockEmbedding(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -229,7 +255,7 @@ class BedrockEmbedding(BaseAWSLLM): returned_response = AmazonTitanG1Config()._transform_response(response_list=response_list, model=model) elif provider == "twelvelabs": returned_response = TwelveLabsMarengoEmbeddingConfig()._transform_response( - response_list=response_list, model=model + response_list=response_list, model=model, batch_data=batch_data ) elif provider == "nova": returned_response = AmazonNovaEmbeddingConfig()._transform_response( @@ -332,7 +358,8 @@ class BedrockEmbedding(BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - prepped = self.get_request_headers( + prepped = await run_aws_signing( + self.get_request_headers, credentials=credentials, aws_region_name=aws_region_name, extra_headers=extra_headers, @@ -474,12 +501,13 @@ class BedrockEmbedding(BaseAWSLLM): elif provider == "twelvelabs": batch_data = [] for i in input: - twelvelabs_request = TwelveLabsMarengoEmbeddingConfig()._transform_request( + twelvelabs_request = TwelveLabsMarengoEmbeddingConfig(model=model)._transform_request( input=i, inference_params=inference_params, async_invoke_route=has_async_invoke, model_id=modelId, output_s3_uri=inference_params.get("output_s3_uri"), + drop_params=drop_params_enabled(litellm_params), ) batch_data.append(twelvelabs_request) elif provider == "nova": @@ -589,9 +617,6 @@ class BedrockEmbedding(BaseAWSLLM): dict: Status response from AWS Bedrock """ - # Get AWS credentials using the same method as other Bedrock methods - credentials, _ = self._load_credentials(kwargs) - # Get the runtime endpoint endpoint_url, _ = self.get_runtime_endpoint( api_base=None, @@ -608,27 +633,13 @@ class BedrockEmbedding(BaseAWSLLM): # Prepare headers for GET request headers: Final = {"Content-Type": "application/json"} - # Use AWSRequest directly for GET requests (get_request_headers hardcodes POST) - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + def sign_status_request() -> AWSPreparedRequest: + credentials, _ = self._load_credentials(kwargs) + return _sign_get_request( + credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name + ) - # Create AWSRequest with GET method and encoded URL - request: Final = AWSRequest( - method="GET", - url=status_url, - data=None, # GET request, no body - headers=headers, - ) - - # Sign the request - SigV4Auth will create canonical string from request URL - sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name) - sigv4.add_auth(request) - - # Prepare the request - prepped: Final = request.prepare() + prepped: Final = await run_aws_signing(sign_status_request) # LOGGING if logging_obj is not None: diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py new file mode 100644 index 00000000000..4aac6f22bae --- /dev/null +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py @@ -0,0 +1,239 @@ +""" +Request builder for Bedrock TwelveLabs Marengo Embed 3.0, whose payload nests the input under a key named after +``inputType`` instead of the flat 2.7 layout. + +Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo-3.html +""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.types.llms.bedrock import ( + TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS, + TWELVELABS_MARENGO_3_EMBEDDING_SCOPES, + TWELVELABS_MARENGO_3_EMBEDDING_TYPES, + TWELVELABS_MARENGO_3_INPUT_TYPES, + TwelveLabsMarengo3AudioRequest, + TwelveLabsMarengo3EmbeddingRequest, + TwelveLabsMarengo3ImageRequest, + TwelveLabsMarengo3MultiInputRequest, + TwelveLabsMarengo3NamedMediaSource, + TwelveLabsMarengo3RequestBase, + TwelveLabsMarengo3Segmentation, + TwelveLabsMarengo3TextImageRequest, + TwelveLabsMarengo3TextRequest, + TwelveLabsMarengo3TimedMediaInput, + TwelveLabsMarengo3TimedMediaOptions, + TwelveLabsMarengo3VideoRequest, + TwelveLabsMediaSource, + TwelveLabsS3Location, +) +from litellm.utils import get_base64_str + +MARENGO_3_MODEL_MARKER: Final = "marengo-embed-3-" +S3_URI_PREFIX: Final = "s3://" +TIMED_MEDIA_OPTION_FIELDS: Final = MappingProxyType( + { + "startSec": True, + "endSec": True, + "segmentation": True, + "embeddingOption": True, + "embeddingType": True, + "embeddingScope": True, + } +) +TIMED_MEDIA_OPTIONS: Final = TypeAdapter(TwelveLabsMarengo3TimedMediaOptions) +TIMED_INPUT_TYPES: Final = frozenset({"video", "audio"}) +MARENGO_2_7_ONLY_PARAMS: Final = ("textTruncate", "lengthSec", "useFixedLengthSec", "minClipSec") +MARENGO_2_7_ONLY_FIELDS: Final = MappingProxyType({name: True for name in MARENGO_2_7_ONLY_PARAMS}) + + +def is_marengo_3_model(model: str | None) -> bool: + return MARENGO_3_MODEL_MARKER in (model or "") + + +class Marengo3Params(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + inputType: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None + input_type: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None + media_source: str | None = None + media_sources: Mapping[str, str] | None = None + bucketOwner: str | None = None + startSec: float | None = None + endSec: float | None = None + segmentation: TwelveLabsMarengo3Segmentation | None = None + embeddingOption: tuple[TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS, ...] | None = None + embeddingType: tuple[TWELVELABS_MARENGO_3_EMBEDDING_TYPES, ...] | None = None + embeddingScope: tuple[TWELVELABS_MARENGO_3_EMBEDDING_SCOPES, ...] | None = None + inferenceId: str | None = None + textTruncate: object = None + lengthSec: object = None + useFixedLengthSec: object = None + minClipSec: object = None + + @property + def resolved_input_type(self) -> TWELVELABS_MARENGO_3_INPUT_TYPES: + return self.inputType or self.input_type or "text" + + def timed_media_options(self) -> TwelveLabsMarengo3TimedMediaOptions: + return TIMED_MEDIA_OPTIONS.validate_python(self.given_timed_media_options()) + + def given_timed_media_options(self) -> dict[str, object]: + return self.model_dump(include=TIMED_MEDIA_OPTION_FIELDS, exclude_none=True) + + def given_2_7_only_params(self) -> dict[str, object]: + return self.model_dump(include=MARENGO_2_7_ONLY_FIELDS, exclude_none=True) + + +def _require_bucket_owner(bucket_owner: str | None) -> str: + if bucket_owner is None: + raise BedrockError( + status_code=400, + message="s3:// media requires the 'bucketOwner' parameter, the account id that owns the bucket", + ) + return bucket_owner + + +def _media_source(media: str, bucket_owner: str | None) -> TwelveLabsMediaSource: + if not media.startswith(S3_URI_PREFIX): + inline: Final[TwelveLabsMediaSource] = {"base64String": get_base64_str(media)} + return inline + s3_location: Final[TwelveLabsS3Location] = {"uri": media, "bucketOwner": _require_bucket_owner(bucket_owner)} + remote: Final[TwelveLabsMediaSource] = {"s3Location": s3_location} + return remote + + +def _named_media_source(name: str, media: str, bucket_owner: str | None) -> TwelveLabsMarengo3NamedMediaSource: + named: Final[TwelveLabsMarengo3NamedMediaSource] = { + "name": name, + "mediaType": "image", + **_media_source(media, bucket_owner), + } + return named + + +def _timed_media_input(media: str, params: Marengo3Params) -> TwelveLabsMarengo3TimedMediaInput: + timed: Final[TwelveLabsMarengo3TimedMediaInput] = { + "mediaSource": _media_source(media, params.bucketOwner), + **params.timed_media_options(), + } + return timed + + +def _describe(error: ValidationError) -> str: + return "; ".join( + f"{'.'.join(str(part) for part in problem['loc'])}: {problem['msg']}" for problem in error.errors() + ) + + +def _validated_params(inference_params: Mapping[str, object]) -> Marengo3Params: + try: + return Marengo3Params.model_validate(inference_params) + except ValidationError as error: + raise BedrockError(status_code=400, message=f"Invalid Marengo 3.0 parameters: {_describe(error)}") from error + + +def _reject_unless_dropped(given: Mapping[str, object], drop_params: bool, reason: str) -> None: + if not given or drop_params: + return + raise BedrockError(status_code=400, message=f"{reason} {', '.join(given)}; set drop_params to drop them") + + +def _require(value: str | None, input_type: str, param_name: str) -> str: + if value is None: + raise BedrockError(status_code=400, message=f"Input type '{input_type}' requires the '{param_name}' parameter") + return value + + +def _require_media_sources(value: Mapping[str, str] | None) -> Mapping[str, str]: + if not value: + raise BedrockError( + status_code=400, + message="Input type 'multi_input' requires a non-empty 'media_sources' mapping of name to media", + ) + return value + + +def _request_base(inference_id: str | None) -> TwelveLabsMarengo3RequestBase: + if inference_id is None: + anonymous: Final[TwelveLabsMarengo3RequestBase] = {} + return anonymous + identified: Final[TwelveLabsMarengo3RequestBase] = {"inferenceId": inference_id} + return identified + + +def build_marengo_3_request( + input: str, inference_params: Mapping[str, object], drop_params: bool = False +) -> TwelveLabsMarengo3EmbeddingRequest: + params: Final = _validated_params(inference_params) + base: Final = _request_base(params.inferenceId) + input_type: Final = params.resolved_input_type + _reject_unless_dropped( + params.given_2_7_only_params(), drop_params, "Marengo 3.0 does not accept the Marengo 2.7 parameters" + ) + if input_type not in TIMED_INPUT_TYPES: + _reject_unless_dropped( + params.given_timed_media_options(), drop_params, f"Input type '{input_type}' does not accept" + ) + match input_type: + case "text": + text_request: Final[TwelveLabsMarengo3TextRequest] = { + **base, + "inputType": "text", + "text": {"inputText": input}, + } + return text_request + case "image": + image_request: Final[TwelveLabsMarengo3ImageRequest] = { + **base, + "inputType": "image", + "image": {"mediaSource": _media_source(input, params.bucketOwner)}, + } + return image_request + case "video": + video_request: Final[TwelveLabsMarengo3VideoRequest] = { + **base, + "inputType": "video", + "video": _timed_media_input(input, params), + } + return video_request + case "audio": + audio_request: Final[TwelveLabsMarengo3AudioRequest] = { + **base, + "inputType": "audio", + "audio": _timed_media_input(input, params), + } + return audio_request + case "text_image": + text_image_request: Final[TwelveLabsMarengo3TextImageRequest] = { + **base, + "inputType": "text_image", + "text_image": { + "inputText": input, + "mediaSource": _media_source( + _require(params.media_source, input_type, "media_source"), params.bucketOwner + ), + }, + } + return text_image_request + case "multi_input": + media_sources: Final = tuple( + _named_media_source(name, media, params.bucketOwner) + for name, media in _require_media_sources(params.media_sources).items() + ) + multi_input_request: Final[TwelveLabsMarengo3MultiInputRequest] = { + **base, + "inputType": "multi_input", + "multi_input": {"inputText": input, "mediaSources": media_sources} + if input + else {"mediaSources": media_sources}, + } + return multi_input_request + case _: + assert_never(input_type) diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py index a39c59b0efd..65ca2be191f 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py @@ -4,19 +4,120 @@ Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Mar Why separate file? Make it easy to see how transformation works Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html +Marengo 3.0 docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo-3.html """ +from collections.abc import Mapping from typing import Final, cast +from pydantic import BaseModel, ConfigDict, TypeAdapter +from typing_extensions import assert_never + +import litellm +from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( + MARENGO_2_7_ONLY_PARAMS, + build_marengo_3_request, + is_marengo_3_model, +) from litellm.types.llms.bedrock import ( TWELVELABS_EMBEDDING_INPUT_TYPES, + TWELVELABS_MARENGO_3_INPUT_TYPES, TwelveLabsAsyncInvokeRequest, + TwelveLabsMarengo3EmbeddingRequest, TwelveLabsMarengoEmbeddingRequest, TwelveLabsOutputDataConfig, TwelveLabsS3Location, TwelveLabsS3OutputDataConfig, ) -from litellm.types.utils import Embedding, EmbeddingResponse, Usage +from litellm.types.utils import Embedding, EmbeddingResponse, PromptTokensDetailsWrapper, Usage + + +class MarengoEmbeddingItem(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + embedding: tuple[float, ...] | None = None + + +class MarengoInvokeResponse(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + data: tuple[MarengoEmbeddingItem, ...] = () + embedding: tuple[float, ...] | None = None + embeddings: tuple[MarengoEmbeddingItem, ...] = () + + def vectors(self) -> tuple[tuple[float, ...], ...]: + if self.data: + return tuple(item.embedding for item in self.data if item.embedding is not None) + if self.embedding is not None: + return (self.embedding,) + return tuple(item.embedding for item in self.embeddings if item.embedding is not None) + + +class MarengoBilledMultiInput(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + inputText: str | None = None + mediaSources: tuple[Mapping[str, object], ...] = () + + +class MarengoBilledRequest(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + inputType: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None + multi_input: MarengoBilledMultiInput | None = None + + +INVOKE_RESPONSES: Final = TypeAdapter(tuple[MarengoInvokeResponse, ...]) +BILLED_REQUESTS: Final = TypeAdapter(tuple[MarengoBilledRequest, ...]) + + +def _billed_units(request: MarengoBilledRequest) -> tuple[int, int]: + input_type: Final = request.inputType + match input_type: + case "text": + return (1, 0) + case "image": + return (0, 1) + case "text_image": + return (1, 1) + case "multi_input": + multi_input: Final = request.multi_input or MarengoBilledMultiInput() + return (1 if multi_input.inputText else 0, len(multi_input.mediaSources)) + case "video" | "audio" | None: + return (0, 0) + case _: + assert_never(input_type) + + +def _billed_usage(batch_data: list[dict] | None) -> Usage: + units: Final = tuple(_billed_units(request) for request in BILLED_REQUESTS.validate_python(batch_data or ())) + query_count: Final = sum(text_requests for text_requests, _ in units) + image_count: Final = sum(images for _, images in units) + details: Final = ( + PromptTokensDetailsWrapper(query_count=query_count or None, image_count=image_count or None) + if query_count or image_count + else None + ) + return Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0, prompt_tokens_details=details) + + +MARENGO_SHARED_PARAMS: Final = ( + "encoding_format", + "embeddingOption", + "startSec", + "input_type", + "endSec", + "segmentation", + "embeddingType", + "embeddingScope", + "inferenceId", + "media_source", + "media_sources", +) + + +def drop_params_enabled(litellm_params: Mapping[str, object]) -> bool: + return litellm.drop_params is True or litellm_params.get("drop_params") is True class TwelveLabsMarengoEmbeddingConfig: @@ -26,28 +127,24 @@ class TwelveLabsMarengoEmbeddingConfig: Supports text, image, video, and audio inputs. - InvokeModel: text and image inputs - StartAsyncInvoke: video, audio, image, and text inputs + + Marengo 3.0 (model ids containing "marengo-embed-3") nests the input under a key named after inputType and + adds the text_image and multi_input input types; that payload is built by build_marengo_3_request. """ - def __init__(self) -> None: - pass + def __init__(self, model: str | None = None) -> None: + self.is_marengo_3: Final = is_marengo_3_model(model) def get_supported_openai_params(self) -> list[str]: - return [ - "encoding_format", - "textTruncate", - "embeddingOption", - "startSec", - "lengthSec", - "useFixedLengthSec", - "minClipSec", - "input_type", - ] + if self.is_marengo_3: + return list(MARENGO_SHARED_PARAMS) + return [*MARENGO_SHARED_PARAMS, *MARENGO_2_7_ONLY_PARAMS] def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): if k == "encoding_format": # TwelveLabs doesn't have encoding_format, but we can map it to embeddingOption - if v == "float": + if v == "float" and not self.is_marengo_3: optional_params["embeddingOption"] = ["visual-text", "visual-image"] elif k == "textTruncate": optional_params["textTruncate"] = v @@ -56,7 +153,19 @@ class TwelveLabsMarengoEmbeddingConfig: elif k == "input_type": # Map input_type to inputType for Bedrock optional_params["inputType"] = v - elif k in ["startSec", "lengthSec", "useFixedLengthSec", "minClipSec"]: + elif k in ( + "startSec", + "lengthSec", + "useFixedLengthSec", + "minClipSec", + "endSec", + "segmentation", + "embeddingType", + "embeddingScope", + "inferenceId", + "media_source", + "media_sources", + ): optional_params[k] = v return optional_params @@ -77,7 +186,8 @@ class TwelveLabsMarengoEmbeddingConfig: async_invoke_route: bool = False, model_id: str | None = None, output_s3_uri: str | None = None, - ) -> TwelveLabsMarengoEmbeddingRequest | TwelveLabsAsyncInvokeRequest: + drop_params: bool = False, + ) -> TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest | TwelveLabsAsyncInvokeRequest: """ Transform OpenAI-style input to TwelveLabs Marengo format/async-invoke format. @@ -87,20 +197,29 @@ class TwelveLabsMarengoEmbeddingConfig: - Video inputs (async-invoke only) - Audio inputs (async-invoke only) - S3 URLs for all media types (async-invoke only) + - Marengo 3.0 only: text_image and multi_input inputs (nested payload) """ - # Get input_type or default to "text" input_type: Final = cast( TWELVELABS_EMBEDDING_INPUT_TYPES, inference_params.get("inputType") or inference_params.get("input_type") or "text", ) - # Validate that async-invoke is used for video/audio if input_type in ["video", "audio"] and not async_invoke_route: raise ValueError( f"Input type '{input_type}' requires async_invoke route. " f"Use model format: 'bedrock/async_invoke/model_id'" ) + if self.is_marengo_3: + marengo_3_request: Final = build_marengo_3_request( + input=input, inference_params=inference_params, drop_params=drop_params + ) + if async_invoke_route and model_id: + return self._wrap_async_invoke_request( + model_input=marengo_3_request, model_id=model_id, output_s3_uri=output_s3_uri + ) + return marengo_3_request + transformed_request: Final[TwelveLabsMarengoEmbeddingRequest] = {"inputType": input_type} if input_type == "text": @@ -154,7 +273,7 @@ class TwelveLabsMarengoEmbeddingConfig: def _wrap_async_invoke_request( self, - model_input: TwelveLabsMarengoEmbeddingRequest, + model_input: TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest, model_id: str, output_s3_uri: str | None = None, ) -> TwelveLabsAsyncInvokeRequest: @@ -188,62 +307,16 @@ class TwelveLabsMarengoEmbeddingConfig: ), ) - def _transform_response(self, response_list: list[dict], model: str) -> EmbeddingResponse: - """ - Transform TwelveLabs response to OpenAI format. - Handles the actual TwelveLabs response format: {"data": [{"embedding": [...]}]} - """ - embeddings: Final[list[Embedding]] = [] - total_tokens = 0 - - for response in response_list: - # TwelveLabs response format has a "data" field containing the embeddings - if "data" in response and isinstance(response["data"], list): - for item in response["data"]: - if "embedding" in item: - # Single embedding response - embedding = Embedding( - embedding=item["embedding"], - index=len(embeddings), - object="embedding", - ) - embeddings.append(embedding) - - # Estimate token count (rough approximation) - if "inputTextTokenCount" in item: - total_tokens += item["inputTextTokenCount"] - else: - # Rough estimate: 1 token per 4 characters for text, or use embedding size - total_tokens += len(item["embedding"]) // 4 - elif "embedding" in response: - # Direct embedding response (fallback for other formats) - embedding = Embedding( - embedding=response["embedding"], - index=len(embeddings), - object="embedding", - ) - embeddings.append(embedding) - - # Estimate token count (rough approximation) - if "inputTextTokenCount" in response: - total_tokens += response["inputTextTokenCount"] - else: - # Rough estimate: 1 token per 4 characters for text - total_tokens += len(response.get("inputText", "")) // 4 - elif "embeddings" in response: - # Multiple embeddings response (from video/audio) - for i, emb in enumerate(response["embeddings"]): - embedding = Embedding( - embedding=emb["embedding"], - index=len(embeddings), - object="embedding", - ) - embeddings.append(embedding) - total_tokens += len(emb["embedding"]) // 4 # Rough estimate - - usage: Final = Usage(prompt_tokens=total_tokens, total_tokens=total_tokens) - - return EmbeddingResponse(data=embeddings, model=model, usage=usage) + def _transform_response( + self, response_list: list[dict], model: str, batch_data: list[dict] | None = None + ) -> EmbeddingResponse: + vectors: Final = tuple( + vector for response in INVOKE_RESPONSES.validate_python(response_list) for vector in response.vectors() + ) + embeddings: Final = [ + Embedding(embedding=list(vector), index=index, object="embedding") for index, vector in enumerate(vectors) + ] + return EmbeddingResponse(data=embeddings, model=model, usage=_billed_usage(batch_data)) def _transform_async_invoke_response(self, response: dict, model: str) -> EmbeddingResponse: """ diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 33b27943ad8..9875ac2b9c3 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -7,13 +7,13 @@ from contextlib import suppress from functools import cache from itertools import chain from types import MappingProxyType -from typing import Any, Final, TypeAlias, TypedDict +from typing import Any, Final, Literal, TypeAlias, TypedDict from urllib.parse import unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from typing_extensions import ReadOnly from litellm._logging import verbose_logger @@ -60,11 +60,12 @@ from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id -# litellm_params key used to hand the SigV4-signed GET headers from -# `transform_file_content_request` to `validate_environment` (the only hook -# the shared file-content HTTP handler exposes for setting request headers). -# Same pattern as the `upload_url` handoff in `transform_create_file_request`. -S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers" +S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" + + +class _S3DeleteContext(BaseModel): + file_id: str = Field(min_length=1) + # litellm_params key carrying the size of the body uploaded to S3, handed from # `transform_create_file_request` to `transform_create_file_response`. @@ -291,7 +292,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) -> dict: result: Final[dict[str, object]] = {} result.update(headers) - signed_headers: Final = litellm_params.pop(S3_SIGNED_GET_HEADERS_PARAM, None) + signed_headers: Final = litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM, None) if isinstance(signed_headers, Mapping): result.update(signed_headers) # any-ok: untyped handoff headers # otherwise no extra headers - AWS credentials are handled by BaseAWSLLM @@ -1187,18 +1188,27 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): def transform_delete_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - raise NotImplementedError("BedrockFilesConfig does not support file deletion") + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: + return self._transform_s3_file_request( + file_id=file_id, method="DELETE", optional_params=optional_params, litellm_params=litellm_params + ) def transform_delete_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> FileDeleted: - raise NotImplementedError("BedrockFilesConfig does not support file deletion") + if raw_response.status_code != 204: + raise BedrockError( + status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, + message=raw_response.text or f"S3 file deletion returned HTTP {raw_response.status_code}", + headers=raw_response.headers, + ) + context: Final = _S3DeleteContext.model_validate(logging_obj.model_call_details.get("additional_args")) + return FileDeleted(id=context.file_id, deleted=True, object="file") def transform_list_files_request( self, @@ -1233,6 +1243,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if not file_id: raise ValueError("file_id is required for Bedrock file content retrieval") + return self._transform_s3_file_request( + file_id=file_id, method="GET", optional_params=optional_params, litellm_params=litellm_params + ) + + def _transform_s3_file_request( + self, + *, + file_id: str, + method: Literal["GET", "DELETE"], + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: s3_uri: Final = extract_s3_uri_from_file_id(file_id) bucket_name, object_key = _validate_file_id_against_configured_buckets( s3_uri=s3_uri, @@ -1240,40 +1262,32 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) - # The shared file-content handler passes optional_params={}, so AWS - # credentials/region arrive via litellm_params here (unlike the upload - # path). s3_region_name wins over aws_region_name, same priority as - # get_complete_file_url above. - merged_params: Final[dict[str, object]] = {} - merged_params.update(litellm_params) - merged_params.update(optional_params) - request_params: Final = _BedrockS3RequestParams.model_validate(merged_params) + request_params: Final = _BedrockS3RequestParams.model_validate({**litellm_params, **optional_params}) region_preference: Final = request_params.s3_region_name or request_params.aws_region_name region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference} aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="") - s3_endpoint_url = ( + s3_endpoint_url: Final = ( request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" ).rstrip("/") url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}" - litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request( + litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = self._sign_s3_request_without_body( api_base=url, aws_region_name=aws_region_name, request_params=request_params, + method=method, ) return url, {} - def _sign_s3_get_request( + def _sign_s3_request_without_body( self, api_base: str, aws_region_name: str, request_params: _BedrockS3RequestParams, + method: Literal["GET", "DELETE"] = "GET", ) -> dict[str, str]: - """ - SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT). - """ try: import hashlib @@ -1297,7 +1311,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): empty_body_hash: Final = hashlib.sha256(b"").hexdigest() aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped - method="GET", + method=method, url=api_base, headers={"x-amz-content-sha256": empty_body_hash}, ) diff --git a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py index 18d47301ee5..acb0cc8dcb7 100644 --- a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py +++ b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py @@ -20,6 +20,7 @@ import httpx from litellm._logging import verbose_logger from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import FileTypes, ImageObject, ImageResponse @@ -228,6 +229,14 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): """ return _supports_nova_canvas_image_edit_from_model_cost(model or "") + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_supported_openai_params(self, model: str) -> list: return [ "n", diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 5c517f2049c..be6489f20ae 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -114,7 +114,12 @@ class BedrockImageEdit(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -156,7 +161,12 @@ class BedrockImageEdit(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index 24e7ba73075..bc9a64f587a 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.stability import ( OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO, @@ -84,6 +85,14 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): return True return False + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_supported_openai_params(self, model: str) -> list: """ Return list of OpenAI params supported by Bedrock Stability. diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index c78e3c147cb..87762b648e0 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -119,7 +119,12 @@ class BedrockImageGeneration(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") ### FORMAT RESPONSE TO OPENAI FORMAT ### @@ -162,7 +167,12 @@ class BedrockImageGeneration(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 6ff9f0155f9..a715d150b4c 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -29,6 +29,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( + BedrockError, apply_bedrock_invoke_structured_output, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, @@ -79,6 +80,14 @@ class AmazonAnthropicClaudeMessagesConfig( BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys()) + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def __init__(self, **kwargs): BaseAnthropicMessagesConfig.__init__(self, **kwargs) AmazonInvokeConfig.__init__(self, **kwargs) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index d0a3c37ffb3..6a120f41cb6 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -1,23 +1,138 @@ import json -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional, cast +import httpx from httpx import Response +from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector +from litellm.types.utils import ModelResponseStream from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo +from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo if TYPE_CHECKING: + from botocore.eventstream import EventStreamMessage from httpx import URL from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder from litellm.types.utils import CostResponseTypes +_TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"}) + + +def _plain_text_delta(chunk: ModelResponseStream) -> str | None: + """Return the delta text when the chunk carries nothing else that stream_chunk_builder reads.""" + if chunk.get("usage") is not None or chunk.provider_specific_fields or len(chunk.choices) != 1: + return None + choice: Final = chunk.choices[0] + if choice.finish_reason or choice.logprobs is not None: + return None + populated: Final = frozenset(key for key, value in choice.delta.model_dump().items() if value is not None) + if not populated <= _TEXT_ONLY_DELTA_FIELDS: + return None + content: Final = choice.delta.get("content") + return content if isinstance(content, str) else None + + +class _CoalescedChunks: + """Retains translated chunks with consecutive text deltas folded into one, so memory tracks the response text, + not the event count.""" + + def __init__(self) -> None: + self._chunks: list[ModelResponseStream] = [] # mutable-ok: instance accumulator for streaming chunks + self._open_text_parts: list[str] = [] # mutable-ok: text deltas pending fold into self._chunks[-1] + + def add(self, chunk: ModelResponseStream) -> None: + text: Final = _plain_text_delta(chunk) + if text is not None and self._open_text_parts: + self._open_text_parts.append(text) + return + self._seal_text_run() + self._chunks.append(chunk) + if text is not None: + self._open_text_parts.append(text) + + def _seal_text_run(self) -> None: + if len(self._open_text_parts) > 1: + self._chunks[-1].choices[0].delta.content = "".join(self._open_text_parts) + self._open_text_parts.clear() + + def chunks(self) -> Sequence[ModelResponseStream]: + self._seal_text_run() + return self._chunks + + +def _translate_message(decoder: "AWSEventStreamDecoder", message: str) -> ModelResponseStream | None: + from litellm.litellm_core_utils.streaming_handler import ( + convert_generic_chunk_to_model_response_stream, + generic_chunk_has_all_required_fields, + ) + from litellm.types.utils import GenericStreamingChunk + + translated_chunk: Final = decoder._chunk_parser(chunk_data=json.loads(message)) + if isinstance(translated_chunk, ModelResponseStream): + return translated_chunk + if generic_chunk_has_all_required_fields(cast(dict, translated_chunk)): + return convert_generic_chunk_to_model_response_stream(cast(GenericStreamingChunk, translated_chunk)) + return None + + +def _build_logged_response( + chunks: Sequence[ModelResponseStream], litellm_logging_obj: "LiteLLMLoggingObj" +) -> Optional["CostResponseTypes"]: + from litellm.main import stream_chunk_builder + + if len(chunks) == 0: + return None + return stream_chunk_builder(chunks=list(chunks), logging_obj=litellm_logging_obj) + + +class BedrockEventStreamCollector: + """Decodes and translates Bedrock event-stream frames as they are relayed instead of buffering the stream.""" + + def __init__( + self, + parse_event: Callable[["EventStreamMessage"], str | None], + decoder: Optional["AWSEventStreamDecoder"], + ) -> None: + from botocore.eventstream import EventStreamBuffer + + self._parse_event = parse_event + self._decoder = decoder + self._event_stream_buffer: Final[EventStreamBuffer] = EventStreamBuffer() + self._chunks: Final = _CoalescedChunks() + + def add(self, chunk: bytes) -> None: + if self._decoder is None: + return + self._event_stream_buffer.add_data(chunk) + for event in self._event_stream_buffer: + self._add_event(self._decoder, event) + + def _add_event(self, decoder: "AWSEventStreamDecoder", event: "EventStreamMessage") -> None: + message: Final = self._parse_event(event) + translated: Final = _translate_message(decoder, message) if message is not None else None + if translated is not None: + self._chunks.add(translated) + + def build_logged_response(self, litellm_logging_obj: "LiteLLMLoggingObj") -> Optional["CostResponseTypes"]: + return _build_logged_response(self._chunks.chunks(), litellm_logging_obj) + + class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig): + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in endpoint @@ -159,87 +274,32 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD return litellm_model_response - def _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]: - from botocore.eventstream import EventStreamBuffer - - all_chunks: Final = [] - event_stream_buffer: Final = EventStreamBuffer() - for chunk in raw_bytes: - event_stream_buffer.add_data(chunk) - for event in event_stream_buffer: - message = self._parse_message_from_event(event) - if message is not None: - all_chunks.append(message) - - return all_chunks - - def handle_logging_collected_chunks( - self, - all_chunks: list[str], - litellm_logging_obj: "LiteLLMLoggingObj", - model: str, - custom_llm_provider: str, - endpoint: str, - ) -> Optional["CostResponseTypes"]: - """ - 1. Convert all_chunks to a ModelResponseStream - 2. combine model_response_stream to model_response - 3. Return the model_response - """ - - from litellm.litellm_core_utils.streaming_handler import ( - convert_generic_chunk_to_model_response_stream, - generic_chunk_has_all_required_fields, + def create_stream_collector( + self, model: str, custom_llm_provider: str, endpoint: str + ) -> PassthroughStreamCollector: + return BedrockEventStreamCollector( + parse_event=self._parse_message_from_event, + decoder=self._get_event_stream_decoder(model=model, endpoint=endpoint), ) + + def _get_event_stream_decoder(self, model: str, endpoint: str) -> Optional["AWSEventStreamDecoder"]: from litellm.llms.bedrock.chat import get_bedrock_event_stream_decoder from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, ) - from litellm.main import stream_chunk_builder - from litellm.types.utils import GenericStreamingChunk, ModelResponseStream - all_translated_chunks: Final = [] if "invoke" in endpoint: invoke_provider: Final = AmazonInvokeConfig.get_bedrock_invoke_provider(model) if invoke_provider is None: - raise ValueError(f"Invalid invoke provider: {invoke_provider}, for model: {model}") - obj = get_bedrock_event_stream_decoder( - invoke_provider=invoke_provider, - model=model, - sync_stream=True, - json_mode=False, - ) - elif "converse" in endpoint: - obj = get_bedrock_event_stream_decoder( - invoke_provider=None, - model=model, - sync_stream=True, - json_mode=False, - ) - else: - return None - - for chunk in all_chunks: - message = json.loads(chunk) - translated_chunk = obj._chunk_parser(chunk_data=message) - - if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields( - cast(dict, translated_chunk) - ): - chunk_obj = convert_generic_chunk_to_model_response_stream( - cast(GenericStreamingChunk, translated_chunk) + verbose_logger.warning( + "Bedrock passthrough spend tracking skipped: no invoke provider for model %s", model ) - elif isinstance(translated_chunk, ModelResponseStream): - chunk_obj = translated_chunk - else: - continue - - all_translated_chunks.append(chunk_obj) - - if len(all_translated_chunks) > 0: - model_response: Final = stream_chunk_builder( - chunks=all_translated_chunks, - logging_obj=litellm_logging_obj, + return None + return get_bedrock_event_stream_decoder( + invoke_provider=invoke_provider, model=model, sync_stream=True, json_mode=False + ) + if "converse" in endpoint: + return get_bedrock_event_stream_decoder( + invoke_provider=None, model=model, sync_stream=True, json_mode=False ) - return model_response return None diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 42fe8941443..ca2370303f2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -21,7 +21,7 @@ from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeE from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig @@ -149,7 +149,8 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model) - credentials: Final = self.get_credentials( + credentials: Final = await run_aws_signing( + self.get_credentials, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, @@ -169,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM): "or configure credentials in the environment" ), ) - frozen_credentials: Final = credentials.get_frozen_credentials() + frozen_credentials: Final = await run_aws_signing(credentials.get_frozen_credentials) # Initialize Bedrock client with aws_sdk_bedrock_runtime config: Final = Config( diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 1f4c81d6491..3b972961940 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -9,12 +9,14 @@ import json import uuid as uuid_lib from typing import Final, cast +import httpx from pydantic import BaseModel from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm from litellm.types.llms.openai import ( OpenAIRealtimeContentPartDone, @@ -121,6 +123,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self._cumulative_usage = BedrockUsageEvent() self._reported_usage = BedrockUsageEvent() + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" return headers diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 4860c99268e..8847381cbc9 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -46,7 +46,12 @@ class BedrockRerankHandler(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -117,7 +122,12 @@ class BedrockRerankHandler(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 920e566c9dd..e7d706c3731 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -39,7 +39,6 @@ from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -380,6 +379,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): raise BedrockError( status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, message=f"AgentCore gateway MCP error: {error}", + headers=raw_response.headers, ) # A failed tools/call is reported in-band, as HTTP 200 with result.isError @@ -389,6 +389,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): raise BedrockError( status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}", + headers=raw_response.headers, ) text_items: Final = tuple( @@ -440,6 +441,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): raise BedrockError( status_code=502, message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}", + headers=raw_response.headers, ) def get_error_class( @@ -448,7 +450,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): status_code: int, headers: dict, # mutable-ok: BaseSearchConfig.get_error_class takes the response headers as a dict ) -> Exception: - return BaseLLMException( + return BedrockError( status_code=status_code, message=error_message, headers=headers, diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 6940077391f..27c90c9d71e 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBContent, BedrockKBResponse, @@ -38,6 +39,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): BaseVectorStoreConfig.__init__(self) BaseAWSLLM.__init__(self) + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: return {} diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 64d7ef2bed6..d91157c3d10 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -13,6 +13,8 @@ Auth: Bearer token (litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the from collections.abc import AsyncIterator, Iterator from typing import Any, Final +import httpx + import litellm from litellm._logging import verbose_logger from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -24,6 +26,8 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams +from ...base_llm.chat.transformation import BaseLLMException +from ...bedrock.common_utils import BedrockError from ...openai_like.chat.transformation import OpenAILikeChatConfig from ..common_utils import mantle_base_segment @@ -45,6 +49,11 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): def get_config(cls): return super().get_config() + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers + ) -> BaseLLMException: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def _get_openai_compatible_provider_info( self, api_base: str | None, diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 850738bc320..ac94d9cb922 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -23,7 +23,7 @@ from botocore.exceptions import ( ProfileNotFound, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS from litellm.secret_managers.main import get_secret_str BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1" @@ -55,7 +55,7 @@ def resolve_mantle_region(params: Mapping[str, object]) -> str: ) -class BedrockMantleAuthMixin: +class BedrockMantleAuthMixin(SignsRequestsWithAWS): _aws_signer: BaseAWSLLM @staticmethod diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 5179c966584..bbbda4d14b6 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -19,11 +19,14 @@ import json from collections.abc import Mapping from typing import Any, Final +import httpx from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock_mantle.common_utils import ( MANTLE_HOST_RE, BedrockMantleAuthMixin, @@ -98,6 +101,11 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK_MANTLE + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers + ) -> BaseLLMException: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_complete_url( self, api_base: str | None, diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 62b631a7671..ddab4f54d57 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -9,6 +9,7 @@ API Reference: https://docs.bfl.ai/ import base64 import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -16,7 +17,7 @@ from httpx._types import RequestFiles import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH -from litellm.litellm_core_utils.url_utils import safe_get +from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -37,6 +38,22 @@ else: LiteLLMLoggingObj = Any +_BFL_REQUEST_PARAMS: Final = ( + "seed", + "output_format", + "safety_tolerance", + "prompt_upsampling", + "aspect_ratio", + "steps", + "guidance", + "grow_mask", + "top", + "bottom", + "left", + "right", +) + + class BlackForestLabsImageEditConfig(BaseImageEditConfig): """ Configuration for Black Forest Labs image editing. @@ -85,34 +102,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): BFL-specific params are passed through directly. """ optional_params: Final[dict[str, object]] = {} - - # Pass through BFL-specific params - bfl_params: Final = [ - "seed", - "output_format", - "safety_tolerance", - "prompt_upsampling", - # Kontext-specific - "aspect_ratio", - # Fill/Inpaint-specific - "steps", - "guidance", - "grow_mask", - # Expand-specific - "top", - "bottom", - "left", - "right", - ] - - # Convert TypedDict to regular dict for access - params_dict: Final = dict(image_edit_optional_params) - - for param in bfl_params: - if param in params_dict: - value = params_dict[param] - if value is not None: - optional_params[param] = value + params: Final[Mapping[str, object]] = image_edit_optional_params + for param in _BFL_REQUEST_PARAMS: + if (value := params.get(param)) is not None: + optional_params[param] = value # Set default output format if "output_format" not in optional_params: @@ -251,23 +244,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): "input_image": b64_image, } - # Add optional params (only BFL-recognized parameters) - bfl_request_params: Final = [ - "seed", - "output_format", - "safety_tolerance", - "prompt_upsampling", - "aspect_ratio", - "steps", - "guidance", - "grow_mask", - "top", - "bottom", - "left", - "right", - ] for key, value in image_edit_optional_request_params.items(): - if key in bfl_request_params and value is not None: + if key in _BFL_REQUEST_PARAMS and value is not None: request_body[key] = value # Handle mask if provided (for inpainting) @@ -277,7 +255,39 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8") # BFL uses JSON, not multipart - return empty files - return request_body, [] + return request_body, () + + async def async_transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | None, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict, RequestFiles]: + downloaded_image: Final = await self._fetch_remote_image(image) + downloaded_mask: Final = await self._fetch_remote_image(image_edit_optional_request_params.get("mask")) + return self.transform_image_edit_request( + model=model, + prompt=prompt, + image=image if downloaded_image is None else downloaded_image, + image_edit_optional_request_params=( + dict(image_edit_optional_request_params) + if downloaded_mask is None + else {**image_edit_optional_request_params, "mask": downloaded_mask} + ), + litellm_params=litellm_params, + headers=dict(headers), + ) + + async def _fetch_remote_image(self, image: object) -> bytes | None: + candidate: Final = image[0] if isinstance(image, list) and image else image + if not isinstance(candidate, str) or not candidate.startswith(("http://", "https://")): + return None + response: Final = await async_safe_get(litellm.module_level_aclient, candidate, timeout=60.0) + response.raise_for_status() + return response.content def transform_image_edit_response( self, diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e1f0fc9e7d3..f4883b57fbc 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -9,6 +9,7 @@ import threading import time from collections.abc import AsyncIterable, Callable, Iterable, Mapping from http.cookiejar import CookieJar, DefaultCookiePolicy +from io import BytesIO from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict, TypeVar @@ -505,6 +506,10 @@ async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> N raise MaskedHTTPStatusError(e, message=_text, text=_text) from None +class HTTPResponseLimitError(ValueError): + pass + + class MaskedHTTPStatusError(httpx.HTTPStatusError): def __init__(self, original_error, message: str | None = None, text: str | None = None): # Create a new error with the masked URL @@ -654,6 +659,7 @@ class AsyncHTTPHandler: headers: dict | None = None, follow_redirects: bool | None = None, timeout: float | httpx.Timeout | None = None, + max_response_bytes: int | None = None, ): # Set follow_redirects to UseClientDefault if None _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT @@ -661,6 +667,16 @@ class AsyncHTTPHandler: params = params or {} params.update(HTTPHandler.extract_query_params(url)) + if max_response_bytes is not None: + return await self._get_with_response_limit( + url, + params=httpx.QueryParams(params), + headers=httpx.Headers(headers), + max_bytes=max_response_bytes, + follow_redirects=self.client.follow_redirects if follow_redirects is None else follow_redirects, + timeout=self.client.timeout if timeout is None else httpx.Timeout(timeout), + ) + response: Final = await self.client.get( url, params=params, @@ -670,6 +686,57 @@ class AsyncHTTPHandler: ) return response + async def _get_with_response_limit( + self, + url: str, + *, + params: httpx.QueryParams, + headers: httpx.Headers, + timeout: httpx.Timeout, + max_bytes: int, + follow_redirects: bool, + ) -> httpx.Response: + request: Final = self.client.build_request( + "GET", + url, + headers=MappingProxyType({**headers, "accept-encoding": "identity"}), + params=params, + timeout=timeout, + ) + response: Final = await self.client.send(request, stream=True, follow_redirects=False) + return await self._read_with_response_limit(response, max_bytes=max_bytes, follow_redirects=follow_redirects) + + async def _read_with_response_limit( + self, response: httpx.Response, *, max_bytes: int, follow_redirects: bool, redirects_remaining: int = 10 + ) -> httpx.Response: + try: + if response.next_request is not None and follow_redirects: + if redirects_remaining == 0: + raise ValueError("Too many redirects") + await response.aclose() + following: Final = await self.client.send( + response.next_request, auth=None, stream=True, follow_redirects=False + ) + return await self._read_with_response_limit( + following, max_bytes=max_bytes, follow_redirects=True, redirects_remaining=redirects_remaining - 1 + ) + if response.is_redirect or response.is_error: + return httpx.Response(response.status_code, headers=response.headers, request=response.request) + if response.headers.get("content-encoding", "identity").lower() != "identity": + raise HTTPResponseLimitError("Response size limits require an uncompressed response") + if int(response.headers.get("content-length", "0")) > max_bytes: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + with BytesIO() as body: + async for chunk in response.aiter_bytes(chunk_size=65536): + if body.tell() + len(chunk) > max_bytes: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + body.write(chunk) + return httpx.Response( + response.status_code, headers=response.headers, content=body.getvalue(), request=response.request + ) + finally: + await response.aclose() + @track_llm_api_timing() async def post( self, @@ -751,7 +818,9 @@ class AsyncHTTPHandler: timeout: float | httpx.Timeout | None = None, stream: bool = False, content: _RequestContent | None = None, + follow_redirects: bool | None = None, ): + _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT try: if timeout is None: timeout = self.timeout @@ -769,22 +838,30 @@ class AsyncHTTPHandler: timeout=timeout, content=request_content, ) - response: Final = await self.client.send(req) + response: Final = await self.client.send(req, follow_redirects=_follow_redirects) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): # Retry the request with a new session if there is a connection error new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks) try: - return await self.single_connection_post_request( - url=url, - client=new_client, - data=data, + retry_data, retry_content = _prepare_request_data_and_content(data, content) + retry: Final = new_client.build_request( + "PUT", + url, + data=retry_data, json=json, params=params, headers=headers, - stream=stream, + timeout=timeout, + content=retry_content, ) + retried: Final = await new_client.send(retry, stream=stream, follow_redirects=_follow_redirects) + try: + retried.raise_for_status() + except httpx.HTTPStatusError as retried_error: + await _raise_masked_async_error(retried_error, stream) + return retried finally: await new_client.aclose() except httpx.TimeoutException as e: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f281c249c72..7109e6942d1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -77,6 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, sign_request_off_loop_if_aws from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -116,6 +117,7 @@ from litellm.types.llms.anthropic_skills import ( Skill, ) from litellm.types.llms.openai import ( + AllMessageValues, CreateBatchRequest, CreateFileRequest, FileContentRequest, @@ -163,13 +165,10 @@ from litellm.utils import ( def _rust_responses_websocket_enabled( custom_llm_provider: str | None, - litellm_params: GenericLiteLLMParams, ) -> bool: from litellm.rust_bridge.configuration import rust_enabled - raw_request_override: Final = litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - return custom_llm_provider == "openai" and rust_enabled(request_override=request_override) + return custom_llm_provider == "openai" and rust_enabled() from .http_handler import get_shared_realtime_ssl_context @@ -488,7 +487,7 @@ class BaseLLMHTTPHandler: def completion( self, model: str, - messages: list, + messages: list[AllMessageValues], api_base: str | None, custom_llm_provider: str, model_response: ModelResponse, @@ -507,7 +506,7 @@ class BaseLLMHTTPHandler: shared_session: Optional["ClientSession"] = None, ): json_mode: Final[bool] = optional_params.pop("json_mode", False) - extra_body: Final[dict | None] = optional_params.pop("extra_body", None) + extra_body: Final[Mapping[str, object] | None] = optional_params.pop("extra_body", None) provider_config = provider_config or ProviderConfigManager.get_provider_chat_config( model=model, provider=litellm.LlmProviders(custom_llm_provider) @@ -522,14 +521,17 @@ class BaseLLMHTTPHandler: ) # get config from model, custom llm provider - headers = provider_config.validate_environment( - api_key=api_key, - headers=headers or {}, - model=model, - messages=messages, - optional_params=optional_params, - api_base=api_base, - litellm_params=litellm_params, + request_headers: Final = cast( # cast-ok: validate_environment is declared as a bare dict + "dict[str, object]", + provider_config.validate_environment( + api_key=api_key, + headers=headers or {}, + model=model, + messages=messages, + optional_params=optional_params, + api_base=api_base, + litellm_params=litellm_params, + ), ) api_base = provider_config.get_complete_url( @@ -541,93 +543,122 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) - data: dict[str, object] = provider_config.transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=headers, - ) - - if extra_body is not None: - data = {**data, **extra_body} - - headers, signed_json_body = provider_config.sign_request( - headers=headers, - optional_params={ - **optional_params, - **_aws_signing_overrides(optional_params, litellm_params), - }, - request_data=data, - api_base=api_base, - api_key=api_key, - stream=stream, - fake_stream=fake_stream, - model=model, - ) - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key=api_key, - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": headers, - }, - ) - - # Check if stream was converted for WebSearch interception - # This is set by the async_pre_request_hook in WebSearchInterceptionLogger - if litellm_params.get("_websearch_interception_converted_stream", False): - logging_obj.model_call_details["websearch_interception_converted_stream"] = True - - if acompletion is True: - if stream is True: - data = self._add_stream_param_to_request_body( - data=data, - provider_config=provider_config, + def sign_and_log( + transformed: dict[str, object], # mutable-ok: async_completion takes dict + ) -> tuple[dict[str, object], dict[str, object], bytes | None]: # mutable-ok: async_completion takes dict + data: Final = {**transformed, **extra_body} if extra_body is not None else transformed + signed: Final = cast( # cast-ok: sign_request is declared as a bare dict + "tuple[dict[str, object], bytes | None]", + provider_config.sign_request( + headers=request_headers, + optional_params={ + **optional_params, + **_aws_signing_overrides(optional_params, litellm_params), + }, + request_data=data, + api_base=api_base, + api_key=api_key, + stream=stream, fake_stream=fake_stream, - ) + model=model, + ), + ) + logging_obj.pre_call( + input=messages, + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": signed[0], + }, + ) + if litellm_params.get("_websearch_interception_converted_stream", False): + logging_obj.model_call_details["websearch_interception_converted_stream"] = True + return data, signed[0], signed[1] + + def dispatch_async( + data: dict[str, object], # mutable-ok: async_completion takes dict + signed_headers: dict[str, object], # mutable-ok: async_completion takes dict + signed_json_body: bytes | None, + ) -> Coroutine[object, object, ModelResponse | CustomStreamWrapper]: + async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None + if stream is True: return self.acompletion_stream_function( model=model, messages=messages, api_base=api_base, - headers=headers, + headers=signed_headers, custom_llm_provider=custom_llm_provider, provider_config=provider_config, timeout=timeout, logging_obj=logging_obj, - data=data, + data=self._add_stream_param_to_request_body( + data=data, + provider_config=provider_config, + fake_stream=fake_stream, + ), fake_stream=fake_stream, - client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), + client=async_client, litellm_params=litellm_params, json_mode=json_mode, optional_params=optional_params, signed_json_body=signed_json_body, ) + return self.async_completion( + custom_llm_provider=custom_llm_provider, + provider_config=provider_config, + api_base=api_base, + headers=signed_headers, + data=data, + timeout=timeout, + model=model, + model_response=model_response, + logging_obj=logging_obj, + api_key=api_key, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + client=async_client, + json_mode=json_mode, + signed_json_body=signed_json_body, + shared_session=shared_session, + ) - else: - return self.async_completion( - custom_llm_provider=custom_llm_provider, - provider_config=provider_config, - api_base=api_base, - headers=headers, - data=data, - timeout=timeout, - model=model, - model_response=model_response, - logging_obj=logging_obj, - api_key=api_key, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - encoding=encoding, - client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), - json_mode=json_mode, - signed_json_body=signed_json_body, - shared_session=shared_session, + if acompletion is True and provider_config.uses_async_transform_request: + + async def transform_then_dispatch() -> ModelResponse | CustomStreamWrapper: + transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict + "dict[str, object]", + await provider_config.async_transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, + ), ) + signed_request: Final = await ( + run_aws_signing(sign_and_log, transformed) + if isinstance(provider_config, SignsRequestsWithAWS) + else asyncio.to_thread(sign_and_log, transformed) + ) + return await dispatch_async(*signed_request) + + return transform_then_dispatch() + + data, signed_headers, signed_json_body = sign_and_log( + provider_config.transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, + ) + ) + + if acompletion is True: + return dispatch_async(data, signed_headers, signed_json_body) if stream is True: data = self._add_stream_param_to_request_body( @@ -641,7 +672,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, messages=messages, @@ -651,7 +682,7 @@ class BaseLLMHTTPHandler: completion_stream, headers = self.make_sync_call( provider_config=provider_config, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, original_data=data, @@ -684,7 +715,7 @@ class BaseLLMHTTPHandler: sync_httpx_client=sync_httpx_client, provider_config=provider_config, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, timeout=timeout, @@ -1716,6 +1747,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + return self._transform_ocr_response( provider_config=provider_config, model=model, @@ -1779,6 +1816,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + # Use async response transform for async operations return await provider_config.async_transform_ocr_response( model=model, @@ -1875,6 +1918,7 @@ class BaseLLMHTTPHandler: url=complete_url, headers=signed_headers, ) + response.raise_for_status() else: # A signed body must be sent verbatim, re-serializing it would break the signature response = client.post( @@ -1884,6 +1928,8 @@ class BaseLLMHTTPHandler: json=data if signed_json_body is None else None, timeout=timeout, ) + except httpx.HTTPStatusError as e: + raise provider_config.get_http_error_class(e) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -1936,7 +1982,9 @@ class BaseLLMHTTPHandler: api_key=api_key, ) - signed_headers, signed_json_body = provider_config.sign_request( + signed_headers, signed_json_body = await sign_request_off_loop_if_aws( + provider_config, + provider_config.sign_request, headers=headers, optional_params=optional_params, request_data=data, @@ -1974,6 +2022,7 @@ class BaseLLMHTTPHandler: url=complete_url, headers=signed_headers, ) + response.raise_for_status() else: # A signed body must be sent verbatim, re-serializing it would break the signature response = await async_httpx_client.post( @@ -1983,6 +2032,8 @@ class BaseLLMHTTPHandler: json=data if signed_json_body is None else None, timeout=timeout, ) + except httpx.HTTPStatusError as e: + raise provider_config.get_http_error_class(e) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -2037,7 +2088,9 @@ class BaseLLMHTTPHandler: max_attempts, ) provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body) - headers, signed_json_body = provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + provider_config, + provider_config.sign_request, headers=headers, optional_params=optional_params_dict, request_data=request_body, @@ -2197,7 +2250,9 @@ class BaseLLMHTTPHandler: stream=stream, ) - headers, signed_json_body = anthropic_messages_provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + anthropic_messages_provider_config, + anthropic_messages_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), # dynamic aws_* params are passed under litellm_params request_data=request_body, @@ -2403,9 +2458,7 @@ class BaseLLMHTTPHandler: return None from litellm.rust_bridge.configuration import rust_enabled - raw_request_override: Final = litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - if not rust_enabled(request_override=request_override): + if not rust_enabled(): return None if has_agentic_hook: return None @@ -2875,7 +2928,9 @@ class BaseLLMHTTPHandler: fake_stream=fake_stream, ) - headers, signed_body = responses_api_provider_config.sign_request( + headers, signed_body = await sign_request_off_loop_if_aws( + responses_api_provider_config, + responses_api_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), request_data=data, @@ -4583,7 +4638,9 @@ class BaseLLMHTTPHandler: ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) - headers, signed_body = responses_api_provider_config.sign_request( + headers, signed_body = await sign_request_off_loop_if_aws( + responses_api_provider_config, + responses_api_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), request_data=data, @@ -6514,7 +6571,7 @@ class BaseLLMHTTPHandler: @asynccontextmanager async def _backend_connection(): - if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params): + if _rust_responses_websocket_enabled(custom_llm_provider): from litellm.rust_bridge import responses_websocket as rust_responses_websocket rust_backend: Final = await rust_responses_websocket.connect( @@ -6759,7 +6816,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - data, files = image_edit_provider_config.transform_image_edit_request( + data, files = await image_edit_provider_config.async_transform_image_edit_request( model=model, image=image, prompt=prompt, @@ -9791,7 +9848,7 @@ class BaseLLMHTTPHandler: vector_store_search_optional_params=vector_store_search_optional_params, api_base=api_base, litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), + litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)), extra_body=extra_body, embedding_executor=embedding_executor, ) @@ -9810,7 +9867,9 @@ class BaseLLMHTTPHandler: ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) - headers, signed_json_body = vector_store_provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + vector_store_provider_config, + vector_store_provider_config.sign_request, headers=headers, optional_params=all_optional_params, request_data=request_body, @@ -9836,6 +9895,12 @@ class BaseLLMHTTPHandler: data=request_data, timeout=timeout, ) + except httpx.TimeoutException: + raise vector_store_provider_config.get_error_class( + error_message="Vector store search exceeded the caller timeout.", + status_code=408, + headers=httpx.Headers(), + ) from None except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -9920,7 +9985,7 @@ class BaseLLMHTTPHandler: vector_store_search_optional_params=vector_store_search_optional_params, api_base=api_base, litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), + litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)), extra_body=extra_body, embedding_executor=embedding_executor, ) @@ -9965,7 +10030,14 @@ class BaseLLMHTTPHandler: url=url, headers=headers, data=request_data, + timeout=timeout, ) + except httpx.TimeoutException: + raise vector_store_provider_config.get_error_class( + error_message="Vector store search exceeded the caller timeout.", + status_code=408, + headers=httpx.Headers(), + ) from None except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -9995,6 +10067,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client + vector_store_provider_config.validate_create_vector_store() + headers: Final = vector_store_provider_config.validate_environment( headers=extra_headers or {}, litellm_params=litellm_params ) @@ -10065,6 +10139,8 @@ class BaseLLMHTTPHandler: else: sync_httpx_client = client + vector_store_provider_config.validate_create_vector_store() + headers: Final = vector_store_provider_config.validate_environment( headers=extra_headers or {}, litellm_params=litellm_params ) diff --git a/litellm/llms/dashscope/common_utils.py b/litellm/llms/dashscope/common_utils.py index b7c97893a15..9ed9c276e43 100644 --- a/litellm/llms/dashscope/common_utils.py +++ b/litellm/llms/dashscope/common_utils.py @@ -2,7 +2,8 @@ Common utilities for the DashScope LLM provider. """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Final +from urllib.parse import urlparse import httpx @@ -16,6 +17,27 @@ if TYPE_CHECKING: ) from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +DASHSCOPE_CHAT_COMPATIBLE_PATH: Final = "/compatible-mode/v1" +DASHSCOPE_RERANK_PATH: Final = "/compatible-api/v1/reranks" + + +def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None: + if api_base is None: + return None + parsed: Final = urlparse(api_base) + host: Final = parsed.hostname or "" + on_aliyun_host: Final = host == "aliyuncs.com" or host.endswith(".aliyuncs.com") + if not on_aliyun_host or parsed.path.rstrip("/") != DASHSCOPE_CHAT_COMPATIBLE_PATH: + return None + return f"{parsed.scheme}://{parsed.netloc}{DASHSCOPE_RERANK_PATH}" + + +def resolve_dashscope_family_rerank_api_base(api_base: str | None, env_var: str, default_rerank_base: str) -> str: + remapped: Final = _rerank_base_for_chat_shaped_base(api_base) + if api_base is not None and remapped is None: + return api_base + return get_secret_str(env_var) or remapped or default_rerank_base + def get_dashscope_family_embedding_config(custom_llm_provider: str) -> "BaseEmbeddingConfig": if custom_llm_provider == "qwencloud": diff --git a/litellm/llms/dashscope/qwen_ai_platform.py b/litellm/llms/dashscope/qwen_ai_platform.py index 9a44eaf574a..6998a57e2b7 100644 --- a/litellm/llms/dashscope/qwen_ai_platform.py +++ b/litellm/llms/dashscope/qwen_ai_platform.py @@ -3,6 +3,7 @@ from typing import Final from litellm.secret_managers.main import get_secret_str from .chat.transformation import DashScopeChatConfig +from .common_utils import resolve_dashscope_family_rerank_api_base from .embed.transformation import DashScopeEmbeddingConfig from .image_generation.transformation import DashScopeImageGenerationConfig from .rerank.transformation import DashScopeRerankConfig @@ -51,7 +52,9 @@ class QwenAIPlatformRerankConfig(DashScopeRerankConfig): return _require_qwen_ai_platform_api_key(api_key) def _resolve_rerank_api_base(self, api_base: str | None) -> str: - return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_RERANK") or QWEN_AI_PLATFORM_RERANK_API_BASE + return resolve_dashscope_family_rerank_api_base( + api_base, "QWEN_AI_PLATFORM_API_BASE_RERANK", QWEN_AI_PLATFORM_RERANK_API_BASE + ) class QwenAIPlatformImageGenerationConfig(DashScopeImageGenerationConfig): diff --git a/litellm/llms/dashscope/qwencloud.py b/litellm/llms/dashscope/qwencloud.py index d8d53e340ef..ac21a48dd3e 100644 --- a/litellm/llms/dashscope/qwencloud.py +++ b/litellm/llms/dashscope/qwencloud.py @@ -3,6 +3,7 @@ from typing import Final from litellm.secret_managers.main import get_secret_str from .chat.transformation import DashScopeChatConfig +from .common_utils import resolve_dashscope_family_rerank_api_base from .embed.transformation import DashScopeEmbeddingConfig from .image_generation.transformation import DashScopeImageGenerationConfig from .rerank.transformation import DashScopeRerankConfig @@ -51,7 +52,9 @@ class QwenCloudRerankConfig(DashScopeRerankConfig): return _require_qwencloud_api_key(api_key) def _resolve_rerank_api_base(self, api_base: str | None) -> str: - return api_base or get_secret_str("QWENCLOUD_API_BASE_RERANK") or QWENCLOUD_RERANK_API_BASE + return resolve_dashscope_family_rerank_api_base( + api_base, "QWENCLOUD_API_BASE_RERANK", QWENCLOUD_RERANK_API_BASE + ) class QwenCloudImageGenerationConfig(DashScopeImageGenerationConfig): diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py index 490757a0948..14ad756ec9c 100644 --- a/litellm/llms/dashscope/rerank/transformation.py +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -12,8 +12,12 @@ Endpoint - https://dashscope.aliyuncs.com/compatible-api/v1/reranks Note: chat/embed live under `/compatible-mode/v1/`, but DashScope's rerank -route is exposed under `/compatible-api/v1/reranks` per the docs. Override -with `DASHSCOPE_API_BASE_RERANK` to point at a different host or path. +route is exposed under `/compatible-api/v1/reranks` per the docs. A chat-shaped +`.aliyuncs.com/compatible-mode/v1` base reaching this config (the chat default +from `get_llm_provider`, or a `DASHSCOPE_API_BASE` env var) is redirected to +the same host's rerank route, since `/compatible-mode/v1/reranks` is a dead +route on every DashScope host. Override with `DASHSCOPE_API_BASE_RERANK` to +point at a different host or path. Empirically, qwen3-rerank accepts `return_documents=true` and echoes `results[].document.text` back, even though the public docs list the flag @@ -40,7 +44,7 @@ from litellm.types.rerank import ( RerankTokens, ) -from ..common_utils import DashScopeError +from ..common_utils import DashScopeError, resolve_dashscope_family_rerank_api_base DEFAULT_RERANK_URL: Final = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks" @@ -67,9 +71,7 @@ class DashScopeRerankConfig(BaseRerankConfig): return resolved_api_key def _resolve_rerank_api_base(self, api_base: str | None) -> str: - if api_base is not None: - return api_base - return get_secret_str("DASHSCOPE_API_BASE_RERANK") or DEFAULT_RERANK_URL + return resolve_dashscope_family_rerank_api_base(api_base, "DASHSCOPE_API_BASE_RERANK", DEFAULT_RERANK_URL) def get_complete_url( self, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index e2e2ea5b553..09aaf970dc5 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _extract_reasoning_content, # pyright: ignore[reportPrivateUsage] # same import as the OpenAI transformation strip_litellm_internal_message_fields, strip_name_from_message, ) @@ -23,7 +24,9 @@ from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.databricks import ( AllDatabricksContentValues, DatabricksChoice, + DatabricksDelta, DatabricksFunction, + DatabricksMessage, DatabricksResponse, DatabricksTool, ) @@ -247,8 +250,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = self._get_api_base(api_base) - complete_url: Final = f"{api_base}/chat/completions" + use_ai_gateway: Final = model.removeprefix("databricks/").count(".") >= 2 + api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway) + url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base + complete_url: Final = f"{url_base}/chat/completions" return complete_url def get_supported_openai_params(self, model: str | None = None) -> list: @@ -534,6 +539,19 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): thinking_blocks.append(thinking_block) return reasoning_content, thinking_blocks + @staticmethod + def extract_top_level_reasoning_content(delta: DatabricksDelta) -> str | None: + return delta.get("reasoning_content") + + @staticmethod + def resolve_reasoning_and_content( + message: DatabricksMessage, block_reasoning_content: str | None + ) -> tuple[str | None, str | None]: + content_str: Final = DatabricksConfig.extract_content_str(message["content"]) + if block_reasoning_content is not None: + return block_reasoning_content, content_str + return _extract_reasoning_content({**message, "content": content_str}) + @staticmethod def extract_citations( content: AllDatabricksContentValues | None, @@ -577,14 +595,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): finish_reason = "stop" if translated_message is None: - ## get the content str - content_str = DatabricksConfig.extract_content_str(choice["message"]["content"]) - - ## get the reasoning content ( - reasoning_content, + block_reasoning_content, thinking_blocks, ) = DatabricksConfig.extract_reasoning_content(choice["message"].get("content")) + reasoning_content, content_str = DatabricksConfig.resolve_reasoning_and_content( + choice["message"], block_reasoning_content + ) citations = DatabricksConfig.extract_citations(choice["message"].get("content")) @@ -738,12 +755,16 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator): # extract the reasoning content ( - reasoning_content, + block_reasoning_content, thinking_blocks, ) = DatabricksConfig.extract_reasoning_content(choice["delta"].get("content")) choice["delta"]["content"] = content_str - choice["delta"]["reasoning_content"] = reasoning_content + choice["delta"]["reasoning_content"] = ( + block_reasoning_content + if block_reasoning_content is not None + else DatabricksConfig.extract_top_level_reasoning_content(choice["delta"]) + ) choice["delta"]["thinking_blocks"] = thinking_blocks translated_choices.append(choice) return ModelResponseStream( diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 7695b1cb35e..a4ec2c5378b 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -177,19 +177,13 @@ class DatabricksBase: # Default: just litellm return f"litellm/{version}" - def _get_api_base(self, api_base: str | None) -> str: - """ - Get the Databricks API base URL. - - If not provided, attempts to get it from the Databricks SDK. - """ + def _get_api_base(self, api_base: str | None, use_ai_gateway: bool = False) -> str: if api_base is None: try: from databricks.sdk import WorkspaceClient databricks_client: Final = WorkspaceClient() api_base = f"{databricks_client.config.host}/serving-endpoints" - return api_base except ImportError: raise DatabricksException( status_code=400, @@ -198,6 +192,18 @@ class DatabricksBase: "or install the databricks-sdk Python library." ), ) + + if not use_ai_gateway: + return api_base + + normalized_api_base: Final = api_base.rstrip("/") + if normalized_api_base.endswith("/ai-gateway/mlflow/v1"): + return normalized_api_base + if normalized_api_base.endswith("/serving-endpoints"): + return f"{normalized_api_base.removesuffix('/serving-endpoints')}/ai-gateway/mlflow/v1" + api_base_parts: Final = urlsplit(normalized_api_base) + if api_base_parts.path in ("", "/"): + return f"{normalized_api_base}/ai-gateway/mlflow/v1" return api_base def _get_oauth_m2m_token( diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index ac934ad0cb5..21a630a76d7 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from httpx import Headers @@ -13,7 +15,7 @@ class FireworksAIException(BaseLLMException): pass -def get_fireworks_session_id(litellm_params: dict) -> str | None: +def get_fireworks_session_id(litellm_params: Mapping[str, object]) -> str | None: """ Session id to send as `x-session-affinity`, or None when the caller gave none. @@ -23,19 +25,39 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None: """ params: Final = litellm_params metadata: Final = params.get("metadata") - if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): + if isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return None for key in ("litellm_session_id", "session_id"): value = params.get(key) if value: return str(value) - if isinstance(metadata, dict): + if isinstance(metadata, Mapping): value = metadata.get("session_id") if value: return str(value) return None +def with_fireworks_session_affinity( + headers: Mapping[str, str], litellm_params: Mapping[str, object] +) -> Mapping[str, str]: + if any(key.lower() == "x-session-affinity" for key in headers): + return headers + session_id: Final = get_fireworks_session_id(litellm_params) + if not session_id: + return headers + return MappingProxyType({**headers, "x-session-affinity": session_id}) + + +def resolve_fireworks_api_key(api_key: str | None) -> str | None: + return api_key or ( + get_secret_str("FIREWORKS_API_KEY") + or get_secret_str("FIREWORKS_AI_API_KEY") + or get_secret_str("FIREWORKSAI_API_KEY") + or get_secret_str("FIREWORKS_AI_TOKEN") + ) + + AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX: Final = "FW-" @@ -63,13 +85,7 @@ class FireworksAIMixin: ) def _get_api_key(self, api_key: str | None) -> str | None: - dynamic_api_key: Final = api_key or ( - get_secret_str("FIREWORKS_API_KEY") - or get_secret_str("FIREWORKS_AI_API_KEY") - or get_secret_str("FIREWORKSAI_API_KEY") - or get_secret_str("FIREWORKS_AI_TOKEN") - ) - return dynamic_api_key + return resolve_fireworks_api_key(api_key) def validate_environment( self, @@ -92,9 +108,5 @@ class FireworksAIMixin: return self._add_session_affinity_header({**auth_headers, **content_type_header}, litellm_params) def _add_session_affinity_header(self, headers: dict, litellm_params: dict) -> dict: - if any(key.lower() == "x-session-affinity" for key in headers): - return headers - session_id: Final = get_fireworks_session_id(litellm_params) - if not session_id: - return headers - return {**headers, "x-session-affinity": session_id} + pinned: Final = with_fireworks_session_affinity(headers, litellm_params) + return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py new file mode 100644 index 00000000000..f7dd774ea18 --- /dev/null +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -0,0 +1,188 @@ +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final +from urllib.parse import unquote + +import httpx +from openai.types.responses import EasyInputMessageParam, ResponseInputContentParam, ResponseInputItemParam + +from litellm.llms.fireworks_ai.common_utils import ( + resolve_fireworks_api_key, + resolve_fireworks_resource_name, + with_fireworks_session_affinity, +) +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ResponseInputParam +from litellm.types.responses.main import DeleteResponseResult +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +FIREWORKS_AI_DEFAULT_API_BASE: Final = "https://api.fireworks.ai/inference/v1" + + +def _session_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object]: + extras: Final[Mapping[str, object]] = litellm_params.model_extra or MappingProxyType({}) + return MappingProxyType( + {"litellm_session_id": extras.get("litellm_session_id"), "metadata": extras.get("litellm_metadata")} + ) + + +_INSTRUCTION_ROLES: Final = frozenset({"system", "developer"}) + + +def _role(item: ResponseInputItemParam) -> str | None: + match item: + case {"role": str(role)}: + return role + case _: + return None + + +def _developer_item_as_system(item: ResponseInputItemParam) -> ResponseInputItemParam: + if "role" not in item or item["role"] != "developer": + return item + return EasyInputMessageParam(role="system", content=item["content"], type="message") + + +def _developer_items_as_system(input: str | ResponseInputParam) -> str | ResponseInputParam: + if isinstance(input, str): + return input + return [_developer_item_as_system(item) for item in input] + + +def _text_part(part: ResponseInputContentParam) -> str | None: + match part: + case {"type": "input_text", "text": str(text)}: + return text + case _: + return None + + +def _text_only_content(item: ResponseInputItemParam) -> str | None: + match item: + case {"role": "system" | "developer", "content": str(text)}: + return text + case {"role": "system" | "developer", "content": [*parts]}: + texts: Final = tuple(map(_text_part, parts)) + return None if any(text is None for text in texts) else "\n\n".join(text for text in texts if text) + case _: + return None + + +def _leading_instruction_block_length(roles: Sequence[str | None]) -> int: + return next((index for index, role in enumerate(roles) if role not in _INSTRUCTION_ROLES), len(roles)) + + +def _closing_instruction_block_start(roles: Sequence[str | None], leading_length: int) -> int: + last_conversation_index: Final = next( + (index for index in range(len(roles) - 1, leading_length - 1, -1) if roles[index] not in _INSTRUCTION_ROLES), + None, + ) + if last_conversation_index is None or roles[last_conversation_index] != "assistant": + return len(roles) + return last_conversation_index + 1 + + +def _hoisted_indices(roles: Sequence[str | None]) -> tuple[int, ...]: + leading_length: Final = _leading_instruction_block_length(roles) + closing_start: Final = _closing_instruction_block_start(roles, leading_length) + return tuple( + index for index, role in enumerate(roles[:closing_start]) if index < leading_length or role == "developer" + ) + + +def _with_instruction_items_folded( + input: str | ResponseInputParam, instructions: str | None +) -> tuple[str | None, str | ResponseInputParam]: + if isinstance(input, str): + return instructions, input + items: Final = tuple(input) + folded: Final = MappingProxyType( + { + index: text + for index in _hoisted_indices(tuple(map(_role, items))) + if (text := _text_only_content(items[index])) is not None + } + ) + joined: Final = "\n\n".join(chunk for chunk in (instructions, *folded.values()) if chunk) + return ( + instructions if not folded else joined or None, + [ # mutable-ok: the base class takes the input items as a list + _developer_item_as_system(item) for index, item in enumerate(items) if index not in folded + ], + ) + + +class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.FIREWORKS_AI + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + litellm_params: GenericLiteLLMParams | None, + ) -> dict: # mutable-ok: overrides the base class signature + params: Final = litellm_params or GenericLiteLLMParams() + api_key: Final = resolve_fireworks_api_key(params.api_key) + if api_key is None: + raise ValueError("FIREWORKS_API_KEY is not set") + authorized: Final = MappingProxyType( + {"Content-Type": "application/json", **headers, "Authorization": f"Bearer {api_key}"} + ) + pinned: Final = with_fireworks_session_affinity(authorized, _session_params(params)) + return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place + + def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/") + return f"{base}/responses" + + def transform_responses_api_request( + self, + model: str, + input: str | ResponseInputParam, + response_api_optional_request_params: dict, # mutable-ok: overrides the base class signature + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: overrides the base class signature + ) -> dict: # mutable-ok: overrides the base class signature + instructions_param: Final[object] = response_api_optional_request_params.get("instructions") + validated_input: Final = self._validate_input_param(input) + instructions, folded_input = ( + _with_instruction_items_folded(validated_input, instructions_param) + if isinstance(instructions_param, str | None) + else (instructions_param, _developer_items_as_system(validated_input)) + ) + instruction_entries: Final = () if instructions is None else (("instructions", instructions),) + folded_params: Final = { # mutable-ok: the base class takes the optional params as a dict + key: value + for key, value in ( + *((key, value) for key, value in response_api_optional_request_params.items() if key != "instructions"), + *instruction_entries, + ) + } + return super().transform_responses_api_request( + model=resolve_fireworks_resource_name(model), + input=folded_input, + response_api_optional_request_params=folded_params, + litellm_params=litellm_params, + headers=headers, + ) + + def transform_delete_response_api_response( + self, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> DeleteResponseResult: + deleted_id: Final = unquote(raw_response.request.url.path.rsplit("/", 1)[-1]) + return DeleteResponseResult(id=deleted_id, object="response", deleted=True) + + def supports_native_websocket(self) -> bool: + return False + + def supports_native_file_search(self) -> bool: + return False diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 1a67b33665b..42c9ef13730 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionFileObject from litellm.types.llms.vertex_ai import ContentType, PartType from litellm.utils import supports_reasoning -from ...vertex_ai.gemini.transformation import _gemini_convert_messages_with_history +from ...vertex_ai.gemini.transformation import GEMINI_FILES_API_URI_PREFIX, _gemini_convert_messages_with_history from ...vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -127,7 +127,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): if element.get("type") == "image_url": img_element = cast(ChatCompletionImageObject, element) # cast-ok: runtime type tag checked _image_url, format, detail = _image_url_fields(img_element) - if _image_url and "https://" in _image_url: + if ( + _image_url + and "https://" in _image_url + and not _image_url.startswith(GEMINI_FILES_API_URI_PREFIX) + ): image_obj = convert_to_anthropic_image_obj(_image_url, format=format) converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj) if detail is not None: @@ -147,7 +151,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): llm_provider="gemini", ) file_id = _file_field.get("file_id") - if file_id and ("http://" in file_id or "https://" in file_id): + if ( + file_id + and ("http://" in file_id or "https://" in file_id) + and not file_id.startswith(GEMINI_FILES_API_URI_PREFIX) + ): # Convert HTTP/HTTPS file URL to base64 data try: base64_data = convert_url_to_base64(file_id) diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/__init__.py b/litellm/llms/gemini/google_genai/guardrail_translation/__init__.py new file mode 100644 index 00000000000..494a72d6999 --- /dev/null +++ b/litellm/llms/gemini/google_genai/guardrail_translation/__init__.py @@ -0,0 +1,20 @@ +"""Google GenAI generateContent guardrail translation handler.""" + +from typing import Final + +from litellm.llms.gemini.google_genai.guardrail_translation.handler import ( + GoogleGenAIGenerateContentHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings: Final = { # mutable-ok: discover_guardrail_translation_mappings only accepts isinstance(mappings, dict) + CallTypes.generate_content: GoogleGenAIGenerateContentHandler, + CallTypes.agenerate_content: GoogleGenAIGenerateContentHandler, + CallTypes.generate_content_stream: GoogleGenAIGenerateContentHandler, + CallTypes.agenerate_content_stream: GoogleGenAIGenerateContentHandler, +} + +__all__ = ( + "GoogleGenAIGenerateContentHandler", + "guardrail_translation_mappings", +) diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py new file mode 100644 index 00000000000..e13e1e63cbb --- /dev/null +++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py @@ -0,0 +1,255 @@ +""" +Google GenAI generateContent handler for Unified Guardrails. + +Extracts text from generateContent requests (systemInstruction.parts[].text +and contents[].parts[].text) and responses (candidates[].content.parts[].text), +applies the guardrail, and +writes the guardrailed text back in place. Requests and responses may be +dicts (wire format) or google-genai SDK objects; streaming chunks may +additionally be raw SSE frames, which are scanned for detection (a blocking +guardrail raises) without rewriting the frames. +""" + +import json +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamTransformSink, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth + +_EMPTY_REQUEST_DATA: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _field(container: object, name: str) -> object | None: + if isinstance(container, dict): + return container.get(name) + return getattr(container, name, None) + + +def _part_text(part: object) -> str | None: + text: Final = _field(part, "text") + if isinstance(text, str) and text: + return text + return None + + +def _write_part_text(part: object, text: str) -> None: + if isinstance(part, dict): + part["text"] = text # rebind-ok: guardrail write-back rewrites the caller's part in place by handler contract + return + setattr(part, "text", text) # noqa: B010 # SDK parts are typed as object here; direct assignment cannot type-check + + +def _content_text_parts(content: object) -> tuple[object, ...]: + parts: Final = _field(content, "parts") + if not isinstance(parts, (list, tuple)): + return () + return tuple(part for part in parts if _part_text(part) is not None) + + +def _system_instruction(data: Mapping[str, object]) -> object | None: + return next( + ( + value + for container in (data, data.get("config")) + if container is not None + for key in ("systemInstruction", "system_instruction") + for value in (_field(container, key),) + if value is not None + ), + None, + ) + + +def _request_text_parts(data: Mapping[str, object]) -> tuple[object, ...]: + contents: Final = data.get("contents") + content_list: Final = ( + (contents,) if isinstance(contents, dict) else tuple(contents) if isinstance(contents, list) else () + ) + return ( + *_content_text_parts(_system_instruction(data)), + *(part for content in content_list for part in _content_text_parts(content)), + ) + + +def _response_text_parts(response: object) -> tuple[object, ...]: + candidates: Final = _field(response, "candidates") + if not isinstance(candidates, (list, tuple)): + return () + return tuple(part for candidate in candidates for part in _content_text_parts(_field(candidate, "content"))) + + +def _part_texts(text_parts: Sequence[object]) -> tuple[str, ...]: + return tuple(text for part in text_parts for text in (_part_text(part),) if text is not None) + + +def _texts_payload( + texts: Sequence[str], +) -> list[str]: # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str] + return list(texts) # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str] + + +def _write_back_texts(text_parts: Sequence[object], guardrailed_texts: Sequence[str] | None) -> None: + if not guardrailed_texts or len(guardrailed_texts) != len(text_parts): + return + for part, text in zip(text_parts, guardrailed_texts): + _write_part_text(part, text) + + +def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None: + try: + parsed: Final = json.loads(payload) + except json.JSONDecodeError: + return None + if isinstance(parsed, dict): + return parsed + return None + + +def _sse_payload_texts(sse_text: str) -> tuple[str, ...]: + return tuple( + text + for line in sse_text.splitlines() + if line.startswith("data:") + for payload in (line[len("data:") :].strip(),) + if payload and payload != "[DONE]" + for parsed in (_parse_json_dict_or_none(payload),) + if parsed is not None + for text in _part_texts(_response_text_parts(parsed)) + ) + + +def _chunk_sse_text(chunk: object) -> str | None: + if isinstance(chunk, bytes): + return chunk.decode("utf-8", errors="replace") + if isinstance(chunk, str): + return chunk + return None + + +def _accumulated_stream_text(responses_so_far: Sequence[object]) -> str: + object_texts: Final = tuple( + text + for chunk in responses_so_far + if _chunk_sse_text(chunk) is None + for text in _part_texts(_response_text_parts(chunk)) + ) + sse_text: Final = "".join(sse for chunk in responses_so_far for sse in (_chunk_sse_text(chunk),) if sse is not None) + return "".join(object_texts) + "".join(_sse_payload_texts(sse_text)) + + +class GoogleGenAIGenerateContentHandler(BaseTranslation): + """ + Guardrail translation for the google genai generateContent surface + (/models/{model}:generateContent, :streamGenerateContent, and the + litellm SDK generate_content call types). + """ + + async def process_input_messages( + self, + data: dict, # mutable-ok: base handler contract passes the proxy's request dict through to apply_guardrail + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> object: + text_parts: Final = _request_text_parts(data) + if not text_parts: + verbose_proxy_logger.debug("Google GenAI guardrail: no request text found, skipping") + return data + model: Final = data.get("model") + inputs: Final = ( + GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)), model=model) + if isinstance(model, str) + else GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts))) + ) + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + _write_back_texts(text_parts, guardrailed_inputs.get("texts")) + return data + + async def process_output_response( + self, + response: object, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: Mapping[str, object] | None = None, + ) -> object: + text_parts: Final = _response_text_parts(response) + if not text_parts: + verbose_proxy_logger.debug("Google GenAI guardrail: no response text found, skipping") + return response + guardrail_request_data: Final = self._merged_request_data( + request_data=request_data, + user_api_key_dict=user_api_key_dict, + context_key="response", + context_value=response, + ) + model: Final = guardrail_request_data.get("model") + inputs: Final = ( + GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)), model=model) + if isinstance(model, str) + else GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts))) + ) + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=guardrail_request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + _write_back_texts(text_parts, guardrailed_inputs.get("texts")) + return response + + async def process_output_streaming_response( + self, + responses_so_far: Sequence[object], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: Mapping[str, object] | None = None, + stream_transform_sink: StreamTransformSink | None = None, + ) -> object: + accumulated_text: Final = _accumulated_stream_text(responses_so_far) + if not accumulated_text: + return responses_so_far + guardrail_request_data: Final = self._merged_request_data( + request_data=request_data, + user_api_key_dict=user_api_key_dict, + context_key="responses_so_far", + context_value=responses_so_far, + ) + _guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=_texts_payload((accumulated_text,))), + request_data=guardrail_request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + def _merged_request_data( + self, + request_data: Mapping[str, object] | None, + user_api_key_dict: Optional["UserAPIKeyAuth"], + context_key: str, + context_value: object, + ) -> dict: # mutable-ok: CustomGuardrail.apply_guardrail requires a plain dict request payload + base: Final = request_data if request_data is not None else _EMPTY_REQUEST_DATA + user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) + context_pairs: Final = ((context_key, context_value),) if context_key not in base else () + metadata_pairs: Final = ( + (("litellm_metadata", user_metadata),) if user_metadata and "litellm_metadata" not in base else () + ) + return dict((*base.items(), *context_pairs, *metadata_pairs)) # mutable-ok: apply_guardrail takes a plain dict diff --git a/litellm/llms/hosted_vllm/image_edit/__init__.py b/litellm/llms/hosted_vllm/image_edit/__init__.py new file mode 100644 index 00000000000..27e005e8a0d --- /dev/null +++ b/litellm/llms/hosted_vllm/image_edit/__init__.py @@ -0,0 +1,9 @@ +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig + +from .transformation import HostedVLLMImageEditConfig + +__all__ = ("HostedVLLMImageEditConfig",) + + +def get_hosted_vllm_image_edit_config(model: str) -> BaseImageEditConfig: + return HostedVLLMImageEditConfig() diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py new file mode 100644 index 00000000000..3b8cc437168 --- /dev/null +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -0,0 +1,43 @@ +from typing import Final + +from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig +from litellm.secret_managers.main import get_secret_str + +PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT: Final = frozenset({"mask", "quality", "input_fidelity"}) + + +class HostedVLLMImageEditConfig(OpenAIImageEditConfig): + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseImageEditConfig contract + return [ # mutable-ok: BaseImageEditConfig returns list + param + for param in super().get_supported_openai_params(model) + if param not in PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT + ] + + def validate_environment( + self, + headers: dict, # mutable-ok: BaseImageEditConfig contract + model: str, + api_key: str | None = None, + litellm_params: dict | None = None, # mutable-ok: BaseImageEditConfig contract + api_base: str | None = None, + ) -> dict: # mutable-ok: BaseImageEditConfig contract + resolved_key: Final = api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" + return {**headers, "Authorization": f"Bearer {resolved_key}"} # mutable-ok: httpx headers are a dict + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, # mutable-ok: BaseImageEditConfig contract + ) -> str: + resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE") + if resolved_api_base is None: + raise ValueError( + "api_base not set for Hosted VLLM images edits API. " + "Set via api_base parameter or HOSTED_VLLM_API_BASE environment variable" + ) + trimmed: Final = resolved_api_base.rstrip("/") + if trimmed.endswith("/v1"): + return f"{trimmed}/images/edits" + return f"{trimmed}/v1/images/edits" diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index a8f3388d092..2a81c38fe34 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -157,9 +157,6 @@ class JinaAIRerankConfig(BaseRerankConfig): billed_units: RerankBilledUnits | None = None, model_info: ModelInfo | None = None, ) -> tuple[float, float]: - """ - Jina AI reranker is priced at $0.000000018 per token. - """ if ( model_info is None or "input_cost_per_token" not in model_info diff --git a/litellm/llms/litellm_proxy/skills/constants.py b/litellm/llms/litellm_proxy/skills/constants.py index a6c88718f11..04a3a7dbc91 100644 --- a/litellm/llms/litellm_proxy/skills/constants.py +++ b/litellm/llms/litellm_proxy/skills/constants.py @@ -16,3 +16,8 @@ DEFAULT_MAX_ITERATIONS: Final[int] = 10 DEFAULT_SANDBOX_TIMEOUT: Final[int] = 120 """Default timeout in seconds for sandbox code execution.""" + +MAX_SKILLS_PER_SEARCH: Final[int] = 5000 +"""Upper bound on how many of the caller's accessible skills a single semantic +search embeds. Ranking runs in memory over this candidate set (no tsvector/DB-side +filtering yet), so this caps worst-case embedding cost per search request.""" diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 73f6ed23092..9b625cb0571 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -6,11 +6,15 @@ Used by the transformation layer and skills injection hook. """ import uuid +from collections.abc import Sequence from typing import Final from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache -from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX +from litellm.llms.litellm_proxy.skills.constants import ( + LITELLM_SKILL_ID_PREFIX, + MAX_SKILLS_PER_SEARCH, +) from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -131,6 +135,19 @@ class LiteLLMSkillsHandler: ) return [_prisma_skill_to_litellm(s) for s in skills] + @staticmethod + async def list_skills_for_search( + user_api_key_dict: UserAPIKeyAuth | None = None, + ) -> Sequence[LiteLLM_SkillsTable]: + """Every skill the caller can access, for ranking. Same owner-scope filter as + ``list_skills``, but unpaginated (up to ``MAX_SKILLS_PER_SEARCH``) since a query + must be scored against the whole accessible set, not one page of it.""" + return await LiteLLMSkillsHandler.list_skills( + limit=MAX_SKILLS_PER_SEARCH, + offset=0, + user_api_key_dict=user_api_key_dict, + ) + @staticmethod async def _load_skill(skill_id: str) -> object | None: """Cache-first read of the Prisma skill row. Owner-scope filtering diff --git a/litellm/llms/litellm_proxy/skills/skill_search.py b/litellm/llms/litellm_proxy/skills/skill_search.py new file mode 100644 index 00000000000..f975c6c4cab --- /dev/null +++ b/litellm/llms/litellm_proxy/skills/skill_search.py @@ -0,0 +1,161 @@ +"""Semantic ranking over the LiteLLM-hosted skill registry, shared by GET /v1/skills?query= and the skill_search MCP tool.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import BaseModel, ConfigDict + +from litellm.llms.litellm_proxy.skills.constants import MAX_SKILLS_PER_SEARCH +from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler +from litellm.proxy.common_utils.semantic_text_index import ( + Embedder, + EmbeddingFailed, + SemanticTextIndex, + router_embedder, +) +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + +DEFAULT_SKILL_SEARCH_TOP_K: Final = 5 +MAX_SKILL_SEARCH_TOP_K: Final = 100 +"""Matches the ``le=100`` bound GET /v1/skills?query= enforces via FastAPI's Query +validation, so the MCP tool can't return a larger payload than the REST endpoint allows.""" +MAX_SKILL_SEARCH_TEXT_CHARS: Final = 4000 +"""Per-skill cap on the title + description + instructions text that gets embedded, so one +search embeds at most ``MAX_SKILLS_PER_SEARCH * MAX_SKILL_SEARCH_TEXT_CHARS`` characters no +matter how long the stored instructions are.""" + + +@dataclass(frozen=True, slots=True) +class SkillSearchHit: + skill: LiteLLM_SkillsTable + score: float + + +@dataclass(frozen=True, slots=True) +class SkillSearchHits: + hits: tuple[SkillSearchHit, ...] + + +@dataclass(frozen=True, slots=True) +class SkillSearchNotConfigured: + reason: str + + +@dataclass(frozen=True, slots=True) +class SkillSearchEmbeddingFailed: + reason: str + + +@dataclass(frozen=True, slots=True) +class SkillSearchUnsupportedProvider: + reason: str + + +SkillSearchOutcome: TypeAlias = SkillSearchHits | SkillSearchNotConfigured | SkillSearchEmbeddingFailed +HostedSkillSearchOutcome: TypeAlias = SkillSearchOutcome | SkillSearchUnsupportedProvider + + +class SkillSearchResult(BaseModel): + model_config = ConfigDict(frozen=True) + + skill_id: str + display_title: str | None + description: str | None + score: float + + +def skill_search_text(skill: LiteLLM_SkillsTable) -> str: + joined: Final = "\n".join(part for part in (skill.display_title, skill.description, skill.instructions) if part) + return joined[:MAX_SKILL_SEARCH_TEXT_CHARS] + + +def skill_search_result(hit: SkillSearchHit) -> SkillSearchResult: + return SkillSearchResult( + skill_id=hit.skill.skill_id, + display_title=hit.skill.display_title, + description=hit.skill.description, + score=hit.score, + ) + + +class SkillSearchIndex: + """Caches one vector per distinct skill text per embedding model, so repeat searches only embed the query.""" + + def __init__(self, max_entries: int = MAX_SKILLS_PER_SEARCH) -> None: + self._index: Final = SemanticTextIndex(max_entries=max_entries) + + async def search( + self, + query: str, + skills: Sequence[LiteLLM_SkillsTable], + top_k: int, + embed: Embedder, + embedding_model: str, + ) -> SkillSearchHits | SkillSearchEmbeddingFailed: + texts: Final = tuple(skill_search_text(skill) for skill in skills) + scores: Final = await self._index.scores(query, texts, embed, embedding_model) + if isinstance(scores, EmbeddingFailed): + return SkillSearchEmbeddingFailed(reason=scores.reason) + ranked: Final = sorted( + (SkillSearchHit(skill=skill, score=score) for skill, score in zip(skills, scores, strict=True)), + key=lambda hit: hit.score, + reverse=True, + ) + return SkillSearchHits(hits=tuple(ranked[:top_k])) + + +global_skill_search_index: Final = SkillSearchIndex() + + +async def search_skills( + query: str, + skills: Sequence[LiteLLM_SkillsTable], + top_k: int, + router: Router | None, + embedding_model: str | None, + index: SkillSearchIndex, + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> SkillSearchOutcome: + if embedding_model is None: + return SkillSearchNotConfigured( + reason="skill search needs litellm_settings.skill_search_embedding_model set to an embedding model from model_list" + ) + if router is None: + return SkillSearchNotConfigured(reason="skill search needs a model_list so the embedding model can be called") + embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj) + return await index.search(query, skills, top_k, embed, embedding_model) + + +async def search_hosted_skills( + custom_llm_provider: str | None, + query: str, + top_k: int, + router: Router | None, + embedding_model: str | None, + index: SkillSearchIndex, + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> HostedSkillSearchOutcome: + """GET /v1/skills?query= for the skills LiteLLM hosts itself: only ``litellm_proxy`` has a registry to rank.""" + if custom_llm_provider != LlmProviders.LITELLM_PROXY.value: + return SkillSearchUnsupportedProvider(reason="query is only supported for custom_llm_provider=litellm_proxy") + skills: Final = await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict=user_api_key_dict) + return await search_skills( + query=query, + skills=skills, + top_k=top_k, + router=router, + embedding_model=embedding_model, + index=index, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py index c972dc349c9..9fc2d2cbb45 100644 --- a/litellm/llms/litellm_proxy/skills/transformation.py +++ b/litellm/llms/litellm_proxy/skills/transformation.py @@ -154,7 +154,7 @@ class LiteLLMSkillsTransformationHandler: user_api_key_dict=user_api_key_dict, ) - return self._db_skill_to_response(db_skill) + return self.db_skill_to_response(db_skill) def list_skills_handler( self, @@ -222,7 +222,9 @@ class LiteLLMSkillsTransformationHandler: user_api_key_dict=user_api_key_dict, ) - skills: Final = [self._db_skill_to_response(s) for s in db_skills] + skills: Final = [ # mutable-ok: ListSkillsResponse.data needs list[Skill]; never mutated after + self.db_skill_to_response(s) for s in db_skills + ] return ListSkillsResponse( data=skills, has_more=len(skills) >= limit, @@ -288,7 +290,7 @@ class LiteLLMSkillsTransformationHandler: skill_id=skill_id, user_api_key_dict=user_api_key_dict, ) - return self._db_skill_to_response(db_skill) + return self.db_skill_to_response(db_skill) def delete_skill_handler( self, @@ -354,7 +356,7 @@ class LiteLLMSkillsTransformationHandler: type=result.get("type", "skill_deleted"), ) - def _db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill: + def db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill: """ Convert a database skill record to Anthropic-compatible Skill response. @@ -375,4 +377,5 @@ class LiteLLMSkillsTransformationHandler: latest_version=db_skill.latest_version, source=db_skill.source or "custom", type="skill", + description=db_skill.description, ) diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py new file mode 100644 index 00000000000..2b3264dc756 --- /dev/null +++ b/litellm/llms/mistral/audio_speech/transformation.py @@ -0,0 +1,210 @@ +""" +Support for Mistral Voxtral text-to-speech via ``/v1/audio/speech``. + +API reference: https://docs.mistral.ai/api/#tag/audio/operation/audio_speech_v1_audio_speech_post +""" + +import base64 +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.text_to_speech.transformation import ( + BaseTextToSpeechConfig, + TextToSpeechRequestData, +) +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import HttpxBinaryResponseContent + + +class MistralTextToSpeechException(BaseLLMException): + pass + + +class MistralTextToSpeechConfig(BaseTextToSpeechConfig): + TTS_BASE_URL: Final[str] = "https://api.mistral.ai/v1" + AUDIO_CONTENT_TYPES: Final[MappingProxyType[str, str]] = MappingProxyType( + { + "mp3": "audio/mpeg", + "wav": "audio/wav", + "pcm": "audio/pcm", + "flac": "audio/flac", + "opus": "audio/ogg", + } + ) + DROPPED_RESPONSE_HEADERS: Final[frozenset[str]] = frozenset( + {"content-encoding", "transfer-encoding", "content-length", "content-type"} + ) + OPENAI_VOICE_ALIASES: Final[MappingProxyType[str, str]] = MappingProxyType( + { + "alloy": "en_paul_neutral", + "echo": "gb_oliver_neutral", + "fable": "en_paul_cheerful", + "onyx": "en_paul_confident", + "nova": "gb_jane_sarcasm", + "shimmer": "gb_jane_sarcasm", + } + ) + + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a plain list + return ["voice", "response_format"] # mutable-ok: base class contract returns a plain list + + def _map_openai_voice(self, voice_id: str) -> str: + return self.OPENAI_VOICE_ALIASES.get(voice_id.lower(), voice_id) + + def _resolve_voice_id(self, voice: object) -> str | None: + if isinstance(voice, str) and voice.strip(): + return self._map_openai_voice(voice.strip()) + if isinstance(voice, Mapping): + candidates: Final = (voice.get(key) for key in ("voice_id", "id", "name")) + resolved: Final = next( + (candidate.strip() for candidate in candidates if isinstance(candidate, str) and candidate.strip()), + None, + ) + return self._map_openai_voice(resolved) if resolved else None + return None + + def map_openai_params( + self, + model: str, + optional_params: Mapping[str, object], + voice: object = None, + drop_params: bool = False, + kwargs: Mapping[str, object] | None = None, + ) -> tuple[str | None, dict]: # mutable-ok: base class contract returns a plain dict + response_format: Final = optional_params.get("response_format") + ref_audio: Final = kwargs.get("ref_audio") if kwargs else None + voice_id_kwarg: Final = kwargs.get("voice_id") if kwargs else None + mapped_voice: Final = self._resolve_voice_id(voice) or self._resolve_voice_id(voice_id_kwarg) + mapped_params: Final = { # mutable-ok: base class contract returns a plain dict + key: value + for key, value in (("response_format", response_format), ("ref_audio", ref_audio)) + if isinstance(value, str) + } + return mapped_voice, mapped_params + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: # mutable-ok: base class contract returns a plain dict + resolved_key: Final = api_key or get_secret_str("MISTRAL_API_KEY") + if resolved_key is None: + raise MistralTextToSpeechException( + status_code=401, + message="Mistral API key is required. Set MISTRAL_API_KEY or pass api_key.", + ) + return { # mutable-ok: base class contract returns a plain dict + **headers, + "Authorization": f"Bearer {resolved_key}", + "Content-Type": "application/json", + } + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: Mapping[str, object], + ) -> str: + configured_base: Final = (api_base or self.TTS_BASE_URL).rstrip("/") + versioned_base: Final = configured_base if configured_base.endswith("/v1") else f"{configured_base}/v1" + return f"{versioned_base}/audio/speech" + + def transform_text_to_speech_request( + self, + model: str, + input: str, + voice: str | None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, str], + ) -> TextToSpeechRequestData: + response_format: Final = optional_params.get("response_format") + ref_audio: Final = optional_params.get("ref_audio") + request_data: Final[TextToSpeechRequestData] = { + "dict_body": { + "model": model, + "input": input, + **({"voice_id": voice} if voice else {}), + **({"response_format": response_format} if isinstance(response_format, str) else {}), + **({"ref_audio": ref_audio} if isinstance(ref_audio, str) else {}), + }, + "headers": {"Content-Type": "application/json"}, + } + return request_data + + def _requested_content_type(self, request: httpx.Request) -> str: + request_body: Final = json.loads(request.content or b"{}") + requested_format: Final = request_body.get("response_format") + if not isinstance(requested_format, str): + return "audio/mpeg" + return self.AUDIO_CONTENT_TYPES.get(requested_format, "audio/mpeg") + + def transform_text_to_speech_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> "HttpxBinaryResponseContent": + from litellm.types.llms.openai import HttpxBinaryResponseContent + + try: + response_json: Final = raw_response.json() + except (json.JSONDecodeError, ValueError): + raise MistralTextToSpeechException( + status_code=raw_response.status_code, + message=f"Non-JSON response from Mistral speech API: {raw_response.text[:500]}", + headers=raw_response.headers, + ) + audio_b64: Final = response_json.get("audio_data") + if not isinstance(audio_b64, str) or not audio_b64: + raise MistralTextToSpeechException( + status_code=500, + message=f"No audio_data in Mistral speech response. Response keys: {tuple(response_json.keys())}", + headers=raw_response.headers, + ) + try: + audio_bytes: Final = base64.b64decode(audio_b64, validate=True) + except ValueError: + raise MistralTextToSpeechException( + status_code=500, + message="Invalid base64 audio_data in Mistral speech response.", + headers=raw_response.headers, + ) + retained_headers: Final = tuple( + (key, value) + for key, value in raw_response.headers.items() + if key.lower() not in self.DROPPED_RESPONSE_HEADERS + ) + response_headers: Final = retained_headers + ( + ("content-length", str(len(audio_bytes))), + ("content-type", self._requested_content_type(raw_response.request)), + ) + binary_response: Final = httpx.Response( + status_code=200, + headers=response_headers, + content=audio_bytes, + request=raw_response.request, + ) + return HttpxBinaryResponseContent(binary_response) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict | httpx.Headers, # mutable-ok: BaseLLMException takes a plain dict or httpx.Headers + ) -> BaseLLMException: + return MistralTextToSpeechException( + message=error_message, + status_code=status_code, + headers=headers, + ) diff --git a/litellm/llms/mongodb/common_utils.py b/litellm/llms/mongodb/common_utils.py deleted file mode 100644 index 02c0b359407..00000000000 --- a/litellm/llms/mongodb/common_utils.py +++ /dev/null @@ -1,303 +0,0 @@ -"""Shared helpers for the MongoDB integrations. pymongo lives in the optional ``mongodb`` extra, -so every import of it is deferred to call time.""" - -import asyncio -import threading -import weakref -from asyncio import AbstractEventLoop -from collections import OrderedDict -from collections.abc import Callable, Mapping -from dataclasses import dataclass -from types import MappingProxyType -from typing import TYPE_CHECKING, Final, TypeAlias, TypeVar - -from litellm.exceptions import BadRequestError, ServiceUnavailableError, Timeout - -if TYPE_CHECKING: - from pymongo import AsyncMongoClient, MongoClient - -PYMONGO_INSTALL_HINT: Final = ( - "The MongoDB vector store requires the 'pymongo' package. " - "Run 'pip install litellm[mongodb]' (or 'pip install pymongo') to install it." -) - -MONGODB_PROVIDER: Final = "mongodb" - - -def config_error(message: str) -> BadRequestError: - """400 rather than the 500 a bare ValueError becomes once litellm.exception_type wraps it.""" - return BadRequestError(message=message, model=None, llm_provider=MONGODB_PROVIDER) - - -def timeout_error(message: str) -> Timeout: - return Timeout(message=message, model=None, llm_provider=MONGODB_PROVIDER) - - -def unavailable_error(message: str) -> ServiceUnavailableError: - """litellm only retries 408, 409, 429 and 5xx, so a 400 here would make a failover permanent.""" - return ServiceUnavailableError(message=message, model=None, llm_provider=MONGODB_PROVIDER) - - -DEFAULT_CONNECT_TIMEOUT_MS: Final = 10_000 -DEFAULT_SOCKET_TIMEOUT_MS: Final = 30_000 -DEFAULT_SERVER_SELECTION_TIMEOUT_MS: Final = 10_000 - -_MAX_CACHED_CLIENTS: Final = 32 - -_APP_NAME: Final = "litellm" - - -@dataclass(frozen=True, slots=True) -class MongoClientKey: - connection_string: str - connect_timeout_ms: int - socket_timeout_ms: int - server_selection_timeout_ms: int - - -SyncClientFactory: TypeAlias = Callable[..., "MongoClient"] -AsyncClientFactory: TypeAlias = Callable[..., "AsyncMongoClient"] - -_K = TypeVar("_K") -_V = TypeVar("_V") - -_AsyncClientCacheKey: TypeAlias = tuple[MongoClientKey, int] -# CPython recycles id() aggressively, so the id alone would hand a new loop a closed loop's client -_AsyncClientEntry: TypeAlias = tuple["weakref.ref[AbstractEventLoop]", "AsyncMongoClient"] - -_SyncClientCache: TypeAlias = "OrderedDict[MongoClientKey, MongoClient]" -_AsyncClientCache: TypeAlias = "OrderedDict[_AsyncClientCacheKey, _AsyncClientEntry]" - -_sync_clients: Final[_SyncClientCache] = OrderedDict() # mutable-ok: process-level client cache -_async_clients: Final[_AsyncClientCache] = OrderedDict() # mutable-ok: same cache, per loop -# async searches reach the sync client through executor threads, so both caches are shared state -_cache_lock: Final = threading.Lock() - - -def _store_bounded(cache: "OrderedDict[_K, _V]", cache_key: "_K", value: "_V") -> None: - """Eviction only drops this cache's reference; an in-flight search keeps its client alive.""" - with _cache_lock: - cache[cache_key] = value # mutable-ok: an LRU cache is mutable state by definition - cache.move_to_end(cache_key) - while len(cache) > _MAX_CACHED_CLIENTS: - cache.popitem(last=False) - - -def _mark_used(cache: "OrderedDict[_K, _V]", cache_key: "_K") -> None: - with _cache_lock: - if cache_key in cache: - cache.move_to_end(cache_key) - - -def import_sync_mongo_client() -> "type[MongoClient]": - try: - from pymongo import MongoClient as SyncMongoClient - except ImportError as e: - raise config_error(PYMONGO_INSTALL_HINT) from e - return SyncMongoClient - - -def import_async_mongo_client() -> "type[AsyncMongoClient]": - try: - from pymongo import AsyncMongoClient as AsyncMongoClientClass - except ImportError as e: - raise config_error(PYMONGO_INSTALL_HINT) from e - return AsyncMongoClientClass - - -def _client_kwargs(key: MongoClientKey) -> Mapping[str, object]: - return MappingProxyType( - { - "connectTimeoutMS": key.connect_timeout_ms, - "socketTimeoutMS": key.socket_timeout_ms, - "serverSelectionTimeoutMS": key.server_selection_timeout_ms, - "appname": _APP_NAME, - } - ) - - -def get_sync_client(key: MongoClientKey, client_class: SyncClientFactory | None = None) -> "MongoClient": - cached: Final = _sync_clients.get(key) - if cached is not None: - _mark_used(_sync_clients, key) - return cached - build: Final = client_class if client_class is not None else import_sync_mongo_client() - client: Final = build(key.connection_string, **_client_kwargs(key)) - _store_bounded(_sync_clients, key, client) - return client - - -def _purge_dead_loops() -> None: - """A cached client holds its loop alive, so a closed loop's entry would pin that client and its - sockets for the life of the process.""" - with _cache_lock: - for stale in tuple( - cache_key - for cache_key, (loop_ref, _) in _async_clients.items() - if (cached_loop := loop_ref()) is None or cached_loop.is_closed() - ): - del _async_clients[stale] - - -def get_async_client(key: MongoClientKey, client_class: AsyncClientFactory | None = None) -> "AsyncMongoClient": - """Async clients bind to the loop that created them, so the cache is keyed per loop.""" - loop: Final = asyncio.get_running_loop() - loop_key: Final = (key, id(loop)) - cached: Final = _async_clients.get(loop_key) - if cached is not None and cached[0]() is loop: - _mark_used(_async_clients, loop_key) - return cached[1] - _purge_dead_loops() - build: Final = client_class if client_class is not None else import_async_mongo_client() - client: Final = build(key.connection_string, **_client_kwargs(key)) - _store_bounded(_async_clients, loop_key, (weakref.ref(loop), client)) - return client - - -def reset_client_cache() -> None: - with _cache_lock: - _sync_clients.clear() - _async_clients.clear() - - -_AUTHENTICATION_FAILED_CODE: Final = 18 -_UNAUTHORIZED_CODE: Final = 13 -# Atlas reports a rejected user as code 8000 "AtlasError" where a self-managed mongod reports 18 -_AUTHENTICATION_MESSAGE_MARKERS: Final = ("bad auth", "authentication failed", "not authorized") -_RESOLUTION_TIMEOUT_MARKERS: Final = ("resolution lifetime expired", "dns operation timed out") -_UNKNOWN_HOSTNAME_MARKERS: Final = ("dns query name does not exist", "name or service not known") -_CREDENTIAL_ESCAPING_MARKERS: Final = ("must be escaped according to rfc 3986", "bad database name") - - -def _index_hint(index_name: str, database: str, collection: str) -> str: - return ( - f"No queryable MongoDB Vector Search index named '{index_name}' was found on " - f"'{database}.{collection}'. Confirm the index exists on that exact collection, that its " - "status is READY rather than still building, and that the vector store id matches the index name." - ) - - -def missing_index_error(index_name: str, database: str, collection: str) -> BadRequestError: - """$vectorSearch against a missing index, database or collection returns zero documents rather - than failing, so an empty result set is checked against the catalogue and reported as this.""" - return config_error( - f"{_index_hint(index_name, database, collection)} A vector search against a database, " - "collection or index that does not exist returns no results rather than an error, so this " - "was reported as an empty result set by MongoDB." - ) - - -def index_not_ready_error(index_name: str, database: str, collection: str, status: str) -> BadRequestError: - return config_error( - f"The MongoDB Vector Search index '{index_name}' on '{database}.{collection}' is not queryable " - f"yet; its status is {status}. Searches against it return no results until the build finishes." - ) - - -def translate_mongo_error(error: Exception, index_name: str, database: str, collection: str) -> Exception: - """Returns the exception to raise, so callers keep the driver error as ``__cause__``.""" - try: - from pymongo.errors import ( - ConfigurationError, - ConnectionFailure, - ExecutionTimeout, - InvalidOperation, - NetworkTimeout, - OperationFailure, - ServerSelectionTimeoutError, - ) - except ImportError: - return error - - if isinstance(error, ServerSelectionTimeoutError): - return timeout_error( - "Could not reach the MongoDB deployment before the timeout. On Atlas this is usually the " - "project's IP access list not containing this host, or a paused cluster. On a self-managed " - "deployment it is usually the host or port in the URI, or a firewall between this process " - f"and mongod. Either way it can also be an unresolvable hostname. Driver detail: {error}" - ) - # ExecutionTimeout subclasses OperationFailure, so it has to be matched before it - if isinstance(error, (NetworkTimeout, ExecutionTimeout)): - return timeout_error( - f"The MongoDB vector search against '{database}.{collection}' timed out before returning. " - f"Driver detail: {error}" - ) - # ServerSelectionTimeoutError and NetworkTimeout also subclass ConnectionFailure, so this only - # sees what those branches left - if isinstance(error, ConnectionFailure): - return unavailable_error( - f"The connection to '{database}.{collection}' was dropped or refused. That is usually a " - "replica set failover or a restarted node, so the search is worth retrying. If it keeps " - "happening: on Atlas the usual cause is a connection string with no username and password, " - "or a TLS failure, so confirm the URI is the one Atlas shows under Connect, Drivers; on a " - "self-managed deployment, check that mongod is listening on the host and port in the URI. " - f"Driver detail: {error}" - ) - if isinstance(error, OperationFailure): - code: Final = error.code - detail: Final = str(error).lower() - if code in (_AUTHENTICATION_FAILED_CODE, _UNAUTHORIZED_CODE) or any( - marker in detail for marker in _AUTHENTICATION_MESSAGE_MARKERS - ): - return config_error( - "MongoDB rejected the credentials in mongodb_connection_string, or the database user " - f"lacks read access to '{database}.{collection}'. Driver detail: {error.details}" - ) - if "dimension" in detail: - return config_error( - "The query embedding does not match the vector dimensions the index was built for. " - "litellm_embedding_model must be the same model that produced the stored vectors. " - f"Driver detail: {error}" - ) - if "is not indexed as vector" in detail: - return config_error( - "mongodb_embedding_field names a field the MongoDB Vector Search index does not cover. " - f"It must match the 'path' the index '{index_name}' was created on. Driver detail: {error}" - ) - if "index" in detail and ("not found" in detail or "does not exist" in detail or "unknown" in detail): - return config_error(f"{_index_hint(index_name, database, collection)} Driver detail: {error}") - return config_error( - f"MongoDB rejected the vector search against '{database}.{collection}' using index " - f"'{index_name}'. Driver detail: {error}" - ) - if isinstance(error, ConfigurationError): - configuration_detail: Final = str(error).lower() - if any(marker in configuration_detail for marker in _RESOLUTION_TIMEOUT_MARKERS): - return timeout_error( - "The DNS lookup for the cluster in mongodb_connection_string did not finish in time. " - "A mongodb+srv:// URI needs an SRV lookup before any connection is attempted, so this " - f"is DNS or the configured timeout, not MongoDB. Driver detail: {error}" - ) - if any(marker in configuration_detail for marker in _UNKNOWN_HOSTNAME_MARKERS): - return config_error( - "The hostname in mongodb_connection_string does not exist in DNS. On Atlas, check the " - "cluster name against the URI shown under Connect, Drivers. On a self-managed deployment, " - f"check that the hostname resolves from this process. Driver detail: {error}" - ) - if any(marker in configuration_detail for marker in _CREDENTIAL_ESCAPING_MARKERS): - return config_error( - "mongodb_connection_string could not be parsed. A username or password containing " - "'@', '/', ':' or '%' has to be percent-encoded per RFC 3986, so 'p@ss/word' becomes " - "'p%40ss%2Fword'. If the credentials are already encoded, check the database name in " - f"the URI path instead. Driver detail: {error}" - ) - return config_error( - f"mongodb_connection_string is not a usable MongoDB connection string. Driver detail: {error}" - ) - if isinstance(error, InvalidOperation): - return config_error(f"The MongoDB client was already closed or is unusable. Driver detail: {error}") - # An unreadable tlsCAFile or tlsCertificateKeyFile raises OSError, not a PyMongoError - if isinstance(error, OSError) and error.filename: - return config_error( - f"'{error.filename}', named by a TLS option in mongodb_connection_string, could not be read. " - "Check that tlsCAFile and tlsCertificateKeyFile point at files this process can open; inside " - f"a container that is the path in the container, not on the host. Driver detail: {error}" - ) - # pymongo raises a plain ValueError, not a PyMongoError, for an unusable port - if isinstance(error, ValueError): - return config_error( - "The host and port in mongodb_connection_string could not be parsed. If the port is a " - "number between 0 and 65535, the cause is usually an unescaped ':' in the password, which " - f"has to be percent-encoded per RFC 3986 as '%3A'. Driver detail: {error}" - ) - return error diff --git a/litellm/llms/mongodb/vector_stores/transformation.py b/litellm/llms/mongodb/vector_stores/transformation.py index 3382c931c96..a59f39d3be8 100644 --- a/litellm/llms/mongodb/vector_stores/transformation.py +++ b/litellm/llms/mongodb/vector_stores/transformation.py @@ -1,37 +1,29 @@ -"""MongoDB Vector Search has no HTTP query API, so this is a direct provider that runs the -``$vectorSearch`` aggregation through pymongo. ``vector_store_id`` is the search index name.""" - -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence +from ipaddress import ip_address +from math import isfinite from types import MappingProxyType -from typing import TYPE_CHECKING, Final, NoReturn +from typing import TYPE_CHECKING, Final, Literal, NoReturn +from urllib.parse import quote, urlsplit import httpx -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from litellm.exceptions import AuthenticationError, BadRequestError, ServiceUnavailableError, Timeout +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.vector_store.transformation import ( - BaseDirectVectorStoreConfig, + BaseQueryEmbeddingVectorStoreConfig, LiteLLMVectorStoreEmbeddingExecutor, VectorStoreEmbeddingExecutor, ) -from litellm.llms.mongodb.common_utils import ( - DEFAULT_CONNECT_TIMEOUT_MS, - DEFAULT_SERVER_SELECTION_TIMEOUT_MS, - DEFAULT_SOCKET_TIMEOUT_MS, - MongoClientKey, - config_error, - get_async_client, - get_sync_client, - index_not_ready_error, - missing_index_error, - translate_mongo_error, -) +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import ( + BaseVectorStoreAuthCredentials, VectorStoreCreateOptionalRequestParams, - VectorStoreResultContent, + VectorStoreIndexEndpoints, VectorStoreSearchOptionalRequestParams, VectorStoreSearchResponse, - VectorStoreSearchResult, ) if TYPE_CHECKING: @@ -39,26 +31,45 @@ if TYPE_CHECKING: DEFAULT_EMBEDDING_FIELD_NAME: Final = "embedding" DEFAULT_TEXT_FIELD_NAME: Final = "text" -SCORE_FIELD_NAME: Final = "score" - DEFAULT_MAX_NUM_RESULTS: Final = 10 MIN_MAX_NUM_RESULTS: Final = 1 MAX_MAX_NUM_RESULTS: Final = 50 - NUM_CANDIDATES_MULTIPLIER: Final = 10 MIN_NUM_CANDIDATES: Final = 100 MAX_NUM_CANDIDATES: Final = 10_000 - MAX_QUERY_CHARACTERS: Final = 32_000 - _EMPTY_EMBEDDING_CONFIG: Final = MappingProxyType({}) - _SEARCH_ONLY_MESSAGE: Final = ( "MongoDB vector store is search-only. Create the collection and its MongoDB Vector Search " "index in MongoDB directly, then register it here by index name." ) +def config_error(message: str) -> BadRequestError: + return BadRequestError(message=message, model=None, llm_provider="mongodb") + + +class _Content(BaseModel): + model_config = ConfigDict(frozen=True, strict=True) + type: Literal["text"] + text: str + + +class _Result(BaseModel): + model_config = ConfigDict(frozen=True, strict=True, allow_inf_nan=False) + score: float | None + content: Sequence[_Content] + file_id: str | None + filename: str | None + + +class _SearchResponse(BaseModel): + model_config = ConfigDict(frozen=True, strict=True) + object: Literal["vector_store.search_results.page"] + search_query: str + data: Sequence[_Result] + + class _MongoDBSearchParams(BaseModel): """Typed view over the vector store's litellm_params; unrelated keys are ignored.""" @@ -66,7 +77,6 @@ class _MongoDBSearchParams(BaseModel): litellm_embedding_model: str | None = None litellm_embedding_config: Mapping[str, object] | None = None - mongodb_connection_string: str | None = None mongodb_database: str | None = None mongodb_collection: str | None = None mongodb_text_field: str | None = None @@ -91,21 +101,6 @@ class _MongoDBSearchParams(BaseModel): ) return self.litellm_embedding_model - def require_connection_string(self) -> str: - if not self.mongodb_connection_string: - raise config_error( - "mongodb_connection_string is required in litellm_params for the MongoDB vector store. " - "Example: mongodb+srv://:@.mongodb.net for Atlas, or " - "mongodb://:@:27017 for a self-managed deployment" - ) - scheme: Final = self.mongodb_connection_string.split("://", 1)[0].lower() - if scheme not in ("mongodb", "mongodb+srv"): - raise config_error( - "mongodb_connection_string must start with 'mongodb://' or 'mongodb+srv://', " - f"got '{self.mongodb_connection_string.split('://', 1)[0]}://'" - ) - return self.mongodb_connection_string - def require_database(self) -> str: if not self.mongodb_database: raise config_error( @@ -127,30 +122,28 @@ _MONGODB_PARAM_PREFIX: Final = "mongodb_" _KNOWN_MONGODB_PARAMS: Final = frozenset( name for name in _MongoDBSearchParams.model_fields if name.startswith(_MONGODB_PARAM_PREFIX) ) +_RESPONSE_ADAPTER: Final = TypeAdapter(VectorStoreSearchResponse) -class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): - def __init__( - self, - embedding_executor: VectorStoreEmbeddingExecutor | None = None, - sync_client_factory: Callable[[MongoClientKey], object] | None = None, - async_client_factory: Callable[[MongoClientKey], object] | None = None, - ) -> None: - super().__init__() - self.embedding_executor: Final[VectorStoreEmbeddingExecutor] = ( - embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor() - ) - self.sync_client_factory: Final[Callable[[MongoClientKey], object]] = ( - sync_client_factory if sync_client_factory is not None else get_sync_client - ) - self.async_client_factory: Final[Callable[[MongoClientKey], object]] = ( - async_client_factory if async_client_factory is not None else get_async_client - ) +class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): + def __init__(self, embedding_executor: VectorStoreEmbeddingExecutor | None = None) -> None: + self.embedding_executor: Final = embedding_executor or LiteLLMVectorStoreEmbeddingExecutor() + + def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials: + return BaseVectorStoreAuthCredentials() + + def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: + return VectorStoreIndexEndpoints(read=[], write=[]) # mutable-ok: the TypedDict declares list fields @staticmethod def _reject_unknown_params(litellm_params: Mapping[str, object]) -> None: """Without this a mistyped mongodb_collection reads as 'mongodb_collection is required', naming a key the reader can see they have set.""" + if litellm_params.get("mongodb_connection_string") is not None: + raise config_error( + "MongoDB vector stores now use the BETA sidecar. Move mongodb_connection_string to " + "MONGODB_CONNECTION_STRING in the sidecar, remove it from LiteLLM, and configure api_base and api_key." + ) unknown: Final = sorted( key for key in litellm_params if key.startswith(_MONGODB_PARAM_PREFIX) and key not in _KNOWN_MONGODB_PARAMS ) @@ -191,239 +184,203 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): return configured return min(max(limit * NUM_CANDIDATES_MULTIPLIER, MIN_NUM_CANDIDATES), MAX_NUM_CANDIDATES) - @staticmethod - def _timeout_ms(timeout: float | httpx.Timeout | None) -> tuple[int, int]: - """The connect and socket budgets pymongo is built with, in that order.""" - if isinstance(timeout, httpx.Timeout): - return ( - int((timeout.connect or DEFAULT_CONNECT_TIMEOUT_MS / 1000) * 1000), - int((timeout.read or DEFAULT_SOCKET_TIMEOUT_MS / 1000) * 1000), + def validate_environment( + self, headers: Mapping[str, object], litellm_params: GenericLiteLLMParams | None + ) -> dict[str, object]: # mutable-ok: the shared HTTP handler requires writable headers + if litellm_params is None: + raise config_error("Configure api_base and api_key for the MongoDB BETA sidecar.") + self._reject_unknown_params(MappingProxyType(dict(litellm_params))) + api_key: Final = litellm_params.api_key or get_secret_str("MONGODB_SIDECAR_API_KEY") + if not api_key: + raise config_error("MongoDB sidecar api_key is required. Set api_key or MONGODB_SIDECAR_API_KEY.") + return { + **headers, + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } # mutable-ok: writable HTTP headers + + def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + if not api_base: + raise config_error("MongoDB sidecar api_base is required, for example http://127.0.0.1:8080.") + try: + parsed: Final = urlsplit(api_base) + valid: Final = parsed.scheme in ("http", "https") and bool(parsed.hostname) and parsed.port != 0 + except ValueError: + raise config_error("MongoDB sidecar api_base must be a valid HTTP or HTTPS URL.") from None + if not valid or parsed.username or parsed.password or parsed.query or parsed.fragment: + raise config_error( + "MongoDB sidecar api_base must be an HTTP or HTTPS URL without credentials, query, or fragment." ) - if timeout is None: - return DEFAULT_CONNECT_TIMEOUT_MS, DEFAULT_SOCKET_TIMEOUT_MS - return min(int(float(timeout) * 1000), DEFAULT_CONNECT_TIMEOUT_MS), int(float(timeout) * 1000) + if parsed.scheme == "http": + try: + loopback: Final = ip_address(parsed.hostname or "").is_loopback + except ValueError: + raise config_error( + "MongoDB sidecar requires HTTPS. HTTP is supported only for a loopback IP such as 127.0.0.1." + ) from None + if not loopback: + raise config_error( + "MongoDB sidecar requires HTTPS. HTTP is supported only for a loopback IP such as 127.0.0.1." + ) + return api_base.rstrip("/") + + @staticmethod + def _timeout_ms(value: object) -> int: + seconds: Final = value.read if isinstance(value, httpx.Timeout) else value + if seconds is None: + return 30_000 + if not isinstance(seconds, (int, float)) or not isfinite(seconds) or seconds <= 0: + raise config_error("MongoDB search timeout must be a positive finite number.") + try: + return max(1, int(seconds * 1000)) + except (ValueError, OverflowError): + raise config_error("MongoDB search timeout must be a positive finite number.") from None @classmethod - def _client_key(cls, params: _MongoDBSearchParams, timeout: float | httpx.Timeout | None) -> MongoClientKey: - connect_ms, socket_ms = cls._timeout_ms(timeout) - return MongoClientKey( - connection_string=params.require_connection_string(), - connect_timeout_ms=connect_ms, - socket_timeout_ms=socket_ms, - server_selection_timeout_ms=min(connect_ms, DEFAULT_SERVER_SELECTION_TIMEOUT_MS), - ) + def _params( + cls, + litellm_params: Mapping[str, object], + optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Mapping[str, object] | None, + ) -> _MongoDBSearchParams: + cls._reject_unknown_params(litellm_params) + if extra_body: + raise config_error("MongoDB vector store does not support extra_body overrides.") + for unsupported in ("filters", "ranking_options", "rewrite_query"): + if optional_params.get(unsupported) is not None: + raise config_error(f"MongoDB vector store does not support the {unsupported} parameter.") + try: + params: Final = _MongoDBSearchParams.model_validate(litellm_params) + except ValidationError: + raise config_error( + "Invalid MongoDB vector-store configuration. Check the database, collection, fields, and candidate count." + ) from None + params.require_database() + params.require_collection() + params.require_embedding_model() + cls._num_candidates(cls._limit(optional_params), params.mongodb_num_candidates) + cls._timeout_ms(litellm_params.get("timeout")) + return params @classmethod - def _pipeline( + def _request( cls, vector_store_id: str, - query_vector: Sequence[float], + query_text: str, params: _MongoDBSearchParams, - vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - ) -> Sequence[Mapping[str, object]]: - if vector_store_search_optional_params.get("filters") is not None: + optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + embedding_response: EmbeddingResponse, + timeout: object, + ) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body + if not embedding_response.data: raise config_error( - "MongoDB vector store does not support the filters parameter yet. " - "Restrict the collection or the MongoDB Vector Search index definition instead." + "The embedding model returned no embedding for the search query. Check litellm_embedding_model." ) - if vector_store_search_optional_params.get("ranking_options") is not None: - raise config_error( - "MongoDB vector store does not support the ranking_options parameter yet. " - "Every result already carries the vectorSearchScore, so filter or re-rank " - "on that rather than having the threshold silently ignored." - ) - if vector_store_search_optional_params.get("rewrite_query") is not None: - raise config_error( - "MongoDB vector store does not support the rewrite_query parameter. The query is " - "embedded exactly as sent; rewrite it before calling if you need that." - ) - limit: Final = cls._limit(vector_store_search_optional_params) - search: Final = MappingProxyType( - { - "index": vector_store_id, - "path": params.embedding_field, - "queryVector": tuple(query_vector), - "numCandidates": cls._num_candidates(limit, params.mongodb_num_candidates), - "limit": limit, - } - ) - projection: Final = MappingProxyType( - {params.text_field: 1, SCORE_FIELD_NAME: MappingProxyType({"$meta": "vectorSearchScore"})} - ) - return [ # mutable-ok: pymongo rejects any non-list pipeline in common.validate_list - MappingProxyType({"$vectorSearch": search}), - MappingProxyType({"$project": projection}), - ] - - @classmethod - def _field_value(cls, document: Mapping[str, object], dotted_path: str) -> str | None: - """None means absent, which is what separates a mistyped field from genuinely empty text.""" - head, _, rest = dotted_path.partition(".") - if head not in document: - return None - value: Final = document[head] - if not rest: - return None if value is None else str(value) - return cls._field_value(value, rest) if isinstance(value, Mapping) else None - - @classmethod - def _to_result(cls, document: Mapping[str, object], text_field: str) -> VectorStoreSearchResult: - document_id: Final = document.get("_id") - identifier: Final = None if document_id is None else str(document_id) - content: Final = [ # mutable-ok: VectorStoreSearchResult declares a list of content parts - VectorStoreResultContent(text=cls._field_value(document, text_field) or "", type="text") - ] - raw_score: Final = document.get(SCORE_FIELD_NAME) - return VectorStoreSearchResult( - score=float(raw_score) if isinstance(raw_score, (int, float)) else None, - content=content, - file_id=identifier, - filename=identifier, + vector: Final = embedding_response.data[0]["embedding"] + if not vector or any(not isinstance(value, (float, int)) or not isfinite(value) for value in vector): + raise config_error("The embedding model must return a non-empty, finite query vector.") + limit: Final = cls._limit(optional_params) + return ( + f"{api_base}/v1/vector_stores/{quote(vector_store_id, safe='')}/search", + { # mutable-ok: JSON transport requires a dict + "query": query_text, + "query_vector": tuple(vector), + "mongodb_database": params.require_database(), + "mongodb_collection": params.require_collection(), + "mongodb_embedding_field": params.embedding_field, + "mongodb_text_field": params.text_field, + "mongodb_num_candidates": cls._num_candidates(limit, params.mongodb_num_candidates), + "max_num_results": limit, + "timeout_ms": cls._timeout_ms(timeout), + }, ) - @classmethod - def _raise_for_missing_text_field( - cls, documents: Sequence[Mapping[str, object]], text_field: str, database: str, collection: str - ) -> None: - """$vectorSearch matches documents carrying no text, so a mistyped mongodb_text_field - returns well-scored results with empty content instead of failing.""" - if documents and all(cls._field_value(document, text_field) is None for document in documents): - raise config_error( - f"None of the {len(documents)} matched documents in '{database}.{collection}' has a " - f"'{text_field}' field, so every result would carry empty text. Set mongodb_text_field " - "to the field holding the readable text; it accepts a dotted path such as metadata.body." - ) - - @classmethod - def _to_response( - cls, documents: Sequence[Mapping[str, object]], query_text: str, text_field: str - ) -> VectorStoreSearchResponse: - return VectorStoreSearchResponse( - object="vector_store.search_results.page", - search_query=query_text, - data=[ # mutable-ok: VectorStoreSearchResponse declares data as a list - cls._to_result(document, text_field) for document in documents - ], - ) - - @staticmethod - def _raise_for_unusable_index( - catalogue: Sequence[Mapping[str, object]], index_name: str, database: str, collection: str - ) -> None: - """mongod returns zero documents both for a query that matched nothing and for a missing - database, collection or index, so the catalogue decides which one happened.""" - if not catalogue: - raise missing_index_error(index_name, database, collection) - entry: Final = catalogue[0] - if not entry.get("queryable"): - raise index_not_ready_error(index_name, database, collection, str(entry.get("status") or "unknown")) - - @staticmethod - def _embedding_vector(embedding_response: EmbeddingResponse) -> Sequence[float]: - data: Final = embedding_response.data - if not data: - raise config_error( - "The embedding model returned no embedding for the search query, so there is nothing " - "to search MongoDB with. Check the embedding deployment named by litellm_embedding_model." - ) - return data[0]["embedding"] - - def execute_search_vector_store_request( + def transform_search_vector_store_request( self, vector_store_id: str, query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, litellm_logging_obj: "LiteLLMLoggingObj", litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, - timeout: float | httpx.Timeout | None = None, - ) -> VectorStoreSearchResponse: - self._reject_unknown_params(litellm_params) - params: Final = _MongoDBSearchParams.model_validate(litellm_params) + ) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body + params: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body) query_text: Final = self._query_text(query) - key: Final = self._client_key(params, timeout) - database: Final = params.require_database() - collection: Final = params.require_collection() - - embedding_response: Final = (embedding_executor or self.embedding_executor).embed( - params.require_embedding_model(), + response: Final = (embedding_executor or self.embedding_executor).embed( + params.require_embedding_model(), query_text, params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG + ) + return self._request( + vector_store_id, query_text, - params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG, - ) - pipeline: Final = self._pipeline( - vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params + params, + vector_store_search_optional_params, + api_base, + response, + litellm_params.get("timeout"), ) - try: - client: Final = self.sync_client_factory(key) - target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted - documents: Final = tuple(target.aggregate(pipeline)) - except Exception as e: - raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e - if not documents: - try: - catalogue: Final = tuple(target.list_search_indexes(vector_store_id)) - except Exception as e: - raise translate_mongo_error( - e, index_name=vector_store_id, database=database, collection=collection - ) from e - self._raise_for_unusable_index(catalogue, vector_store_id, database, collection) - self._raise_for_missing_text_field(documents, params.text_field, database, collection) - return self._to_response(documents, query_text, params.text_field) - - async def aexecute_search_vector_store_request( + async def atransform_search_vector_store_request( self, vector_store_id: str, query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, litellm_logging_obj: "LiteLLMLoggingObj", litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, - timeout: float | httpx.Timeout | None = None, - ) -> VectorStoreSearchResponse: - self._reject_unknown_params(litellm_params) - params: Final = _MongoDBSearchParams.model_validate(litellm_params) + ) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body + params: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body) query_text: Final = self._query_text(query) - key: Final = self._client_key(params, timeout) - database: Final = params.require_database() - collection: Final = params.require_collection() - - embedding_response: Final = await (embedding_executor or self.embedding_executor).aembed( - params.require_embedding_model(), + response: Final = await (embedding_executor or self.embedding_executor).aembed( + params.require_embedding_model(), query_text, params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG + ) + return self._request( + vector_store_id, query_text, - params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG, - ) - pipeline: Final = self._pipeline( - vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params + params, + vector_store_search_optional_params, + api_base, + response, + litellm_params.get("timeout"), ) + def transform_search_vector_store_response( + self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj" + ) -> VectorStoreSearchResponse: try: - client: Final = self.async_client_factory(key) - target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted - cursor: Final = await target.aggregate(pipeline) - documents: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly - document async for document in cursor - ] - except Exception as e: - raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e - if not documents: - try: - index_cursor: Final = await target.list_search_indexes(vector_store_id) - catalogue: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly - entry async for entry in index_cursor - ] - except Exception as e: - raise translate_mongo_error( - e, index_name=vector_store_id, database=database, collection=collection - ) from e - self._raise_for_unusable_index(catalogue, vector_store_id, database, collection) - self._raise_for_missing_text_field(documents, params.text_field, database, collection) - return self._to_response(documents, query_text, params.text_field) + validated: Final = _SearchResponse.model_validate_json(response.content) + return _RESPONSE_ADAPTER.validate_python(validated.model_dump()) + except ValidationError: + raise ServiceUnavailableError( + message="MongoDB sidecar returned an invalid search response. Check the sidecar version and deployment.", + model=None, + llm_provider="mongodb", + ) from None + + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, object] | httpx.Headers + ) -> BaseLLMException: + if status_code == 400: + raise config_error(error_message) + if status_code == 401: + raise AuthenticationError(message="MongoDB sidecar rejected api_key.", model=None, llm_provider="mongodb") + if status_code == 408: + raise Timeout(message=error_message, model=None, llm_provider="mongodb") + raise ServiceUnavailableError( + message="MongoDB sidecar is unavailable. Check its address, health, and logs.", + model=None, + llm_provider="mongodb", + ) + + def validate_create_vector_store(self) -> NoReturn: + raise config_error(_SEARCH_ONLY_MESSAGE) def transform_create_vector_store_request( - self, - vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, - api_base: str, + self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str ) -> NoReturn: raise config_error(_SEARCH_ONLY_MESSAGE) diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 6e9bb83b0a0..1b494ebad47 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -290,19 +290,14 @@ def handle_cohere_stream_chunk( ) -> ModelResponseStream: """Parse a single Cohere SSE chunk into a LiteLLM ModelResponseStream. - ``prior_tool_calls_emitted`` lets the caller signal whether tool calls - were already emitted in earlier chunks of the same stream. When set, the - terminal consolidation chunk's tool calls are suppressed (they would - duplicate prior deltas); otherwise they are passed through so a stream - that delivers tool calls only on the terminal chunk doesn't silently - drop them. - - ``prior_text_emitted`` plays the analogous role for the ``text`` field: - when set, the terminal consolidation chunk's ``text`` is suppressed - (it would re-emit the full assembled response on top of prior deltas); - when unset (e.g. a degenerate stream that delivers the entire response - in a single SSE event carrying both ``chatHistory`` and ``finishReason``), - the text is passed through so the response content isn't silently lost. + OCI Cohere streams the answer as single-token ``text`` deltas, then restates + the whole assembled ``text`` on every chunk that carries ``toolCalls`` or + ``chatHistory`` (the tool-calls event and the terminal event). Once the + caller reports that earlier chunks already emitted text + (``prior_text_emitted``), those restatements are dropped so the client does + not see the answer twice; a stream whose only text lives on such a chunk + keeps it. ``prior_tool_calls_emitted`` plays the same role for the tool + calls the terminal ``chatHistory`` chunk repeats. """ try: typed_chunk: Final = CohereStreamChunk.model_validate(dict_chunk) @@ -315,33 +310,10 @@ def handle_cohere_stream_chunk( if typed_chunk.index is None: typed_chunk.index = 0 - # OCI Cohere's terminal SSE event re-sends the full assembled response in - # `text` alongside a populated `chatHistory` and a non-null `finishReason`. - # Emitting that text would concatenate the whole response onto the - # already-streamed deltas. We require both signals to be present so that a - # future API change which adds `chatHistory` to intermediate chunks (or a - # rare early-populated case) doesn't silently drop legitimate token deltas. - is_terminal_consolidation: Final = typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None - # On non-terminal text-free chunks (e.g. tool-call-only or keep-alive - # chunks) emit ``content=None`` rather than ``content=""`` so downstream - # stream-mergers that distinguish "no text in this delta" from "an - # explicitly empty text delta" behave correctly. - # - # We only suppress the terminal chunk's ``text`` when the caller has - # confirmed that text deltas were already emitted earlier — otherwise - # (e.g. a degenerate stream that delivers the whole response in a - # single SSE event), passing it through is the only chance to surface it. - text: Final[str | None] = None if (is_terminal_consolidation and prior_text_emitted) else typed_chunk.text - - # Tool calls on the terminal consolidation chunk (whether from - # `typed_chunk.toolCalls` or from `chatHistory`) typically restate what - # was already streamed in intermediate chunks. Re-emitting them would - # mint fresh `uuid4` IDs and cause downstream consumers to execute each - # tool call twice. We only suppress when the caller has confirmed that - # tool calls were already emitted earlier — otherwise (e.g. a short - # response that delivers tool calls exclusively on the terminal chunk), - # passing them through is the only chance to surface them. - cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls + restates_text: Final = typed_chunk.chatHistory is not None or typed_chunk.toolCalls is not None + restates_tool_calls: Final = typed_chunk.chatHistory is not None + text: Final[str | None] = None if (restates_text and prior_text_emitted) else typed_chunk.text + cohere_tool_calls: Final = None if (restates_tool_calls and prior_tool_calls_emitted) else typed_chunk.toolCalls tool_calls: list[dict[str, object]] | None = None if cohere_tool_calls: diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 98e23a59eea..8e4e41b4ac1 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -745,13 +745,25 @@ class OCIStreamWrapper(CustomStreamWrapper): # single-event case (terminal chunk carries the only copy of the text). self._cohere_text_emitted = False - def chunk_creator(self, chunk: Any) -> ModelResponseStream: + def _emit_chunk(self, parsed: ModelResponseStream) -> ModelResponseStream: + for choice in parsed.choices: + if getattr(choice.delta, "tool_calls", None): + self.tool_call = True + if choice.finish_reason is not None: + self.received_finish_reason = choice.finish_reason + self.sent_last_chunk = True + return self.model_response_creator(chunk={"choices": parsed.choices}) + + def chunk_creator(self, chunk: Any) -> ModelResponseStream | None: if not isinstance(chunk, str): raise ValueError(f"Chunk is not a string: {chunk}") if not chunk.startswith("data:"): raise ValueError(f"Chunk does not start with 'data:': {chunk}") + payload: Final = chunk[5:].strip() + if payload == "[DONE]": + return None try: - dict_chunk: Final = json.loads(chunk[5:]) + dict_chunk: Final = json.loads(payload) except json.JSONDecodeError as e: raise OCIError( status_code=500, @@ -774,8 +786,8 @@ class OCIStreamWrapper(CustomStreamWrapper): if getattr(choice.delta, "content", None): self._cohere_text_emitted = True break - return result - return handle_generic_stream_chunk(dict_chunk) + return self._emit_chunk(result) + return self._emit_chunk(handle_generic_stream_chunk(dict_chunk)) __all__ = [ diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index dccc83efed4..a1340ba1952 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -16,6 +16,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( custom_prompt, ollama_pt, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import ( + async_inline_remote_media, + inline_remote_image_urls, +) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock @@ -344,6 +348,26 @@ class OllamaConfig(BaseConfig): ) return model_response + @property + def uses_async_transform_request(self) -> bool: + return True + + async def async_transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, object], # mutable-ok: BaseConfig signature + litellm_params: dict[str, object], # mutable-ok: BaseConfig signature + headers: dict[str, object], # mutable-ok: BaseConfig signature + ) -> dict[str, object]: # mutable-ok: BaseConfig signature + return self.transform_request( + model=model, + messages=await async_inline_remote_media(messages, should_inline=inline_remote_image_urls), + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + def transform_request( self, model: str, diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 9afc6331d96..9b410cf073e 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -19,10 +19,12 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, drop_tool_reference_parts_from_tool_messages, + flatten_combinators_and_drop_non_python_regex_patterns, get_tool_call_names, hoist_images_from_tool_messages, - tool_with_flattened_parameters, + tool_with_sanitized_parameters, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, @@ -432,7 +434,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): custom_llm_provider, api_base ) - def _flattened_tools_update_for_openai( + def _sanitized_tools_update_for_openai( self, optional_params: Mapping[str, object], litellm_params: Mapping[str, object], @@ -440,22 +442,26 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ OpenAI's chat completions validator rejects tool `parameters` carrying 'oneOf'/'anyOf'/'allOf'/'enum'/'const'/'not' at the top level for every - model family, unlike the Responses API, where GPT-5+ accepts them. + model family, unlike the Responses API, where GPT-5+ accepts them, and + a `pattern` Python's `re` cannot compile for every model family on both. + A custom api_base on the `openai` provider is usually a proxy in front of + the same validator, so regexes are dropped there too, while the lossier + combinator flattening stays limited to api.openai.com hosts. """ tools: Final = optional_params.get("tools") - if not isinstance(tools, list): - return _NO_TOOLS_UPDATE provider: Final = litellm_params.get("custom_llm_provider") - raw_api_base: Final = litellm_params.get("api_base") - if not self._targets_openai_hosted_endpoint( - provider if isinstance(provider, str) else None, - raw_api_base if isinstance(raw_api_base, str) else None, - ): + if not isinstance(tools, list) or provider != "openai": return _NO_TOOLS_UPDATE - flattened: Final = [ # mutable-ok: request tools are a JSON list - tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + raw_api_base: Final = litellm_params.get("api_base") + sanitize: Final = ( + flatten_combinators_and_drop_non_python_regex_patterns + if self._targets_openai_hosted_endpoint(provider, raw_api_base if isinstance(raw_api_base, str) else None) + else drop_non_python_regex_patterns + ) + sanitized: Final = [ # mutable-ok: request tools are a JSON list + tool_with_sanitized_parameters(tool, sanitize) if isinstance(tool, dict) else tool for tool in tools ] - return MappingProxyType({"tools": flattened}) + return MappingProxyType({"tools": sanitized}) def transform_request( self, @@ -489,7 +495,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "model": model, "messages": messages, **optional_params, - **self._flattened_tools_update_for_openai(optional_params, litellm_params), + **self._sanitized_tools_update_for_openai(optional_params, litellm_params), } async def async_transform_request( @@ -521,7 +527,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "model": model, "messages": transformed_messages, **optional_params, - **self._flattened_tools_update_for_openai(optional_params, litellm_params), + **self._sanitized_tools_update_for_openai(optional_params, litellm_params), } else: ## allow for any object specific behaviour to be handled diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index d41c8557d72..58ff03e6a0d 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -49,6 +49,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import coerce_stream_holdback_value, ) from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + ChatCompletionMessageToolCall, Choices, GenericGuardrailAPIInputs, ModelResponse, @@ -78,6 +80,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + delivers_ended_stream_rewrites = True + assembles_streamed_response = True + def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ Convert chat completions request data to OpenAI-spec structured messages. @@ -453,6 +458,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """ Process output streaming responses by applying guardrails to text content. @@ -467,6 +473,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): accumulated text (``responses_so_far`` is left untouched so it stays a correct raw accumulator across rounds) and the guardrailed text plus requested holdback are reported per choice on the sink. + deliver_ended_stream_rewrites: When True and the buffered stream has + ended, guardrail text rewrites are written back across + ``responses_so_far`` (full rewritten text in each choice's first + content-carrying chunk, the rest blanked) instead of discarded. Returns: The (unmodified) list of responses. @@ -492,6 +502,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, request_data=request_data, + deliver_ended_stream_rewrites=deliver_ended_stream_rewrites, ) async def _process_streaming_block_only( @@ -502,27 +513,23 @@ class OpenAIChatCompletionsHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", request_data: dict | None, + deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """Block-only streaming path: run the guardrail so an in-flight BLOCK can terminate the stream. Text rewrites are not propagated to the client here - (see ``_process_streaming_transform`` for the incremental_diff path).""" + (see ``_process_streaming_transform`` for the incremental_diff path) unless + ``deliver_ended_stream_rewrites`` opts the ended-stream branch in.""" has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) if has_stream_ended: - # convert to model response - model_response: Final = cast( - ModelResponse, - stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), - ) - # run process_output_response - await self.process_output_response( - response=model_response, + await self._process_ended_stream( + responses_so_far=responses_so_far, guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, request_data=request_data, + deliver_ended_stream_rewrites=deliver_ended_stream_rewrites, ) - return responses_so_far # Step 0: Check if any response has text content to process @@ -595,6 +602,48 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return responses_so_far + async def _process_ended_stream( + self, + *, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None", + user_api_key_dict: "UserAPIKeyAuth | None", + request_data: dict[str, object] | None, # mutable-ok: same request-payload shape the hooks take + deliver_ended_stream_rewrites: bool, + ) -> None: + """Ended-stream path: rebuild the full response, run the non-streaming + output guardrail against it, and (when opted in) write any text or + tool-call rewrite back across the buffered chunks.""" + model_response: Final = cast( + ModelResponse, + stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), + ) + pre_guardrail_texts: Final = self._string_choice_contents(model_response) + pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response) + await self.process_output_response( + response=model_response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + if not deliver_ended_stream_rewrites: + return + guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown" + await self._write_ended_stream_text_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_texts=pre_guardrail_texts, + guardrail_name=guardrail_name, + ) + self._write_ended_stream_tool_call_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_name, + ) + def build_stream_error_items( self, exc: "HTTPException", @@ -745,8 +794,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ combined_texts: Final[dict[tuple[int, int | None], str]] = {} - for response_idx, response in enumerate(responses_so_far): - for choice_idx, choice in enumerate(response.choices): + for response in responses_so_far: + for choice in response.choices: if isinstance(choice, litellm.StreamingChoices): content = choice.delta.content elif isinstance(choice, litellm.Choices): @@ -759,7 +808,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - accumulate for this choice - str_key: tuple[int, int | None] = (choice_idx, None) + str_key: tuple[int, int | None] = (choice.index, None) if str_key not in combined_texts: combined_texts[str_key] = "" combined_texts[str_key] += content @@ -770,7 +819,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): text_str = content_item.get("text") if text_str: list_key: tuple[int, int | None] = ( - choice_idx, + choice.index, content_idx, ) if list_key not in combined_texts: @@ -960,6 +1009,117 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "name" in func_dict: existing_tool_call.function.name = func_dict["name"] + @staticmethod + def _string_choice_contents(response: "ModelResponse") -> tuple[str | None, ...]: + return tuple( + choice.message.content if isinstance(choice.message.content, str) else None for choice in response.choices + ) + + async def _write_ended_stream_text_rewrites( + self, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrailed_response: "ModelResponse", + pre_guardrail_texts: tuple[str | None, ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail text rewrites back across the buffered + chunks: the full rewritten text lands in the choice's first + content-carrying chunk and the rest are blanked, the same shape the + in-flight write-back uses. Chunks carrying only finish_reason or usage + stay untouched. A rewrite on a stream carrying more than one distinct + choice index is reported as undeliverable, so the pipeline executor + discards it and releases the original chunks.""" + post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response) + changed: Final = tuple( + after + for before, after in zip(pre_guardrail_texts, post_guardrail_texts) + if before is not None and after is not None and after != before + ) + if not changed: + return + stream_choice_indices: Final = frozenset( + choice.index for response in responses_so_far for choice in response.choices + ) + if len(stream_choice_indices) != 1: + # stream_chunk_builder collapses every choice into one index-0 + # choice, so a rewrite of the rebuilt response cannot be attributed + # back to a single choice on an n>1 stream: report it undeliverable + # rather than deliver the rewrite on the wrong choice + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + target_choice_index: Final = next(iter(stream_choice_indices)) + await self._apply_guardrail_responses_to_output_streaming( + responses=responses_so_far, + guardrailed_texts=list(changed), # mutable-ok: callee takes lists + task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists + ) + + @staticmethod + def _function_tool_call_shapes(response: "ModelResponse") -> tuple[tuple[str | None, str], ...]: + return tuple( + (tool_call.function.name, tool_call.function.arguments) + for choice in response.choices + for tool_call in choice.message.tool_calls or () + if isinstance(tool_call, ChatCompletionMessageToolCall) + ) + + @staticmethod + def _function_tool_call_fragments( + responses_so_far: Sequence["ModelResponseStream"], + ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]: + """Group the stream's function tool-call fragments by their tool-call index, in + the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping + only the indices the builder keeps (an id and a name somewhere in the stream).""" + fragments: Final = tuple( + tool_call + for response in responses_so_far + for choice in response.choices + for tool_call in choice.delta.tool_calls or () + if isinstance(tool_call, ChatCompletionDeltaToolCall) + ) + identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id) + named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name) + return tuple( + tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named) + ) + + def _write_ended_stream_tool_call_rewrites( + self, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrailed_response: "ModelResponse", + pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail tool-call rewrites back across the buffered + chunks: the rewritten name and full arguments land in the tool call's first + fragment and the arguments of its later fragments are blanked, mirroring the + text write-back. A rewrite on a stream carrying more than one distinct choice + index, or whose fragments do not line up with the rebuilt tool calls, is + reported as undeliverable, so the pipeline executor discards it and releases + the original chunks.""" + post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response) + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + stream_choice_indices: Final = frozenset( + choice.index for response in responses_so_far for choice in response.choices + ) + fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) + if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + for before, (name, arguments), fragments in zip( + pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call + ): + if (name, arguments) == before: + continue + head, *tail = fragments + head.function.name = name + head.function.arguments = arguments + for fragment in tail: + fragment.function.arguments = "" + async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], @@ -975,7 +1135,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Args: responses: List of ModelResponseStream objects to modify guardrailed_texts: List of guardrailed text responses (combined from all chunks) - task_mappings: List of tuples (choice_idx, content_idx) + task_mappings: List of tuples (choice_idx, content_idx), where choice_idx + is the choice's ``index`` field, not its position in a chunk's list Override this method to customize how responses are applied to streaming responses. """ @@ -991,9 +1152,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Key: (choice_idx, content_idx), Value: boolean (True if already set) already_set: Final[dict[tuple[int, int | None], bool]] = {} - # Iterate through all responses and update content - for response_idx, response in enumerate(responses): - for choice_idx_in_response, choice in enumerate(response.choices): + # Iterate through all responses and update content, matching each chunk's + # choice by its index field: on n>1 streams a chunk usually carries one + # choice at list position 0 whose index names the logical choice. + for response in responses: + for choice in response.choices: if isinstance(choice, litellm.StreamingChoices): content = choice.delta.content elif isinstance(choice, litellm.Choices): @@ -1006,7 +1169,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - str_key: tuple[int, int | None] = (choice_idx_in_response, None) + str_key: tuple[int, int | None] = (choice.index, None) if str_key in guardrail_map: if str_key not in already_set: # First chunk - set the complete guardrailed text @@ -1027,7 +1190,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for content_idx, content_item in enumerate(content): if "text" in content_item: list_key: tuple[int, int | None] = ( - choice_idx_in_response, + choice.index, content_idx, ) if list_key in guardrail_map: diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index edc8d64d9c2..5f04ebe0c01 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -797,6 +797,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): "api_base": openai_client._base_url._uri_reference, "acompletion": acompletion, "complete_input_dict": data, + "openai_sdk": True, }, ) @@ -938,6 +939,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): "api_base": openai_aclient._base_url._uri_reference, "acompletion": True, "complete_input_dict": data, + "openai_sdk": True, }, ) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2dec2b3f178..33e0a0a923f 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -33,22 +33,23 @@ import time import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass -from itertools import accumulate +from itertools import accumulate, chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast -from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, OpenAiResponsesToChatCompletionStreamIterator, + tool_call_dict_from_output_item, ) from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, + StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, @@ -83,7 +84,6 @@ from litellm.types.llms.openai import ( ) from litellm.types.responses.main import ( GenericResponseOutputItem, - OutputFunctionToolCall, OutputText, ) from litellm.types.utils import GenericGuardrailAPIInputs @@ -100,6 +100,72 @@ if TYPE_CHECKING: from litellm.types.llms.openai import ResponseInputParam +class _ToolCallShape(NamedTuple): + name: str | None + arguments: str + + +class _ToolCallFunctionFields(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str | None = None + arguments: str = "" + + +class _ToolCallFields(BaseModel): + model_config = ConfigDict(frozen=True) + + function: _ToolCallFunctionFields + + +def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]: + return tuple( + _ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", "")) + for tool_call in tool_calls + ) + + +def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None: + payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call + try: + fields: Final = _ToolCallFields.model_validate(payload) + except ValidationError: + return None + return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments) + + +def _post_guardrail_tool_call_shapes( + returned_tool_calls: Sequence[object] | None, + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str | None, +) -> tuple[_ToolCallShape, ...]: + if not pre_guardrail_tool_calls: + return pre_guardrail_tool_calls + if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls): + verbose_proxy_logger.warning( + "OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, " + "leaving the tool call output items unchanged", + guardrail_name, + "no" if returned_tool_calls is None else len(returned_tool_calls), + len(pre_guardrail_tool_calls), + ) + return pre_guardrail_tool_calls + returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls) + validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None) + if len(validated_shapes) != len(returned_shapes): + verbose_proxy_logger.warning( + "OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, " + "leaving the tool call output items unchanged", + guardrail_name, + ) + return pre_guardrail_tool_calls + return validated_shapes + + +def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape: + return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments) + + class ResponseOutputEnvelope(TypedDict, total=False): """Dict form of a Responses API response, as far as guardrail write-back reads it.""" @@ -118,6 +184,29 @@ class ResponsesStreamChunk(TypedDict, total=False): content_index: ReadOnly[int] +_TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( + { + ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, + ResponsesAPIStreamEvents.RESPONSE_FAILED.value, + ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, + } +) + + +_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"}) +_TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"function_call": "arguments", "custom_tool_call": "input"} +) +_TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: Final = frozenset( + {"response.function_call_arguments.delta", "response.custom_tool_call_input.delta"} +) +_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"response.function_call_arguments.done": "arguments", "response.custom_tool_call_input.done": "input"} +) +_TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | frozenset( + _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS +) +_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) _PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( {"function_call_output": "output", "message": "content"} ) @@ -154,8 +243,20 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts -def _is_function_call_item(item: object) -> bool: - return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call") +def _is_tool_call_item(item: object) -> bool: + return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES + + +def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None: + if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES: + return None + if isinstance(item, Mapping): + return cast("Mapping[str, object]", item) # cast-ok: output items are str-keyed JSON objects + return item.model_dump() if isinstance(item, BaseModel) else None + + +def _is_tool_call_output_item(item: object) -> bool: + return _tool_call_output_item_mapping(item) is not None def _last_message_role(messages: Sequence[object]) -> str | None: @@ -179,7 +280,7 @@ def _provenance_unit_bounds( start_indexes: Final = tuple( index for index in range(len(raw_input)) - if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") + if index == 0 or not (_is_tool_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") ) return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input)))) @@ -330,6 +431,9 @@ class OpenAIResponsesHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + delivers_ended_stream_rewrites = True + assembles_streamed_response = True + def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ Convert Responses API request data to OpenAI-spec structured messages. @@ -575,7 +679,7 @@ class OpenAIResponsesHandler(BaseTranslation): - response.output is a list of output items - Each output item can be: * GenericResponseOutputItem with a content list of OutputText objects - * ResponseFunctionToolCall with tool call data + * ResponseFunctionToolCall or CustomToolCallOutputItem with tool call data - Each OutputText object has a text field """ @@ -640,6 +744,7 @@ class OpenAIResponsesHandler(BaseTranslation): if response_model: inputs["model"] = response_model + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -648,6 +753,11 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name, + ) # Step 3: Map guardrail responses back to original response structure await self._apply_guardrail_responses_to_output( @@ -655,6 +765,11 @@ class OpenAIResponsesHandler(BaseTranslation): responses=guardrailed_texts, task_mappings=task_mappings, ) + self._write_tool_call_rewrites_to_output( + tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=post_guardrail_tool_calls, + ) verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response) @@ -667,6 +782,8 @@ class OpenAIResponsesHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, + stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> list[Any]: """ Process output streaming response by applying guardrails to text content. @@ -675,10 +792,18 @@ class OpenAIResponsesHandler(BaseTranslation): chunk, apply the guardrail, then write the result back in-place so the caller sees the modified content (e.g. PII tokens replaced). - For ``response.completed`` events (the normal end-of-stream signal) we - use the same per-item extraction + task-mapping approach as - ``process_output_response`` so that unmasking / blocking works correctly - for every output item. + For terminal envelope events (``response.completed``, and equally + ``response.incomplete`` / ``response.failed``, whose envelopes carry the + partial output) we use the same per-item extraction + task-mapping + approach as ``process_output_response`` so that unmasking / blocking + works correctly for every output item. With + ``deliver_ended_stream_rewrites`` the earlier text-carrying events + (``response.output_text.delta`` / ``.done``, + ``response.content_part.done``, ``response.output_item.done``) are synced + to the rewritten envelope too, so a client reading deltas sees the + rewrite instead of the raw model output; a rewrite observed where no + write-back is possible is reported as undeliverable, so the pipeline + executor discards it and releases the original events. """ if not responses_so_far: return responses_so_far @@ -690,14 +815,16 @@ class OpenAIResponsesHandler(BaseTranslation): return responses_so_far # ------------------------------------------------------------------ # - # Case 1: response.completed — full response is available in the # - # final chunk; iterate output items, apply guardrail, write back. # + # Case 1: terminal envelope events (completed/incomplete/failed). # + # the accumulated response is available in the final chunk; iterate # + # output items, apply guardrail, write back. Falls through to the # + # string fallback when the envelope yields nothing to check. # # ------------------------------------------------------------------ # - if final_chunk.get("type") == "response.completed": + if final_chunk.get("type") in _TERMINAL_ENVELOPE_EVENT_TYPES: response_obj: Final[ResponseOutputEnvelope] = final_chunk.get("response") or {} - if not hasattr(response_obj, "get"): - return responses_so_far - outputs: Final[Sequence[object]] = response_obj.get("output") or [] + outputs: Final[Sequence[object]] = ( + (response_obj.get("output") or []) if hasattr(response_obj, "get") else [] + ) texts_to_check: Final[list[str]] = [] tool_calls_to_check: Final[list[ChatCompletionToolCallChunk]] = [] @@ -730,6 +857,7 @@ class OpenAIResponsesHandler(BaseTranslation): if response_model: inputs["model"] = response_model + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -738,6 +866,11 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name, + ) # Write guardrailed texts back into the output items in-place. # final_chunk is a reference into responses_so_far so this @@ -747,11 +880,32 @@ class OpenAIResponsesHandler(BaseTranslation): responses=guardrailed_texts, task_mappings=task_mappings, ) - - return responses_so_far + if deliver_ended_stream_rewrites: + rewrites_by_position: Final = MappingProxyType( + { + task_mappings[task_idx]: rewritten + for task_idx, rewritten in enumerate(guardrailed_texts) + if task_idx < len(texts_to_check) and rewritten != texts_to_check[task_idx] + } + ) + if rewrites_by_position: + self._sync_stream_events_with_rewrites( + stream_events=responses_so_far[:-1], + rewrites_by_position=rewrites_by_position, + ) + self._deliver_ended_stream_tool_call_rewrites( + responses_so_far=responses_so_far, + outputs=outputs, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=post_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) + return responses_so_far # ------------------------------------------------------------------ # - # Case 2: response.output_item.done — extract tool calls only. # + # Case 2: response.output_item.done — extract tool calls only, then # + # fall through to the text fallback when a caller expects rewrites # + # delivered, so a truncated buffer still reports text undeliverable. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": model_response_stream: Final = ( @@ -769,12 +923,14 @@ class OpenAIResponsesHandler(BaseTranslation): input_type="response", logging_obj=litellm_logging_obj, ) - return responses_so_far + if not deliver_ended_stream_rewrites: + return responses_so_far # ------------------------------------------------------------------ # # Fallback: apply guardrail to the accumulated text string. # # No structured write-back is possible here; guardrails that only # - # need to block/flag (not rewrite) still work correctly. # + # need to block/flag (not rewrite) still work correctly, and a # + # rewrite a caller expects delivered is reported undeliverable. # # ------------------------------------------------------------------ # string_so_far: Final = self.get_streaming_string_so_far(responses_so_far) if string_so_far: @@ -784,28 +940,225 @@ class OpenAIResponsesHandler(BaseTranslation): ) if response_model: fallback_inputs["model"] = response_model - await guardrail_to_apply.apply_guardrail( + fallback_outputs: Final = await guardrail_to_apply.apply_guardrail( inputs=fallback_inputs, request_data=request_data if request_data is not None else {}, input_type="response", logging_obj=litellm_logging_obj, ) + fallback_texts: Final = fallback_outputs.get("texts") + if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") return responses_so_far + @staticmethod + def _write_event_field(event: object, field: str, value: str) -> None: + if isinstance(event, dict): + event[field] = value # rebind-ok: delivering the rewrite means editing the buffered event in place + else: + setattr(event, field, value) + + def _sync_stream_events_with_rewrites( + self, + stream_events: Sequence[Any], + rewrites_by_position: Mapping[tuple[int, int], str], + ) -> None: + """Sync pre-completion stream events with the rewritten completed + response, keyed by ``(output_index, content_index)``: the first + ``output_text.delta`` for a rewritten item carries the full rewritten + text and the rest are blanked, while ``output_text.done``, + ``content_part.done``, and ``output_item.done`` events carry the full + rewritten text, so every event a client may read agrees with the + rewritten ``response.completed`` payload.""" + delta_replacements: Final = MappingProxyType( + {position: chain((rewritten,), repeat("")) for position, rewritten in rewrites_by_position.items()} + ) + for event in stream_events: + if not (isinstance(event, dict) or hasattr(event, "get")): + continue + event_type = event.get("type") + output_index = event.get("output_index") + content_index = event.get("content_index") + if event_type == "response.output_item.done" and isinstance(output_index, int): + self._sync_output_item_done_event(event.get("item"), output_index, rewrites_by_position) + continue + if not isinstance(output_index, int) or not isinstance(content_index, int): + continue + position = (output_index, content_index) + if event_type == "response.output_text.delta" and position in delta_replacements: + self._write_event_field(event, "delta", next(delta_replacements[position])) + elif event_type == "response.output_text.done" and position in rewrites_by_position: + self._write_event_field(event, "text", rewrites_by_position[position]) + elif event_type == "response.content_part.done" and position in rewrites_by_position: + part = event.get("part") + if isinstance(part, dict) or hasattr(part, "text"): + self._write_event_field(part, "text", rewrites_by_position[position]) + + @staticmethod + def _sync_output_item_done_event( + item: object, + output_index: int, + rewrites_by_position: Mapping[tuple[int, int], str], + ) -> None: + content: Final = item.get("content") if isinstance(item, dict) else getattr(item, "content", None) + if not isinstance(content, list): + return + for (item_idx, content_idx), rewritten in rewrites_by_position.items(): + if item_idx != output_index or content_idx >= len(content): + continue + OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten) + + def _deliver_ended_stream_tool_call_rewrites( + self, + responses_so_far: Sequence[object], + outputs: Sequence[object], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail tool-call rewrites into the completed + envelope's ``function_call`` and ``custom_tool_call`` items and sync the + earlier stream events, keyed by ``call_id``. The guardrail sees the + envelope's tool calls in output order, which is how a rewritten call + finds its ``call_id``; the stream events find their call through the + ``call_id`` on ``output_item`` events and the ``item_id`` on argument + and custom-input events, since an + event's ``output_index`` need not match the envelope's (the chat bridge + numbers tool calls from 1 while the envelope lists them after the + message). A rewrite whose calls do not line up with the envelope, or + whose events cannot be found, is reported as undeliverable, so the + pipeline executor discards it and releases the original events.""" + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + tool_call_items: Final = tuple(output_item for output_item in outputs if _is_tool_call_output_item(output_item)) + call_ids: Final = tuple( + call_id + for output_item in tool_call_items + if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id + ) + stream_events: Final = responses_so_far[:-1] + call_id_by_item_id: Final = self._tool_call_ids_by_item_id(stream_events) + event_call_ids: Final = tuple( + self._tool_call_event_call_id(event, call_id_by_item_id) for event in stream_events + ) + rewrites_by_call_id: Final = MappingProxyType( + { + call_id: _tool_call_rewrite(before, after) + for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if after != before + } + ) + unresolved_argument_event: Final = any( + call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES + for event, call_id in zip(stream_events, event_call_ids) + ) + if ( + len(call_ids) != len(tool_call_items) + or len(frozenset(call_ids)) != len(call_ids) + or len(call_ids) != len(post_guardrail_tool_calls) + or unresolved_argument_event + or not rewrites_by_call_id.keys() <= frozenset(event_call_ids) + ): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + for output_item, rewrite in ( + (output_item, rewrites_by_call_id[call_id]) + for output_item, call_id in zip(tool_call_items, call_ids) + if call_id in rewrites_by_call_id + ): + self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments) + delta_replacements: Final = MappingProxyType( + {call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()} + ) + for event, call_id in zip(stream_events, event_call_ids): + if call_id not in rewrites_by_call_id: + continue + match stream_item_field(event, "type"): + case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: + self._write_event_field(event, "delta", next(delta_replacements[call_id])) + case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: + self._write_event_field( + event, _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS[event_type], rewrites_by_call_id[call_id].arguments + ) + case "response.output_item.added": + self._write_tool_call_item( + stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None + ) + case "response.output_item.done": + self._write_tool_call_item( + stream_item_field(event, "item"), + rewrites_by_call_id[call_id].name, + rewrites_by_call_id[call_id].arguments, + ) + case _: + pass + + def _write_tool_call_rewrites_to_output( + self, + tool_call_items: Sequence[object], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + ) -> None: + if len(tool_call_items) != len(post_guardrail_tool_calls): + return + for output_item, rewrite in ( + (output_item, _tool_call_rewrite(before, after)) + for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if after != before + ): + self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments) + + @staticmethod + def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]: + items: Final = tuple( + stream_item_field(event, "item") + for event in stream_events + if stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES + ) + return MappingProxyType( + { + item_id: call_id + for item in items + if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES + and isinstance(item_id := stream_item_field(item, "id"), str) + and isinstance(call_id := stream_item_field(item, "call_id"), str) + } + ) + + @staticmethod + def _tool_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None: + event_type: Final = stream_item_field(event, "type") + if event_type in _TOOL_CALL_PAYLOAD_EVENT_TYPES: + item_id: Final = stream_item_field(event, "item_id") + return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None + if event_type not in _OUTPUT_ITEM_EVENT_TYPES: + return None + item: Final = stream_item_field(event, "item") + call_id: Final = stream_item_field(item, "call_id") + return ( + call_id if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(call_id, str) else None + ) + + @staticmethod + def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> None: + if item is None: + return + if name is not None: + OpenAIResponsesHandler._write_event_field(item, "name", name) + item_type: Final = stream_item_field(item, "type") + if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS: + OpenAIResponsesHandler._write_event_field(item, _TOOL_CALL_PAYLOAD_FIELDS[item_type], payload) + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if the streaming has ended. """ if not responses_so_far: return False - terminal_types: Final = frozenset( - ( - ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, - ResponsesAPIStreamEvents.RESPONSE_FAILED.value, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, - ) - ) - return stream_item_field(responses_so_far[-1], "type") in terminal_types + return stream_item_field(responses_so_far[-1], "type") in _TERMINAL_ENVELOPE_EVENT_TYPES def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: if not responses_so_far or not hasattr(responses_so_far[-1], "get"): @@ -825,7 +1178,7 @@ class OpenAIResponsesHandler(BaseTranslation): def _completed_response_scan_key(response: object) -> StreamingScanKey: output_items: Final = stream_item_items(response, "output") message_items: Final = tuple( - item for item in output_items if stream_item_field(item, "type") != "function_call" + item for item in output_items if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES ) return StreamingScanKey( texts=tuple( @@ -837,7 +1190,7 @@ class OpenAIResponsesHandler(BaseTranslation): tool_calls=tuple( stream_item_fingerprint(item) for item in output_items - if stream_item_field(item, "type") == "function_call" + if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES ), stream_ended=True, ) @@ -948,34 +1301,10 @@ class OpenAIResponsesHandler(BaseTranslation): Override this method to customize text/image/tool extraction logic. """ - # Check if this is a tool call (OutputFunctionToolCall) - if isinstance(output_item, OutputFunctionToolCall) or ( - isinstance(output_item, BaseModel) - and hasattr(output_item, "type") - and getattr(output_item, "type") == "function_call" - ): + tool_call_item: Final = _tool_call_output_item_mapping(output_item) + if tool_call_item is not None: if tool_calls_to_check is not None: - tool_call_dict = ( - LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=output_item, - index=output_idx, - ) - ) - tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict)) - return - elif isinstance(output_item, dict) and output_item.get("type") == "function_call": - # Handle dict representation of tool call - if tool_calls_to_check is not None: - # Convert dict to ResponseFunctionToolCall for processing - try: - tool_call_obj: Final = ResponseFunctionToolCall(**output_item) - tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=tool_call_obj, - index=output_idx, - ) - tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict)) - except Exception: - pass + tool_calls_to_check.append(tool_call_dict_from_output_item(tool_call_item, output_idx)) return # Handle both GenericResponseOutputItem and dict diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index d7c2fcace09..833ae206024 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,18 +1,24 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence +from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints import httpx from openai.types.responses import ResponseReasoningItem -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + flatten_combinators_and_drop_non_python_regex_patterns, +) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig @@ -38,10 +44,34 @@ else: _NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) _MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4") -_PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) +_PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) +class _ReasoningSupportEntry(BaseModel): + litellm_provider: str | None = None + supports_reasoning: bool | None = None + + +_BUNDLED_COST_MAP: Final = TypeAdapter(dict[str, _ReasoningSupportEntry]) + + +@lru_cache(maxsize=1) +def _bundled_openai_reasoning_models() -> frozenset[str]: + """OpenAI models the cost map shipped with this release flags as reasoning models. + + The live map can lag this release (a pinned mirror, or a proxy on newer code than the + map it fetches), and a lagging entry must never strip `reasoning` from a model this + release knows accepts it. + """ + bundled: Final = _BUNDLED_COST_MAP.validate_json(GetModelCostMap.read_local_model_cost_map_text()) + return frozenset( + name + for name, entry in bundled.items() + if entry.litellm_provider == LlmProviders.OPENAI.value and entry.supports_reasoning is True + ) + + class _DeleteResponseBody(TypedDict): """Decoded body of the Responses API delete call.""" @@ -80,6 +110,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): def supports_native_file_search(self) -> bool: return True + def supports_encrypted_agent_messages(self) -> bool: + return self.custom_llm_provider in (LlmProviders.OPENAI, LlmProviders.AZURE) + @staticmethod def _is_gpt_5_model(model: str) -> bool: """Return True only for actual OpenAI GPT-5 models. @@ -95,13 +128,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): @staticmethod def _supports_reasoning_effort_none(model: str) -> bool: """Return True if the model supports reasoning.effort='none'.""" - from litellm.utils import _supports_factory + from litellm.utils import supports_none_reasoning_effort - return _supports_factory( - model=model, - custom_llm_provider=None, - key="supports_none_reasoning_effort", - ) + return supports_none_reasoning_effort(model=model, custom_llm_provider=None) @staticmethod def _effort_resolves_to_none(model: str, effort: str | None) -> bool: @@ -117,6 +146,28 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return OpenAIGPT5Config.effort_resolves_to_none(model, effort) + @staticmethod + def _supports_reasoning_param(model: str) -> bool: + from litellm.utils import _get_model_info_helper + + try: + info: Final = _get_model_info_helper( + model=model.split("/")[-1], custom_llm_provider=LlmProviders.OPENAI.value + ) + except Exception: + return True + declared: Final = info.get("supports_reasoning") + if declared is not None: + return declared + return info["key"] in _bundled_openai_reasoning_models() + + @staticmethod + def _requests_reasoning_effort(reasoning: object) -> bool: + effort: Final = ( + reasoning.get("effort") if isinstance(reasoning, Mapping) else getattr(reasoning, "effort", None) + ) + return effort is not None + @staticmethod def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None": """Raise sub-minimum max_output_tokens up to the OpenAI Responses API minimum. @@ -166,6 +217,23 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): if "max_output_tokens" in params: params["max_output_tokens"] = self._enforce_min_max_output_tokens(params.get("max_output_tokens")) + if ( + self.custom_llm_provider == LlmProviders.OPENAI + and self._requests_reasoning_effort(params.get("reasoning")) + and not self._supports_reasoning_param(model=model) + ): + if drop_params or litellm.drop_params: + params.pop("reasoning", None) + else: + raise litellm.UnsupportedParamsError( + message=( + f"{model} doesn't support `reasoning.effort` " + "(its model cost map entry lacks `supports_reasoning`). " + "To drop unsupported params set `litellm.drop_params = True`" + ), + status_code=400, + ) + if self._is_gpt_5_model(model=model): temperature: Final = params.get("temperature") if temperature is not None and temperature != 1: @@ -232,7 +300,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): model=model, input=validated_input, tools=tools ) object_schema_tools: Final = self._tools_with_object_parameters(model=model, tools=stripped_tools) - sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai( + sanitized_tools: Final = self._sanitized_tool_schemas_for_openai( model=model, tools=object_schema_tools, litellm_params=litellm_params ) return self._drop_foreign_tool_call_item_ids(stripped_input), sanitized_tools @@ -317,35 +385,35 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return item return {key: value for key, value in item.items() if key != "id"} # mutable-ok: outgoing JSON request item - def _flatten_tool_schema_combinators_for_openai( + def _sanitized_tool_schemas_for_openai( self, model: str, tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, litellm_params: GenericLiteLLMParams, ) -> Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None: - """Flatten top-level schema combinators only where OpenAI's validator rejects them. + """Rewrite tool schemas only where OpenAI's validator rejects them. - OpenAI-compatible backends reusing this config (and the ChatGPT backend - Codex talks to natively) accept them, and so do GPT-5 and later models, - which also call tools better with the union intact. Codex wraps MCP tools - inside namespace entries, so nested ``tools`` arrays are walked too. - Azure OpenAI shares the validator but names deployments arbitrarily, so - the router's declared ``model_info.base_model`` wins over the deployment - name and an unrecognized name without one is left untouched. + Every model family refuses a ``pattern`` Python's ``re`` cannot compile, + while top-level schema combinators are flattened only for the families + whose validator rejects them: OpenAI-compatible backends reusing this + config (and the ChatGPT backend Codex talks to natively) accept them, + and so do GPT-5 and later models, which also call tools better with the + union intact. Codex wraps MCP tools inside namespace entries, so nested + ``tools`` arrays are walked too. Azure OpenAI shares the validator but + names deployments arbitrarily, so the router's declared + ``model_info.base_model`` wins over the deployment name and an + unrecognized name without one keeps its combinators. """ - if tools is None or self.custom_llm_provider not in _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: + if tools is None or self.custom_llm_provider not in _PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: return tools gate_model: Final = self._combinator_gate_model(model=model, litellm_params=litellm_params) - if not self._rejects_top_level_schema_combinators(gate_model): - return tools - flattened: Final = [ # mutable-ok: request tools are a JSON list - self._flattened_tool_or_passthrough(tool) for tool in tools - ] - return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", flattened) # cast-ok: spread keeps each tool's shape - - @staticmethod - def _flattened_tool_or_passthrough(tool: object) -> object: - return OpenAIResponsesAPIConfig._flattened_tool_entry(tool) if isinstance(tool, dict) else tool + sanitize: Final = ( + flatten_combinators_and_drop_non_python_regex_patterns + if self._rejects_top_level_schema_combinators(gate_model) + else drop_non_python_regex_patterns + ) + sanitized: Final = self._sanitized_tools(tools, sanitize) + return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", sanitized) # cast-ok: spread keeps each tool's shape @staticmethod def _rejects_top_level_schema_combinators(model: str) -> bool: @@ -360,35 +428,42 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return base_model if isinstance(base_model, str) and base_model else model @staticmethod - def _flattened_tool_entry( + def _sanitized_tool_entry( entry: Mapping[str, object], - ) -> dict[str, object]: # mutable-ok: request tools are JSON dicts - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - flatten_top_level_schema_combinators, - ) - + sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], + ) -> Mapping[str, object]: parameters: Final = entry.get("parameters") nested_tools: Final = entry.get("tools") + sanitized_parameters: Final = sanitize(parameters) if isinstance(parameters, dict) else parameters + sanitized_nested_tools: Final = ( + OpenAIResponsesAPIConfig._sanitized_tools(nested_tools, sanitize) + if isinstance(nested_tools, list) + else nested_tools + ) parameters_update: Final = ( - MappingProxyType({"parameters": flatten_top_level_schema_combinators(parameters)}) - if isinstance(parameters, dict) + MappingProxyType({"parameters": sanitized_parameters}) + if sanitized_parameters is not parameters else _NO_TOOL_UPDATE ) tools_update: Final = ( - MappingProxyType({"tools": OpenAIResponsesAPIConfig._flattened_nested_tools(nested_tools)}) - if isinstance(nested_tools, list) + MappingProxyType({"tools": sanitized_nested_tools}) + if sanitized_nested_tools is not nested_tools else _NO_TOOL_UPDATE ) + if not parameters_update and not tools_update: + return entry return {**entry, **parameters_update, **tools_update} # mutable-ok: request tools are JSON dicts @staticmethod - def _flattened_nested_tools( - nested_tools: Sequence[object], - ) -> list[object]: # mutable-ok: namespace tools are a JSON list - return [ # mutable-ok: namespace tools are a JSON list - OpenAIResponsesAPIConfig._flattened_tool_entry(item) if isinstance(item, dict) else item - for item in nested_tools + def _sanitized_tools( + tools: Sequence[object], + sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], + ) -> Sequence[object]: + sanitized: Final = [ # mutable-ok: request tools are a JSON list + OpenAIResponsesAPIConfig._sanitized_tool_entry(item, sanitize) if isinstance(item, dict) else item + for item in tools ] + return tools if all(new is old for new, old in zip(sanitized, tools, strict=True)) else sanitized def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam: """ @@ -478,7 +553,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): processed_headers: Final = process_response_headers(raw_response_headers) try: response = ResponsesAPIResponse.model_validate(raw_response_json) - except Exception: + except ValidationError: verbose_logger.debug( "Error constructing ResponsesAPIResponse: %s, using model_construct", raw_response_json ) @@ -559,15 +634,20 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return event_pydantic_model.model_construct(**parsed_chunk) @staticmethod - def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + def parse_terminal_event_from_stream_chunks(all_chunks: Sequence[str]) -> ResponsesTerminalEvent | None: for chunk_str in reversed(all_chunks): for event_model in (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent): try: - return event_model.model_validate_json(chunk_str.removeprefix("data: ")).response + return event_model.model_validate_json(chunk_str.removeprefix("data: ")) except ValueError: continue return None + @staticmethod + def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks) + return None if terminal_event is None else terminal_event.response + @staticmethod def get_event_model_class(event_type: str) -> type[BaseLiteLLMOpenAIResponseObject]: """ @@ -870,7 +950,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): try: response = ResponsesAPIResponse.model_validate(raw_response_json) - except Exception: + except ValidationError: verbose_logger.debug( "Error constructing ResponsesAPIResponse: %s, using model_construct", raw_response_json ) diff --git a/litellm/llms/sagemaker/chat/handler.py b/litellm/llms/sagemaker/chat/handler.py index 3f62b7276df..38978300c52 100644 --- a/litellm/llms/sagemaker/chat/handler.py +++ b/litellm/llms/sagemaker/chat/handler.py @@ -36,6 +36,7 @@ class SagemakerChatHandler(BaseAWSLLM): aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) ### SET REGION NAME ### if aws_region_name is None: @@ -63,6 +64,7 @@ class SagemakerChatHandler(BaseAWSLLM): aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) return credentials, aws_region_name diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index fb8074d3682..fad0a460647 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -59,6 +59,7 @@ class SagemakerLLM(BaseAWSLLM): aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) ### SET REGION NAME ### if aws_region_name is None: @@ -86,6 +87,7 @@ class SagemakerLLM(BaseAWSLLM): aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) return credentials, aws_region_name diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index c64fc583edc..f65b0876202 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( create_anthropic_image_param, select_anthropic_content_block_type_for_file, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload @@ -421,6 +422,21 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body) return self._transform_request_openai(model, messages, optional_params, stream, extra_body) + @property + def uses_async_transform_request(self) -> bool: + return True + + async def async_transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, object], # mutable-ok: BaseConfig signature + litellm_params: dict[str, object], # mutable-ok: BaseConfig signature + headers: dict[str, object], # mutable-ok: BaseConfig signature + ) -> dict[str, object]: # mutable-ok: BaseConfig signature + inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages + return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers) + def _transform_request_openai( self, model: str, diff --git a/litellm/llms/tinyfish/search/transformation.py b/litellm/llms/tinyfish/search/transformation.py index b688dc2cd01..460394c6f2d 100644 --- a/litellm/llms/tinyfish/search/transformation.py +++ b/litellm/llms/tinyfish/search/transformation.py @@ -247,6 +247,13 @@ class TinyfishSearchConfig(BaseSearchConfig): hidden["additional_headers"] = process_response_headers(raw_headers) return parsed + def get_http_error_class(self, error: httpx.HTTPStatusError) -> Exception: + return self._wrap_error( + error_message=error.response.text, + status_code=error.response.status_code, + headers=dict(error.response.headers), # mutable-ok: existing error wrapper requires dict headers + ) + def _wrap_error( self, error_message: str, @@ -256,8 +263,7 @@ class TinyfishSearchConfig(BaseSearchConfig): """ Build an attributed ``BaseLLMException`` from a TinyFish error body. - Used only at the call sites we control inside - ``transform_search_response`` (non-2xx, JSONDecodeError, ValidationError). + Used for HTTP status errors and response transformation errors. Not an override of ``BaseSearchConfig.get_error_class``: that path is left to inherit from the base so it auto-picks-up any future LiteLLM improvements. Trade-off: network failures (routed through LiteLLM diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 377cd9f3437..a15ea4d845b 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,6 +1,7 @@ import json -from collections.abc import Coroutine +from collections.abc import Coroutine, Sequence from typing import TYPE_CHECKING, Final, Protocol +from urllib.parse import urlparse import httpx from typing_extensions import ReadOnly, TypedDict @@ -12,11 +13,13 @@ from litellm.litellm_core_utils.url_utils import ( safe_get, ) from litellm.llms.custom_httpx.http_handler import ( + HTTPHandler, _get_httpx_client, get_async_httpx_client, ) from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.vertex_llm_base import _graft_default_vertex_path from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( VERTEX_CREDENTIALS_TYPES, @@ -55,6 +58,20 @@ class _FetchedResponseView(TypedDict): response: ReadOnly[httpx.Response] +class _VertexEndpointDeployedModel(TypedDict, total=False): + model: ReadOnly[str] + + +class _VertexEndpointResponse(TypedDict, total=False): + deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]] + + +class _VertexEndpointPayloadView(TypedDict): + """Holds one decoded GET endpoints/ response so the payload reads back typed.""" + + payload: ReadOnly[_VertexEndpointResponse] + + def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse: return response.json() @@ -78,7 +95,17 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location: str | None, timeout: float | httpx.Timeout, max_retries: int | None, + custom_endpoint: bool | None = None, ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: + if custom_endpoint: + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. " + "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; " + "use a publisher model or fine-tuned Gemini endpoint deployment instead." + ), + ) sync_handler: Final = _get_httpx_client() access_token, project_id = self._ensure_access_token( @@ -87,6 +114,26 @@ class VertexAIBatchPrediction(VertexLLM): custom_llm_provider="vertex_ai", ) + headers: Final = { + "Content-Type": "application/json; charset=utf-8", + "Authorization": f"Bearer {access_token}", + } + + transformed_batch_request: Final[VertexAIBatchPredictionJob] = ( + VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( + request=create_batch_data, + vertex_project=vertex_project or project_id, + vertex_location=vertex_location or "us-central1", + ) + ) + vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model( + vertex_batch_request=transformed_batch_request, + headers=headers, + sync_handler=sync_handler, + api_base=api_base, + vertex_location=vertex_location or "us-central1", + ) + default_api_base: Final = self.create_vertex_batch_url( vertex_location=vertex_location or "us-central1", vertex_project=vertex_project or project_id, @@ -111,17 +158,6 @@ class VertexAIBatchPrediction(VertexLLM): vertex_api_version="v1", ) - headers: Final = { - "Content-Type": "application/json; charset=utf-8", - "Authorization": f"Bearer {access_token}", - } - - vertex_batch_request: Final[VertexAIBatchPredictionJob] = ( - VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( - request=create_batch_data - ) - ) - if _is_async is True: return self._async_create_batch( vertex_batch_request=vertex_batch_request, @@ -142,6 +178,77 @@ class VertexAIBatchPrediction(VertexLLM): ) return vertex_batch_response + @staticmethod + def _build_endpoint_resolution_url(api_base: str | None, model: str, vertex_location: str) -> str: + """ + Builds the GET url for resolving an endpoint resource (`projects/../endpoints/`). + + A custom `api_base` replaces the Google host: its `/v1`/`/v1beta1` path swallows the + version segment (matching `_check_custom_proxy`'s grafting), any other path is kept as a + mount prefix in front of the full default path. The `:operation` suffix convention from + `_check_custom_proxy` does not apply to a plain resource GET. + """ + default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}" + if not api_base: + return default_endpoint_url + api_base_path: Final = urlparse(api_base).path.rstrip("/") + if api_base_path in ("/v1", "/v1beta1"): + return _graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url) + return api_base.rstrip("/") + urlparse(default_endpoint_url).path + + def _resolve_fine_tuned_endpoint_model( + self, + vertex_batch_request: VertexAIBatchPredictionJob, + headers: dict[str, str], # mutable-ok: HTTPHandler.get only accepts dict headers + sync_handler: HTTPHandler, + api_base: str | None, + vertex_location: str, + ) -> VertexAIBatchPredictionJob: + """ + A fine-tuned Gemini deployment is configured by its endpoint id, but the v1 batch API only + accepts Model resources, so swap the endpoint resource for its deployed tuned model + (`projects/../locations/../models/`) read from GET endpoints/. + """ + model: Final = vertex_batch_request.get("model", "") + if "/endpoints/" not in model: + return vertex_batch_request + + endpoint_url: Final = self._build_endpoint_resolution_url( + api_base=api_base, + model=model, + vertex_location=vertex_location, + ) + # ``api_base`` can come from caller-supplied request kwargs, so wrap the + # fetch in ``safe_get``: it rejects DNS-rebind / private / cloud-metadata + # targets before the bearer token leaves the process (mirrors retrieve_batch). + fetched: Final[_FetchedResponseView] = { + "response": safe_get( + sync_handler, + endpoint_url, + headers=headers, + ) + } + response: Final = fetched["response"] + if response.status_code != 200: + raise VertexAIError( + status_code=response.status_code, + message=f"Failed to resolve fine-tuned Vertex endpoint '{model}': {response.text}", + ) + + payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()} + deployed_models: Final = payload_view["payload"].get("deployedModels") or () + deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else "" + if not deployed_model: + raise VertexAIError( + status_code=400, + message=( + f"Vertex endpoint '{model}' has no deployed model, so there is no tuned model " + "resource to run batch predictions against" + ), + ) + resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model} + return resolved_request + async def _async_create_batch( self, vertex_batch_request: VertexAIBatchPredictionJob, diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index f284b47292b..e63c80dd3cf 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -22,6 +22,8 @@ class VertexAIBatchTransformation: def transform_openai_batch_request_to_vertex_ai_batch_request( cls, request: CreateBatchRequest, + vertex_project: str | None = None, + vertex_location: str | None = None, ) -> VertexAIBatchPredictionJob: """ Transforms OpenAI Batch requests to Vertex AI Batch requests @@ -31,7 +33,11 @@ class VertexAIBatchTransformation: if input_file_id is None: raise ValueError("input_file_id is required, but not provided") input_config: InputConfig = InputConfig(gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl") - model: Final[str] = cls._get_model_from_gcs_file(input_file_id) + model: Final[str] = cls._get_batch_job_model( + input_file_id=input_file_id, + vertex_project=vertex_project, + vertex_location=vertex_location, + ) output_config: Final[OutputConfig] = OutputConfig( predictionsFormat="jsonl", gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)), @@ -188,6 +194,33 @@ class VertexAIBatchTransformation: path_parts: Final = input_file_id.rsplit("/", 1) return path_parts[0] + @classmethod + def _get_batch_job_model( + cls, + input_file_id: str, + vertex_project: str | None, + vertex_location: str | None, + ) -> str: + """ + Returns the `model` for the batchPredictionJobs request: the publisher model path as-is, or + the full `projects/../locations/../endpoints/` resource name for a fine-tuned endpoint. + + The v1 batch API only accepts Model resources, so the handler resolves an endpoint resource + to its deployed tuned model (`projects/../locations/../models/`) before sending the job. + """ + parsed_model: Final = cls._get_model_from_gcs_file(input_file_id) + if not parsed_model.startswith("endpoints/"): + return parsed_model + if not vertex_project: + raise VertexAIError( + status_code=400, + message=( + f"Vertex AI batch jobs against a fine-tuned endpoint ('{parsed_model}') require " + "`vertex_project` to build the endpoint resource name" + ), + ) + return f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}/{parsed_model}" + @classmethod def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str: """ @@ -202,6 +235,9 @@ class VertexAIBatchTransformation: gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" + Fine-tuned Gemini endpoints are stored as `endpoints/` in the uri and returned + in that form. + Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ model: Final = cls._parse_model_from_gcs_file(gcs_file_uri) @@ -210,11 +246,13 @@ class VertexAIBatchTransformation: status_code=400, message=( "Vertex AI batch creation requires the model to be part of `input_file_id`, but " - f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " + f"'{gcs_file_uri}' contains no 'publishers//models/' or " + "'endpoints/' path segment. " "Either upload the input file through LiteLLM (POST /v1/files with " "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " "pass a uri of the form " - "gs:////publishers//models//" + "gs:////publishers//models// " + "(or gs:////endpoints// for fine-tuned models)" ), ) return model @@ -222,18 +260,26 @@ class VertexAIBatchTransformation: @classmethod def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: """ - Returns the `publishers//models/` path from a gcs uri, or None if the uri - does not contain one. + Returns the `publishers//models/` or `endpoints/` path from a + gcs uri, or None if the uri does not contain one. + + A publisher path wins over an `endpoints/` segment, and the last `endpoints/` occurrence is + used, so a user-configured bucket prefix that happens to contain `endpoints/` cannot + override the model path LiteLLM appended after it. """ - _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") - if not separator: - return None + unquoted_uri: Final = unquote(gcs_file_uri) + _, separator, model_path = unquoted_uri.partition("publishers/") + if separator: + parts: Final = model_path.split("/") + if len(parts) >= 3 and parts[1] == "models" and parts[2]: + return f"publishers/{'/'.join(parts[:3])}" - parts: Final = model_path.split("/") - if len(parts) < 3 or parts[1] != "models" or not parts[2]: - return None + _, endpoint_separator, endpoint_path = unquoted_uri.rpartition("endpoints/") + endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else "" + if endpoint_id.isdigit(): + return f"endpoints/{endpoint_id}" - return f"publishers/{'/'.join(parts[:3])}" + return None @classmethod def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 970759479fe..14aebcaabaf 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -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, @@ -310,6 +370,19 @@ def get_vertex_base_model_name(model: str) -> str: return model +def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None: + """ + Fine-tuned Gemini deployments are addressed by a numeric endpoint id, + configured as `vertex_ai/` or `vertex_ai/gemini/`. + + Returns the endpoint id, or None when `model` is a regular publisher model. + Mirrors the online chat path in `_get_vertex_url`, which sends numeric + models to `endpoints/{id}` instead of `publishers/google/models/{model}`. + """ + candidate: Final = model.split("/")[-1] if "gemini/" in model else model + return candidate if candidate.isdigit() else None + + def validate_vertex_location(vertex_location: str | None) -> str: """ Validate a Vertex AI location before interpolating it into a request host or diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index b6ad9fbcc04..263956efc9f 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -39,6 +39,7 @@ from litellm.llms.base_llm.files.transformation import ( ) from litellm.llms.vertex_ai.common_utils import ( _convert_vertex_datetime_to_openai_datetime, + get_vertex_ai_fine_tuned_endpoint_id, ) from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -707,20 +708,39 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, openai_jsonl_content: list[dict[str, Any]], + deployment_model: str | None = None, ) -> str: """ Gets a unique GCS object name for the VertexAI batch prediction job named as: litellm-vertex-{model}-{uuid} + + The stored model path decides which Vertex model the batch job later executes against, so + `deployment_model` (the deployment's own configured model) wins over the user-supplied + JSONL `body.model`; the JSONL value is only a fallback for direct SDK calls that carry no + deployment config. + + Fine-tuned Gemini deployments (numeric endpoint ids) are stored under + `endpoints/` so the batch transformation can round-trip them into a + `projects/../locations/../endpoints/` batch job model instead of a + nonexistent publisher model. """ - _model = openai_jsonl_content[0].get("body", {}).get("model", "") - if "publishers/google/models" not in _model: - _model = f"publishers/google/models/{_model}" - safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model") + raw_model: Final = ( + deployment_model.removeprefix("vertex_ai/") + if deployment_model + else openai_jsonl_content[0].get("body", {}).get("model", "") + ) + endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) + model_path: Final = ( + f"endpoints/{endpoint_id}" + if endpoint_id is not None + else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") + ) + safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name - def get_object_name(self, file_data: FileTypes, purpose: str) -> str: + def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str: """ Get the object name for the request. @@ -728,10 +748,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): upload is never materialized just to derive the GCS object name. """ if purpose == "batch": - ## 1. If jsonl, derive the object name from the first entry's model + ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None) if first_entry is not None: - return self._get_gcs_object_name_from_batch_jsonl([first_entry]) + return self._get_gcs_object_name_from_batch_jsonl([first_entry], deployment_model=deployment_model) ## 2. If not jsonl, store under a server-generated managed object name filename, _ = extract_file_metadata(file_data) @@ -761,6 +781,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Get the complete url for the request """ + if data.get("purpose") == "batch" and litellm_params.get("custom_endpoint"): + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. " + "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; " + "remove this deployment from the batch request (e.g. `target_model_names`) or " + "use a publisher model / fine-tuned Gemini endpoint instead." + ), + ) bucket_name = self._get_configured_bucket_name(litellm_params) bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) file_data: Final = data.get("file") @@ -769,7 +799,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - object_name = self.get_object_name(file_data, purpose) + configured_model: Final = litellm_params.get("model") + object_name = self.get_object_name( + file_data, + purpose, + deployment_model=configured_model if isinstance(configured_model, str) else None, + ) if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name: Final = encode_gcs_object_name_for_url(object_name) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index e2d62be6a69..13e2238fdf6 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -7,6 +7,7 @@ Why separate file? Make it easy to see how transformation works import json import os import re +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -27,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, response_schema_prompt, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.files import ( @@ -68,6 +70,7 @@ _GCS_METADATA_VERTEX_BASE: Any | None = None # Shared sync client for GCS JSON API metadata reads so proxy/SSL settings # from litellm's HTTP stack apply (see Greptile review on PR #27278). _GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None +GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/" _GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = { "image/jpg": "image/jpeg", } @@ -556,7 +559,7 @@ def _process_gemini_media( file_data = FileDataType(mime_type=mime_type, file_uri=image_url) part: PartType = {"file_data": file_data} return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata) - elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"): + elif image_url.startswith(GEMINI_FILES_API_URI_PREFIX): # Gemini Files API URIs — the file is already uploaded to Google's # servers; pass the URI through as file_data without fetching it. # These URLs return 403 when accessed directly, so we must not try @@ -1307,6 +1310,23 @@ def sync_transform_request_body( ) +def _explicit_mime_type(fields: Mapping[str, object]) -> str | None: + hint: Final = fields.get("format") or fields.get("mime_type") or fields.get("content_type") + return hint if isinstance(hint, str) else None + + +def _ai_studio_inlines(media: RemoteMedia) -> bool: + return not media.url.startswith(GEMINI_FILES_API_URI_PREFIX) + + +def _vertex_inlines(media: RemoteMedia) -> bool: + if media.url.startswith(GEMINI_FILES_API_URI_PREFIX): + return False + return media.url.startswith("http://") or ( + _explicit_mime_type(media.fields) is None and _get_image_mime_type_from_url(media.url) is None + ) + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1348,13 +1368,17 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) - if _openai_messages_may_need_sync_gcs_metadata_fetch(messages): + inlined_messages: Final = await async_inline_remote_media( + messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + ) + + if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): # _transform_request_body may issue a sync httpx.get (up to 5s timeout) # via _get_gcs_object_content_type to fetch GCS object metadata. Run the # whole sync transformation on a worker thread so it does not block the # async event loop. return await asyncify(_transform_request_body)( - messages=messages, + messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, @@ -1363,7 +1387,7 @@ async def async_transform_request_body( ) return _transform_request_body( - messages=messages, + messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 69fe5678de9..d113b2b4f6b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -23,6 +23,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE, DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO, ) +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator from litellm.litellm_core_utils.prompt_templates.factory import ( _encode_tool_call_id_with_signature, @@ -108,6 +109,21 @@ else: StreamingChoices = Any +SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "none", "disable") + + +def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsError: + return UnsupportedParamsError( + message=( + f"Invalid `reasoning_effort`: {reasoning_effort!r}. " + f"Must be one of: {', '.join(repr(effort) for effort in SUPPORTED_REASONING_EFFORTS)}. " + "To drop this param, set `litellm.drop_params = True` or pass in `(.., drop_params=True)` " + "in the request - https://docs.litellm.ai/docs/completion/drop_params" + ), + status_code=400, + ) + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -842,7 +858,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "includeThoughts": False, } else: - raise ValueError(f"Invalid reasoning effort: {reasoning_effort}") + raise _unsupported_reasoning_effort(reasoning_effort) @staticmethod def _map_reasoning_effort_to_thinking_level( @@ -890,7 +906,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): else: return {"thinkingLevel": "low", "includeThoughts": False} else: - raise ValueError(f"Invalid reasoning effort: {reasoning_effort}") + raise _unsupported_reasoning_effort(reasoning_effort) @staticmethod def _is_thinking_budget_zero(thinking_budget: int | None) -> bool: diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index d7b4ad22a01..d382f43495f 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -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}), + ) + ) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 7579bc8c02e..508f68b3eca 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final import httpx import litellm +from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, inline_remote_image_urls from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -51,6 +52,9 @@ class VertexAIAnthropicConfig(AnthropicConfig): def custom_llm_provider(self) -> str | None: return "vertex_ai" + def inlines_remote_media(self, media: RemoteMedia) -> bool: + return inline_remote_image_urls(media) + def should_strip_billing_metadata(self) -> bool: return True diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index ec77ce91d33..870de8756bb 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -3,6 +3,7 @@ This module is used to transform the request and response for the Voyage context This would be used for all the contextualized embeddings models in Voyage. """ +from collections.abc import Mapping from typing import Final import httpx @@ -24,7 +25,10 @@ class VoyageError(BaseLLMException): ): self.status_code = status_code self.message = message - self.request = httpx.Request(method="POST", url="https://api.voyageai.com/v1/contextualizedembeddings") + self.request = httpx.Request( + method="POST", + url="https://api.voyageai.com/v1/contextualizedembeddings", + ) self.response = httpx.Response(status_code=status_code, request=self.request) super().__init__( status_code=status_code, @@ -56,16 +60,16 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): return api_base return "https://api.voyageai.com/v1/contextualizedembeddings" - def get_supported_openai_params(self, model: str) -> list: + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class signature return ["encoding_format", "dimensions"] def map_openai_params( self, - non_default_params: dict, - optional_params: dict, + non_default_params: dict, # mutable-ok: base class signature + optional_params: dict, # mutable-ok: base class signature model: str, drop_params: bool, - ) -> dict: + ) -> dict: # mutable-ok: base class signature """ Map OpenAI params to Voyage params @@ -79,7 +83,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): def validate_environment( self, - headers: dict, + headers: dict, # mutable-ok: base class signature model: str, messages: list[AllMessageValues], optional_params: dict, @@ -97,6 +101,8 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): "Authorization": f"Bearer {api_key}", } + AUTO_CHUNK_SIZE: Final = 32000 + def transform_embedding_request( self, model: str, @@ -105,11 +111,27 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): headers: dict, ) -> dict: return { - "inputs": input, + "inputs": [input] if isinstance(input, str) else input, "model": model, + **self._auto_chunk_params(input, optional_params), **optional_params, } + @classmethod + def _auto_chunk_params( + cls, + input: AllEmbeddingInputValues | list[list[str]], + optional_params: Mapping[str, object], + ) -> Mapping[str, object]: + is_flat: Final = isinstance(input, str) or all(isinstance(item, str) for item in input) + if not is_flat or optional_params.get("input_type") == "query": + return {} + return { + "enable_auto_chunking": True, + "chunk_size": cls.AUTO_CHUNK_SIZE, + "input_type": "document", + } + def transform_embedding_response( self, model: str, @@ -124,9 +146,11 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): try: raw_response_json: Final = raw_response.json() except Exception: - raise VoyageError(message=raw_response.text, status_code=raw_response.status_code) + raise VoyageError( + message=raw_response.text, + status_code=raw_response.status_code, + ) - # model_response.usage model_response.model = raw_response_json.get("model") model_response.data = raw_response_json.get("data") model_response.object = raw_response_json.get("object") diff --git a/litellm/llms/wandb/chat/transformation.py b/litellm/llms/wandb/chat/transformation.py index a477c00f643..fdd6644f03d 100644 --- a/litellm/llms/wandb/chat/transformation.py +++ b/litellm/llms/wandb/chat/transformation.py @@ -6,10 +6,17 @@ This is OpenAI compatible - no translation needed / occurs from typing import Final +import litellm from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig class WandbConfig(OpenAIGPTConfig): + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract + supported_params: Final = super().get_supported_openai_params(model) + if litellm.supports_reasoning(model=model, custom_llm_provider="wandb"): + return supported_params + ["reasoning_effort"] # mutable-ok: inherited contract + return supported_params + def map_openai_params( self, non_default_params: dict, diff --git a/litellm/llms/watsonx/chat/transformation.py b/litellm/llms/watsonx/chat/transformation.py index f9e71f9116e..cc616ab5f9f 100644 --- a/litellm/llms/watsonx/chat/transformation.py +++ b/litellm/llms/watsonx/chat/transformation.py @@ -157,11 +157,9 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): @staticmethod async def aapply_prompt_template(model: str, messages: list[dict[str, str]]) -> str | None: """Apply prompt template (async version)""" - import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( ahf_chat_template, custom_prompt, - hf_chat_template, ibm_granite_pt, mistral_instruct_pt, ) @@ -179,11 +177,7 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): else: hf_model = model try: - # Use sync if cached, async if not - if hf_model in litellm.known_tokenizer_config: - result = hf_chat_template(model=hf_model, messages=messages) - else: - result = await ahf_chat_template(model=hf_model, messages=messages) + result = await ahf_chat_template(model=hf_model, messages=messages) # Return result if it's truthy (not None and not empty string) # The caller (_aconvert_watsonx_messages_core) will handle None/empty by falling back to default if result: diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index 0b4c9ae917a..2be007336b4 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -16,6 +16,7 @@ from ..common_utils import ( IBMWatsonXMixin, WatsonXAIError, _get_api_params, + aconvert_watsonx_messages_to_prompt, convert_watsonx_messages_to_prompt, ) @@ -236,7 +237,11 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): **watsonx_auth_payload, } - async def atransform_request( + @property + def uses_async_transform_request(self) -> bool: + return True + + async def async_transform_request( self, model: str, messages: list[AllMessageValues], @@ -244,11 +249,6 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): litellm_params: dict, headers: dict, ) -> dict: - """Async version of transform_request""" - from litellm.llms.watsonx.common_utils import ( - aconvert_watsonx_messages_to_prompt, - ) - provider: Final = model.split("/")[0] prompt: Final = await aconvert_watsonx_messages_to_prompt( model=model, messages=messages, provider=provider, custom_prompt_dict={} diff --git a/litellm/main.py b/litellm/main.py index 2929790f2bd..17edafcdfca 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -19,7 +19,7 @@ import random import sys import time import traceback -from collections.abc import AsyncIterator, Coroutine, Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mapping, Sequence from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy @@ -827,6 +827,35 @@ async def _sleep_for_timeout_async(timeout: float | str | httpx.Timeout): await asyncio.sleep(timeout.connect) +class _AdmissionReservation(BaseModel): + input_tokens: int | None = None + + +class _AdmissionMetadata(BaseModel): + user_api_key_budget_reservation: _AdmissionReservation | None = None + + +def admission_input_tokens(kwargs: Mapping[str, object]) -> int | None: + reservations: Final = ( + _AdmissionMetadata.model_validate(kwargs.get(key) or {}).user_api_key_budget_reservation + for key in ("litellm_metadata", "metadata") + ) + return next( + ( + reservation.input_tokens + for reservation in reservations + if reservation and reservation.input_tokens is not None + ), + None, + ) + + +def admitted_prompt_token_counter(prompt_tokens: int | None) -> Callable[[], int] | None: + if prompt_tokens is None: + return None + return lambda: prompt_tokens + + def mock_completion( model: str, messages: list, @@ -838,6 +867,7 @@ def mock_completion( logging=None, custom_llm_provider=None, timeout: float | str | httpx.Timeout | None = None, + prompt_tokens: int | None = None, **kwargs, ): """ @@ -911,23 +941,26 @@ def mock_completion( if stream is True: model_response = ModelResponseStream() + count_prompt_tokens: Final = admitted_prompt_token_counter(prompt_tokens) # don't try to access stream object, if kwargs.get("acompletion", False) is True: return CustomStreamWrapper( completion_stream=async_mock_completion_streaming_obj( - model_response, mock_response=mock_response, model=model, n=n + model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens ), model=model, custom_llm_provider="openai", logging_obj=logging, + count_prompt_tokens=count_prompt_tokens, ) return CustomStreamWrapper( completion_stream=mock_completion_streaming_obj( - model_response, mock_response=mock_response, model=model, n=n + model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens ), model=model, custom_llm_provider="openai", logging_obj=logging, + count_prompt_tokens=count_prompt_tokens, ) if isinstance(mock_response, litellm.MockException): raise mock_response @@ -953,13 +986,16 @@ def mock_completion( ChatCompletionMessageToolCall(**tool_call) for tool_call in mock_tool_calls ] + usage_prompt_tokens: Final = ( + prompt_tokens if prompt_tokens is not None else DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + ) setattr( model_response, "usage", Usage( - prompt_tokens=DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, + prompt_tokens=usage_prompt_tokens, completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, - total_tokens=DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + total_tokens=usage_prompt_tokens + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, ), ) @@ -4474,7 +4510,7 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client = _dispatch_client_http(ctx) + injected_client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4486,11 +4522,11 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR shared_session: Final = ctx.shared_session stream: Final = ctx.stream timeout: Final = ctx.timeout + client: Final = ( + injected_client if injected_client is not None else (HTTPHandler(timeout=timeout) if stream is False else None) + ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible try: - client = ( - HTTPHandler(timeout=timeout) if stream is False else None - ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible response: Final = base_llm_http_handler.completion( model=model, messages=messages, @@ -5398,6 +5434,14 @@ def completion( if dynamic_api_key is not None: api_key = dynamic_api_key # check if user passed in any of the OpenAI optional params + bridges_to_responses_api: Final = ( + responses_api_model_info.get("mode") == "responses" and not skip_responses_api_bridge + ) + allowed_openai_params: Final[list[str] | None] = ( + [*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"] + if bridges_to_responses_api + else kwargs.get("allowed_openai_params") + ) optional_param_args: Final = { "functions": functions, "function_call": function_call, @@ -5442,7 +5486,7 @@ def completion( "service_tier": service_tier, "store": store, "prompt_cache_key": prompt_cache_key, - "allowed_openai_params": kwargs.get("allowed_openai_params"), + "allowed_openai_params": allowed_openai_params, "base_model": base_model, } optional_params = get_optional_params(**optional_param_args, **non_default_params) @@ -5542,6 +5586,9 @@ def completion( custom_llm_provider=custom_llm_provider, mock_timeout=mock_timeout, timeout=timeout, + prompt_tokens=admission_input_tokens( + cast(Mapping[str, object], kwargs) # cast-ok: completion's **kwargs is untyped + ), ) ## RESPONSES API BRIDGE LOGIC ## - check if model has 'mode: responses' in litellm.model_cost map @@ -6545,7 +6592,7 @@ def embedding( client=client, timeout=timeout, aembedding=aembedding, - litellm_params={}, + litellm_params=litellm_params_dict, api_base=api_base, print_verbose=print_verbose, extra_headers=headers, @@ -7805,6 +7852,7 @@ def transcription( azure_ad_token=azure_ad_token, max_retries=max_retries, litellm_params=litellm_params_dict, + custom_llm_provider=custom_llm_provider, ) elif custom_llm_provider == "openai" or (custom_llm_provider in litellm.openai_compatible_providers): api_base = ( @@ -8247,6 +8295,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 +8320,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) @@ -8384,6 +8437,34 @@ def speech( client=client, _is_async=aspeech or False, ) + elif custom_llm_provider == "mistral": + from litellm.llms.mistral.audio_speech.transformation import ( + MistralTextToSpeechConfig, + ) + + mistral_tts_config: Final = text_to_speech_provider_config or MistralTextToSpeechConfig() + + if api_base is not None: + litellm_params_dict["api_base"] = api_base + if api_key is not None: + litellm_params_dict["api_key"] = api_key + + mistral_voice: Final[str | None] = voice if isinstance(voice, str) else None + + response = base_llm_http_handler.text_to_speech_handler( + model=model, + input=input, + voice=mistral_voice, + text_to_speech_provider_config=mistral_tts_config, + text_to_speech_optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params_dict, + logging_obj=logging_obj, + timeout=timeout, + extra_headers=extra_headers, + client=client, + _is_async=aspeech or False, + ) elif custom_llm_provider == "aws_polly": from litellm.llms.aws_polly.text_to_speech.transformation import ( AWSPollyTextToSpeechConfig, @@ -8553,7 +8634,7 @@ def config_completion(**kwargs): ) -def stream_chunk_builder_text_completion(chunks: list, messages: list | None = None) -> TextCompletionResponse: +def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse: id: Final = chunks[0]["id"] object: Final = chunks[0]["object"] created: Final = chunks[0]["created"] @@ -8670,10 +8751,11 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o def stream_chunk_builder( chunks: list, - messages: list | None = None, + messages: Sequence | None = None, start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, + count_prompt_tokens: Callable[[], int] | None = None, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -8747,6 +8829,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=0, + count_prompt_tokens=count_prompt_tokens, ) setattr(response, "usage", usage) @@ -8924,6 +9007,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=reasoning_tokens, + count_prompt_tokens=count_prompt_tokens, ) setattr(response, "usage", usage) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f607c1e298e..aadd9bd3028 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -364,7 +364,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -380,6 +381,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -399,6 +401,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -416,6 +419,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -435,6 +439,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -452,6 +457,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -471,6 +477,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -488,6 +495,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -507,6 +515,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -537,7 +546,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -551,7 +561,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -650,7 +661,10 @@ }, "twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, "litellm_provider": "bedrock", "max_input_tokens": 77, "max_tokens": 77, @@ -662,7 +676,7 @@ }, "us.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -677,7 +691,7 @@ }, "eu.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -690,6 +704,48 @@ "supports_embedding_image_input": true, "supports_image_input": true }, + "twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "us.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "eu.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, "twelvelabs.pegasus-1-2-v1:0": { "input_cost_per_video_per_second": 0.00049, "output_cost_per_token": 7.5e-06, @@ -711,6 +767,14 @@ "mode": "chat", "supports_video_input": true }, + "global.twelvelabs.pegasus-1-2-v1:0": { + "input_cost_per_video_per_second": 0.00049, + "output_cost_per_token": 7.5e-06, + "litellm_provider": "bedrock", + "mode": "chat", + "supports_video_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "amazon.titan-text-express-v1": { "input_cost_per_token": 1.3e-06, "litellm_provider": "bedrock", @@ -2831,7 +2895,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.7e-08, @@ -2843,7 +2908,8 @@ "output_cost_per_token": 1.48e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "apac.amazon.nova-pro-v1:0": { "input_cost_per_token": 8.4e-07, @@ -2857,7 +2923,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.anthropic.claude-3-5-sonnet-20240620-v1:0": { "deprecation_date": "2026-07-30", @@ -3485,6 +3552,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, @@ -3532,6 +3648,79 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure_ai/codex-mini": { + "cache_read_input_token_cost": 3.75e-07, + "deprecation_date": "2026-11-15", + "input_cost_per_token": 1.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "responses", + "output_cost_per_token": 6e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/whisper": { + "deprecation_date": "2026-12-15", + "input_cost_per_second": 0.0001, + "litellm_provider": "azure_ai", + "mode": "audio_transcription", + "output_cost_per_second": 0.0001, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" + }, "azure_ai/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, @@ -3935,13 +4124,29 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/model_router": { + "deprecation_date": "2027-05-20", "input_cost_per_token": 1.4e-07, "output_cost_per_token": 0, "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "chat", "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, + "azure_ai/model-router": { + "deprecation_date": "2027-05-20", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 0, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", + "comment": "Catalog-name twin of azure_ai/model_router: the flat $0.14 per M input tokens is the router's own fee, the routed model is priced on top of it" + }, "azure/eu/gpt-4o-2024-08-06": { "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.375e-06, @@ -7189,7 +7394,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 +7660,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, @@ -7881,7 +8086,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -7925,7 +8131,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -7969,7 +8176,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -10253,6 +10461,18 @@ "/v1/ocr" ] }, + "azure_ai/cohere-command-a": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 8182, + "max_tokens": 8182, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/", + "supports_function_calling": true, + "supports_tool_choice": true + }, "azure_ai/doc-intelligence/prebuilt-read": { "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.0015, @@ -10604,6 +10824,41 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/grok-4-20-reasoning": { + "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-04-06", + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_reasoning": true + }, + "azure_ai/grok-4-20-non-reasoning": { + "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-04-06", + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure_ai/grok-4-fast-non-reasoning": { "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, @@ -12087,7 +12342,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -12266,7 +12522,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "input_cost_per_token": 4.2e-08, @@ -12278,7 +12535,8 @@ "output_cost_per_token": 1.68e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "input_cost_per_token": 9.6e-07, @@ -12292,7 +12550,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -12859,6 +13118,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "cerebras/qwen-3.8-27b": { + "input_cost_per_token": 9.9e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.49e-06, + "source": "https://api.cerebras.ai/public/v1/models/qwen-3.8-27b", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "chatdolphin": { "input_cost_per_token": 5e-07, "litellm_provider": "nlp_cloud", @@ -13743,7 +14018,8 @@ "max_output_tokens": 3072, "max_tokens": 3072, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/meta/llama-2-7b-chat-int8": { "input_cost_per_token": 1.923e-06, @@ -13752,7 +14028,8 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.1": { "input_cost_per_token": 1.923e-06, @@ -13761,7 +14038,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@hf/thebloke/codellama-7b-instruct-awq": { "input_cost_per_token": 1.923e-06, @@ -13770,7 +14048,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/openai/gpt-oss-120b": { "input_cost_per_token": 3.5e-07, @@ -13780,6 +14059,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13790,7 +14070,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-3b-instruct": { "input_cost_per_token": 5.09e-08, @@ -13799,7 +14080,8 @@ "max_output_tokens": 80000, "max_tokens": 80000, "mode": "chat", - "output_cost_per_token": 3.35e-07 + "output_cost_per_token": 3.35e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-guard-3-8b": { "input_cost_per_token": 4.84e-07, @@ -13808,7 +14090,8 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-08 + "output_cost_per_token": 3e-08, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { "input_cost_per_token": 0.0, @@ -13817,7 +14100,8 @@ "max_output_tokens": 15000, "max_tokens": 15000, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 1.9e-07, @@ -13828,6 +14112,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13839,6 +14124,7 @@ "max_tokens": 80000, "mode": "chat", "output_cost_per_token": 4.881e-06, + "rpm": 300, "supports_reasoning": true }, "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { @@ -13848,7 +14134,8 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 2.87e-07 + "output_cost_per_token": 2.87e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-1b-instruct": { "input_cost_per_token": 2.7e-08, @@ -13857,7 +14144,8 @@ "max_output_tokens": 60000, "max_tokens": 60000, "mode": "chat", - "output_cost_per_token": 2.01e-07 + "output_cost_per_token": 2.01e-07, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.6": { "cache_read_input_token_cost": 1.6e-07, @@ -13868,6 +14156,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13879,6 +14168,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13889,7 +14179,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { "input_cost_per_token": 2.93e-07, @@ -13899,6 +14190,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 2.253e-06, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { @@ -13909,6 +14201,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1.12e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { @@ -13918,7 +14211,8 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1e-06 + "output_cost_per_token": 1e-06, + "rpm": 300 }, "cloudflare/@cf/zai-org/glm-5.2": { "cache_read_input_token_cost": 2.6e-07, @@ -13929,6 +14223,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4.4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13940,6 +14235,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13950,7 +14246,8 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.55e-07 + "output_cost_per_token": 5.55e-07, + "rpm": 300 }, "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { "input_cost_per_token": 5.09e-08, @@ -13960,6 +14257,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 3.35e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13970,7 +14268,8 @@ "max_output_tokens": 3500, "max_tokens": 3500, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/google/gemma-4-26b-a4b-it": { "input_cost_per_token": 1e-07, @@ -13980,6 +14279,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13991,6 +14291,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5.55e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { @@ -14001,6 +14302,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.76e-07, + "rpm": 300, "supports_vision": true }, "cloudflare/@cf/openai/gpt-oss-20b": { @@ -14011,6 +14313,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14022,6 +14325,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 8.5e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwq-32b": { @@ -14032,6 +14336,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 1e-06, + "rpm": 300, "supports_reasoning": true }, "codestral/codestral-2405": { @@ -14158,6 +14463,28 @@ "output_vector_size": 1536, "supports_embedding_image_input": true }, + "us.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, + "global.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "cohere", @@ -20729,7 +21056,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.amazon.nova-micro-v1:0": { "input_cost_per_token": 4.6e-08, @@ -20741,7 +21069,8 @@ "output_cost_per_token": 1.84e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "eu.amazon.nova-pro-v1:0": { "input_cost_per_token": 1.05e-06, @@ -20756,24 +21085,25 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.anthropic.claude-3-5-haiku-20241022-v1:0": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 8e-07, "litellm_provider": "bedrock", "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 4e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 2.5e-08, - "cache_creation_input_token_cost": 3.125e-07, + "cache_read_input_token_cost": 8e-08, + "cache_creation_input_token_cost": 1e-06, "prompt_cache_min_tokens": 2048 }, "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -23780,9 +24110,9 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2e-06, @@ -23805,12 +24135,12 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_url_context": true, + "supports_url_context": false, "supports_vision": true, "supports_web_search": true, "search_context_cost_per_query": { @@ -27602,6 +27932,70 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-5.5": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.5", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-luna": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-luna", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-sol": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-sol", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-terra": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-terra", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "chatgpt/gpt-5.4": { "litellm_provider": "chatgpt", "max_input_tokens": 1050000, @@ -28208,6 +28602,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -29032,7 +29427,12 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + } }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29049,9 +29449,9 @@ "output_cost_per_token_priority": 1e-06, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29150,9 +29550,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29177,9 +29577,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29290,9 +29690,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29317,9 +29717,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29401,6 +29801,66 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2.5-flare": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-flare-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, "low/1024-x-1024/gpt-image-1.5": { "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, @@ -29851,7 +30311,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": false, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": false }, "gpt-5.1-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, @@ -29896,7 +30356,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": false, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": false }, "gpt-5.1-chat-latest": { "cache_read_input_token_cost": 1.25e-07, @@ -29988,7 +30448,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": false }, "gpt-5.2-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, @@ -30034,7 +30494,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": false }, "gpt-5.2-chat-latest": { "cache_read_input_token_cost": 1.75e-07, @@ -30742,6 +31202,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], "source": "https://developers.openai.com/api/docs/models/chat-latest", "supported_endpoints": [ "/v1/chat/completions", @@ -30759,6 +31222,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -30880,8 +31344,6 @@ "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-pro": { - "cache_read_input_token_cost": 3e-06, - "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, @@ -30930,8 +31392,6 @@ "supports_low_reasoning_effort": false }, "gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 3e-06, - "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, @@ -31026,7 +31486,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 @@ -31078,14 +31538,12 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-pro": { - "cache_read_input_token_cost": 3e-06, - "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, @@ -31130,13 +31588,11 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-pro-2026-03-05": { - "cache_read_input_token_cost": 3e-06, - "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, @@ -31181,7 +31637,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, @@ -32628,6 +33084,7 @@ "supports_tool_choice": true }, "groq/gemma-7b-it": { + "deprecation_date": "2024-12-18", "input_cost_per_token": 5e-08, "litellm_provider": "groq", "max_input_tokens": 8192, @@ -33261,13 +33718,14 @@ "supports_tool_choice": true }, "jina-reranker-v2-base-multilingual": { - "input_cost_per_token": 1.8e-08, + "input_cost_per_token": 5e-08, "litellm_provider": "jina_ai", "max_input_tokens": 1024, "max_output_tokens": 1024, "max_tokens": 1024, "mode": "rerank", - "output_cost_per_token": 1.8e-08 + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" }, "jp.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.125e-06, @@ -33441,6 +33899,20 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "inception/mercury-2.5": { + "input_cost_per_token": 2e-07, + "litellm_provider": "inception", + "max_input_tokens": 260000, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://docs.inceptionlabs.ai/get-started/models", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "text-completion-inception/mercury-edit-2": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, @@ -34898,9 +35370,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-latest": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" @@ -37176,6 +37648,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37216,6 +37689,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37427,6 +37901,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37467,6 +37942,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -38673,19 +39149,21 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-chat-v3.1": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 2.5e-07, "input_cost_per_token_cache_hit": 2e-08, "litellm_provider": "openrouter", "max_input_tokens": 163840, "max_output_tokens": 163840, "max_tokens": 163840, "mode": "chat", - "output_cost_per_token": 8e-07, + "output_cost_per_token": 9.5e-07, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.3e-07, + "source": "https://openrouter.ai/deepseek/deepseek-chat-v3.1" }, "openrouter/deepseek/deepseek-v3.2": { "input_cost_per_token": 2.69e-07, @@ -38748,36 +39226,56 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 8.59908e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.719816e-06, "source": "https://openrouter.ai/deepseek/deepseek-v4-pro", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 7.1659e-08 + }, + "openrouter/deepseek/deepseek-v4.1-flash": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 3e-09, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "source": "https://openrouter.ai/deepseek/deepseek-v4.1-flash", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": false, + "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.7948e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.73844e-06, "source": "https://openrouter.ai/deepseek/deepseek-v4-pro-0813", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.9316e-08 }, "openrouter/google/gemini-2.0-flash-001": { "deprecation_date": "2026-06-01", @@ -39517,6 +40015,67 @@ "supports_tool_choice": true, "supports_vision": true }, + "openrouter/openai/gpt-5.6-sol": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "default_reasoning_effort": "medium", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://openrouter.ai/openai/gpt-5.6-sol", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/openai/gpt-5.6-sol-pro": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2e-07, + "cache_creation_input_token_cost": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-5.6-sol-pro", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/openai/gpt-oss-120b": { "input_cost_per_token": 3.7e-08, "litellm_provider": "openrouter", @@ -39645,13 +40204,13 @@ "supports_tool_choice": true }, "openrouter/qwen/qwen3-235b-a22b-2507": { - "input_cost_per_token": 8.75e-08, + "input_cost_per_token": 2.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 3.5e-07, + "output_cost_per_token": 8.8e-07, "source": "https://openrouter.ai/qwen/qwen3-235b-a22b-2507", "supports_function_calling": true, "supports_tool_choice": true @@ -39684,7 +40243,7 @@ "supports_vision": true }, "openrouter/qwen/qwen3.5-35b-a3b": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 3.125e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, @@ -39695,7 +40254,8 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.5625e-07 }, "openrouter/qwen/qwen3.5-27b": { "input_cost_per_token": 1.95e-07, @@ -39712,13 +40272,13 @@ "supports_vision": true }, "openrouter/qwen/qwen3.5-122b-a10b": { - "input_cost_per_token": 2.9e-07, + "input_cost_per_token": 2.6e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 2.4e-06, + "output_cost_per_token": 2.08e-06, "source": "https://openrouter.ai/qwen/qwen3.5-122b-a10b", "supports_function_calling": true, "supports_reasoning": true, @@ -39805,18 +40365,19 @@ "supports_web_search": true }, "openrouter/z-ai/glm-4.6": { - "input_cost_per_token": 5.5e-07, + "input_cost_per_token": 4.3e-07, "litellm_provider": "openrouter", "max_input_tokens": 202800, "max_output_tokens": 131000, "max_tokens": 131000, "mode": "chat", - "output_cost_per_token": 2.2e-06, + "output_cost_per_token": 1.75e-06, "source": "https://openrouter.ai/z-ai/glm-4.6", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8e-08 }, "openrouter/z-ai/glm-4.6:exacto": { "input_cost_per_token": 4.5e-07, @@ -41473,6 +42034,28 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "rerank-v4.0-fast": { + "input_cost_per_query": 0.002, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, + "rerank-v4.0-pro": { + "input_cost_per_query": 0.0025, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -42619,7 +43202,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.2e-06, + "max_input_tokens": 131072, + "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo": { "litellm_provider": "together_ai", @@ -42627,7 +43214,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 3e-07, + "max_input_tokens": 32768, + "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": { "deprecation_date": "2026-07-10", @@ -42839,7 +43430,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2e-07, + "max_input_tokens": 32768, + "source": "https://api.together.xyz/v1/models" }, "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": { "deprecation_date": "2026-04-02", @@ -42847,7 +43442,11 @@ "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, + "max_input_tokens": 32768, + "source": "https://api.together.xyz/v1/models" }, "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": { "deprecation_date": "2026-04-16", @@ -43433,7 +44032,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, @@ -43445,7 +44045,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -43474,7 +44075,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, @@ -45653,8 +46255,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -45678,8 +46280,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -47141,6 +47743,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", @@ -47482,27 +48177,29 @@ "supports_tool_choice": true }, "vertex_ai/mistral-small-2503": { - "input_cost_per_token": 1e-06, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-mistral_models", "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-06, + "output_cost_per_token": 3e-07, "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "vertex_ai/mistral-small-2503@001": { - "input_cost_per_token": 1e-06, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-mistral_models", "max_input_tokens": 32000, "max_output_tokens": 8191, "max_tokens": 8191, "mode": "chat", - "output_cost_per_token": 3e-06, + "output_cost_per_token": 3e-07, "supports_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "vertex_ai/mistral-ocr-2505": { "litellm_provider": "vertex_ai", @@ -47552,23 +48249,24 @@ "supports_reasoning": true }, "vertex_ai/openai/gpt-oss-20b-maas": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 7e-08, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 3e-07, - "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas", - "supports_reasoning": true + "output_cost_per_token": 2.5e-07, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_reasoning": true, + "cache_read_input_token_cost": 7e-09 }, "vertex_ai/xai/grok-4.1-fast-non-reasoning": { "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -47582,9 +48280,9 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -47597,14 +48295,17 @@ }, "vertex_ai/xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -47613,14 +48314,17 @@ }, "vertex_ai/xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -47628,6 +48332,44 @@ "supports_vision": true, "supports_web_search": true }, + "vertex_ai/xai/grok-4.3": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "vertex_ai/xai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -47886,6 +48628,16 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-multilingual-2": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-3-large": { "input_cost_per_token": 1.8e-07, "litellm_provider": "voyage", @@ -47991,6 +48743,7 @@ "output_cost_per_token": 0.0 }, "wandb/openai/gpt-oss-120b": { + "supports_reasoning": true, "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, @@ -48001,6 +48754,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/openai/gpt-oss-20b": { + "supports_reasoning": true, "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, @@ -48011,6 +48765,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/zai-org/GLM-4.5": { + "supports_reasoning": true, "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, @@ -48039,6 +48794,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3-235B-A22B-Thinking-2507": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, @@ -48095,6 +48851,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/deepseek-ai/DeepSeek-V3.1": { + "supports_reasoning": true, "max_tokens": 128000, "max_input_tokens": 161000, "max_output_tokens": 128000, @@ -48105,6 +48862,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/deepseek-ai/DeepSeek-R1-0528": { + "supports_reasoning": true, "max_tokens": 161000, "max_input_tokens": 161000, "max_output_tokens": 161000, @@ -51959,7 +52717,8 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 8e-07, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/openai/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -51998,7 +52757,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-01" }, "scaleway/hcompany/holo2-30b-a3b": { "input_cost_per_token": 3e-07, @@ -52009,7 +52769,8 @@ "mode": "chat", "output_cost_per_token": 7e-07, "supports_reasoning": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-09" }, "scaleway/mistralai/mistral-medium-3.5-128b": { "input_cost_per_token": 1.5e-06, @@ -52032,7 +52793,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 2e-06, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/voxtral-small-24b-2507": { "input_cost_per_audio_token": 1.5e-07, @@ -52043,7 +52805,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.5e-07, - "supports_audio_input": true + "supports_audio_input": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/mistral-small-3.2-24b-instruct-2506": { "input_cost_per_token": 1.5e-07, @@ -52065,7 +52828,8 @@ "mode": "chat", "output_cost_per_token": 2e-07, "supports_vision": true, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/BAAI/bge-multilingual-gemma2": { "input_cost_per_token": 1e-07, @@ -54515,7 +55279,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-20b": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 7e-08, "output_cost_per_token": 3e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, @@ -54549,8 +55313,8 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 65536, @@ -54596,11 +55360,14 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-terra": { @@ -54635,11 +55402,14 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-cyber": { @@ -54664,12 +55434,49 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, @@ -54702,11 +55509,14 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "us.openai.gpt-5.6-sol": { @@ -54731,8 +55541,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "global.openai.gpt-5.6-sol": { @@ -54757,8 +55570,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "us.openai.gpt-5.6-terra": { @@ -54783,8 +55599,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "global.openai.gpt-5.6-terra": { @@ -54809,8 +55628,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "us.openai.gpt-5.6-luna": { @@ -54835,8 +55657,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "global.openai.gpt-5.6-luna": { @@ -54861,10 +55686,115 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "us.openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "global.openai.gpt-6-astra": { + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "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, + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, @@ -54894,11 +55824,13 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.4": { @@ -54930,11 +55862,13 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "bedrock_mantle/google.gemma-4-31b": { @@ -55666,17 +56600,17 @@ "supports_reasoning": true, "source": "https://serverless.tensormesh.ai/v1/models/openrouter" }, - "deepseek-v4-flash": { + "deepseek-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -55690,19 +56624,45 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, - "deepseek-v4-flash-vision-exp": { + "deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "deepseek-v4-flash-vision-exp": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "deepseek", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -55744,17 +56704,17 @@ "supports_tool_choice": true, "supports_vision": false }, - "deepseek/deepseek-v4-flash": { + "deepseek/deepseek-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -55768,19 +56728,45 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, - "deepseek/deepseek-v4-flash-vision-exp": { + "deepseek/deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "deepseek/deepseek-v4-flash-vision-exp": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "deepseek", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56136,6 +57122,23 @@ ], "supports_audio_input": true }, + "gpt-live-1": { + "input_cost_per_second": 0.0008333333333333334, + "litellm_provider": "openai", + "mode": "realtime", + "source": "https://developers.openai.com/api/docs/models/gpt-live-1", + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true + }, "gpt-realtime-translate": { "input_cost_per_second": 0.0005666666666666667, "litellm_provider": "openai", @@ -56374,9 +57377,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-2603": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" @@ -56459,6 +57462,14 @@ "model_info": { "supports_mid_conversation_system": true } + }, + { + "name": "wandb-reasoning-baseline", + "pattern": "^wandb/", + "description": "Any Weights & Biases Inference model id, anchored to the wandb/ namespace so only that provider's ids match. W&B's serverless catalog is reasoning-first and grows faster than this registry names it, so an id the map has not described yet is treated as reasoning-capable and keeps the caller's reasoning_effort instead of dropping it or raising UnsupportedParamsError. Rules lose to exact entries, so a mapped non-reasoning model such as wandb/meta-llama/Llama-3.1-8B-Instruct is unaffected. Carries no mode and no pricing, so cost stays on the standard unpriced behavior and the deployment does not read as catalog-mapped to the router's reasoning-effort resolver.", + "model_info": { + "supports_reasoning": true + } } ] }, @@ -56530,8 +57541,8 @@ "rpm": 10 }, "vertex_ai/gemini-3.5-transcribe-preview": { - "input_cost_per_audio_token": 2.5e-06, - "input_cost_per_token": 2.5e-06, + "input_cost_per_audio_token": 2e-06, + "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai", "mode": "audio_transcription", "output_cost_per_token": 1.2e-05, @@ -56566,6 +57577,27 @@ ], "supports_audio_input": true }, + "vertex_ai/gemini-3.5-live-translate-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 2.1e-05, + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "perplexity/pplx-embed-context-v1-0.6b": { "input_cost_per_token": 8e-09, "litellm_provider": "perplexity", @@ -57816,6 +58848,7 @@ "supports_vision": false }, "wandb/deepseek-ai/DeepSeek-V4-Flash": { + "supports_reasoning": true, "max_tokens": 1048576, "max_input_tokens": 1048576, "input_cost_per_token": 1.4e-07, @@ -57828,6 +58861,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/deepseek-ai/DeepSeek-V4-Flash-0731": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 1.3e-07, @@ -57840,6 +58874,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/deepseek-ai/DeepSeek-V4-Pro": { + "supports_reasoning": true, "max_tokens": 1048576, "max_input_tokens": 1048576, "input_cost_per_token": 1.15e-06, @@ -57852,6 +58887,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/google/gemma-4-31B-it": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 1e-07, @@ -57892,6 +58928,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/MiniMaxAI/MiniMax-M3": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 2.3e-07, @@ -57904,6 +58941,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/moonshotai/Kimi-K2.7-Code": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 7.1e-07, @@ -57916,6 +58954,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/moonshotai/Kimi-K2.6": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 6.5e-07, @@ -57928,6 +58967,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 1e-07, @@ -57940,6 +58980,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 7.5e-07, @@ -57962,6 +59003,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.8-27B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 4e-07, @@ -57974,6 +59016,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.6-35B-A3B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 2.5e-07, @@ -57984,6 +59027,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.6-27B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 6e-07, @@ -57996,6 +59040,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.5-35B-A3B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 2.5e-07, @@ -58015,7 +59060,28 @@ "supports_vision": false, "source": "https://wandb.ai/site/pricing/tokens/" }, + "wandb/deepseek-ai/DeepSeek-V4-Pro-0813": { + "litellm_provider": "wandb", + "mode": "chat", + "supports_reasoning": true, + "input_cost_per_token": 0.00000131, + "output_cost_per_token": 0.00000396, + "cache_read_input_token_cost": 0.000000044, + "supports_prompt_caching": true, + "source": "https://wandb.ai/site/pricing/tokens/" + }, + "wandb/ibm-granite/granite-4.2-8b": { + "litellm_provider": "wandb", + "mode": "chat", + "supports_reasoning": true, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.00000015, + "cache_read_input_token_cost": 0.00000005, + "supports_prompt_caching": true, + "source": "https://wandb.ai/site/pricing/tokens/" + }, "wandb/zai-org/GLM-5.2": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 7.6e-07, @@ -59177,6 +60243,7 @@ ] }, "xai/grok-imagine-image-quality": { + "deprecation_date": "2026-11-02", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -59193,6 +60260,7 @@ ] }, "xai/grok-imagine-image-quality-20260403": { + "deprecation_date": "2026-11-02", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -59209,6 +60277,7 @@ ] }, "xai/grok-imagine-image-quality-latest": { + "deprecation_date": "2026-11-02", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -59257,6 +60326,77 @@ "image" ] }, + "xai/grok-imagine-video": { + "input_cost_per_image": 0.002, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.05, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_720p": 0.07, + "source": "https://docs.x.ai/docs/models/grok-imagine-video", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-2026-05-30": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-preview": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", @@ -59657,6 +60797,113 @@ "output_cost_per_token": 4.7e-07, "source": "https://docs.together.ai/docs/serverless-models" }, + "together_ai/moonshotai/Kimi-K2.6": { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 4.5e-06, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/moonshotai/Kimi-K2.5-fp4": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.8e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/MiniMaxAI/MiniMax-M2.7": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 196608, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/zai-org/GLM-5": { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3.2e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 202752, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/zai-org/GLM-5.1": { + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 202752, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-R1-0528": { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 7e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 163840, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/Qwen/Qwen3-Coder-Next-FP8": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/Qwen/Qwen3-VL-32B-Instruct": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.5e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/Qwen/Qwen3-VL-8B-Instruct": { + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 6.8e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/mistralai/Ministral-3-14B-Instruct-2512": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { + "input_cost_per_token": 6e-08, + "output_cost_per_token": 2.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/mistralai/Mistral-7B-Instruct-v0.3": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/Qwen/QwQ-32B": { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, "cerebras/gemma-4-31b": { "input_cost_per_token": 9.9e-07, "litellm_provider": "cerebras", @@ -60160,10 +61407,12 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true, "input_cost_per_token": 2.64e-06, "input_cost_per_token_above_272k_tokens": 5.28e-06, @@ -60193,10 +61442,12 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true, "input_cost_per_token": 2.64e-07, "input_cost_per_token_above_272k_tokens": 5.28e-07, @@ -60225,10 +61476,12 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true, "input_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 3.3e-07, @@ -60387,10 +61640,12 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true, "input_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 3.3e-07, @@ -60524,6 +61779,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60534,6 +61790,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60589,6 +61846,31 @@ "supports_web_search": false, "output_cost_per_image": 0.08 }, + "gemini/lyria-3.5": { + "input_cost_per_token": 0, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 0, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "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, + "output_cost_per_image": 0.08 + }, "perplexity/anthropic/claude-fable-5": { "litellm_provider": "perplexity", "mode": "responses", @@ -60900,6 +62182,7 @@ "mode": "responses", "supports_web_search": true, "supports_function_calling": true, + "supports_reasoning": true, "input_cost_per_token": 1.15e-08, "output_cost_per_token": 1.7e-07, "cache_read_input_token_cost": 1.15e-09, @@ -60910,6 +62193,7 @@ "mode": "responses", "supports_web_search": true, "supports_function_calling": true, + "supports_reasoning": true, "input_cost_per_token": 2.5e-07, "output_cost_per_token": 2.5e-06, "cache_read_input_token_cost": 2.5e-07, @@ -61273,6 +62557,28 @@ "cache_read_input_token_cost": 2e-08, "supports_prompt_caching": true }, + "openrouter/openai/gpt-5.6-luna-pro": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 2e-08, + "cache_creation_input_token_cost": 2.5e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-5.6-luna-pro", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/openai/gpt-5.6-terra": { "input_cost_per_token": 2e-06, "output_cost_per_token": 1.2e-05, @@ -61292,6 +62598,28 @@ "cache_read_input_token_cost": 2e-07, "supports_prompt_caching": true }, + "openrouter/openai/gpt-5.6-terra-pro": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1.2e-05, + "cache_read_input_token_cost": 2e-07, + "cache_creation_input_token_cost": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-5.6-terra-pro", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/openai/o3": { "input_cost_per_token": 2e-06, "output_cost_per_token": 8e-06, @@ -61525,6 +62853,28 @@ "supports_pdf_input": true, "supports_prompt_caching": true }, + "openrouter/openai/gpt-6-astra-pro": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_creation_input_token_cost": 1.25e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-6-astra-pro", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/qwen/qwen3.8-flash": { "input_cost_per_token": 1.5e-07, "output_cost_per_token": 4.7e-07, @@ -61544,9 +62894,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 2.5e-07, - "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 3e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -61661,6 +63011,25 @@ "supports_vision": true, "supports_prompt_caching": true }, + "openrouter/qwen/qwen3.8-max-0902": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost": 2.5e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "source": "https://openrouter.ai/qwen/qwen3.8-max-0902", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": false, + "supports_prompt_caching": true + }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 6.5e-08, "output_cost_per_token": 1.8e-07, @@ -61732,9 +63101,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 3e-06, - "output_cost_per_token": 1.5e-05, - "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 2.1e-06, + "output_cost_per_token": 1.053e-05, + "cache_read_input_token_cost": 2.35e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -61831,9 +63200,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.2": { - "input_cost_per_token": 9.66e-07, - "output_cost_per_token": 3.036e-06, - "cache_read_input_token_cost": 1.932e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2e-06, + "cache_read_input_token_cost": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 131072, @@ -61864,9 +63233,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 6.6e-07, - "output_cost_per_token": 3.4e-06, - "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 7.1e-07, + "output_cost_per_token": 3.5e-06, + "cache_read_input_token_cost": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 235929, @@ -62114,10 +63483,28 @@ "supports_vision": true, "supports_pdf_input": true }, + "openrouter/openai/gpt-chat-latest": { + "input_cost_per_token": 5e-06, + "output_cost_per_token": 3e-05, + "cache_read_input_token_cost": 5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-chat-latest", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": false, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.778e-08, - "output_cost_per_token": 1.7556e-07, - "cache_read_input_token_cost": 1.7556e-08, + "input_cost_per_token": 8.54e-08, + "output_cost_per_token": 1.708e-07, + "cache_read_input_token_cost": 1.708e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -62150,8 +63537,8 @@ "supports_prompt_caching": true }, "openrouter/google/gemma-4-26b-a4b-it": { - "input_cost_per_token": 7e-08, - "output_cost_per_token": 3.4e-07, + "input_cost_per_token": 4.2e-08, + "output_cost_per_token": 2.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -62833,7 +64220,7 @@ "supports_vision": false }, "openrouter/qwen/qwen3-next-80b-a3b-instruct": { - "input_cost_per_token": 1e-07, + "input_cost_per_token": 9e-08, "output_cost_per_token": 1.1e-06, "cache_read_input_token_cost": 7e-08, "litellm_provider": "openrouter", @@ -62959,8 +64346,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-30b-a3b-instruct-2507": { - "input_cost_per_token": 4.815e-08, - "output_cost_per_token": 1.9305e-07, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32000, @@ -63158,8 +64545,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-14b": { - "input_cost_per_token": 1.2e-07, - "output_cost_per_token": 2.4e-07, + "input_cost_per_token": 2.275e-07, + "output_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 16384, diff --git a/litellm/models/__init__.py b/litellm/models/__init__.py index 07d1ffa743d..50eb6f8af4f 100644 --- a/litellm/models/__init__.py +++ b/litellm/models/__init__.py @@ -3,6 +3,7 @@ Domain models for LiteLLM backend. """ from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.models.autorouter_session import LiteLLM_AutoRouterSession from litellm.models.budget import ( LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, @@ -40,6 +41,7 @@ __all__ = [ "CredentialBase", "CredentialItem", "LiteLLM_AccessGroupTable", + "LiteLLM_AutoRouterSession", "LiteLLM_BudgetTable", "LiteLLM_BudgetTableFull", "LiteLLM_Config", diff --git a/litellm/models/autorouter_session.py b/litellm/models/autorouter_session.py new file mode 100644 index 00000000000..c7126236ec3 --- /dev/null +++ b/litellm/models/autorouter_session.py @@ -0,0 +1,39 @@ +""" +Auto-router per-session rollup model. + +Canonical definition for ``litellm_autoroutersession``, the row the spend flush +maintains per (api_key, session_id, router_name). +""" + +from collections.abc import Mapping +from datetime import datetime + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_AutoRouterSession(LiteLLMPydanticObjectBase): + api_key: str + session_id: str + router_name: str + router_type: str + first_turn_at: datetime + last_turn_at: datetime + last_model: str + turns: int + spend: float + saved_spend: float + classifier_cost: float + tier_turns: Mapping[str, int] + baseline_models: Mapping[str, int] + + @property + def baseline_model(self) -> str | None: + """The baseline most of this session's turns were priced against, or None when no turn recorded one. + + A router reconfigured mid-session leaves turns priced against two baselines; the row keeps both + counts, and the label is the one that priced the most money-carrying turns rather than whatever the + router is configured with now. + """ + if not self.baseline_models: + return None + return max(self.baseline_models, key=lambda model: (self.baseline_models[model], model)) diff --git a/litellm/models/managed_files.py b/litellm/models/managed_files.py index 23d70ef5c48..c90f9b535ea 100644 --- a/litellm/models/managed_files.py +++ b/litellm/models/managed_files.py @@ -32,6 +32,7 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse created_by: str | None = None team_id: str | None = None + org_id: str | None = None class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py index 6cc4a765e46..7ccff9434a7 100644 --- a/litellm/models/mcp_server.py +++ b/litellm/models/mcp_server.py @@ -6,10 +6,12 @@ Canonical definition for ``litellm_mcpservertable``. Re-exported from """ import enum +from collections.abc import Mapping from datetime import datetime +from types import MappingProxyType from typing import Literal -from pydantic import Field +from pydantic import Field, ValidationInfo, field_validator from litellm.types.llms.base import LiteLLMPydanticObjectBase from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPTransportType @@ -98,6 +100,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 @@ -114,3 +117,12 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): submitted_at: datetime | None = None reviewed_at: datetime | None = None review_notes: str | None = None + + @field_validator("static_headers", "env", mode="before") + @classmethod + def decode_stored_secret_map(cls, value: object, info: ValidationInfo) -> Mapping[str, str] | None: + from litellm.proxy.common_utils.encrypt_decrypt_utils import decode_secret_map + + if value is None and info.field_name == "env": + return MappingProxyType({}) + return decode_secret_map(value, key=info.field_name or "secret map") diff --git a/litellm/models/object_permission.py b/litellm/models/object_permission.py index a09d50ddc33..f178c5ad47f 100644 --- a/litellm/models/object_permission.py +++ b/litellm/models/object_permission.py @@ -23,3 +23,4 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): blocked_tools: list[str] | None = [] search_tools: list[str] | None = [] mcp_tool_search_enabled: bool | None = None + skills: list[str] | None = None diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 6c68971f8d5..56bfd98895d 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -10,6 +10,7 @@ import re from collections.abc import Callable, Coroutine, Mapping from dataclasses import dataclass from io import IOBase +from types import MappingProxyType from typing import Any, Final, cast import httpx @@ -19,6 +20,7 @@ from litellm._logging import verbose_logger from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure_ai.ocr.common_utils import ( + is_azure_cohere_parse_model, is_azure_document_intelligence_model, ) from litellm.llms.base_llm.ocr.transformation import ( @@ -29,6 +31,8 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.rust_bridge.bindings import native_exception_types +from litellm.rust_bridge.configuration import rust_enabled from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -50,21 +54,32 @@ class _PreparedOCRRequest: litellm_params: dict[str, object] effective_timeout: float | httpx.Timeout litellm_logging_obj: LiteLLMLoggingObj + caller_supplied_api_key: bool = True + caller_supplied_api_base: bool = True -@dataclass -class _PreparedRustOCRCall: - api_key: str | None - api_base: str | None - headers: dict[str, object] - optional_params: dict[str, object] - - -_RUST_OCR_PROVIDERS: Final = { - "mistral", - "azure_ai", - "vertex_ai", -} +_RUST_OCR_PROVIDERS: Final = frozenset({"mistral", "azure_ai", "vertex_ai"}) +_RUST_OCR_CONFIG_FIELDS: Final = frozenset( + { + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_scope", + "azure_authority_host", + "azure_credential", + "azure_federated_token_file", + "vertex_credentials", + "vertex_ai_credentials", + "vertex_project", + "vertex_ai_project", + "vertex_location", + "vertex_ai_location", + } +) +_RUST_OCR_SECRET_FIELDS: Final = frozenset( + {"azure_ad_token", "client_secret", "azure_federated_token_file", "vertex_credentials", "vertex_ai_credentials"} +) def _prepare_ocr_request( @@ -92,6 +107,7 @@ def _prepare_ocr_request( if doc_type not in ["document_url", "image_url"]: raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'") + caller_supplied_api_key: Final = api_key is not None caller_supplied_api_base: Final = api_base is not None ( @@ -185,155 +201,256 @@ def _prepare_ocr_request( litellm_params=dict(litellm_params), effective_timeout=effective_timeout, litellm_logging_obj=litellm_logging_obj, + caller_supplied_api_key=caller_supplied_api_key, + caller_supplied_api_base=caller_supplied_api_base, ) -def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: - if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native": - return False - if not prepared_request.provider_config.supports_rust_bridge(): - return False - return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS - - -def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool: - raw_request_override: Final = prepared_request.litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override) - - -def _rust_bridge_optional_params( - prepared_request: _PreparedOCRRequest, - resolve_secret: Callable[[str], str | None], -) -> dict[str, object]: - optional_params: Final = dict(prepared_request.optional_params) - if prepared_request.custom_llm_provider == "vertex_ai": - vertex_project: Final = ( - prepared_request.litellm_params.get("vertex_project") - or prepared_request.litellm_params.get("vertex_ai_project") - or litellm.vertex_project - or resolve_secret("VERTEXAI_PROJECT") - ) - vertex_location: Final = ( - prepared_request.litellm_params.get("vertex_location") - or prepared_request.litellm_params.get("vertex_ai_location") - or litellm.vertex_location - or resolve_secret("VERTEXAI_LOCATION") - or resolve_secret("VERTEX_LOCATION") - ) - if vertex_project is not None: - optional_params["vertex_project"] = vertex_project - if vertex_location is not None: - optional_params["vertex_location"] = vertex_location - return optional_params - - -def _rust_bridge_api_base( - prepared_request: _PreparedOCRRequest, - resolve_secret: Callable[[str], str | None], -) -> str | None: - if prepared_request.api_base is not None: - return prepared_request.api_base - if prepared_request.custom_llm_provider == "azure_ai": - if is_azure_document_intelligence_model(prepared_request.model): - return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") - return resolve_secret("AZURE_AI_API_BASE") +def _rust_ocr_provider(request: rust_ocr_bridge.LiteLLMOcrRequest) -> str | None: + if request.custom_llm_provider is not None: + return request.custom_llm_provider + prefix: Final = request.model.partition("/")[0] + if prefix in _RUST_OCR_PROVIDERS: + return prefix + if request.model.startswith("mistral-ocr"): + return "mistral" return None -def _prepare_rust_ocr_call( - prepared_request: _PreparedOCRRequest, - resolve_api_key: Callable[[str], str | None], -) -> _PreparedRustOCRCall: - provider_config: Final = prepared_request.provider_config - api_key_env_var: Final = provider_config.get_api_key_env_var() - resolved_api_key: Final = prepared_request.api_key or ( - resolve_api_key(api_key_env_var) if api_key_env_var is not None else None +def _rust_ocr_supported(request: rust_ocr_bridge.LiteLLMOcrRequest) -> bool: + provider: Final = _rust_ocr_provider(request) + if provider not in _RUST_OCR_PROVIDERS or request.kwargs.get(OCR_REQUEST_FORMAT_PARAM) == "native": + return False + if provider == "azure_ai": + return ( + not is_azure_cohere_parse_model(request.model) + and not callable(request.kwargs.get("azure_ad_token_provider")) + and request.kwargs.get("azure_username") is None + and request.kwargs.get("azure_password") is None + ) + return True + + +def _rust_bridge_optional_params( + request: rust_ocr_bridge.LiteLLMOcrRequest, + resolve_secret: Callable[[str], str | None], +) -> Mapping[str, object]: + optional_params: Final = MappingProxyType( + { + name: value + for name, value in request.kwargs.items() + if (name not in GenericLiteLLMParams.model_fields or name in _RUST_OCR_CONFIG_FIELDS) + and name not in {"litellm_logging_obj", "aocr", "litellm_call_id", "proxy_server_request"} + } ) - resolved_headers: Final = provider_config.validate_environment( - headers=prepared_request.extra_headers or {}, - model=prepared_request.model, - api_key=resolved_api_key, - api_base=prepared_request.api_base, - litellm_params=prepared_request.litellm_params, + provider: Final = _rust_ocr_provider(request) + if provider == "azure_ai" and litellm.enable_azure_ad_token_refresh is True: + return MappingProxyType({**optional_params, "enable_azure_ad_token_refresh": True}) + if provider != "vertex_ai": + return optional_params + project: Final = ( + request.kwargs.get("vertex_project") + or request.kwargs.get("vertex_ai_project") + or litellm.vertex_project + or resolve_secret("VERTEXAI_PROJECT") ) - resolved_complete_url: Final = provider_config.get_complete_url( - api_base=prepared_request.api_base, - model=prepared_request.model, - optional_params=prepared_request.optional_params, - litellm_params=prepared_request.litellm_params, + location: Final = ( + request.kwargs.get("vertex_location") + or request.kwargs.get("vertex_ai_location") + or litellm.vertex_location + or resolve_secret("VERTEXAI_LOCATION") + or resolve_secret("VERTEX_LOCATION") ) - rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key) - rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key) - prepared_request.litellm_logging_obj.pre_call( + credentials: Final = ( + request.kwargs.get("vertex_credentials") + or request.kwargs.get("vertex_ai_credentials") + or resolve_secret("VERTEXAI_CREDENTIALS") + ) + vertex_params: Final = MappingProxyType( + { + name: value + for name, value in ( + ("vertex_project", project), + ("vertex_location", location), + ("vertex_credentials", credentials), + ) + if value is not None + } + ) + return MappingProxyType({**optional_params, **vertex_params}) + + +def _rust_bridge_input_sources( + request: rust_ocr_bridge.LiteLLMOcrRequest, + optional_params: Mapping[str, object], +) -> Mapping[str, str]: + proxy_request: Final = request.kwargs.get("proxy_server_request") + if not isinstance(proxy_request, Mapping): + return MappingProxyType({}) + proxy_request_mapping: Final = cast( # cast-ok: runtime Mapping check loses generic key and value types + Mapping[object, object], proxy_request + ) + body_value: Final = proxy_request_mapping.get("body") + if not isinstance(body_value, Mapping): + return MappingProxyType({}) + body: Final = cast( # cast-ok: runtime Mapping check loses generic key and value types + Mapping[object, object], body_value + ) + credential_fields_value: Final = proxy_request_mapping.get("credential_fields", ()) + credential_fields: Final = ( + frozenset(name for name in credential_fields_value if isinstance(name, str)) + if isinstance(credential_fields_value, (list, tuple, set, frozenset)) + else frozenset() + ) + names: Final = frozenset(optional_params) | frozenset({"api_key", "api_base", "extra_headers"}) + request_sources: Final = MappingProxyType( + {name: "request" for name in names if name in body or name in credential_fields} + ) + if litellm.enable_azure_ad_token_refresh is True and "enable_azure_ad_token_refresh" in optional_params: + return MappingProxyType({**request_sources, "enable_azure_ad_token_refresh": "deployment"}) + return request_sources + + +def _marshal_rust_ocr_request( + request: rust_ocr_bridge.LiteLLMOcrRequest, + resolve_secret: Callable[[str], str | None], +) -> rust_ocr_bridge.LiteLLMOcrRequest: + if not isinstance(request.document, dict): + raise TypeError(f"document must be a dict with 'type' and URL/file field, got {type(request.document)}") + document: Final = ( + convert_file_document_to_url_document(request.document) + if request.document.get("type") == "file" + else request.document + ) + provider: Final = _rust_ocr_provider(request) + api_key: Final = request.api_key or resolve_secret("MISTRAL_API_KEY") if provider == "mistral" else request.api_key + optional_params: Final = _rust_bridge_optional_params(request, resolve_secret) + input_sources: Final = _rust_bridge_input_sources(request, optional_params) + logged_optional_params: Final = MappingProxyType( + {name: "****" if name in _RUST_OCR_SECRET_FIELDS else value for name, value in optional_params.items()} + ) + logged_kwargs: Final = MappingProxyType( + { + name: "****" if name in _RUST_OCR_SECRET_FIELDS else value + for name, value in request.kwargs.items() + if name != "proxy_server_request" + } + ) + logging_obj: Final = cast( # cast-ok: bridge kwargs carry the prepared logging object + LiteLLMLoggingObj, request.kwargs["litellm_logging_obj"] + ) + logging_obj.update_from_kwargs( + kwargs=dict(logged_kwargs), # mutable-ok: logging API requires an owned dict + model=request.model, + optional_params=dict(logged_optional_params), # mutable-ok: logging API requires an owned dict + litellm_params={ + "litellm_call_id": request.kwargs.get("litellm_call_id"), + "api_base": request.api_base, + }, # mutable-ok: legacy logging requires a concrete params dict + custom_llm_provider=provider, + ) + logging_obj.pre_call( input="OCR document processing", - api_key=resolved_api_key, - additional_args={ + api_key=api_key, + additional_args={ # mutable-ok: pre_call mutates the additional_args dict "complete_input_dict": { - "model": prepared_request.model, - "document": prepared_request.document, - **rust_optional_params, - }, - "api_base": resolved_complete_url, - "headers": resolved_headers, + "model": request.model, + "document": document, + **logged_optional_params, + }, # mutable-ok: callbacks consume a JSON-serializable request dict + "api_base": request.api_base or "", + "headers": request.extra_headers or {}, # mutable-ok: logging callbacks consume a concrete headers dict }, ) - return _PreparedRustOCRCall( - api_key=resolved_api_key, - api_base=rust_api_base, - headers=cast(dict[str, object], resolved_headers), - optional_params=rust_optional_params, + return rust_ocr_bridge.LiteLLMOcrRequest( + model=request.model, + document=document, + api_key=api_key, + api_base=request.api_base, + timeout=request.timeout if request.timeout is not None else request_timeout, + custom_llm_provider=request.custom_llm_provider, + extra_headers=request.extra_headers, + kwargs=optional_params, + input_sources=input_sources, ) +def _map_rust_ocr_error( + error: Exception, + request: rust_ocr_bridge.LiteLLMOcrRequest, + exception_types: tuple[type[BaseException], type[BaseException]] | None, +) -> Exception: + if exception_types is None or not isinstance(error, exception_types[1]): + return error + provider: Final = _rust_ocr_provider(request) + if provider is None: + return error + provider_config: Final = ProviderConfigManager.get_provider_ocr_config( + model=request.model.removeprefix(f"{provider}/"), provider=litellm.LlmProviders(provider) + ) + if provider_config is None: + return error + error_args: Final = cast( # cast-ok: Python exceptions expose positional args as a tuple + tuple[object, ...], error.args + ) + status: Final = error_args[0] if error_args and isinstance(error_args[0], int) else 500 + message: Final = str(error_args[1]) if len(error_args) > 1 else str(error) + error_factory: Final = cast( # cast-ok: provider configs expose heterogeneous exception factories + Callable[..., Exception], provider_config.get_error_class + ) + return error_factory( + error_message=message, status_code=status or 500, headers={} + ) # mutable-ok: provider error factories require a concrete headers dict + + def _run_rust_ocr( - prepared_request: _PreparedOCRRequest, + request: rust_ocr_bridge.LiteLLMOcrRequest, resolve_api_key: Callable[[str], str | None], ) -> OCRResponse | None: if rust_ocr_bridge.load_rust_ocr() is None: return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) - rust_response: Final = rust_ocr_bridge.ocr( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, - ) - if rust_response is None: - return None - return OCRResponse.model_validate(rust_response) + marshalled: Final = _marshal_rust_ocr_request(request, resolve_api_key) + input_sources: Final = marshalled.input_sources + try: + response: Final = rust_ocr_bridge.ocr( + model=marshalled.model, + document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict + api_key=marshalled.api_key, + api_base=marshalled.api_base, + custom_llm_provider=marshalled.custom_llm_provider, + extra_headers=marshalled.extra_headers, + optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict + input_sources=input_sources, + timeout=marshalled.timeout, + ) + except Exception as error: + raise _map_rust_ocr_error(error, request, native_exception_types()) from error + return OCRResponse.model_validate(response) if response is not None else None async def _run_rust_aocr( - prepared_request: _PreparedOCRRequest, + request: rust_ocr_bridge.LiteLLMOcrRequest, resolve_api_key: Callable[[str], str | None], ) -> OCRResponse | None: if rust_ocr_bridge.load_rust_aocr() is None: return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) - rust_response: Final = await rust_ocr_bridge.aocr( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, - ) - if rust_response is None: - return None - return OCRResponse.model_validate(rust_response) + marshalled: Final = _marshal_rust_ocr_request(request, resolve_api_key) + input_sources: Final = marshalled.input_sources + try: + response: Final = await rust_ocr_bridge.aocr( + model=marshalled.model, + document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict + api_key=marshalled.api_key, + api_base=marshalled.api_base, + custom_llm_provider=marshalled.custom_llm_provider, + extra_headers=marshalled.extra_headers, + optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict + input_sources=input_sources, + timeout=marshalled.timeout, + ) + except Exception as error: + raise _map_rust_ocr_error(error, request, native_exception_types()) from error + return OCRResponse.model_validate(response) if response is not None else None @client @@ -415,7 +532,29 @@ async def aocr( "extra_headers": extra_headers, "kwargs": kwargs, } + request: Final = rust_ocr_bridge.LiteLLMOcrRequest( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) try: + if rust_enabled() and _rust_ocr_supported(request): + from litellm.secret_managers.main import get_secret_str + + rust_response: Final = await _run_rust_aocr( + request=request, + resolve_api_key=get_secret_str, + ) + if rust_response is None: + verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") + else: + return rust_response + prepared: Final = _prepare_ocr_request( model=model, document=document, @@ -430,18 +569,6 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = await _run_rust_aocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) - if rust_response is None: - verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response - response = base_llm_http_handler.ocr( model=prepared.model, document=prepared.document, @@ -465,9 +592,11 @@ async def aocr( return response except Exception as e: + error_provider: Final = custom_llm_provider or _rust_ocr_provider(request) + error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model raise litellm.exception_type( - model=model, - custom_llm_provider=custom_llm_provider, + model=error_model, + custom_llm_provider=error_provider, original_exception=e, completion_kwargs=completion_kwargs, extra_kwargs=kwargs, @@ -685,9 +814,31 @@ def ocr( "extra_headers": extra_headers, "kwargs": kwargs, } + request: Final = rust_ocr_bridge.LiteLLMOcrRequest( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) try: _is_async: Final = kwargs.pop("aocr", False) is True completion_kwargs["aocr"] = _is_async + if rust_enabled() and _rust_ocr_supported(request): + from litellm.secret_managers.main import get_secret_str + + rust_response: Final = _run_rust_ocr( + request=request, + resolve_api_key=get_secret_str, + ) + if rust_response is None: + verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") + else: + return rust_response + prepared: Final = _prepare_ocr_request( model=model, document=document, @@ -702,18 +853,6 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = _run_rust_ocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) - if rust_response is None: - verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response - response: Final = base_llm_http_handler.ocr( model=prepared.model, document=prepared.document, @@ -731,9 +870,11 @@ def ocr( return response except Exception as e: + error_provider: Final = custom_llm_provider or _rust_ocr_provider(request) + error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model raise litellm.exception_type( - model=model, - custom_llm_provider=custom_llm_provider, + model=error_model, + custom_llm_provider=error_provider, original_exception=e, completion_kwargs=completion_kwargs, extra_kwargs=kwargs, diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 7076683f294..73d8bab686b 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -17,7 +17,7 @@ from httpx._types import CookieTypes, QueryParamTypes, RequestContent, RequestFi from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.passthrough.utils import CommonUtils @@ -36,6 +36,35 @@ def _as_generator(iterable: Iterator[bytes]) -> Generator[bytes, bytes, None]: yield from iterable +class _SpendCollection: + """Feeds relayed chunks to the provider's stream collector without letting spend tracking break the relay.""" + + def __init__(self, provider_config: BasePassthroughConfig, litellm_logging_obj: LiteLLMLoggingObj) -> None: + self.collector: Final[PassthroughStreamCollector] = provider_config.create_stream_collector( + model=litellm_logging_obj.model, + custom_llm_provider=litellm_logging_obj.model_call_details.get("custom_llm_provider", ""), + endpoint=litellm_logging_obj.model_call_details.get("endpoint", ""), + ) + self.chunk_count = 0 + self._failed = False + + def add(self, chunk: bytes) -> None: + self.chunk_count += 1 + if self._failed: + return + try: + self.collector.add(chunk) + except Exception as e: # noqa: BLE001 # Safe catch-all: spend tracking must never break the relayed stream + self._failed = True + verbose_logger.exception( + "Passthrough spend-tracking collector failed; spend dropped for this stream: %s", e + ) + + @property + def should_flush(self) -> bool: + return self.chunk_count > 0 and not self._failed + + class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): def __init__( self, @@ -50,8 +79,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): self._response: httpx.Response self._iterator: AsyncGenerator[bytes, bytes] self._litellm_logging_obj = litellm_logging_obj - self._provider_config = provider_config - self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks + self._spend = _SpendCollection(provider_config, litellm_logging_obj) self._flush_scheduled = False self._background_tasks: set[asyncio.Task] = set() # mutable-ok: instance set for background task tracking self._hidden_params: dict[str, object] = {} # mutable-ok: router attaches response headers here in place @@ -101,16 +129,13 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): return _init().__await__() def _start_flush(self) -> None: - if self._flush_scheduled or not self._raw_bytes: + if self._flush_scheduled or not self._spend.should_flush: return self._flush_scheduled = True try: task: Final = asyncio.create_task( - self._litellm_logging_obj.async_flush_passthrough_collected_chunks( - raw_bytes=self._raw_bytes, - provider_config=self._provider_config, - ) + self._litellm_logging_obj.async_flush_passthrough_collected_chunks(collector=self._spend.collector) ) self._background_tasks.add(task) @@ -118,8 +143,8 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): task.add_done_callback(self._background_tasks.discard) except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging verbose_logger.exception( - "Failed to schedule passthrough spend-tracking flush; %d buffered chunks dropped: %s", - len(self._raw_bytes), + "Failed to schedule passthrough spend-tracking flush; %d collected chunks dropped: %s", + self._spend.chunk_count, e, ) @@ -134,7 +159,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): await self # pyright: ignore[reportGeneralTypeIssues] # structural type check misses __await__ try: chunk: Final = await anext(self._iterator) - self._raw_bytes.append(chunk) + self._spend.add(chunk) except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic self._start_flush() try: @@ -181,13 +206,12 @@ class PassthroughStreamingResponse(Generator[bytes, bytes, None]): self.headers = response.headers self.status_code = response.status_code self._litellm_logging_obj = litellm_logging_obj - self._provider_config = provider_config self._iterator: Generator[bytes, bytes, None] = _as_generator(response.iter_bytes()) - self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks + self._spend = _SpendCollection(provider_config, litellm_logging_obj) self._flush_scheduled = False def _start_flush(self) -> None: - if self._flush_scheduled or not self._raw_bytes: + if self._flush_scheduled or not self._spend.should_flush: return self._flush_scheduled = True @@ -195,14 +219,12 @@ class PassthroughStreamingResponse(Generator[bytes, bytes, None]): try: executor.submit( - self._litellm_logging_obj.flush_passthrough_collected_chunks, - raw_bytes=self._raw_bytes, - provider_config=self._provider_config, + self._litellm_logging_obj.flush_passthrough_collected_chunks, collector=self._spend.collector ) except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging verbose_logger.exception( - "Failed to schedule passthrough spend-tracking flush; %d buffered chunks dropped: %s", - len(self._raw_bytes), + "Failed to schedule passthrough spend-tracking flush; %d collected chunks dropped: %s", + self._spend.chunk_count, e, ) @@ -212,7 +234,7 @@ class PassthroughStreamingResponse(Generator[bytes, bytes, None]): def __next__(self) -> bytes: try: chunk: Final = next(self._iterator) - self._raw_bytes.append(chunk) + self._spend.add(chunk) except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic self._start_flush() try: diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index b9bfb062ec7..7e6de474b0b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -179,16 +179,7 @@ def _gateway_dcr_challenge_target( mcp_servers: list[str] | None, client_ip: str | None, ) -> str | None: - """The single path-named server this request targets, iff it resolves to a - gateway-managed oauth2 server — the one per-server shape the gateway's own keyless - DCR flow serves end to end, so the 401 challenge may advertise the per-server - protected-resource metadata (whose ``authorization_servers`` names the gateway). - - Multi-server CSV paths, header/path mismatches, unknown names, and every - client-forwarded or delegated mode return ``None``: those cells keep their existing - challenge (or absence of one), and a challenge is never emitted for a name the - public discovery routes would 404, so this reveals exactly the server set the - per-server protected-resource metadata already reveals.""" + """Resolve a single path target whose sign-in metadata advertises the gateway.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -197,7 +188,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] @@ -217,7 +208,7 @@ def _is_gateway_dcr_challenge_scope( the caller is not a cold-start DCR client), on the scopes the gateway's keyless flow serves: the aggregate ``/mcp`` endpoint, an ``x-mcp-servers``-scoped request (the resource the client configured is still ``/mcp``), or a per-server path whose - single target is a gateway-managed oauth2 server. Every other named target keeps + single target advertises gateway-owned sign-in. Every other named target keeps its existing behavior, failing closed to the original admission error.""" if not _is_litellm_auth_admission_error(exc): return False @@ -236,7 +227,7 @@ def _gateway_dcr_challenge( ) -> HTTPException: """The RFC 9728 challenge pointing the client at the protected-resource metadata matching the scope it requested: the per-server document (same URL spelling the - request arrived on) when the single target is a gateway-managed oauth2 server, + request arrived on) when the single target advertises gateway-owned sign-in, else the gateway's aggregate document. Either way the client discovers the gateway as its authorization server and starts the same sign-in flow. @@ -3108,15 +3099,17 @@ class MCPRequestHandler: @staticmethod async def _get_allowed_mcp_servers_for_agent( user_api_key_auth: UserAPIKeyAuth | None = None, - agent_object_permission=None, + agent_object_permission: LiteLLM_ObjectPermissionTable | None = None, ) -> list[str]: """ Get allowed MCP servers for an agent (from the agent's object_permission). - Returns the MCP servers from the agent's object_permission. - If agent has no object_permission, returns [] (no extra restriction). An entitlement the - agent LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here so the - resolver denies. + Returns the agent's direct servers, the servers in its access groups, and the servers reached + through its toolsets, exactly as the key, team, and org levels count theirs. If agent has no + object_permission, returns [] (no extra restriction). An entitlement the agent LINKS but that + cannot be read, or a declared toolset that resolves to no grants, raises + ``UnloadableEntitlementError`` out of here so the resolver denies instead of reading the + agent as unrestricted. Args: user_api_key_auth: User auth with agent_id @@ -3126,31 +3119,30 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return [] - obj_perm = agent_object_permission - if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + obj_perm: Final = ( + agent_object_permission + if agent_object_permission is not None + else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + ) if obj_perm is None: return [] try: - direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or [] - if isinstance(direct_mcp_servers, str): - direct_mcp_servers = [] - mcp_access_groups = getattr(obj_perm, "mcp_access_groups", None) or [] - if isinstance(mcp_access_groups, str): - mcp_access_groups = [] - - # Permission entries may be server_ids OR names/aliases — expand to ids. from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(list(direct_mcp_servers)) - - access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups) - all_servers: Final = expanded_direct_servers + access_group_servers - return list(set(all_servers)) + expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list( + obj_perm.mcp_servers or [] + ) + access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups( + obj_perm.mcp_access_groups or [] + ) + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(obj_perm) + return list({*expanded_direct_servers, *access_group_servers, *toolset_grants}) except Exception as e: + if isinstance(e, UnloadableEntitlementError): + raise verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e) return [] @@ -3158,13 +3150,15 @@ class MCPRequestHandler: async def _get_agent_tool_permissions_for_server( server_id: str, user_api_key_auth: UserAPIKeyAuth | None = None, - agent_object_permission=None, + agent_object_permission: LiteLLM_ObjectPermissionTable | None = None, ) -> list[str] | None: """ - Get allowed tool names for a server from the agent's object_permission. - Returns None if agent has no tool restrictions for this server. An entitlement the agent - LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here, which the - tool resolver turns into deny-all for the server rather than an unrestricted tool list. + Get allowed tool names for a server from the agent's object_permission: the union of its + direct tool permissions and the tools its toolsets grant on that server, mirroring the key and + team levels. Returns None if agent has no tool restrictions for this server. An entitlement the + agent LINKS but that cannot be read, or a declared toolset that resolves to no grants, raises + ``UnloadableEntitlementError`` out of here, which the tool resolver turns into deny-all for the + server rather than an unrestricted tool list. Args: server_id: Server ID to check permissions for @@ -3175,24 +3169,30 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return None - obj_perm = agent_object_permission - if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + obj_perm: Final = ( + agent_object_permission + if agent_object_permission is not None + else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + ) if obj_perm is None: return None try: - mcp_tool_permissions: Final = getattr(obj_perm, "mcp_tool_permissions", None) - if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict): - return None - # Dict keys may be server_ids OR names/aliases; normalize before lookup. from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - tools: Final = global_mcp_server_manager.expand_tool_permissions(mcp_tool_permissions).get(server_id) - return list(tools) if tools else None + direct_tools: Final = ( + global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id) + if obj_perm.mcp_tool_permissions + else None + ) + toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(obj_perm, server_id) + agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools) + return list(agent_tools) if agent_tools else None except Exception as e: + if isinstance(e, UnloadableEntitlementError): + raise verbose_logger.warning("Failed to get agent tool permissions for server: %s", e) return None diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 41d0b78b555..789b2ffaef4 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -4,12 +4,20 @@ import hashlib import json from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast + +from fastapi import HTTPException +from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + RefreshTokenPresented, + credential_binding_matches, + enforce_oauth_identity_binding, +) from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -23,8 +31,11 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + SecretMapDecodeError, _get_salt_key, + decode_secret_map, decrypt_value_helper, + encrypt_secret_map, encrypt_value_helper, ) from litellm.proxy.utils import PrismaClient @@ -114,6 +125,7 @@ class _OAuthCredentialAccessToken(TypedDict): class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False): + identity_binding_proof: ReadOnly[str] type: str refresh_token: str expires_at: str @@ -122,6 +134,9 @@ class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False): server_id: str +OAuthGrantState = Literal["valid", "refreshable", "absent"] + + class _OAuthTokenRefreshResponse(TypedDict, total=False): access_token: str refresh_token: str @@ -360,7 +375,7 @@ def _prepare_mcp_server_data( # exclude_unset filter is respected. Reading back from ``data`` would # reintroduce defaults (e.g. ``env={}``) for fields the caller never set. if data_dict.get("static_headers") is not None: - data_dict["static_headers"] = safe_dumps(data_dict["static_headers"]) + data_dict["static_headers"] = encrypt_secret_map(data_dict["static_headers"]) # env_vars is read from ``data_dict`` (not ``data``) like every other JSON # column so the exclude_unset filter is respected: a partial update that @@ -376,7 +391,7 @@ def _prepare_mcp_server_data( data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"]) if data_dict.get("env") is not None: - data_dict["env"] = safe_dumps(data_dict["env"]) + data_dict["env"] = encrypt_secret_map(data_dict["env"]) if "tool_name_to_display_name" in data_dict: data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {}) @@ -589,6 +604,19 @@ def decrypt_credentials( return credentials +def _readable_mcp_servers( + rows: Iterable["prisma_db_models.LiteLLM_MCPServerTable"], +) -> Iterable[LiteLLM_MCPServerTable]: + for row in rows: + try: + table = LiteLLM_MCPServerTable.model_validate(row.model_dump()) + except SecretMapDecodeError: + verbose_proxy_logger.warning("Skipping MCP server %s: cannot decrypt secret map", row.server_id) + continue + decrypt_global_env_var_values(table.env_vars) + yield table + + async def get_all_mcp_servers( prisma_client: PrismaClient, approval_status: str | None = None, @@ -609,10 +637,7 @@ async def get_all_mcp_servers( ) mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where) - tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers] - for table in tables: - decrypt_global_env_var_values(table.env_vars) - return tables + return list(_readable_mcp_servers(mcp_servers)) async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None: @@ -638,13 +663,7 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str] "server_id": {"in": server_ids}, } ) - final_mcp_servers: Final[list[LiteLLM_MCPServerTable]] = [] - for _mcp_server in _mcp_servers: - table = LiteLLM_MCPServerTable.model_validate(_mcp_server.model_dump()) - decrypt_global_env_var_values(table.env_vars) - final_mcp_servers.append(table) - - return final_mcp_servers + return list(_readable_mcp_servers(_mcp_servers)) async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> list[str]: @@ -852,12 +871,10 @@ async def create_mcp_server( data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - new_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.create( - data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable - ) + new_mcp_server: Final = await MCPServerRepository(prisma_client).table.create(data=data_dict) _decrypt_env_vars_on_returned_row(new_mcp_server) - return new_mcp_server + return LiteLLM_MCPServerTable.model_validate(new_mcp_server.model_dump()) async def create_draft_mcp_server( @@ -1066,13 +1083,13 @@ async def update_mcp_server( data_dict["credentials"] = Json(None) - updated_mcp_server: Final[LiteLLM_MCPServerTable | None] = await MCPServerRepository(prisma_client).table.update( + updated_mcp_server: Final = await MCPServerRepository(prisma_client).table.update( where={"server_id": data.server_id}, - data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable + data=data_dict, ) _decrypt_env_vars_on_returned_row(updated_mcp_server) - return updated_mcp_server + return LiteLLM_MCPServerTable.model_validate(updated_mcp_server.model_dump()) if updated_mcp_server else None async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None: @@ -1144,6 +1161,13 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, if rotated_env_vars is not None: update_data["env_vars"] = safe_dumps(rotated_env_vars) + for field in ("static_headers", "env"): + try: + if secret_map := decode_secret_map(getattr(mcp_server, field, None), key=field): + update_data[field] = encrypt_secret_map(secret_map, new_encryption_key=new_master_key) + except SecretMapDecodeError: + verbose_proxy_logger.warning("Cannot rotate MCP %s for server %s", field, mcp_server.server_id) + if not update_data: continue @@ -1378,6 +1402,7 @@ async def store_user_oauth_credential( expires_in: int | None = None, scopes: list[str] | None = None, skip_byok_guard: bool = False, + identity_binding_proof: str | None = None, ) -> None: """Persist an OAuth2 access token for a user+server pair. @@ -1394,6 +1419,7 @@ async def store_user_oauth_credential( "type": "oauth2", "access_token": access_token, "connected_at": datetime.now(timezone.utc).isoformat(), + **({"identity_binding_proof": identity_binding_proof} if identity_binding_proof else {}), } if refresh_token: payload["refresh_token"] = refresh_token @@ -1453,6 +1479,15 @@ def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: in return False +def oauth_grant_state(cred: OAuthCredentialPayload | None) -> OAuthGrantState: + """Classify local grant readiness without attempting a refresh or checking upstream revocation.""" + if not cred or not cred.get("access_token"): + return "absent" + if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS): + return "valid" + return "refreshable" if cred.get("refresh_token") else "absent" + + async def get_user_oauth_credential( prisma_client: PrismaClient, user_id: str, @@ -1604,6 +1639,11 @@ async def refresh_user_oauth_token( warning and returns ``None`` — the caller is responsible for clearing the stale credential and triggering re-authentication. """ + binding: Final = server.oauth_identity_binding + if binding is not None and binding.mode == "enforce": + if not await credential_binding_matches(binding, user_id, server.server_id, cred): + return None + refresh_token: Final[str | None] = cred.get("refresh_token") token_url: Final[str | None] = getattr(server, "effective_token_url", None) or getattr(server, "token_url", None) server_id: Final[str] = getattr(server, "server_id", "") @@ -1653,6 +1693,19 @@ async def refresh_user_oauth_token( ) return None + try: + binding_proof: Final = await enforce_oauth_identity_binding( + server=server, + token_response=body, + litellm_user_id=user_id, + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented(refresh_token), + ) + except HTTPException as exc: + if exc.status_code != 403: + raise + return None + access_token: Final[str | None] = body.get("access_token") if not access_token: verbose_proxy_logger.warning( @@ -1685,6 +1738,7 @@ async def refresh_user_oauth_token( refresh_token=new_refresh_token, expires_in=expires_in, scopes=scopes, + identity_binding_proof=binding_proof, skip_byok_guard=True, # Row is already OAuth2; skip the extra find_unique check ) @@ -1715,12 +1769,15 @@ async def resolve_valid_user_oauth_token( dict it already holds. ``prisma_client`` is fetched lazily and only when a refresh actually happens, so the valid-token path never requires a DB handle. """ - if not cred or not cred.get("access_token"): + grant: Final = oauth_grant_state(cred) + if cred is None or grant == "absent": return None - if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS): + binding: Final = server.oauth_identity_binding + if binding is not None and binding.mode == "enforce": + if not await credential_binding_matches(binding, user_id, server.server_id, cred): + return None + if grant == "valid": return cred - if not cred.get("refresh_token"): - return None if prisma_client is None: from litellm.proxy.utils import get_prisma_client_or_throw @@ -1759,7 +1816,17 @@ async def resolve_user_oauth_access_token( mcp_per_user_token_cache, ) - if prefetched_creds is None: + binding: Final = server.oauth_identity_binding + enforce_binding: Final = binding is not None and binding.mode == "enforce" + if prefetched_creds is None and enforce_binding and binding is not None: + bound_token: Final = await mcp_per_user_token_cache.get_token(user_id, server_id) + if bound_token is not None: + if await credential_binding_matches( + binding, user_id, server_id, {"identity_binding_proof": bound_token.identity_binding_proof} + ): + return bound_token.access_token + await mcp_per_user_token_cache.delete(user_id, server_id) + if prefetched_creds is None and not enforce_binding: cached_token: Final = await mcp_per_user_token_cache.get(user_id, server_id) if cached_token is not None: return cached_token @@ -1793,7 +1860,9 @@ async def resolve_user_oauth_access_token( access_token: Final[str] = cred["access_token"] if prefetched_creds is None: ttl: Final = _compute_per_user_token_ttl(server, _remaining_token_seconds(cred.get("expires_at"))) - await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl) + await mcp_per_user_token_cache.set( + user_id, server_id, access_token, ttl, identity_binding_proof=cred.get("identity_binding_proof") + ) return access_token except Exception as e: verbose_proxy_logger.warning( @@ -1894,9 +1963,7 @@ async def get_mcp_submissions( order={"submitted_at": "desc"}, take=500, # safety cap; paginate if needed in a future iteration ) - items: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in rows] - for item in items: - decrypt_global_env_var_values(item.env_vars) + items: Final = list(_readable_mcp_servers(rows)) pending: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review) active: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 3da5950ce7f..42fbe82531c 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -43,9 +43,11 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + VendorCredentialState, aggregate_authorize, aggregate_token, complete_connect_flow, + describe_connect_flow, introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, @@ -55,6 +57,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( relative_request_url, revoke_refresh_token, ) +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + RefreshOwnershipProven, + RefreshTokenPresented, + enforce_oauth_identity_binding, +) from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, build_upstream_oauth2_token_request, @@ -137,6 +144,7 @@ def encode_state_with_base_url( dcr_client_id: str | None = None, dcr_client_secret: str | None = None, dcr_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None, + oauth_nonce: str | None = None, ) -> str: """ Encode the base_url, original state, and PKCE parameters using encryption. @@ -147,9 +155,8 @@ def encode_state_with_base_url( code_challenge: PKCE code challenge from client code_challenge_method: PKCE code challenge method from client client_redirect_uri: Original redirect_uri from client - litellm_user_id: The SSO-authenticated litellm user captured at the bridge authorize - (interactive dcr_bridge oauth_delegate only); the callback seals it into the gateway - authorization code so the token mint can bind the envelope to this user + litellm_user_id: The authenticated user captured for bridge or identity-bound per-user OAuth; + the callback seals this credential owner into the authorization code mcp_server_id: The server the flow targets, sealed alongside litellm_user_id (bridge) or dcr_client_id (ephemeral mint) so the gateway code cannot be replayed against another server @@ -167,6 +174,7 @@ def encode_state_with_base_url( An encrypted string that encodes all values """ state_data: Final = { + "oauth_nonce": oauth_nonce, "base_url": base_url, "original_state": original_state, "code_challenge": code_challenge, @@ -208,10 +216,10 @@ _BRIDGE_AUTH_CODE_PREFIX: Final = "llm_bcode_" class _BridgeAuthorizationCode(BaseModel): - """The identity and upstream code the gateway seals into the authorization code it hands a DCR - client for an interactive dcr_bridge oauth_delegate sign-in, recovered at the token endpoint.""" + """Authenticated caller and upstream code sealed for bridge or identity-bound per-user OAuth.""" model_config = ConfigDict(frozen=True) + oauth_nonce: str | None = None upstream_code: str = Field(min_length=1) litellm_user_id: str = Field(min_length=1) mcp_server_id: str = Field(min_length=1) @@ -223,7 +231,12 @@ def is_bridge_authorization_code(code: str) -> bool: return code.startswith(_BRIDGE_AUTH_CODE_PREFIX) -def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp_server_id: str) -> str: +def seal_bridge_authorization_code( + upstream_code: str, + litellm_user_id: str, + mcp_server_id: str, + oauth_nonce: str | None = None, +) -> str: """Seal the upstream authorization code and the SSO-captured litellm user into a gateway authorization code. The DCR client only echoes this opaque value back at the token endpoint; the gateway decrypts it there to recover the user (to bind the envelope) and the upstream code (to @@ -232,7 +245,12 @@ def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp authenticated symmetric helper (the same family the OAuth state uses), so the client can neither read nor forge it.""" payload: Final = json.dumps( - {"upstream_code": upstream_code, "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id}, + { + "upstream_code": upstream_code, + "litellm_user_id": litellm_user_id, + "mcp_server_id": mcp_server_id, + "oauth_nonce": oauth_nonce, + }, sort_keys=True, ) return _BRIDGE_AUTH_CODE_PREFIX + encrypt_value_helper(payload) @@ -545,6 +563,7 @@ async def _store_per_user_token_server_side( server: MCPServer, user_id: str, token_response: dict[str, Any], + identity_binding_proof: str | None = None, ) -> None: """Persist the OAuth token server-side and warm the Redis cache. @@ -586,6 +605,7 @@ async def _store_per_user_token_server_side( refresh_token=refresh_token, expires_in=expires_in, scopes=scopes, + identity_binding_proof=identity_binding_proof, ) verbose_logger.info( "_store_per_user_token_server_side: stored token for user=%s server=%s", @@ -614,6 +634,7 @@ async def _store_per_user_token_server_side( server_id=server.server_id, access_token=access_token, ttl=ttl, + identity_binding_proof=identity_binding_proof, ) @@ -798,21 +819,7 @@ def _bridge_access_denied_redirect(redirect_uri: str, state: str, mcp_server: MC return RedirectResponse(_append_query_params(redirect_uri, params), status_code=302) -async def _bridge_authorize_access_denial( - litellm_user_id: str, - mcp_server: MCPServer, - redirect_uri: str, - state: str, -) -> RedirectResponse | None: - """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed. - - Admits the user exactly as MCP egress will (the same ``reload_admitted_user`` constructor and the - same ``get_allowed_mcp_servers`` resolver), so an envelope is minted only when the resulting - session can actually list and call the server's tools. Without this gate the flow completes, the - client shows connected, and every tool request fail-closes to an empty list with nothing telling - the operator why. An availability fault (5xx, e.g. a DB outage's 503) propagates; an unknown or - deactivated user denies like a missing grant, fail closed. - """ +async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool: from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) @@ -821,13 +828,22 @@ async def _bridge_authorize_access_denial( ) try: - admitted: Final = await MCPRequestHandler.reload_admitted_user(litellm_user_id) + admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id) except HTTPException as exc: if exc.status_code >= 500: raise - return _bridge_access_denied_redirect(redirect_uri, state, mcp_server) - allowed_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(admitted) - if mcp_server.server_id in allowed_server_ids: + return False + return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted) + + +async def _bridge_authorize_access_denial( + litellm_user_id: str, + mcp_server: MCPServer, + redirect_uri: str, + state: str, +) -> RedirectResponse | None: + """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed.""" + if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id): return None return _bridge_access_denied_redirect(redirect_uri, state, mcp_server) @@ -857,6 +873,11 @@ async def authorize_with_server( ), ) + binding: Final = resolved_server.oauth_identity_binding + enforce_binding: Final = binding is not None and binding.mode == "enforce" + if enforce_binding: + _require_s256_pkce(code_challenge, code_challenge_method) + if resolved_server.is_dcr_bridge: # Enforce S256 PKCE on both bridge arms. The relay arm forwards the validated, # now-non-optional pair to the upstream authorize; the short-circuit arm keeps @@ -887,19 +908,16 @@ async def authorize_with_server( base_url: Final = urlunparse(parsed._replace(query="")) request_base_url: Final = get_request_base_url(request) - # Interactive dcr_bridge oauth_delegate sign-in: this arm runs the gateway /callback and /token in - # the loop, so the gateway can capture the litellm user here (from the browser's UI session) and - # carry it to the back-channel token mint. Seal the SSO user and the target server into the state; - # the callback reads them back to mint the gateway authorization code. A DCR client cannot present a - # litellm key, so the browser session is the only identity source; without one there is nothing to - # bind, so send the user through login first. Every other oauth2 server keeps the identity-less state. + # Seal the authenticated caller into state so the token exchange cannot select another credential owner. litellm_user_id: str | None = None - if resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate: + if enforce_binding or (resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate): from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import _user_id_from_session_cookie, ) - litellm_user_id = _user_id_from_session_cookie(request) + litellm_user_id = ( + await _extract_user_id_from_request(request) if enforce_binding else None + ) or _user_id_from_session_cookie(request) if litellm_user_id is None: return _redirect_to_litellm_login(request) denial: Final = await _bridge_authorize_access_denial( @@ -911,9 +929,11 @@ async def authorize_with_server( if denial is not None: return denial + oauth_nonce: Final = secrets.token_urlsafe(32) if enforce_binding else None encoded_state: Final = encode_state_with_base_url( base_url=base_url, original_state=state, + oauth_nonce=oauth_nonce, code_challenge=code_challenge, code_challenge_method=code_challenge_method, client_redirect_uri=redirect_uri, @@ -933,11 +953,16 @@ async def authorize_with_server( "state": relay_state, "response_type": response_type or "code", } + if oauth_nonce: + params["nonce"] = oauth_nonce if scope: params["scope"] = scope elif resolved_server.scopes: params["scope"] = " ".join(resolved_server.scopes) + if enforce_binding and "openid" not in params.get("scope", "").split(): + params["scope"] = f"openid {params.get('scope', '')}".strip() + if code_challenge: params["code_challenge"] = code_challenge if code_challenge_method: @@ -1018,6 +1043,12 @@ async def exchange_token_with_server( except TokenEndpointAuthConfigError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc + request_user_id: Final = ( + await _extract_user_id_from_request(request) + if resolved_server.needs_user_oauth_token or resolved_server.oauth_identity_binding is not None + else None + ) + bridge_identity: _BridgeAuthorizationCode | None = None bridge_mint_ready: _BridgeMintReady | None = None bridge_upstream_refresh: SecretStr | None = None @@ -1054,7 +1085,13 @@ async def exchange_token_with_server( refresh_request_scope = scope or bridge_upstream_scope if refresh_request_scope: token_data["scope"] = refresh_request_scope + refresh_ownership = ( # rebind-ok: grant-specific branches assign one ownership value + RefreshOwnershipProven() + if bridge_upstream_refresh is not None + else RefreshTokenPresented(upstream_refresh_token) + ) else: + refresh_ownership = None # rebind-ok: grant-specific branches assign one ownership value if not code: raise HTTPException( status_code=400, @@ -1073,6 +1110,14 @@ async def exchange_token_with_server( detail="Authorization code was issued for a different MCP server", ) code = bridge_identity.upstream_code + binding: Final = resolved_server.oauth_identity_binding + if binding is not None and binding.mode == "enforce": + if bridge_identity is None or not bridge_identity.oauth_nonce: + raise HTTPException(status_code=403, detail={"error": "oauth_identity_binding_failed"}) + if request_user_id is not None and request_user_id != bridge_identity.litellm_user_id: + raise HTTPException(status_code=403, detail={"error": "oauth_principal_mismatch"}) + if not code_verifier: + raise HTTPException(status_code=403, detail={"error": "oauth_identity_binding_failed"}) bridge_token_relay: Final = _dcr_bridge_relays_client_registration(resolved_server) if bridge_token_relay and not redirect_uri: raise HTTPException( @@ -1100,6 +1145,16 @@ async def exchange_token_with_server( return _bridge_mint_error_response(prepared) bridge_mint_ready = prepared + refresh_binding: Final = resolved_server.oauth_identity_binding + if grant_type == "refresh_token" and refresh_binding is not None and refresh_binding.mode == "enforce": + await enforce_oauth_identity_binding( + server=resolved_server, + token_response={}, + litellm_user_id=request_user_id, + grant_type=grant_type, + refresh_ownership=refresh_ownership, + ) + async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) try: response: Final = await async_client.post( @@ -1140,17 +1195,34 @@ async def exchange_token_with_server( server_id=resolved_server.server_id, ) + # Bind the exchanged token to the LiteLLM caller BEFORE it is returned, stored, or cached, so a + # token minted for a different upstream principal never becomes usable under the caller's user_id. + resolved_user_id: Final = bridge_identity.litellm_user_id if bridge_identity else request_user_id + binding_proof: Final = ( + await enforce_oauth_identity_binding( + server=resolved_server, + token_response=token_response, + litellm_user_id=resolved_user_id, + grant_type=grant_type, + refresh_ownership=refresh_ownership, + expected_nonce=bridge_identity.oauth_nonce if bridge_identity else None, + ) + if isinstance(token_response, dict) + else None + ) + # Store server-side when the server is configured for per-user OAuth and # the calling client has provided a valid LiteLLM identity. # Errors are non-fatal: the token is still returned to the client. if resolved_server.needs_user_oauth_token: - user_id: Final = await _extract_user_id_from_request(request) + user_id: Final = resolved_user_id if user_id: try: await _store_per_user_token_server_side( server=resolved_server, user_id=user_id, token_response=token_response, + identity_binding_proof=binding_proof, ) except Exception as exc: verbose_logger.warning( @@ -1481,11 +1553,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", ) @@ -1902,6 +1982,38 @@ async def token_endpoint( ) +async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState: + """Whether the gateway itself can see a live vendor credential for this user and server. + + The one reading of "authorized" the connect page displays and the finish step enforces, so + the button a user sees and the grant they get cannot disagree. A read fault is neither, and + fails the scoped grant closed.""" + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # circular import at module load + get_user_oauth_credential, + oauth_grant_state, + ) + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # circular import at module load + + if prisma_client is None: + return "unavailable" + try: + credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id) + except Exception: # noqa: BLE001 # a credential-read fault must fail the scoped grant closed + return "unavailable" + return "absent" if oauth_grant_state(credential) == "absent" else "present" + + +@router.get("/authorize/flow") +async def authorize_flow(request: Request, flow: str) -> Response: + return await describe_connect_flow( + request=request, + flow_handle=flow, + session_user_id=_session_cookie_user_id(request), + lookup_vendor_credential=_vendor_credential_state, + lookup_server_reachability=_user_can_reach_mcp_server, + ) + + @router.post("/authorize/complete") async def authorize_complete( request: Request, @@ -1926,6 +2038,8 @@ async def authorize_complete( delivery=delivery, team_id=team_id, decision=decision, + lookup_vendor_credential=_vendor_credential_state, + lookup_server_reachability=_user_can_reach_mcp_server, ) @@ -2095,7 +2209,10 @@ async def callback( forwarded_code = code if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id: forwarded_code = seal_bridge_authorization_code( - upstream_code=code, litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id + upstream_code=code, + litellm_user_id=litellm_user_id, + mcp_server_id=mcp_server_id, + oauth_nonce=state_data.get("oauth_nonce"), ) elif isinstance(dcr_client_id, str) and dcr_client_id and isinstance(mcp_server_id, str) and mcp_server_id: forwarded_code = seal_passthrough_authorization_code( @@ -2271,8 +2388,7 @@ async def _build_oauth_protected_resource_response( it. Only the legacy ``is_oauth_passthrough`` opt-in rewrites ``resource`` to the gateway's own URL so clients present the bearer token back to the gateway. - An explicitly named gateway-managed oauth2 server (interactive with - gateway-vaulted per-user tokens, or M2M) advertises the gateway's own + An explicitly named server with gateway-owned sign-in advertises the gateway's own authorization server (``{base}/mcp``): a keyless DCR client that configured the per-server URL completes the same sign-in flow the aggregate ``/mcp`` endpoint supports and is admitted with a gateway session bearer. The per-server relay @@ -2362,18 +2478,16 @@ async def _build_oauth_protected_resource_response( if obo_response is not None: return obo_response - # An OBO server with no configured issuer falls through to the gateway default so discovery still - # returns metadata; every other non-oauth2 named server 404s to avoid enumeration. - 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, "scopes_supported": (mcp_server.scopes if mcp_server.scopes else []), } + 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") + return { "authorization_servers": [ (f"{request_base_url}/{mcp_server_name}" if mcp_server_name else f"{request_base_url}") diff --git a/litellm/proxy/_experimental/mcp_server/faults/__init__.py b/litellm/proxy/_experimental/mcp_server/faults/__init__.py index 1b9ee77d795..de7ee5c866a 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/__init__.py +++ b/litellm/proxy/_experimental/mcp_server/faults/__init__.py @@ -22,6 +22,7 @@ from litellm.proxy._experimental.mcp_server.faults.types import ( GatewayRejected, UpstreamOAuthFault, UpstreamProtocolFault, + UpstreamRegistrationRefused, UpstreamReportedFault, ) @@ -31,6 +32,7 @@ __all__ = [ "GatewayRejected", "UpstreamOAuthFault", "UpstreamProtocolFault", + "UpstreamRegistrationRefused", "UpstreamReportedFault", "classify_upstream_dcr_rejection", "classify_upstream_token_rejection", diff --git a/litellm/proxy/_experimental/mcp_server/faults/classify.py b/litellm/proxy/_experimental/mcp_server/faults/classify.py index 2162d078c09..bb41436b495 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/classify.py +++ b/litellm/proxy/_experimental/mcp_server/faults/classify.py @@ -21,6 +21,7 @@ from litellm.proxy._experimental.mcp_server.faults.types import ( GatewayRejected, UpstreamOAuthFault, UpstreamProtocolFault, + UpstreamRegistrationRefused, UpstreamReportedFault, ) @@ -122,11 +123,13 @@ def classify_upstream_dcr_rejection(response: httpx.Response, log_context: str) """Classify a dynamic-client-registration rejection. RFC 7591 §3.2.2 errors carry ``error`` / ``error_description`` and go through the same blame assignment as token errors (registration sends no client credentials, so credential codes stay caller-actionable); anything - without a usable ``error`` field is an upstream protocol fault.""" + without a usable ``error`` field is a registration refusal for 401/403 and a protocol fault otherwise.""" parsed: Final = _safe_json(response) fields: Final = parsed if isinstance(parsed, dict) else {} code: Final = _bounded_field(fields.get("error")) if code is None: + if response.status_code == 401 or response.status_code == 403: + return UpstreamRegistrationRefused(status_code=response.status_code) _log_out_of_contract("registration", response, log_context) return UpstreamProtocolFault(note=f"upstream registration failed with HTTP {response.status_code}") return _classify_oauth_error_code( diff --git a/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py b/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py index 64d14140a5b..3ecf5310482 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py +++ b/litellm/proxy/_experimental/mcp_server/faults/render_oauth.py @@ -11,7 +11,7 @@ from typing import Final from fastapi.responses import JSONResponse from typing_extensions import assert_never -from litellm.proxy._experimental.mcp_server.faults.types import UpstreamOAuthFault +from litellm.proxy._experimental.mcp_server.faults.types import CallerRejected, UpstreamOAuthFault from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS @@ -35,6 +35,24 @@ def _upstream_reported_status_and_description(code: str) -> tuple[int, str]: return 502, "the upstream authorization server reported an internal error" +def _registration_refused_description(status_code: int) -> str: + return ( + f"the upstream authorization server refused dynamic client registration (HTTP {status_code}). " + "This provider may require a pre-registered OAuth client. Configure client_id and, if required " + "by the provider, client_secret for this MCP server to skip dynamic registration" + ) + + +def _render_caller_rejected(fault: CallerRejected) -> JSONResponse: + content: Final = { + "error": fault.code, + **({"error_description": fault.description} if fault.description else {}), + **({"error_uri": fault.error_uri} if fault.error_uri else {}), + } + status_code: Final = 401 if fault.code == "invalid_client" else 400 + return JSONResponse(status_code=status_code, content=content, headers=TOKEN_NO_CACHE_HEADERS) + + def render_token_fault(fault: UpstreamOAuthFault) -> JSONResponse: """RFC 6749 §5.2 response for a token-endpoint fault. Caller-actionable rejections relay the upstream's code on the status that code implies (401 for invalid_client per §5.2, else 400); @@ -42,13 +60,7 @@ def render_token_fault(fault: UpstreamOAuthFault) -> JSONResponse: blamed for, or shown the internals of, a failure only the operator can fix.""" match fault.tag: case "caller_rejected": - content: Final = { - "error": fault.code, - **({"error_description": fault.description} if fault.description else {}), - **({"error_uri": fault.error_uri} if fault.error_uri else {}), - } - status_code = 401 if fault.code == "invalid_client" else 400 - return JSONResponse(status_code=status_code, content=content, headers=TOKEN_NO_CACHE_HEADERS) + return _render_caller_rejected(fault) case "gateway_rejected": return JSONResponse( status_code=502, @@ -65,6 +77,13 @@ def render_token_fault(fault: UpstreamOAuthFault) -> JSONResponse: content={"error": fault.code, "error_description": description}, headers=TOKEN_NO_CACHE_HEADERS, ) + case "upstream_registration_refused": + return _render_caller_rejected( + CallerRejected( + code="unauthorized_client", + description=_registration_refused_description(fault.status_code), + ) + ) case "upstream_protocol_fault": return JSONResponse( status_code=502, @@ -78,7 +97,7 @@ def render_token_fault(fault: UpstreamOAuthFault) -> JSONResponse: def dcr_fault_detail(fault: UpstreamOAuthFault) -> tuple[int, str]: """Status and detail string for a registration fault, raised as HTTPException by the caller. RFC 7591 §3.2.2 defines registration errors as 400, so a contract-conformant rejection is 400 - regardless of the status the upstream chose; everything else is a 502 upstream fault.""" + regardless of the upstream status; a bare 401/403 is a registration refusal rendered as 403.""" match fault.tag: case "caller_rejected": detail: Final = f"{fault.code}: {fault.description}" if fault.description else fault.code @@ -87,6 +106,8 @@ def dcr_fault_detail(fault: UpstreamOAuthFault) -> tuple[int, str]: return 502, _gateway_rejected_description(fault.code) case "upstream_reported_fault": return _upstream_reported_status_and_description(fault.code) + case "upstream_registration_refused": + return 403, _registration_refused_description(fault.status_code) case "upstream_protocol_fault": return 502, fault.note case _: diff --git a/litellm/proxy/_experimental/mcp_server/faults/types.py b/litellm/proxy/_experimental/mcp_server/faults/types.py index 4b9505ad801..d081d9d735e 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/types.py +++ b/litellm/proxy/_experimental/mcp_server/faults/types.py @@ -77,4 +77,12 @@ class UpstreamProtocolFault(BaseModel): note: str -UpstreamOAuthFault: TypeAlias = CallerRejected | GatewayRejected | UpstreamReportedFault | UpstreamProtocolFault +class UpstreamRegistrationRefused(BaseModel): + model_config = ConfigDict(frozen=True) + tag: Literal["upstream_registration_refused"] = "upstream_registration_refused" + status_code: Literal[401, 403] + + +UpstreamOAuthFault: TypeAlias = ( + CallerRejected | GatewayRejected | UpstreamReportedFault | UpstreamProtocolFault | UpstreamRegistrationRefused +) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index c7b0045dde5..f3fdd54b39d 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -122,13 +122,13 @@ _USED_CODE_CACHE_PREFIX: Final = "mcp_gateway_dcr_code_used:" _USED_FLOW_CACHE_PREFIX: Final = "mcp_gateway_dcr_flow_used:" _USED_REFRESH_CACHE_PREFIX: Final = "mcp_gateway_dcr_refresh_used:" -MAX_REDIRECT_URIS: Final = 3 +MAX_REDIRECT_URIS: Final = 4 MAX_REDIRECT_URI_LENGTH: Final = 256 MAX_CLIENT_ID_LENGTH: Final = 2048 """Registration bounds. They exist to bound the sealed client_id, which rides inside -every session-token claim set: 3 URIs of 256 bytes seal to roughly 1.2KB, comfortably -under this cap and under the session token's own 4KB ceiling. Claude Desktop and MCP -Inspector register one or two redirect URIs.""" +every session-token claim set. Four 256-character ASCII URIs seal to roughly 1.5KB; +the encoded client_id is checked against its own cap before registration succeeds. +VS Code registers four callbacks for its web and desktop environments.""" MAX_STATE_LENGTH: Final = 1024 """Bound on the client ``state`` sealed into the flow cookie and echoed on the auth-code @@ -152,10 +152,8 @@ _AUTH_CODE_DEBUG_KEY: Final = "gateway_authorization_code" ReloadUserFailure = Literal["unresolvable", "unavailable", "faulted", "no_active_key"] ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]] -"""Injected live-user revalidation (the token endpoint's mirror of admission): -``None`` means the user is active; ``unavailable`` is a retryable DB outage; ``faulted`` is -a DB fault retrying will not clear (still 503, worded so nobody just waits); anything else -fails the grant closed.""" +VendorCredentialState = Literal["present", "absent", "unavailable"] +"""The per-user vendor credential read has three outcomes: present, absent, or unavailable.""" _DB_UNAVAILABLE_DESCRIPTION: Final = "the gateway database is unavailable; retry" _DB_FAULTED_DESCRIPTION: Final = ( @@ -195,6 +193,16 @@ class ConsentTeam(BaseModel): team_alias: str | None = None +class LookupVendorCredential(Protocol): + """Injected read of a user's vendor credential for one server.""" + + def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ... + + +class LookupServerReachability(Protocol): + def __call__(self, user_id: str, server_id: str, /) -> Awaitable[bool]: ... + + class LookupConsentTeams(Protocol): """Injected lookup of the teams a signed-in user may bind a proxy-API credential to.""" @@ -205,6 +213,14 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr return "unresolvable" +async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState: + return "unavailable" + + +async def _unreachable_server(user_id: str, server_id: str) -> bool: + return False + + class GatewayDcrClient(BaseModel): """The registration record sealed into a gateway DCR ``client_id``. @@ -395,7 +411,7 @@ def relative_request_url(request: Request) -> str: def resolve_scoped_resource_server(request: Request, resource: str | None) -> MCPServer | None: - """Resolve an RFC 8707 ``resource`` value to the single gateway-managed oauth2 server it + """Resolve an RFC 8707 ``resource`` value to the single gateway-owned server it names, or ``None`` for every other shape: absent, the aggregate resource, a foreign host, an unparseable value, a multi-server path, an unknown name, or any server mode the keyless gateway flow does not serve (whose protected-resource metadata never directs a @@ -427,7 +443,7 @@ def resolve_scoped_resource_server(request: Request, resource: str | None) -> MC if len(names) != 1: return None server: Final = global_mcp_server_manager.get_mcp_server_by_name(names[0]) - if server is None or not server.is_gateway_managed_oauth2: + if server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server): return None return server @@ -449,7 +465,10 @@ def aggregate_authorize( A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the flow to that one server: the scope is sealed into the flow, carried into the code, and - bound into the session token, while the connect page interlude runs exactly as before. + bound into the session token. The connect URL carries only the flow handle; the page + learns the client origin, the scoped server, and whether its vendor OAuth is done from + :func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry + steers which server the page authorizes or names on the confirmation. Validation failures respond directly with 400 and never redirect: per RFC 6749 section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and @@ -474,10 +493,7 @@ def aggregate_authorize( resource_server_id=scoped_server.server_id if scoped_server is not None else None, audience=None, ) - connect_url: Final = _append_query_params( - f"{base_url}/ui/connect", - (("connect_flow", handle), ("connect_client", _origin_only(redirect_uri))), - ) + connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) response: Final = RedirectResponse(connect_url, status_code=303) _set_flow_cookie(response, request, handle, flow) return response @@ -684,6 +700,103 @@ def _origin_only(url: str) -> str: return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else "" +def _open_flow_for( + request: Request, flow_handle: str, session_user_id: str | None, now: datetime +) -> _ConnectFlow | Response: + sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle)) + if sealed_flow is None: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) + if flow is None or now.timestamp() >= flow.exp: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + if session_user_id is None: + return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") + if session_user_id != flow.user_id: + return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + return flow + + +async def _flow_target( + flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability +) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]: + if flow.resource_server_id is None: + return "unscoped", None + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle + MCPServerManager, + global_mcp_server_manager, + ) + + server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id) + if ( + server is None + or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server) + or not await lookup_server_reachability(flow.user_id, server.server_id) + ): + return "stale", None + state: Final = ( + "interactive" + if server.is_gateway_managed_oauth2 and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" + else "m2m" + ) + return state, server + + +class ConnectFlowDescription(TypedDict): + """What the connect page is allowed to know about one in-flight flow.""" + + state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]] + client_origin: ReadOnly[str] + server_id: ReadOnly[str | None] + server_name: ReadOnly[str | None] + connected: ReadOnly[bool | None] + + +async def _describe_opened_flow( + flow: _ConnectFlow, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> ConnectFlowDescription | Response: + state, server = await _flow_target(flow, lookup_server_reachability) + if state == "interactive" and server is not None: + credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id) + if credential == "unavailable": + return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) + interactive_description: Final[ConnectFlowDescription] = { + "state": state, + "client_origin": _origin_only(flow.redirect_uri), + "server_id": server.server_id, + "server_name": server.server_name or server.alias or server.name, + "connected": credential == "present", + } + return interactive_description + described: Final[ConnectFlowDescription] = { + "state": state, + "client_origin": _origin_only(flow.redirect_uri), + "server_id": None if server is None else server.server_id, + "server_name": None if server is None else (server.server_name or server.alias or server.name), + "connected": state == "m2m" or None, + } + return described + + +async def describe_connect_flow( + request: Request, + flow_handle: str, + session_user_id: str | None, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> Response: + opened: Final = _open_flow_for(request, flow_handle, session_user_id, datetime.now(timezone.utc)) + if isinstance(opened, Response): + return opened + described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability) + return ( + described + if isinstance(described, Response) + else JSONResponse(content=described, headers=TOKEN_NO_CACHE_HEADERS) + ) + + async def complete_connect_flow( request: Request, flow_handle: str, @@ -692,56 +805,34 @@ async def complete_connect_flow( delivery: str | None = None, team_id: str | None = None, decision: str | None = None, + lookup_vendor_credential: LookupVendorCredential = _unavailable_vendor_credential, + lookup_server_reachability: LookupServerReachability = _unreachable_server, ) -> Response: - """The deliberate finish step of the connect flow: mint the gateway authorization - code and send the browser back to the client. + """Mint the code only after a deliberate POST by the sealed user. - Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly - per-flow cookie plus an exact match between the signed-in user and the user sealed - into the flow: a link crafted by another party dies here with ``access_denied`` - instead of minting a code for the victim's identity. The flow is single-use (an atomic - claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in. - - ``delivery`` chooses how the code reaches the client. Default (absent or - ``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"`` - renders the callback URL on a page instead, for a client whose redirect URI is a - loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box, - container): the 303 would dereference the browser machine's loopback and the code - would never arrive, so the user carries it over by pasting the URL into the client or - fetching it from the client machine's terminal. Manual delivery is honored only for - loopback redirect URIs; a routable redirect URI works from any browser by - construction, so those flows always redirect. The user who sees the page is exactly - the user the 303 would have carried the code to, and the same user already sees the - code today in the dead redirect's address bar, so the page exposes the code to no new - party. Unknown ``delivery`` values are rejected rather than defaulted: a client that - asked for manual delivery and got a dead redirect instead would silently lose its - code. - - ``decision`` and ``team_id`` come from the native-client consent page. ``"deny"`` - burns the flow and sends the client ``error=access_denied`` so it stops waiting; - ``team_id`` is sealed into the code only for proxy-API flows, where it picks which of - the user's teams the minted credential is attributed to. + A scoped flow additionally requires its sealed server to have a live vendor credential + before a code can be minted. The check happens before the single-use claim, so a + premature submit can be retried after authorization; denial deliberately bypasses it. """ if delivery not in (None, "redirect", "manual"): return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'") if decision not in (None, "approve", "deny"): return _oauth_error(400, "invalid_request", "decision must be 'approve' or 'deny'") - sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle)) - if sealed_flow is None: - return _oauth_error(400, "invalid_request", "unknown or expired connect flow") - flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) - if flow is None: - return _oauth_error(400, "invalid_request", "unknown or expired connect flow") now: Final = datetime.now(timezone.utc) - if now.timestamp() >= flow.exp: - return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection") - if session_user_id is None: - return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") - if session_user_id != flow.user_id: - return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + opened: Final = _open_flow_for(request, flow_handle, session_user_id, now) + if isinstance(opened, Response): + return opened + if decision != "deny": + described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability) + if isinstance(described, Response): + return described + if described["state"] == "stale": + return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available") + if described["connected"] is False: + return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing") flow_refusal: Final = _claim_refusal( await _SingleUseGuard(cache).claim( - f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS + f"{_USED_FLOW_CACHE_PREFIX}{opened.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS ), replayed=_oauth_error( 400, "invalid_request", "this connect flow was already completed; restart the connection" @@ -750,7 +841,7 @@ async def complete_connect_flow( if flow_refusal is not None: return flow_refusal response: Final = ( - _denied_flow_response(flow) if decision == "deny" else _approved_flow_response(flow, delivery, team_id, now) + _denied_flow_response(opened) if decision == "deny" else _approved_flow_response(opened, delivery, team_id, now) ) path, secure = _cookie_path_and_secure(request) response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_context.py b/litellm/proxy/_experimental/mcp_server/mcp_context.py index 7936edad753..74cc0c900d9 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_context.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_context.py @@ -21,3 +21,6 @@ _mcp_gateway_initialize_instructions: Final[ContextVar[str | None]] = ContextVar # Per-request scoped server name; set in MCP HTTP/SSE handlers when the path # identifies exactly one upstream server. Never populated from client-supplied headers. _mcp_gateway_server_name: Final[ContextVar[str | None]] = ContextVar("_mcp_gateway_server_name", default=None) + +# Set server-side by the /mcp/proxy route. Never populated from client-supplied headers. +_mcp_proxy_mode: Final[ContextVar[bool]] = ContextVar("_mcp_proxy_mode", default=False) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 6aaa00ae415..1f157aefdc3 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -22,7 +22,16 @@ Response headers returned (all values are masked for safety): x-mcp-debug-auth-resolution Which auth priority was used for the outbound MCP call: ``per-request-header``, ``m2m-client-credentials``, ``static-token``, - ``oauth2-passthrough``, or ``no-auth``. + ``oauth2-passthrough``, ``stored-user-token``, ``token-exchange``, + ``id-jag``, ``aws-sigv4``, ``extra-headers``, or ``no-auth``. + ``unresolved`` means no outcome was available before the first response + frame; ``multiple`` means several servers resolved credentials; + ``not-applicable`` covers stdio; ``resolution-failed`` is a resolver error. + + x-mcp-debug-auth-resolutions + For multiple servers, a JSON map of server IDs to resolution labels. + At most 32 entries are included; x-mcp-debug-auth-resolutions-truncated + is true when additional servers were omitted. No credentials are included. x-mcp-debug-outbound-url The upstream MCP server URL that will receive the request. @@ -58,10 +67,16 @@ header is free for OAuth2 discovery:: Symptom: ``x-mcp-debug-oauth2-token`` shows ``(none)`` and ``x-mcp-debug-auth-resolution`` shows ``no-auth``. -This means the client didn't go through the OAuth2 flow. Check that: -1. The ``Authorization`` header is NOT set as a static header in the client config. -2. The ``.well-known/oauth-protected-resource`` endpoint returns valid metadata. -3. The MCP server in LiteLLM config has ``auth_type: oauth2``. +``no-auth`` means the resolved upstream client carries no authentication. +An absent inbound OAuth2 token does not imply the user skipped OAuth: the gateway +can retrieve a stored per-user token, reported as ``stored-user-token``. +``unresolved`` is used when a stream starts before credential resolution, or a +request (such as initialization or a cached tool listing) resolves no credential. +Debug reporting does not fetch credentials or delay a streaming frame to resolve them. +``extra-headers`` identifies supplied headers that won over the resolver or were +the only headers supplied; their values are never inspected to guess a scheme. +``per-request-header`` denotes a legacy credential override, including a BYOK +credential supplied by the gateway; it does not imply a caller-supplied token. **Common issue: M2M token used instead of user token** @@ -69,8 +84,8 @@ Symptom: ``x-mcp-debug-auth-resolution`` shows ``m2m-client-credentials``. This means the server has ``client_id``/``client_secret``/``token_url`` configured and LiteLLM is fetching a machine-to-machine token instead of -using the per-user OAuth2 token. If you want per-user tokens, remove the -client credentials from the server config. +using the per-user OAuth2 token. For gateway-stored per-user tokens, +configure ``oauth2_flow: authorization_code``. Usage from Claude Code:: @@ -85,14 +100,26 @@ Usage with curl:: http://localhost:4000/mcp/atlassian_mcp """ -from typing import TYPE_CHECKING, Final +import asyncio +import base64 +import io +import json +import re +from collections.abc import AsyncIterator, Callable, Mapping +from http.cookies import CookieError, SimpleCookie +from itertools import islice +from types import MappingProxyType +from typing import Final +from urllib.parse import parse_qsl, quote, quote_plus, unquote_plus, urlencode +import httpx +from pydantic import JsonValue, TypeAdapter +from starlette.requests import HTTPConnection from starlette.types import Message, Send +from litellm.litellm_core_utils.secret_redaction import REDACTED, redact_string from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker - -if TYPE_CHECKING: - from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution # Header the client sends to opt into debug mode MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug" @@ -101,6 +128,83 @@ MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug" _RESPONSE_HEADER_PREFIX: Final = "x-mcp-debug" +MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: Final = "litellm.mcp.auth_diagnostics" + + +def record_auth_resolution(server_id: str, source: AuthResolution) -> None: + from mcp.server.lowlevel.server import request_ctx + + context: Final[object] = request_ctx.get(None) + request: Final[object] = getattr(context, "request", None) + if isinstance(request, HTTPConnection): + diagnostics: Final[object] = request.scope.get(MCP_AUTH_DIAGNOSTICS_SCOPE_KEY) + if isinstance(diagnostics, MCPAuthDiagnostics): + diagnostics.record(server_id, source) + + +class MCPAuthDiagnostics: + def __init__(self) -> None: + self._outcomes: tuple[tuple[str, AuthResolution], ...] = () + + def record(self, server_id: str, resolution: AuthResolution) -> None: + self._outcomes = tuple(item for item in self._outcomes if item[0] != server_id) + ((server_id, resolution),) + + def resolution(self) -> str: + match self._outcomes: + case (): + return AuthResolution.unresolved.value + case ((_, source),): + return source.value + case _: + return AuthResolution.multiple.value + + def headers(self) -> Mapping[str, str]: + if len(self._outcomes) <= 1: + return MappingProxyType({"x-mcp-debug-auth-resolution": self.resolution()}) + return MappingProxyType( + { + "x-mcp-debug-auth-resolution": AuthResolution.multiple.value, + "x-mcp-debug-auth-resolutions": json.dumps( + { + server_id: source.value for server_id, source in self._outcomes[:32] + }, # mutable-ok: JSON encoder requires a concrete dict + separators=(",", ":"), + ensure_ascii=True, + ), + **( + MappingProxyType({"x-mcp-debug-auth-resolutions-truncated": "true"}) + if len(self._outcomes) > 32 + else MappingProxyType({}) + ), + } + ) + + +class _DiagnosticSend: + def __init__(self, send: Send, headers: Mapping[str, str], resolution: Callable[[], Mapping[str, str]]) -> None: + self._send = send + self._headers = headers + self._resolution = resolution + self._start: Message | None = None + + async def __call__(self, message: Message) -> None: + if message["type"] == "http.response.start": + self._start = message + return + if self._start is not None: + start: Final = self._start + self._start = None + headers: Final = MappingProxyType({**self._headers, **self._resolution()}) + await self._send( + { # mutable-ok: ASGI send consumes a mutable message mapping + **start, + "headers": tuple(start.get("headers", ())) + + tuple((key.encode(), value.encode()) for key, value in headers.items()), + } + ) + await self._send(message) + + class MCPDebug: """ Static helper class for MCP OAuth2 debug headers. @@ -127,6 +231,10 @@ class MCPDebug: @staticmethod def _mask(value: str | None) -> str: """Mask a single value for safe display in headers.""" + return MCPDebug.mask_secret(value) + + @staticmethod + def mask_secret(value: str | None) -> str: if not value: return "(none)" return MCPDebug._masker._mask_value(value) @@ -144,37 +252,6 @@ class MCPDebug: return val.strip().lower() in ("true", "1", "yes") return False - @staticmethod - def resolve_auth_resolution( - server: "MCPServer", - mcp_auth_header: str | None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, - oauth2_headers: dict[str, str] | None, - ) -> str: - """ - Determine which auth priority will be used for the outbound MCP call. - - Returns one of: ``per-request-header``, ``m2m-client-credentials``, - ``static-token``, ``oauth2-passthrough``, or ``no-auth``. - """ - from litellm.types.mcp import MCPAuth - - has_server_specific: Final = bool( - mcp_server_auth_headers - and ( - mcp_server_auth_headers.get(server.alias or "") or mcp_server_auth_headers.get(server.server_name or "") - ) - ) - if has_server_specific or mcp_auth_header: - return "per-request-header" - if server.has_client_credentials: - return "m2m-client-credentials" - if server.authentication_token: - return "static-token" - if oauth2_headers and server.auth_type == MCPAuth.oauth2: - return "oauth2-passthrough" - return "no-auth" - @staticmethod def build_debug_headers( *, @@ -244,12 +321,21 @@ class MCPDebug: return debug @staticmethod - def wrap_send_with_debug_headers(send: Send, debug_headers: dict[str, str]) -> Send: + def wrap_send_with_debug_headers( + send: Send, + debug_headers: Mapping[str, str], + resolution: Callable[[], Mapping[str, str]] | None = None, + *, + request_method: str | None = None, + ) -> Send: """ Return a new ASGI ``send`` callable that injects *debug_headers* into the ``http.response.start`` message. """ + if resolution is not None and request_method == "POST": + return _DiagnosticSend(send, debug_headers, resolution) + async def _send_with_debug(message: Message) -> None: if message["type"] == "http.response.start": headers: Final = list(message.get("headers", [])) @@ -266,8 +352,6 @@ class MCPDebug: raw_headers: dict[str, str] | None, scope: dict, mcp_servers: list[str] | None, - mcp_auth_header: str | None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, oauth2_headers: dict[str, str] | None, client_ip: str | None, ) -> dict[str, str]: @@ -288,16 +372,13 @@ class MCPDebug: server_url: str | None = None server_auth_type: str | None = None - auth_resolution = "no-auth" + auth_resolution: Final = AuthResolution.unresolved.value for server_name in mcp_servers or []: server = global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip) if server: server_url = server.url server_auth_type = server.auth_type - auth_resolution = MCPDebug.resolve_auth_resolution( - server, mcp_auth_header, mcp_server_auth_headers, oauth2_headers - ) break scope_headers: Final = MCPRequestHandler._safe_get_headers_from_scope(scope) @@ -311,3 +392,230 @@ class MCPDebug: server_url=server_url, server_auth_type=server_auth_type, ) + + +_BODY_PREVIEW_CHARS: Final = 512 +_BODY_CAPTURE_BYTES: Final = 16384 +_CAPTURE_TIMEOUT_SECONDS: Final = 1.0 +_CAPTURE_EXTENSION: Final = "litellm_mcp_error_preview" +_SAFE_HEADER_NAMES: Final = frozenset({"content-type", "content-length", "accept"}) +_PUBLIC_HEADER_NAMES: Final = _SAFE_HEADER_NAMES | frozenset(("host", "user-agent", "accept-encoding", "connection")) +_JSON_BODY: Final = TypeAdapter(JsonValue) +_LOG_MASKER: Final = SensitiveDataMasker(visible_prefix=0, visible_suffix=0) + + +def _safe_text(value: str, limit: int = _BODY_PREVIEW_CHARS) -> str: + escaped: Final = "".join(json.dumps(char)[1:-1] if ord(char) < 32 or ord(char) == 127 else char for char in value) + return escaped if len(escaped) <= limit else f"{escaped[:limit]}...(truncated)" + + +def safe_upstream_url(url: httpx.URL) -> str: + return _safe_text(str(url.copy_with(username="", password="", path="/", query=None, fragment=None))) + + +def _sensitive_field(key: str) -> bool: + normalized: Final = re.sub(r"[^a-z0-9]", "", key.casefold()) + return normalized in ("code", "clientassertion") or any( + pattern in normalized for pattern in _LOG_MASKER.sensitive_patterns + ) + + +def _redact_object( + fields: Mapping[str, JsonValue], +) -> dict[str, JsonValue]: # mutable-ok: the standard JSON encoder requires dict objects + return { # mutable-ok: construct the JSON object once for the standard parser and encoder + key: REDACTED if _sensitive_field(key) else value for key, value in fields.items() + } + + +def _header_secret_values(name: str, value: str) -> tuple[str, ...]: + if name == "cookie": + cookie: Final = SimpleCookie[str]() + try: + cookie.load(value) + except CookieError: + return (value,) + return (value, *(item.value for item in cookie.values())) + if name not in ("authorization", "proxy-authorization"): + return (value,) + scheme, _, credential = value.partition(" ") + if scheme.lower() != "basic": + return (value, credential) + try: + decoded: Final = base64.b64decode(credential, validate=True).decode("utf-8") + except ValueError: + return (value, credential) + password: Final = decoded.partition(":")[2] + return (value, credential, decoded, password, unquote_plus(password)) + + +def _body_secret_values(request: httpx.Request) -> tuple[str, ...] | None: + try: + raw: Final = request.content + except httpx.RequestNotRead: + return None + if not raw: + return () + if len(raw) > _BODY_CAPTURE_BYTES: + return None + if request.headers.get("content-type", "").split(";", 1)[0].strip().lower() == "application/x-www-form-urlencoded": + return tuple(value for key, value in parse_qsl(raw.decode("utf-8", errors="replace")) if _sensitive_field(key)) + try: + body: Final = _JSON_BODY.validate_json(raw) + except ValueError: + return None + from litellm.proxy._experimental.mcp_server.utils import ( # noqa: PLC0415 # MCP utils imports clients; inspect bodies only after initialization + json_string_leaves, + ) + + leaves: Final = json_string_leaves(body) + if leaves is None: + return None + return tuple( + value + for path, value in leaves + if not path or any(isinstance(part, str) and _sensitive_field(part) for part in path) + ) + + +def _request_secret_values(request: httpx.Request) -> tuple[str, ...] | None: + body_values: Final = _body_secret_values(request) + if body_values is None: + return None + values: Final = ( + *body_values, + request.url.password, + *(value for _, value in request.url.params.multi_items()), + *( + secret + for name, value in request.headers.items() + if name not in _PUBLIC_HEADER_NAMES + for secret in _header_secret_values(name, value) + ), + ) + return tuple(sorted(frozenset(value for value in values if value), key=len, reverse=True)) + + +def _mask_known_values(value: str, secrets: tuple[str, ...]) -> str: + variants: Final = tuple( + sorted( + frozenset( + variant + for secret in secrets + for variant in (secret, json.dumps(secret)[1:-1], quote(secret, safe=""), quote_plus(secret)) + ), + key=len, + reverse=True, + ) + ) + return re.sub("|".join(re.escape(secret) for secret in variants), REDACTED, value) if variants else value + + +def _preview(raw: bytes, content_type: str = "", secrets: tuple[str, ...] = ()) -> str: + if not raw: + return "(empty)" + if len(raw) > _BODY_CAPTURE_BYTES: + return "(omitted: body exceeds capture limit)" + try: + parsed: Final = _JSON_BODY.validate_python(json.loads(raw, object_hook=_redact_object)) + except (ValueError, RecursionError): + text: Final = raw.decode("utf-8", errors="replace") + if ( + content_type.split(";", 1)[0].strip().lower() != "application/x-www-form-urlencoded" + or "=" not in text + or any(char in text for char in "<>\n\r") + ): + return "(omitted: unstructured body)" + fields: Final = parse_qsl(text, keep_blank_values=True) + return _safe_text( + _mask_known_values( + urlencode(tuple((key, REDACTED if _sensitive_field(key) else value) for key, value in fields)), secrets + ) + ) + if not isinstance(parsed, (dict, list)): + return "(omitted: unstructured body)" + return _safe_text(redact_string(_mask_known_values(json.dumps(parsed, separators=(",", ":")), secrets))) + + +def _masked_headers(headers: httpx.Headers) -> str: + return _safe_text(", ".join(f"{name}={value}" for name, value in headers.items() if name in _SAFE_HEADER_NAMES)) + + +def _request_body_preview(request: httpx.Request, secrets: tuple[str, ...] | None) -> str: + try: + return _preview(request.content, request.headers.get("content-type", ""), secrets or ()) + except httpx.RequestNotRead: + return "(streamed, not captured)" + + +def _response_body_preview(response: httpx.Response, secrets: tuple[str, ...] | None) -> str: + if secrets is None: + return "(omitted: request credentials unavailable)" + captured: Final = response.extensions.get(_CAPTURE_EXTENSION) + if isinstance(captured, str): + return captured + try: + return _preview(response.content, response.headers.get("content-type", ""), secrets) + except httpx.ResponseNotRead: + return "(not read)" + + +async def _read_error_prefix(chunks: AsyncIterator[bytes], limit: int) -> bytes: + buffer: Final = io.BytesIO() + async for chunk in chunks: + buffer.write(chunk[: limit - buffer.tell()]) + if buffer.tell() >= limit: + break + return buffer.getvalue() + + +async def capture_upstream_error_response(response: httpx.Response) -> None: + if not response.is_error: + return + try: + prefix: Final = await asyncio.wait_for( + _read_error_prefix(response.aiter_bytes(chunk_size=4096), _BODY_CAPTURE_BYTES + 1), + timeout=_CAPTURE_TIMEOUT_SECONDS, + ) + response._content = prefix # pyright: ignore[reportPrivateUsage] # rebind-ok: httpx has no public setter to retain consumed bytes for auth retries + secrets: Final = _request_secret_values(response.request) + preview: Final = ( + _preview(prefix, response.headers.get("content-type", ""), secrets) + if secrets is not None + else "(omitted: request credentials unavailable)" + ) + except (asyncio.TimeoutError, httpx.HTTPError, httpx.StreamError): + response._content = b"" # pyright: ignore[reportPrivateUsage] # rebind-ok: httpx auth retries must survive diagnostic read failures + response.extensions[_CAPTURE_EXTENSION] = ( + "(unavailable: error body read failed)" # rebind-ok: httpx response hooks communicate through extensions + ) + return + response.extensions[_CAPTURE_EXTENSION] = preview # rebind-ok: httpx response hooks communicate through extensions + + +def describe_upstream_response(response: httpx.Response) -> str: + try: + request: Final = response.request + except RuntimeError: + return f"HTTP {response.status_code} | request unavailable" + secrets: Final = _request_secret_values(request) + return ( + f"{_safe_text(request.method)} {safe_upstream_url(request.url)} -> HTTP {response.status_code}" + f" | request headers: {_masked_headers(request.headers)}" + f" | request body: {_request_body_preview(request, secrets)}" + f" | response body: {_response_body_preview(response, secrets)}" + ) + + +def describe_upstream_http_failure(exc: BaseException) -> str | None: + from litellm.proxy._experimental.mcp_server.faults.traversal import ( # noqa: PLC0415 # fault package initialization imports the credential resolver + iter_exception_tree, + ) + + lines: Final = tuple( + describe_upstream_response(response) + for current in islice(iter_exception_tree(exc), 16) + for response in (getattr(current, "response", None),) + if isinstance(response, httpx.Response) + ) + return " | ".join(lines) or None diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index dc1e8db1628..cc7ea0c1ea5 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -25,6 +25,7 @@ from collections.abc import ( ) from contextlib import asynccontextmanager from dataclasses import dataclass, replace +from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast from urllib.parse import ParseResult, urlparse @@ -80,6 +81,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( raise_classified_list_failure, upstream_auth_challenge, ) +from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure, record_auth_resolution from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, mcp_per_user_token_cache, @@ -108,12 +110,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import ( LazyPerUserOAuthTokenStore, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import ( build_token_exchanger, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( DEFAULT_CREDENTIAL_HEADER, AuthorizationCodeConfig, + AuthResolution, ClientCredentialsConfig, CredError, IdJagConfig, @@ -155,6 +159,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 @@ -162,7 +167,10 @@ from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( id_jag_assertion_capture_gap_at_startup, ) -from litellm.proxy.utils import PrismaClient, ProxyLogging, get_server_root_path +from litellm.proxy.middleware.per_request_root_path_middleware import ( + get_request_root_path, +) +from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.table_repositories import MCPServerRepository from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import ( @@ -247,6 +255,7 @@ _TRUE_ENV_VALUES: Final = frozenset(("1", "true", "yes", "on")) _OAUTH_DISCOVERY_RETRY_DELAYS_SECONDS: Final = (0.05, 0.15) _OAUTH_DISCOVERY_RETRY_BASE_SECONDS: Final = 30.0 _OAUTH_DISCOVERY_RETRY_MAX_SECONDS: Final = 900.0 +_OAUTH_TEMPORARY_DISCOVERY_TTL_SECONDS: Final = 300.0 def _oauth_discovery_now() -> float: @@ -344,6 +353,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 +424,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. @@ -850,6 +885,53 @@ def _sanitized_error_text(exc: Exception) -> str: return re.sub(r"https?://\S+", "", str(exc))[:200] +async def _openapi_spec_health( + spec_path: str, *, timeout: float +) -> tuple[Literal["healthy", "unhealthy", "unknown"], str | None]: + """Check specification availability, not upstream operations or user credentials.""" + from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + if not spec_path.startswith(("http://", "https://")): + return "unknown", "OpenAPI servers have no protocol-level health probe" + try: + await asyncio.wait_for(load_openapi_spec_async(spec_path, max_bytes=10 * 1024 * 1024), timeout=timeout) + except asyncio.TimeoutError: + return "unhealthy", f"OpenAPI specification check timed out after {timeout} seconds" + except HTTPStatusError as exc: + return "unhealthy", f"OpenAPI specification request failed (HTTP {exc.response.status_code})" + except HTTPResponseLimitError as exc: + return "unknown", f"OpenAPI specification probe refused: {exc}" + except (httpx.RequestError, ValueError, OSError) as exc: + return "unhealthy", f"OpenAPI specification could not be loaded ({type(exc).__name__})" + return "healthy", None + + +class _OpenAPIHealthProbe: + def __init__(self, spec_path: str, clock: Callable[[], float] = time.monotonic) -> None: + self.spec_path = spec_path + self.clock = clock + self.lock = asyncio.Lock() + self.checked_at = float("-inf") + self.result: tuple[Literal["healthy", "unhealthy", "unknown"], str | None, datetime.datetime] | None = None + + async def check(self) -> tuple[Literal["healthy", "unhealthy", "unknown"], str | None, datetime.datetime]: + async with self.lock: + if self.result is not None and self.clock() - self.checked_at < 30.0: + return self.result + try: + status, error = await _openapi_spec_health(self.spec_path, timeout=MCP_HEALTH_CHECK_TIMEOUT) + except asyncio.CancelledError: + return ( + "unknown", + "OpenAPI specification check was cancelled", + datetime.datetime.now(datetime.timezone.utc), + ) + self.result = (status, error, datetime.datetime.now(datetime.timezone.utc)) + self.checked_at = self.clock() + return self.result + + def _discovery_failure_leaves_needs_unresolved( *, needs_authorization_url: bool, @@ -1323,6 +1405,11 @@ def _extract_upstream_auth_failure( return upstream_auth_challenge(exc) +def _upstream_failure_suffix(exc: BaseException) -> str: + detail: Final = describe_upstream_http_failure(exc) + return f"\n upstream exchange: {detail}" if detail else "" + + def _obo_retry_applies(server: MCPServer, subject_token: str | None) -> bool: """Whether an upstream 401/403 should invalidate the minted credential and retry once. @@ -1716,6 +1803,7 @@ class MCPServerManager: token_exchanger=build_token_exchanger(), ) self.registry: dict[str, MCPServer] = {} + self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe) self.config_mcp_servers: dict[str, MCPServer] = {} """ eg. @@ -1865,6 +1953,10 @@ class MCPServerManager: slot: Final = self._oauth_discovery_slot(server_id) return slot is not None and slot.generation == generation + def _expire_temporary_oauth_discovery(self, server_id: str, generation: int) -> None: + if self._oauth_discovery_slot_is_current(server_id, generation): + self._remove_oauth_discovery_slot(server_id) + def _publish_resolved_oauth_server( self, server: MCPServer, @@ -1877,7 +1969,13 @@ class MCPServerManager: elif server.server_id in self.config_mcp_servers: self.config_mcp_servers[server.server_id] = server else: - return None + asyncio.get_running_loop().call_later( + _OAUTH_TEMPORARY_DISCOVERY_TTL_SECONDS, + self._expire_temporary_oauth_discovery, + server.server_id, + generation, + ) + return server self._remove_oauth_discovery_slot(server.server_id) return server @@ -1969,6 +2067,12 @@ class MCPServerManager: if slot.task is not None: if not slot.task.done() or _oauth_discovery_now() < slot.retry_not_before: return slot.task, slot.generation + if ( + not slot.task.cancelled() + and slot.task.exception() is None + and isinstance(slot.task.result(), _OAuthDiscoveryResolved) + ): + return slot.task, slot.generation task: Final = asyncio.create_task( self._run_oauth_metadata_resolution(self._registered_server(server), slot.generation) ) @@ -1998,7 +2102,7 @@ class MCPServerManager: if should_defer != has_slot: self._set_oauth_discovery_deferred(server.server_id, should_defer) - async def ensure_oauth_metadata_discovered(self, server: MCPServer) -> MCPServer: + async def ensure_oauth_metadata_discovered(self, server: MCPServer, *, _retry_stale: bool = True) -> MCPServer: """Join the bounded discovery task and return the resolved server. Concurrent callers share one task per server. A failed attempt remains @@ -2025,13 +2129,13 @@ class MCPServerManager: outcome: Final = await asyncio.shield(task) except asyncio.CancelledError: if task.cancelled() and not self._oauth_discovery_slot_is_current(server.server_id, generation): - return await self.ensure_oauth_metadata_discovered(server) + return await self._rejoin_oauth_metadata_discovery(server, retry_stale=_retry_stale) raise match outcome: case _OAuthDiscoveryResolved(resolved_server): return resolved_server case _OAuthDiscoveryStale(): - return await self.ensure_oauth_metadata_discovered(server) + return await self._rejoin_oauth_metadata_discovery(server, retry_stale=_retry_stale) case _OAuthDiscoveryFailed(timed_out=timed_out): current: Final = self._registered_server(server) if current.is_client_forwarded_token: @@ -2043,6 +2147,14 @@ class MCPServerManager: detail=f"OAuth metadata discovery {reason} for MCP server {server_ref!r}", ) + async def _rejoin_oauth_metadata_discovery(self, server: MCPServer, *, retry_stale: bool) -> MCPServer: + if retry_stale: + return await self.ensure_oauth_metadata_discovered(server, _retry_stale=False) + current: Final = self._registered_server(server) + if not _oauth_endpoints_unresolved(current) or current.is_client_forwarded_token: + return current + raise HTTPException(status_code=503, detail="OAuth metadata discovery changed repeatedly; retry shortly") + def _remember_upstream_initialize_instructions(self, server: MCPServer, client: MCPClient) -> None: raw: Final[str | None] = getattr(client, "_last_initialize_instructions", None) if raw and str(raw).strip(): @@ -2307,6 +2419,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 +2493,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), @@ -2407,6 +2523,8 @@ class MCPServerManager: allow_elicitation=bool(server_config.get("allow_elicitation", False)), timeout=server_config.get("timeout", None), max_concurrent_requests=server_config.get("max_concurrent_requests", None), + token_validation=server_config.get("token_validation", None), + oauth_identity_binding=server_config.get("oauth_identity_binding", None), ) self._assign_unique_short_prefix(new_server) _warn_internal_delegate_pkce_if_applicable(new_server, source="config") @@ -2903,6 +3021,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)), @@ -3797,13 +3916,21 @@ class MCPServerManager: (authorization_code's browser-OAuth 401, token_exchange's RFC 9728 challenge) or maps any other ``CredError`` onto its public HTTP status; it never returns an error as a value. """ - match await provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): - case Ok(auth): + match await resolve_credentials_with_source(provider, to_subject(user_api_key_auth, subject_token), spec): + case Ok(credential): + auth: Final = credential.auth # NoOpAuth has no header_name and so never conflicts. header_name: Final[str | None] = getattr(auth, "header_name", None) if header_name is None or not extra_headers: + source: Final = ( + AuthResolution.extra_headers + if credential.source == AuthResolution.no_auth and extra_headers + else credential.source + ) + record_auth_resolution(server.server_id, source) return auth, extra_headers if not has_header(extra_headers, header_name): + record_auth_resolution(server.server_id, credential.source) return auth, extra_headers if isinstance( spec.config, @@ -3818,15 +3945,18 @@ class MCPServerManager: # one-shot 401 refetch is lost with it). Drop only the header the resolved # credential is about to occupy, so a static credential the operator aimed at a # DIFFERENT header still reaches upstream. + record_auth_resolution(server.server_id, credential.source) return auth, without_header(extra_headers, header_name) # Other modes: an Authorization already supplied via extra_headers (a forwarded caller # header or static_headers) is intentional and wins; v1 applies those last. + record_auth_resolution(server.server_id, AuthResolution.extra_headers) return None, extra_headers case Error(err): + record_auth_resolution(server.server_id, AuthResolution.failed) if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig): # authorization_code's missing per-user token -> the per-server browser-OAuth # challenge, built here where the full MCPServer is in hand. - raise_user_oauth_challenge(server, root_path=get_server_root_path()) + raise_user_oauth_challenge(server, root_path=get_request_root_path()) if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig): # token_exchange (OBO): a missing/rejected subject token -> the RFC 9728 challenge # pointing at the IdP the client must SSO with to obtain one, rather than an opaque @@ -3834,7 +3964,7 @@ class MCPServerManager: # Access) threads its claims blob into the challenge for the client to satisfy. raise_token_exchange_challenge( server, - root_path=get_server_root_path(), + root_path=get_request_root_path(), claims=err.unauthorized.claims, ) raise_public(err) @@ -3882,7 +4012,7 @@ class MCPServerManager: if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): return if subject_token is None and isinstance(spec.config, TokenExchangeConfig): - raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path()) + raise_token_exchange_challenge(resolved_server, root_path=get_request_root_path()) match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): case Ok(_): return @@ -3890,7 +4020,7 @@ class MCPServerManager: if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig): raise_token_exchange_challenge( resolved_server, - root_path=get_server_root_path(), + root_path=get_request_root_path(), claims=err.unauthorized.claims, ) raise_public(err) @@ -3925,6 +4055,7 @@ class MCPServerManager: Returns: Configured MCP client instance. """ + record_auth_resolution(server.server_id, AuthResolution.unresolved) resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) transport: Final = resolved_server.transport or MCPTransport.sse spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(resolved_server) @@ -3997,6 +4128,7 @@ class MCPServerManager: env=resolved_env, ) + record_auth_resolution(server.server_id, AuthResolution.not_applicable) return MCPClient( server_url="", # Not used for stdio transport_type=transport, @@ -4051,6 +4183,20 @@ class MCPServerManager: aws_session_name=resolved_server.aws_session_name, ) + legacy_source: Final = ( + AuthResolution.aws_sigv4 + if aws_auth is not None + else AuthResolution.extra_headers + if extra_headers and has_header(extra_headers, auth_header_name or "Authorization") + else AuthResolution.per_request_header + if mcp_auth_header + else AuthResolution.static_token + if auth_value + else AuthResolution.extra_headers + if extra_headers + else AuthResolution.no_auth + ) + record_auth_resolution(server.server_id, legacy_source) return MCPClient( server_url=server_url, transport_type=transport, @@ -4224,7 +4370,9 @@ class MCPServerManager: except MCPServerListError: raise except Exception as e: - verbose_logger.warning("Failed to get tools from server %s: %s", server.name, e) + verbose_logger.warning( + "Failed to get tools from server %s: %s%s", server.name, type(e).__name__, _upstream_failure_suffix(e) + ) raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge) async def get_prompts_from_server( @@ -4969,7 +5117,9 @@ class MCPServerManager: verbose_logger.warning("Connection error while listing tools from %s: %s", server_name, e) raise MCPServerListError(ServerListFault(tag="unreachable"), server_name) from e except Exception as e: - verbose_logger.warning("Error listing tools from %s: %s", server_name, e) + verbose_logger.warning( + "Error listing tools from %s: %s%s", server_name, type(e).__name__, _upstream_failure_suffix(e) + ) raise_classified_list_failure(e, server_name) _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 @@ -6237,8 +6387,7 @@ class MCPServerManager: ] } ) - db_mcp_servers: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in raw_rows] - verbose_logger.info("Found %s MCP servers in database", len(db_mcp_servers)) + verbose_logger.info("Found %s MCP servers in database", len(raw_rows)) previous_registry: Final = self.registry new_registry: Final[dict[str, MCPServer]] = {} @@ -6246,8 +6395,9 @@ class MCPServerManager: # Stage one: build every server. Stage two assigns short prefixes # against the *full* set so dedup is deterministic regardless of # iteration order. - for server in db_mcp_servers: + for row in raw_rows: try: + server = LiteLLM_MCPServerTable.model_validate(row.model_dump()) existing_server = previous_registry.get(server.server_id) if ( @@ -6285,8 +6435,8 @@ class MCPServerManager: except Exception as e: verbose_logger.exception( "Skipping MCP server %s (%s) during DB reload: %s", - server.server_id, - getattr(server, "alias", None), + getattr(row, "server_id", None), + getattr(row, "alias", None), e, ) @@ -6612,6 +6762,18 @@ class MCPServerManager: last_health_check=datetime.now(), ) + if server.spec_path: + spec_status, spec_error, spec_checked_at = await self._openapi_health_probes(server.spec_path).check() + return self._build_mcp_server_table(server).model_copy( + update=MappingProxyType( + { + "status": spec_status, + "health_check_error": spec_error, + "last_health_check": spec_checked_at, + } + ) + ) + status: Literal["healthy", "unhealthy", "unknown"] = "unknown" health_check_error = None @@ -6692,6 +6854,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 +6973,7 @@ class MCPServerManager: delegate_auth_to_upstream=server.delegate_auth_to_upstream, oauth_passthrough=getattr(server, "oauth_passthrough", False), dcr_bridge=server.dcr_bridge, + per_server_oauth_discovery=server.per_server_oauth_discovery, is_byok=server.is_byok, byok_description=server.byok_description, byok_api_key_help_url=server.byok_api_key_help_url, diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index a4ef970b87a..42edc2999ab 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -28,6 +28,8 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( build_upstream_oauth2_token_request, resolve_upstream_resource, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import OAuthTokenCacheCodec from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -233,8 +235,17 @@ class MCPPerUserTokenCache: def _cache_key(self, user_id: str, server_id: str) -> str: return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}" + def _codec(self) -> OAuthTokenCacheCodec: + return OAuthTokenCacheCodec( + encrypt_value_helper, + lambda blob: decrypt_value_helper(blob, key="mcp_per_user_token", exception_type="debug"), + ) + async def get(self, user_id: str, server_id: str) -> str | None: - """Return the plaintext access_token, or None on miss/error.""" + token: Final = await self.get_token(user_id, server_id) + return token.access_token if token is not None else None + + async def get_token(self, user_id: str, server_id: str) -> OAuthToken | None: try: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 @@ -242,12 +253,7 @@ class MCPPerUserTokenCache: encrypted: Final = await user_api_key_cache.async_get_cache(key) if encrypted is None: return None - plaintext: Final = decrypt_value_helper( - encrypted, - key="mcp_per_user_token", - exception_type="debug", - ) - return plaintext or None + return self._codec().decode(encrypted) except Exception as exc: verbose_logger.debug( "MCPPerUserTokenCache.get failed for user=%s server=%s: %s", @@ -263,13 +269,16 @@ class MCPPerUserTokenCache: server_id: str, access_token: str, ttl: int, + identity_binding_proof: str | None = None, ) -> None: """Store NaCl-encrypted access_token in Redis with the given TTL.""" try: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 key: Final = self._cache_key(user_id, server_id) - encrypted: Final = encrypt_value_helper(access_token) + encrypted: Final = self._codec().encode( + OAuthToken(access_token=access_token, identity_binding_proof=identity_binding_proof) + ) await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl) verbose_logger.debug( "MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds", diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py new file mode 100644 index 00000000000..f02e6c85d9b --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -0,0 +1,423 @@ +"""Per-user OAuth identity binding: verify the upstream OIDC principal matches the LiteLLM caller. + +Closes the confused-deputy gap where a browser authenticated upstream as one principal produces a +token that the relay stores under a different, LiteLLM-authenticated principal: before the token +endpoint returns, stores, or caches an exchanged token for an identity-bound server, the id_token +is validated (signature via the pinned issuer's JWKS, issuer, audience, expiry) and its principal +claim is compared to the caller's trusted LiteLLM identity. Mismatches fail closed in enforce mode +and are logged in audit mode. +""" + +import hashlib +import hmac +import json +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Final, Literal, Protocol, TypeAlias + +import jwt +from fastapi import HTTPException +from jwt.types import Options +from typing_extensions import assert_never + +from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.llms.custom_http import httpxSpecialProvider +from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + +_ALLOWED_ID_TOKEN_ALGORITHMS: Final = ( + "RS256", + "RS384", + "RS512", + "ES256", + "ES384", + "ES512", + "PS256", + "PS384", + "PS512", +) +_JWKS_CACHE_TTL_SECONDS: Final = 3600 +_jwks_cache: Final = InMemoryCache(default_ttl=_JWKS_CACHE_TTL_SECONDS) + +JwksFetcher: TypeAlias = Callable[ + [MCPOAuthIdentityBinding], # mutable-ok: Callable parameter syntax requires a list + Awaitable[Sequence[Mapping[str, object]]], +] +CallerPrincipalLoader: TypeAlias = Callable[ + [str, MCPOAuthIdentityBinding], # mutable-ok: Callable parameter syntax requires a list + Awaitable[str | None], +] + + +@dataclass(frozen=True, slots=True) +class VerifiedRefreshToken: + refresh_token: str + binding_proof: str + + +StoredRefreshTokenLoader: TypeAlias = Callable[ + [str, str, MCPOAuthIdentityBinding], # mutable-ok: Callable parameter syntax requires a list + Awaitable[VerifiedRefreshToken | None], +] + +_RejectionCode: TypeAlias = Literal["oauth_principal_mismatch", "oauth_identity_binding_failed"] + + +@dataclass(frozen=True, slots=True) +class _BindingRejection: + code: _RejectionCode + description: str + + +@dataclass(frozen=True, slots=True) +class RefreshOwnershipProven: + """The gateway itself unwrapped the upstream refresh token from a sealed per-user envelope.""" + + +@dataclass(frozen=True, slots=True) +class RefreshTokenPresented: + refresh_token: str + + +RefreshOwnership: TypeAlias = RefreshOwnershipProven | RefreshTokenPresented | None + + +class BindingValidator(Protocol): + async def __call__( + self, + *, + server: MCPServer, + token_response: Mapping[str, object], + litellm_user_id: str | None, + grant_type: str, + refresh_ownership: RefreshOwnership, + ) -> str | None: ... + + +async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> Sequence[Mapping[str, object]]: + jwks_url: Final[str] = binding.jwks_url or await _discover_jwks_url(binding.issuer) + cached: Final = await _jwks_cache.async_get_cache(jwks_url) + if isinstance(cached, list): + return cached + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + response: Final = await client.get(jwks_url) + response.raise_for_status() + document: Final = response.json() + keys: Final = document.get("keys") if isinstance(document, dict) else None + if not isinstance(keys, list): + raise TypeError(f"JWKS document at {jwks_url} has no 'keys' array") + await _jwks_cache.async_set_cache(jwks_url, keys, ttl=_JWKS_CACHE_TTL_SECONDS) + return keys + + +async def _discover_jwks_url(issuer: str) -> str: + discovery_url: Final = f"{issuer.rstrip('/')}/.well-known/openid-configuration" + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + response: Final = await client.get(discovery_url) + response.raise_for_status() + metadata: Final = response.json() + jwks_uri: Final = metadata.get("jwks_uri") if isinstance(metadata, dict) else None + if not isinstance(jwks_uri, str) or not jwks_uri: + raise ValueError(f"OIDC discovery at {discovery_url} returned no jwks_uri") + return jwks_uri + + +def _select_signing_key(id_token: str, keys: Sequence[Mapping[str, object]]) -> "jwt.PyJWK | _BindingRejection": + header: Final = jwt.get_unverified_header(id_token) + kid: Final = header.get("kid") + for key in keys: + if kid is None or key.get("kid") == kid: + return jwt.PyJWK(dict(key)) # mutable-ok: PyJWT requires a concrete JWK dictionary + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"id_token signing key (kid={kid!r}) not found in the issuer's JWKS", + ) + + +def _decode_id_token( + id_token: str, + binding: MCPOAuthIdentityBinding, + signing_key: "jwt.PyJWK", +) -> "Mapping[str, object] | _BindingRejection": + try: + decode_options: Final[Options] = {"require": ("iss", "exp", "aud", "sub", "iat")} + return jwt.decode( + id_token, + signing_key.key, + algorithms=_ALLOWED_ID_TOKEN_ALGORITHMS, + issuer=binding.issuer, + audience=binding.audiences, + options=decode_options, + ) + except jwt.InvalidTokenError as exc: + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"id_token validation failed: {exc}", + ) + + +def _upstream_principal( + claims: Mapping[str, object], + binding: MCPOAuthIdentityBinding, +) -> "str | _BindingRejection": + principal: Final = claims.get(binding.principal_claim) + if not isinstance(principal, str) or not principal: + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"id_token has no usable '{binding.principal_claim}' claim", + ) + if ( + binding.principal_claim == "email" + and binding.require_email_verified + and claims.get("email_verified") is not True + ): + return _BindingRejection( + code="oauth_identity_binding_failed", + description="id_token email is not verified (email_verified is not true)", + ) + return principal + + +async def _load_caller_principal(litellm_user_id: str, binding: MCPOAuthIdentityBinding) -> str | None: + if binding.caller_field == "user_id": + return litellm_user_id + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # inline import avoids a module-load circular import + load_active_user_by_id, + ) + + loaded: Final = await load_active_user_by_id(litellm_user_id) + if isinstance(loaded, str): + return None + return loaded.user_email + + +async def _load_stored_refresh_token( + litellm_user_id: str, server_id: str, binding: MCPOAuthIdentityBinding +) -> VerifiedRefreshToken | None: + try: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # keep database imports lazy + get_user_oauth_credential, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # keep database imports lazy + + prisma_client: Final = get_prisma_client_or_throw( + "Database not connected. Cannot verify OAuth refresh token ownership." + ) + cred: Final = await get_user_oauth_credential( + prisma_client=prisma_client, + user_id=litellm_user_id, + server_id=server_id, + ) + if not cred or not await credential_binding_matches(binding, litellm_user_id, server_id, cred): + return None + refresh_token: Final = cred.get("refresh_token") + proof: Final = cred.get("identity_binding_proof") + return VerifiedRefreshToken(refresh_token, proof) if refresh_token and proof else None + except Exception: # noqa: BLE001 # a credential lookup failure must fail closed + return None + + +async def current_binding_proof( + binding: MCPOAuthIdentityBinding, + user_id: str, + server_id: str, + caller_principal_loader: CallerPrincipalLoader = _load_caller_principal, +) -> str | None: + principal: Final = await caller_principal_loader(user_id, binding) + if not principal: + return None + return _binding_proof(binding, user_id, server_id, principal) + + +def _binding_proof(binding: MCPOAuthIdentityBinding, user_id: str, server_id: str, principal: str) -> str: + payload: Final = json.dumps( + ("oidc-nonce-v1", server_id, user_id, principal, binding.model_dump(mode="json")), + sort_keys=True, + ) + return hashlib.sha256(payload.encode()).hexdigest() + + +async def credential_binding_matches( + binding: MCPOAuthIdentityBinding, + user_id: str, + server_id: str, + credential: Mapping[str, object], + caller_principal_loader: CallerPrincipalLoader = _load_caller_principal, +) -> bool: + stored: Final = credential.get("identity_binding_proof") + if not isinstance(stored, str) or not stored: + return False + expected: Final = await current_binding_proof(binding, user_id, server_id, caller_principal_loader) + return expected is not None and hmac.compare_digest(stored, expected) + + +def _principals_match(upstream: str, caller: str, binding: MCPOAuthIdentityBinding) -> bool: + if binding.principal_claim == "email" or binding.caller_field == "user_email": + return upstream.strip().casefold() == caller.strip().casefold() + return upstream == caller + + +async def _evaluate_refresh_ownership( + binding: MCPOAuthIdentityBinding, + litellm_user_id: str | None, + server_id: str, + refresh_ownership: RefreshOwnership, + stored_refresh_token_loader: StoredRefreshTokenLoader, +) -> _BindingRejection | str: + match refresh_ownership: + case RefreshOwnershipProven(): + return _BindingRejection( + code="oauth_identity_binding_failed", + description="an identity envelope alone does not prove upstream principal binding", + ) + case None: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="refresh_token grant without an id_token carries no refresh token to prove ownership of", + ) + case RefreshTokenPresented(refresh_token): + if not litellm_user_id: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the request carries no resolvable LiteLLM user identity to bind the credential to", + ) + stored: Final = await stored_refresh_token_loader(litellm_user_id, server_id, binding) + if stored is None or not hmac.compare_digest(stored.refresh_token, refresh_token): + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the presented refresh_token is not the caller's stored credential for this server", + ) + return stored.binding_proof + assert_never(refresh_ownership) # pragma: no cover + + +async def _evaluate_binding( + binding: MCPOAuthIdentityBinding, + token_response: Mapping[str, object], + litellm_user_id: str | None, + grant_type: str, + server_id: str, + refresh_ownership: RefreshOwnership, + jwks_fetcher: JwksFetcher, + caller_principal_loader: CallerPrincipalLoader, + stored_refresh_token_loader: StoredRefreshTokenLoader, + expected_nonce: str | None, +) -> _BindingRejection | str: + id_token: Final = token_response.get("id_token") + if not isinstance(id_token, str) or not id_token: + if grant_type != "refresh_token": + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the upstream token response carries no id_token to bind the credential to a principal", + ) + return await _evaluate_refresh_ownership( + binding, + litellm_user_id, + server_id, + refresh_ownership, + stored_refresh_token_loader, + ) + if not litellm_user_id: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the request carries no resolvable LiteLLM user identity to bind the credential to", + ) + try: + keys: Final = await jwks_fetcher(binding) + except Exception as exc: # noqa: BLE001 # a JWKS fetch failure must fail closed, not surface as a 500 + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"could not fetch the issuer's JWKS: {exc}", + ) + try: + signing_key: Final = _select_signing_key(id_token, keys) + except (jwt.PyJWTError, ValueError, TypeError): + return _BindingRejection( + code="oauth_identity_binding_failed", + description="invalid id_token header or issuer signing key", + ) + if isinstance(signing_key, _BindingRejection): + return signing_key + claims: Final = _decode_id_token(id_token, binding, signing_key) + if isinstance(claims, _BindingRejection): + return claims + if grant_type == "authorization_code" and (binding.mode == "enforce" or expected_nonce is not None): + nonce: Final = claims.get("nonce") + if not expected_nonce or not isinstance(nonce, str) or not hmac.compare_digest(nonce, expected_nonce): + return _BindingRejection( + code="oauth_identity_binding_failed", + description="id_token nonce does not match the authenticated authorization transaction", + ) + upstream: Final = _upstream_principal(claims, binding) + if isinstance(upstream, _BindingRejection): + return upstream + caller: Final = await caller_principal_loader(litellm_user_id, binding) + if not caller: + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"the LiteLLM user has no '{binding.caller_field}' to compare the upstream principal against", + ) + if not _principals_match(upstream, caller, binding): + return _BindingRejection( + code="oauth_principal_mismatch", + description="The browser account does not match the selected credential owner.", + ) + return _binding_proof(binding, litellm_user_id, server_id, caller) + + +async def enforce_oauth_identity_binding( + server: MCPServer, + token_response: Mapping[str, object], + litellm_user_id: str | None, + grant_type: str, + refresh_ownership: RefreshOwnership, + jwks_fetcher: JwksFetcher = _fetch_issuer_jwks, + caller_principal_loader: CallerPrincipalLoader = _load_caller_principal, + stored_refresh_token_loader: StoredRefreshTokenLoader = _load_stored_refresh_token, + expected_nonce: str | None = None, +) -> str | None: + """Validate the exchanged token's upstream principal against the LiteLLM caller. + + No-op when the server has no binding or it is disabled. In enforce mode a failure raises 403 + before the caller returns, stores, or caches the token; in audit mode failures are logged only. + A refresh_token grant without an id_token is allowed only when the presented refresh token matches + the caller's stored credential and that credential was previously identity-validated. + """ + binding: Final = server.oauth_identity_binding + if binding is None or binding.mode not in ("audit", "enforce"): + return + rejection: Final = await _evaluate_binding( + binding=binding, + token_response=token_response, + litellm_user_id=litellm_user_id, + grant_type=grant_type, + server_id=server.server_id, + refresh_ownership=refresh_ownership, + jwks_fetcher=jwks_fetcher, + caller_principal_loader=caller_principal_loader, + stored_refresh_token_loader=stored_refresh_token_loader, + expected_nonce=expected_nonce, + ) + if isinstance(rejection, str): + return rejection if binding.mode == "enforce" else None + if binding.mode == "audit": + verbose_logger.warning( + "oauth_identity_binding audit: server=%s user=%s grant=%s rejected=%s (%s)", + server.server_id, + litellm_user_id, + grant_type, + rejection.code, + rejection.description, + ) + return + raise HTTPException( + status_code=403, + detail={ + "error": rejection.code, + "error_description": rejection.description, + "server_id": server.server_id, + "credential_owner": "caller", + "credential_stored": False, + }, + ) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 16f58ef5b76..d115eb8b3c1 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -163,10 +163,14 @@ def load_openapi_spec(filepath: str) -> dict[str, Any]: return asyncio.run(load_openapi_spec_async(filepath)) -async def load_openapi_spec_async(filepath: str) -> dict[str, Any]: +async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None) -> dict[str, Any]: if filepath.startswith("http://") or filepath.startswith("https://"): client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) - r: Final[httpx.Response] = await async_safe_get(client, filepath) + r: Final[httpx.Response] = ( + await async_safe_get(client, filepath) + if max_bytes is None + else await async_safe_get(client, filepath, max_response_bytes=max_bytes) + ) r.raise_for_status() return r.json() diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index ea2318bd6f1..77979a15199 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -12,6 +12,7 @@ every other mode so the caller defers to v1 (parity-safe); it grows one branch p from __future__ import annotations import base64 +import os from typing import TYPE_CHECKING, Final, Literal, NoReturn from fastapi import HTTPException @@ -310,13 +311,30 @@ def raise_public(error: CredError) -> NoReturn: def oauth_protected_resource_path(root_path: str, server: MCPServer) -> str: """The server's RFC 9728 Protected Resource Metadata path, the shared anchor of both challenges. - ``root_path`` is the proxy's ``SERVER_ROOT_PATH``, resolved by the caller (the imperative shell) - so this stays a pure function of its inputs; ``"/"`` and ``""`` both mean no prefix. The path is - relative, so it resolves against the caller's own host (correct even behind a reverse proxy). + ``root_path`` is the prefix the request was routed under, resolved by the caller (the imperative + shell); ``"/"`` and ``""`` both mean no prefix. The path is relative, so it resolves against the + caller's own host (correct even behind a reverse proxy). + + URL structure depends on how the prefix is served: + + - The scalar ``SERVER_ROOT_PATH`` deployment registers the well-known routes with the prefix + *inserted* into the path (via :func:`well_known_root_suffix` at import time), matching RFC 8414 + §3 well-known path insertion. When ``root_path`` equals ``SERVER_ROOT_PATH`` the URL must use + the same insertion or a client fetching it 404s. + - The per-request ``SERVER_ROOT_PATHS`` deployment can't register routes per prefix (the prefix + set is dynamic and could contain many entries); the middleware strips the prefix from + ``scope["path"]`` and the router matches the un-inserted well-known route. The URL must place + the prefix *before* ``.well-known`` so the strip leaves a matching path. + + Picking the wrong form 404s the client's discovery fetch — the discovery document and the 401 + challenge would then disagree on where the resource metadata lives. """ prefix: Final = "" if root_path == "/" else root_path name: Final = server.alias or server.server_name or server.name or server.server_id - return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}" + scalar_env: Final = os.getenv("SERVER_ROOT_PATH", "").rstrip("/") + if not prefix or (scalar_env and prefix == scalar_env): + return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}" + return f"{prefix}/.well-known/oauth-protected-resource/mcp/{name}" def raise_user_oauth_challenge(server: MCPServer, *, root_path: str) -> NoReturn: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py index 92bd30694af..af1e82eab82 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py @@ -14,10 +14,17 @@ import time from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Final, Protocol +from fastapi import HTTPException + from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( TokenEndpointAuthConfigError, ) +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + BindingValidator, + RefreshTokenPresented, + enforce_oauth_identity_binding, +) from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, @@ -39,6 +46,7 @@ class CredentialPersist(Protocol): refresh_token: str | None, expires_in: int | None, scopes: tuple[str, ...] | None, + identity_binding_proof: str | None = None, ) -> None: ... @@ -78,13 +86,23 @@ class AuthorizationCodeRefresher: persist: CredentialPersist, *, clock: Callable[[], float] = time.time, + identity_validator: BindingValidator = enforce_oauth_identity_binding, ) -> None: self._server_lookup = server_lookup self._token_endpoint = token_endpoint self._persist = persist self._clock = clock + self._identity_validator = identity_validator async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: + try: + return await self._refresh(user_id, server_id, token) + except HTTPException as exc: + if exc.status_code != 403: + raise + return None + + async def _refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: if token.refresh_token is None: return None server: Final = self._server_lookup(server_id) @@ -104,6 +122,15 @@ class AuthorizationCodeRefresher: except TokenEndpointAuthConfigError as exc: verbose_logger.warning("MCP OAuth refresh misconfigured for server %s: %s", server_id, exc) return None + binding: Final = server.oauth_identity_binding + if binding is not None and binding.mode == "enforce": + await self._identity_validator( + server=server, + token_response={}, + litellm_user_id=user_id, + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented(token.refresh_token), + ) form: Final = { "grant_type": "refresh_token", "refresh_token": token.refresh_token, @@ -116,15 +143,34 @@ class AuthorizationCodeRefresher: if not isinstance(access_token, str) or not access_token: return None + binding_proof: Final = await self._identity_validator( + server=server, + token_response=body, + litellm_user_id=user_id, + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented(token.refresh_token), + ) rotated: Final = body.get("refresh_token") new_refresh: Final = rotated if isinstance(rotated, str) and rotated else token.refresh_token expires_in: Final = _parse_expires_in(body.get("expires_in")) scopes: Final = _parse_scopes(body.get("scope")) or token.scopes - await self._persist(user_id, server_id, access_token, new_refresh, expires_in, scopes or None) + if binding_proof is not None: + await self._persist( + user_id, + server_id, + access_token, + new_refresh, + expires_in, + scopes or None, + identity_binding_proof=binding_proof, + ) + else: + await self._persist(user_id, server_id, access_token, new_refresh, expires_in, scopes or None) return OAuthToken( access_token=access_token, expires_at=self._clock() + expires_in if expires_in is not None else None, refresh_token=new_refresh, scopes=scopes, + identity_binding_proof=binding_proof, ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index ad18d1bb10f..43d97abe4db 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -37,6 +37,7 @@ import httpx from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError from typing_extensions import assert_never +from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( InMemoryTokenCacheBackend, OAuthToken, @@ -101,6 +102,11 @@ async def post_client_credentials_grant( from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 # defer heavy handler import to call time get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # handler factory params are coarsely typed ) + from litellm.proxy._experimental.mcp_server.mcp_debug import ( # noqa: PLC0415 # diagnostics import credential enums through this package + describe_upstream_http_failure, + describe_upstream_response, + safe_upstream_url, + ) from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 # deferred with the handler import try: @@ -110,15 +116,28 @@ async def post_client_credentials_grant( ) except httpx.HTTPStatusError as status_err: status_code: Final = status_err.response.status_code + verbose_logger.warning( + "OAuth2 client_credentials token request denied:\n upstream exchange: %s", + describe_upstream_http_failure(status_err), + ) return TokenEndpointDenied(status_code=status_code, detail=f"token endpoint returned HTTP {status_code}") except Exception as exc: # noqa: BLE001 # any transport failure is the same outcome: unreachable - return TokenEndpointUnreachable(detail=str(exc)) + verbose_logger.warning( + "OAuth2 client_credentials POST %s failed: %s", safe_upstream_url(httpx.URL(url)), type(exc).__name__ + ) + return TokenEndpointUnreachable(detail=type(exc).__name__) try: body: Final = _TOKEN_BODY_ADAPTER.validate_json(response.content) except ValidationError: + verbose_logger.warning("OAuth2 client_credentials invalid response: %s", describe_upstream_response(response)) return TokenEndpointDenied( status_code=response.status_code, detail="token endpoint returned a non-JSON-object body" ) + access_token: Final = body.get("access_token") + if not isinstance(access_token, str) or not access_token: + verbose_logger.warning( + "OAuth2 client_credentials response has no access token | %s", describe_upstream_response(response) + ) return TokenEndpointSuccess(body=body) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py index e0cd9e8e5ed..0ac4296f498 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py @@ -41,6 +41,7 @@ class OAuthToken: expires_at: float | None = None refresh_token: str | None = None scopes: tuple[str, ...] = () + identity_binding_proof: str | None = None def __repr__(self) -> str: has_refresh: Final = self.refresh_token is not None diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index 5e28396dcfb..2897c3e8e4a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -12,9 +12,11 @@ from __future__ import annotations import asyncio from collections.abc import Callable, Mapping +from functools import partial from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import credential_binding_matches from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import ( AuthorizationCodeRefresher, ) @@ -69,6 +71,7 @@ async def _persist_credential( refresh_token: str | None, expires_in: int | None, scopes: tuple[str, ...] | None, + identity_binding_proof: str | None = None, ) -> None: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 store_user_oauth_credential, @@ -86,6 +89,7 @@ async def _persist_credential( expires_in=expires_in, scopes=list(scopes) if scopes else None, skip_byok_guard=True, + identity_binding_proof=identity_binding_proof, ) @@ -142,12 +146,26 @@ def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, Refres return backend, coordinator, True +async def _read_bound_credential( + server_lookup: ServerLookup, user_id: str, server_id: str +) -> Mapping[str, object] | None: + credential: Final = await _read_credential(user_id, server_id) + server: Final = server_lookup(server_id) + binding: Final = server.oauth_identity_binding if server else None + if credential is not None and binding is not None and binding.mode == "enforce": + if not await credential_binding_matches(binding, user_id, server_id, credential): + return None + return credential + + def _build_per_user_oauth_token_store( server_lookup: ServerLookup, ) -> tuple[CachedOAuthTokenStore, bool]: backend, coordinator, uses_redis = _runtime_backend_and_coordinator() refresher: Final = AuthorizationCodeRefresher(server_lookup, _post_token_endpoint, _persist_credential) - refreshing: Final = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher, coordinator=coordinator) + refreshing: Final = RefreshingTokenStore( + V2PerUserTokenStore(partial(_read_bound_credential, server_lookup)), refresher, coordinator=coordinator + ) return CachedOAuthTokenStore(refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS, backend=backend), uses_redis @@ -182,6 +200,18 @@ class LazyPerUserOAuthTokenStore: self._local_fetches = 0 async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + token: Final = await self._fetch_token(user_id, server_id) + server: Final = self._server_lookup(server_id) + binding: Final = server.oauth_identity_binding if server else None + if token is not None and binding is not None and binding.mode == "enforce": + if not await credential_binding_matches( + binding, user_id, server_id, {"identity_binding_proof": token.identity_binding_proof} + ): + await self.invalidate(user_id, server_id) + return None + return token + + async def _fetch_token(self, user_id: str, server_id: str) -> OAuthToken | None: if self._uses_redis: store = self._store if store is not None: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 8328aae01ab..85c7f68719d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -65,6 +65,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, + AuthResolution, AuthSpecKind, AwsSigV4Config, Byok, @@ -76,6 +77,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, + ResolvedCredential, ServerSpec, SharedKey, Subject, @@ -448,3 +450,32 @@ def _client_auth_fingerprint(client_auth: ClientAuth) -> str: def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet")) + + +async def resolve_credentials_with_source( + provider: UpstreamCredentialProvider, subject: Subject, server: ServerSpec +) -> Result[ResolvedCredential, CredError]: + match await provider.resolve_credentials(subject, server): + case Error(err): + return Error(err) + case Ok(auth): + if isinstance(auth, NoOpAuth): + return Ok(ResolvedCredential(auth, AuthResolution.no_auth)) + match server.config: + case NoneConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.no_auth)) + case ApiKeyConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.static_token)) + case PassthroughConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.oauth2_passthrough)) + case ClientCredentialsConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.client_credentials)) + case TokenExchangeConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.token_exchange)) + case IdJagConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.id_jag)) + case AuthorizationCodeConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.stored_user_token)) + case AwsSigV4Config(): + return Ok(ResolvedCredential(auth, AuthResolution.aws_sigv4)) + assert_never(server.config) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py index 14cc309b685..56dc3f19cdd 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py @@ -1,24 +1,26 @@ """Serialize + encrypt boundary for caching an OAuth token in a shared (Redis) cache. -A cross-replica cache must serialize the token, and a plaintext bearer in Redis is a leak, so this -encrypts the value (NaCl in production via the injected ``encrypt``, identity in tests). It caches -**only** the ``access_token``: the hot path needs just the bearer, expiry is carried by the cache -entry's TTL (set from the token's ``expires_at`` by the cache), and the long-lived refresh_token stays -in the DB - the refresh path is always a cache miss that re-reads it - so it never reaches Redis. A -decoded token therefore carries only the bearer (``expires_at`` and ``refresh_token`` both None); the -TTL, not the value, bounds its life. An empty/undecryptable blob (e.g. master-key rotation) is a miss. +Shared cache values contain an encrypted access token and optional identity-binding proof. +Refresh tokens remain in the database; cache TTL bounds the access token's lifetime. +Legacy bearer-only entries decode without proof and cannot satisfy identity enforcement. """ from __future__ import annotations +import json from collections.abc import Callable from dataclasses import dataclass from typing import Final +from pydantic import TypeAdapter, ValidationError + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, ) +_BOUND_PREFIX: Final = "litellm-bound-oauth-v1:" +_BOUND_PAYLOAD: Final = TypeAdapter(dict[str, str]) + @dataclass(frozen=True, slots=True) class OAuthTokenCacheCodec: @@ -26,10 +28,30 @@ class OAuthTokenCacheCodec: decrypt: Callable[[str], str | None] def encode(self, token: OAuthToken) -> str: + if token.identity_binding_proof is not None: + return self.encrypt( + _BOUND_PREFIX + + json.dumps( + { + "access_token": token.access_token, + "identity_binding_proof": token.identity_binding_proof, + } + ) + ) return self.encrypt(token.access_token) def decode(self, blob: str) -> OAuthToken | None: access_token: Final = self.decrypt(blob) if not access_token: return None + if access_token.startswith(_BOUND_PREFIX): + try: + payload: Final = _BOUND_PAYLOAD.validate_json(access_token[len(_BOUND_PREFIX) :]) + except ValidationError: + return None + bearer: Final = payload.get("access_token") + proof: Final = payload.get("identity_binding_proof") + if not bearer or not proof: + return None + return OAuthToken(access_token=bearer, identity_binding_proof=proof) return OAuthToken(access_token=access_token, refresh_token=None) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 632dc57dcf6..d186724fd9f 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -26,10 +26,11 @@ union (see `result.py`), not `expression.Result`. from __future__ import annotations from collections.abc import Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum from typing import Annotated, Final, Literal +import httpx from expression import case, tag, tagged_union from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator from typing_extensions import assert_never @@ -46,6 +47,29 @@ from litellm.types.mcp import ( ) +class AuthResolution(str, Enum): + no_auth = "no-auth" + stored_user_token = "stored-user-token" + static_token = "static-token" + per_request_header = "per-request-header" + oauth2_passthrough = "oauth2-passthrough" + client_credentials = "m2m-client-credentials" + token_exchange = "token-exchange" + id_jag = "id-jag" + aws_sigv4 = "aws-sigv4" + extra_headers = "extra-headers" + not_applicable = "not-applicable" + unresolved = "unresolved" + failed = "resolution-failed" + multiple = "multiple" + + +@dataclass(frozen=True, slots=True) +class ResolvedCredential: + auth: httpx.Auth = field(repr=False) + source: AuthResolution + + class AuthSpecKind(str, Enum): """The server's statically-declared upstream-auth mode — derived from its `config`. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py index 4d88ec8b025..0f18931d118 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py @@ -46,11 +46,13 @@ def _to_oauth_token(payload: Mapping[str, object]) -> OAuthToken | None: return None refresh_token: Final = payload.get("refresh_token") expires_at: Final = payload.get("expires_at") + binding_proof: Final = payload.get("identity_binding_proof") return OAuthToken( access_token=access_token, expires_at=_iso_to_epoch(expires_at) if isinstance(expires_at, str) else None, refresh_token=refresh_token if isinstance(refresh_token, str) else None, scopes=_to_scopes(payload.get("scopes")), + identity_binding_proof=binding_proof if isinstance(binding_proof, str) else None, ) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 5fbfad54a39..7a97e995570 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,13 +1,18 @@ import asyncio import importlib from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime +from traceback import walk_tb from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal +from uuid import uuid4 import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status +from pydantic import ValidationError +from starlette.datastructures import Headers from litellm._logging import verbose_logger from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT @@ -28,6 +33,8 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( list_fault_http_status, outcome_wire_value, ) +from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree +from litellm.proxy._experimental.mcp_server.oauth_utils import _redact_mcp_resource_url from litellm.proxy._experimental.mcp_server.ui_session_utils import ( acting_user_auth, build_effective_auth_contexts, @@ -76,11 +83,39 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + reference: Final = uuid4().hex + verbose_logger.error( + "MCP connection test failed (reference=%s): %s", + reference, + tuple( + ( + type(cause).__name__, + tuple( + (frame.f_code.co_filename, lineno, frame.f_code.co_name) + for frame, lineno in walk_tb(cause.__traceback__) + ), + ) + for cause in iter_exception_tree(exc) + ), + ) + return next( + ( + message + for cause in iter_exception_tree(exc) + if (message := _known_connection_error_message(cause, url, timeout_seconds)) is not None + ), + "An unexpected error occurred while testing the MCP connection. " + f"Retry; if it persists, share reference {reference} with your gateway administrator.", + ) + + +def _known_connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str | None: if isinstance(exc, MCPServerURLCredentialsError): return str(exc.detail) if isinstance(exc, TimeoutError): return ( - f"Failed to connect to MCP server: no response from {url or 'the server'} " + "Failed to connect to MCP server: no valid MCP response received from " + f"{_redact_mcp_resource_url(url) or 'the server'} " f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." ) @@ -97,13 +132,48 @@ def _connection_error_message(exc: BaseException, url: str | None, timeout_secon return "Failed to connect to MCP server: the connection timed out." if isinstance(exc, httpx.HTTPStatusError): return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}." - return "Failed to connect to MCP server. Check proxy logs for details." + if isinstance(exc, (httpx.NetworkError, httpx.RemoteProtocolError, ConnectionError)): + return ( + "Failed to connect to MCP server: the connection was interrupted. " + "Check the server and network connection, then retry." + ) + if isinstance(exc, ValueError) and str(exc).startswith("Unexpected content type:"): + return ( + "Failed to connect to MCP server: the endpoint returned an unsupported content type. " + "Check that the URL is an MCP endpoint, not a web page, and matches the selected transport." + ) + if isinstance(exc, ValidationError) and exc.title in ("JSONRPCMessage", "InitializeResult", "ListToolsResult"): + return ( + "Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response. " + "Check the MCP endpoint URL and the server's protocol implementation." + ) + if MCP_AVAILABLE and isinstance(exc, McpError): + if exc.error.code == -32000 and exc.error.message == "Connection closed": + return ( + "Failed to connect to MCP server: the connection was closed before the request completed. " + "Check that the server stays running and returns a complete MCP response, then retry." + ) + if exc.error.code == 32600 and exc.error.message == "Session terminated": + return ( + "Failed to connect to MCP server: the MCP session was terminated. " + "Check that the URL points to an MCP endpoint and matches the selected transport, " + "then retry to start a new session." + ) + return ( + f"Failed to connect to MCP server: the MCP request failed (JSON-RPC code {exc.error.code}). " + "Check that the endpoint supports MCP initialization and tool listing, and check the upstream server logs." + ) + return None if MCP_AVAILABLE: + from mcp.shared.exceptions import McpError from mcp.types import Tool as MCPTool - from litellm.experimental_mcp_client.client import MCPClient + from litellm.experimental_mcp_client.client import MCPClient, as_mcp_read_timeout + from litellm.llms.litellm_proxy.skills.skill_search import ( + DEFAULT_SKILL_SEARCH_TOP_K, + ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, global_mcp_server_manager, @@ -121,6 +191,7 @@ if MCP_AVAILABLE: execute_mcp_tool, filter_tools_by_allowed_tools, filter_tools_by_key_team_permissions, + fire_mcp_tool_call_failure_logging, ) ######################################################## @@ -162,6 +233,20 @@ if MCP_AVAILABLE: return result return outcome + async def _safe_fire_mcp_tool_call_failure_logging( + logging_obj: "LiteLLMLoggingObj | None", + exception: Exception, + start_time: datetime, + user_api_key_auth: UserAPIKeyAuth, + request_data: Mapping[str, object], + ) -> None: + try: + await fire_mcp_tool_call_failure_logging( + logging_obj, exception, start_time, user_api_key_auth, request_data + ) + except Exception as logging_error: + verbose_logger.warning("MCP tool call failure logging failed (continuing): %s", logging_error) + def _relay_upstream_auth_http_exception(e: MCPUpstreamAuthError, request: Request) -> HTTPException: """Convert a client-forwarded pass-through upstream 401 into an HTTPException that preserves the upstream WWW-Authenticate, so a standards-compliant MCP client can run the upstream OAuth flow @@ -188,10 +273,12 @@ if MCP_AVAILABLE: AGENT_SEARCH_TOOL_NAME, DEFAULT_AGENT_SEARCH_TOP_K, MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, coerce_top_k, handle_agent_search, handle_mcp_tool_call, handle_mcp_tool_search, + handle_skill_search, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj @@ -210,6 +297,14 @@ if MCP_AVAILABLE: ), user_api_key_dict=user_api_key_dict, ) + if tool_name == SKILL_SEARCH_TOOL_NAME: + return await handle_skill_search( + query=str(tool_arguments.get("query", "")), + top_k=coerce_top_k( + tool_arguments.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K + ), + user_api_key_dict=user_api_key_dict, + ) rest_client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) ( virtual_mcp_auth_header, @@ -230,26 +325,39 @@ if MCP_AVAILABLE: ) # MCP_TOOL_CALL_TOOL_NAME: run the same pre-call pipeline as the normal path so the tool # execution is spend-logged and guardrail-checked. - (_, virtual_logging_obj) = await ProxyBaseLLMRequestProcessing(data=data).common_processing_pre_call_logic( - request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, - route_type=CallTypes.call_mcp_tool.value, - proxy_logging_obj=proxy_logging_obj, - general_settings=general_settings, - ) - _tool_start_time: Final = datetime.now() - result: Final = await handle_mcp_tool_call( - tool_name=tool_arguments.get("tool_name", ""), - arguments=tool_arguments.get("arguments") or {}, - user_api_key_dict=user_api_key_dict, - client_ip=rest_client_ip, - mcp_auth_header=virtual_mcp_auth_header, - mcp_server_auth_headers=virtual_mcp_server_auth_headers, - oauth2_headers=virtual_oauth2_headers, - raw_headers=virtual_raw_headers, - litellm_logging_obj=virtual_logging_obj, - ) + virtual_processor: Final = ProxyBaseLLMRequestProcessing(data=data) + _request_start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the tool start time below + try: + (_, virtual_logging_obj) = await virtual_processor.common_processing_pre_call_logic( + request=request, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + route_type=CallTypes.call_mcp_tool.value, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + ) + _tool_start_time: Final = datetime.now() + result: Final = await handle_mcp_tool_call( + tool_name=tool_arguments.get("tool_name", ""), + arguments=tool_arguments.get("arguments") or {}, + user_api_key_dict=user_api_key_dict, + client_ip=rest_client_ip, + mcp_auth_header=virtual_mcp_auth_header, + mcp_server_auth_headers=virtual_mcp_server_auth_headers, + oauth2_headers=virtual_oauth2_headers, + raw_headers=virtual_raw_headers, + litellm_logging_obj=virtual_logging_obj, + ) + except Exception as e: + virtual_request_data: Final = virtual_processor.data + await _safe_fire_mcp_tool_call_failure_logging( + virtual_request_data.get("litellm_logging_obj"), + e, + _request_start_time, + user_api_key_dict, + virtual_request_data, + ) + raise return await _safe_fire_mcp_tool_call_logging( virtual_logging_obj, result, @@ -885,7 +993,9 @@ if MCP_AVAILABLE: apply_tool_filters=apply_tool_filters, ) except Exception as e: - verbose_logger.exception("Error getting tools from %s: %s", server.name, e) + verbose_logger.warning( + "Error getting tools from %s: %s", server.name, classify_list_exception(e).tag + ) return (), classify_list_exception(e) return tools_result, ServerListOk(tool_count=len(tools_result)) @@ -999,65 +1109,73 @@ if MCP_AVAILABLE: ) proxy_base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) - ( - data, - logging_obj, - ) = await proxy_base_llm_response_processor.common_processing_pre_call_logic( - request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, - route_type=CallTypes.call_mcp_tool.value, - proxy_logging_obj=proxy_logging_obj, - general_settings=general_settings, - ) + _request_start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the tool start time below + try: + ( + data, + logging_obj, + ) = await proxy_base_llm_response_processor.common_processing_pre_call_logic( + request=request, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + route_type=CallTypes.call_mcp_tool.value, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + ) - # Extract MCP auth headers from request and add to data dict - ( - mcp_auth_header, - mcp_server_auth_headers, - raw_headers_from_request, - ) = _extract_mcp_headers_from_request(request, MCPRequestHandler) - if mcp_auth_header: - data["mcp_auth_header"] = mcp_auth_header - if mcp_server_auth_headers: - data["mcp_server_auth_headers"] = mcp_server_auth_headers - data["raw_headers"] = raw_headers_from_request + # Extract MCP auth headers from request and add to data dict + ( + mcp_auth_header, + mcp_server_auth_headers, + raw_headers_from_request, + ) = _extract_mcp_headers_from_request(request, MCPRequestHandler) + if mcp_auth_header: + data["mcp_auth_header"] = mcp_auth_header + if mcp_server_auth_headers: + data["mcp_server_auth_headers"] = mcp_server_auth_headers + data["raw_headers"] = raw_headers_from_request - # Extract user_api_key_auth from metadata and add to top level - # call_mcp_tool expects user_api_key_auth as a top-level parameter - if "metadata" in data and "user_api_key_auth" in data["metadata"]: - data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"] + # Extract user_api_key_auth from metadata and add to top level + # call_mcp_tool expects user_api_key_auth as a top-level parameter + if "metadata" in data and "user_api_key_auth" in data["metadata"]: + data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"] - # Resolve allowed MCP servers with IP filtering - ( - allowed_mcp_servers, - canonical_server_id, - ) = await _resolve_allowed_mcp_servers_with_ip_filter(request, user_api_key_dict, server_id) + # Resolve allowed MCP servers with IP filtering + ( + allowed_mcp_servers, + canonical_server_id, + ) = await _resolve_allowed_mcp_servers_with_ip_filter(request, user_api_key_dict, server_id) - # Look up per-user OAuth headers for this server (mirrors list_tool_rest_api). - user_oauth_extra_headers: dict[str, str] | None = None - target_server: Final = next( - (s for s in allowed_mcp_servers if s.server_id == canonical_server_id), - None, - ) - if target_server is not None: - user_oauth_extra_headers = await _get_user_oauth_extra_headers(target_server, user_api_key_dict) + # Look up per-user OAuth headers for this server (mirrors list_tool_rest_api). + user_oauth_extra_headers: dict[str, str] | None = None + target_server: Final = next( + (s for s in allowed_mcp_servers if s.server_id == canonical_server_id), + None, + ) + if target_server is not None: + user_oauth_extra_headers = await _get_user_oauth_extra_headers(target_server, user_api_key_dict) - # Call execute_mcp_tool directly (permission checks already done) - _tool_start_time: Final = datetime.now() - result: Final = await execute_mcp_tool( - name=tool_name, - arguments=tool_arguments, - allowed_mcp_servers=allowed_mcp_servers, - start_time=_tool_start_time, - user_api_key_auth=data.get("user_api_key_auth"), - mcp_auth_header=data.get("mcp_auth_header"), - mcp_server_auth_headers=data.get("mcp_server_auth_headers"), - oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"), - raw_headers=data.get("raw_headers"), - litellm_logging_obj=data.get("litellm_logging_obj"), - requested_server_id=canonical_server_id, - ) + # Call execute_mcp_tool directly (permission checks already done) + _tool_start_time: Final = datetime.now() + result: Final = await execute_mcp_tool( + name=tool_name, + arguments=tool_arguments, + allowed_mcp_servers=allowed_mcp_servers, + start_time=_tool_start_time, + user_api_key_auth=data.get("user_api_key_auth"), + mcp_auth_header=data.get("mcp_auth_header"), + mcp_server_auth_headers=data.get("mcp_server_auth_headers"), + oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"), + raw_headers=data.get("raw_headers"), + litellm_logging_obj=data.get("litellm_logging_obj"), + requested_server_id=canonical_server_id, + ) + except Exception as e: + request_data: Final = proxy_base_llm_response_processor.data + await _safe_fire_mcp_tool_call_failure_logging( + request_data.get("litellm_logging_obj"), e, _request_start_time, user_api_key_dict, request_data + ) + raise return await _safe_fire_mcp_tool_call_logging( logging_obj, result, @@ -1153,6 +1271,68 @@ if MCP_AVAILABLE: scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None return client_id, client_secret, scopes + _STAGED_AUTH_VALUE_AUTH_TYPES: Final = frozenset( + (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization, MCPAuth.token) + ) + + @dataclass(frozen=True, slots=True) + class _StagedServerTest: + request: NewMCPServerRequest + mcp_auth_header: str | None + oauth2_headers: dict[str, str] | None + + def _preview_origin(url: str | None) -> tuple[str, str, int | None] | None: + if not url: + return None + try: + parsed: Final = httpx.URL(url) + except httpx.InvalidURL: + return None + if parsed.scheme not in ("http", "https") or not parsed.host: + return None + return parsed.scheme, parsed.host, parsed.port + + def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest: + """ + Resolve the credentials a not-yet-saved server config carries for a preview call. + + Both preview endpoints (``/test/connection`` and ``/test/tools/list``) must hand the + temporary client the same credentials, or a server that the saved connection reaches + fine fails one of them. + """ + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + saved_server: Final = ( + global_mcp_server_manager.get_mcp_server_by_id(new_mcp_server_request.server_id) + if new_mcp_server_request.server_id + else None + ) + saved_origin: Final = _preview_origin(saved_server.url) if saved_server else None + preview_origin: Final = _preview_origin(new_mcp_server_request.url) + may_inherit: Final = new_mcp_server_request.auth_type not in _STAGED_AUTH_VALUE_AUTH_TYPES or ( + saved_origin is not None and saved_origin == preview_origin + ) + request: Final = ( + _inherit_credentials_from_existing_server(new_mcp_server_request) if may_inherit else new_mcp_server_request + ) + mcp_auth_header: Final = ( + request.credentials.get("auth_value") + if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict) + else None + ) + # Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY): + # when the primary x-litellm-api-key header is absent, the Authorization value is the + # caller's LiteLLM key, not an upstream token, and must never be forwarded upstream. + oauth2_headers: Final = ( + MCPRequestHandler._get_oauth2_headers_from_headers(headers) + if request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + and headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY) + else None + ) + return _StagedServerTest(request=request, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers) + async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None: with anyio.move_on_after(deadline): return await client.list_tools(raise_on_error=True) @@ -1199,8 +1379,15 @@ if MCP_AVAILABLE: if _oauth2_flow == "client_credentials" and not request.token_url: _oauth2_flow = None + # Static previews inherit credentials before this step, but must not resolve back to + # the saved record during client creation and discard the edited connection settings. + preview_server_id: Final = ( + "" + if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES or request.auth_type in (None, MCPAuth.none) + else request.server_id or "" + ) server_model: Final = MCPServer( - server_id=request.server_id or "", + server_id=preview_server_id, name=request.alias or request.server_name or "", url=request.url, transport=request.transport, @@ -1288,11 +1475,18 @@ if MCP_AVAILABLE: except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise except BaseException as e: - verbose_logger.error("Error in MCP operation: %s", e, exc_info=True) + effective_timeout: Final = ( + min(request.timeout if request.timeout is not None else MCP_CLIENT_TIMEOUT, timeout_seconds) + if any( + isinstance(cause, McpError) and as_mcp_read_timeout(cause) is not None + for cause in iter_exception_tree(e) + ) + else timeout_seconds + ) return { "status": "error", "error": True, - "message": _connection_error_message(e, request.url, timeout_seconds), + "message": _connection_error_message(e, request.url, effective_timeout), } async def _preview_openapi_tools(spec_path: str) -> dict: @@ -1374,6 +1568,8 @@ if MCP_AVAILABLE: }, ) + staged: Final = _stage_server_test(new_mcp_server_request, request.headers) + async def _test_connection_operation(client): async def _noop(session): return "ok" @@ -1382,8 +1578,10 @@ if MCP_AVAILABLE: return {"status": "ok"} return await _execute_with_mcp_client( - new_mcp_server_request, + staged.request, _test_connection_operation, + mcp_auth_header=staged.mcp_auth_header, + oauth2_headers=staged.oauth2_headers, raw_headers=_safe_get_request_headers(request), ) @@ -1404,37 +1602,11 @@ if MCP_AVAILABLE: }, ) - new_mcp_server_request = _inherit_credentials_from_existing_server(new_mcp_server_request) + staged: Final = _stage_server_test(new_mcp_server_request, request.headers) # For OpenAPI spec servers, generate tools from the spec directly - if new_mcp_server_request.spec_path: - return await _preview_openapi_tools(new_mcp_server_request.spec_path) - - from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCPRequestHandler, - ) - - headers: Final = request.headers - - mcp_auth_header: str | None = None - if new_mcp_server_request.auth_type in { - MCPAuth.api_key, - MCPAuth.bearer_token, - MCPAuth.basic, - MCPAuth.authorization, - }: - credentials: Final = getattr(new_mcp_server_request, "credentials", None) - if isinstance(credentials, dict): - mcp_auth_header = credentials.get("auth_value") - - # Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY): - # when the primary x-litellm-api-key header is absent, the Authorization value is the - # caller's LiteLLM key, not an upstream token, and must never be forwarded upstream. - oauth2_headers: dict[str, str] | None = None - if new_mcp_server_request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and headers.get( - MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY - ): - oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) + if staged.request.spec_path: + return await _preview_openapi_tools(staged.request.spec_path) async def _list_tools_operation(client): # Bound the whole pagination walk: without this the preview is limited only by the @@ -1465,9 +1637,9 @@ if MCP_AVAILABLE: } return await _execute_with_mcp_client( - new_mcp_server_request, + staged.request, _list_tools_operation, - mcp_auth_header=mcp_auth_header, - oauth2_headers=oauth2_headers, + mcp_auth_header=staged.mcp_auth_header, + oauth2_headers=staged.oauth2_headers, raw_headers=_safe_get_request_headers(request), ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 26f5d6e7c8c..43a7f894db8 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -15,14 +15,15 @@ import types import uuid from collections.abc import AsyncIterator, Callable, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol import httpx from fastapi import FastAPI, HTTPException -from pydantic import AnyUrl, ConfigDict +from pydantic import AnyUrl, ConfigDict, TypeAdapter, ValidationError from starlette.requests import Request as StarletteRequest from starlette.responses import JSONResponse from starlette.types import Message, Receive, Scope, Send +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG @@ -46,8 +47,13 @@ from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_active_toolset_id, _mcp_gateway_initialize_instructions, _mcp_gateway_server_name, + _mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode +) +from litellm.proxy._experimental.mcp_server.mcp_debug import ( + MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, + MCPAuthDiagnostics, + MCPDebug, ) -from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, get_passthrough_www_authenticate, @@ -107,9 +113,9 @@ _MAX_STATEFUL_SESSIONS_PER_OWNER: Final = 100 # prevents an authenticated client from forcing the proxy to buffer an # arbitrarily large body just to make a routing decision. _MCP_ROUTING_PEEK_MAX_BYTES: Final = 4096 -# ASGI scope key holding the tracing span of the request carrying an MCP -# message, written on the request task and read back by the message handler. +# ASGI scope keys carrying OTel request state into a stateful MCP message handler. _MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span" +_MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations" def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: @@ -327,18 +333,17 @@ def _otel_publish_transport_span_on_scope(scope: Scope) -> None: scope[_MCP_TRANSPORT_SPAN_SCOPE_KEY] = span -def _otel_transport_span_from_message(req_ctx: object) -> object: - """The tracing span of the HTTP request that carried this MCP message. - - Read off that request's ASGI scope, reached through the ``Request`` the - streamable-HTTP transport attaches to each message, so it is this message's - transport and not whichever request happens to have touched the session last. - Returns whatever the scope holds; the otel plumbing validates it.""" +def _otel_value_from_message_scope(req_ctx: object, key: str) -> object: request: Final = getattr(req_ctx, "request", None) scope: Final = getattr(request, "scope", None) if not isinstance(scope, Mapping): return None - return scope.get(_MCP_TRANSPORT_SPAN_SCOPE_KEY) + return scope.get(key) + + +def _otel_transport_span_from_message(req_ctx: object) -> object: + """The tracing span of the HTTP request that carried this MCP message.""" + return _otel_value_from_message_scope(req_ctx, _MCP_TRANSPORT_SPAN_SCOPE_KEY) def _otel_set_mcp_transport_span(span: object) -> object: @@ -371,6 +376,44 @@ def _otel_reset_mcp_transport_span(token: object) -> None: return +def _otel_publish_request_destinations_on_scope(scope: Scope) -> None: + try: + from litellm.integrations.otel.plumbing.context import request_destinations + + scope[_MCP_DESTINATIONS_SCOPE_KEY] = request_destinations() + except ImportError: + return + + +def _otel_set_mcp_request_destinations(req_ctx: object) -> object: + destinations: Final = _otel_value_from_message_scope(req_ctx, _MCP_DESTINATIONS_SCOPE_KEY) + if not isinstance(destinations, tuple): + return None + try: + from litellm.integrations.otel.model.destination import OtelDestination + from litellm.integrations.otel.plumbing.context import set_request_destinations + + destination_adapter: Final[TypeAdapter[tuple[OtelDestination, ...]]] = TypeAdapter( + tuple[OtelDestination, ...], + config=ConfigDict(revalidate_instances="always"), + ) + validated_destinations: Final = destination_adapter.validate_python(destinations, strict=True) + return set_request_destinations(validated_destinations) + except (ImportError, ValidationError): + return None + + +def _otel_reset_mcp_request_destinations(token: object) -> None: + if token is None: + return + try: + from litellm.integrations.otel.plumbing.context import reset_request_destinations + + reset_request_destinations(token) + except ImportError: + return + + def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: """Map a ``ProxyException`` to an ``HTTPException`` that preserves its real status code and headers. @@ -499,11 +542,22 @@ if MCP_AVAILABLE: notification_options: NotificationOptions | None = None, experimental_capabilities: dict[str, dict[str, object]] | None = None, ) -> InitializationOptions: - opts: Final = Server.create_initialization_options( + base_options: Final = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) + opts: Final = ( + base_options.model_copy( + update={ # mutable-ok: Pydantic update payload + "capabilities": base_options.capabilities.model_copy( + update={"prompts": None, "resources": None} # mutable-ok: Pydantic update payload + ) + } + ) + if _mcp_proxy_mode.get() + else base_options + ) updates: Final[dict[str, str]] = {} merged: Final = _mcp_gateway_initialize_instructions.get() if merged is not None: @@ -717,12 +771,12 @@ if MCP_AVAILABLE: _stateful_auth_context_cleanup_task.cancel() with contextlib.suppress(asyncio.CancelledError): await _stateful_auth_context_cleanup_task - if _session_manager_cm: - await _session_manager_cm.__aexit__(None, None, None) - if _session_manager_stateful_cm: - await _session_manager_stateful_cm.__aexit__(None, None, None) if _sse_session_manager_cm: await _sse_session_manager_cm.__aexit__(None, None, None) + if _session_manager_stateful_cm: + await _session_manager_stateful_cm.__aexit__(None, None, None) + if _session_manager_cm: + await _session_manager_cm.__aexit__(None, None, None) except Exception as e: verbose_logger.exception("Error during session manager shutdown: %s", e) @@ -762,10 +816,12 @@ if MCP_AVAILABLE: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None _transport_token = None + _destinations_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) _transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx)) + _destinations_token = _otel_set_mcp_request_destinations(req_ctx) # Get user authentication from context variable ( user_api_key_auth, @@ -782,17 +838,20 @@ if MCP_AVAILABLE: "MCP list_tools - MCP server auth headers: %s", list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, ) + from mcp.types import Tool + + from litellm.proxy._experimental.mcp_server.tool_search import ( + get_mcp_proxy_tool_definitions, + get_virtual_tool_definitions, + ) + + if _mcp_proxy_mode.get(): + return [Tool.model_validate(d) for d in get_mcp_proxy_tool_definitions()] # mutable-ok: MCP SDK list if getattr( getattr(user_api_key_auth, "object_permission", None), "mcp_tool_search_enabled", False, ): - from mcp.types import Tool - - from litellm.proxy._experimental.mcp_server.tool_search import ( - get_virtual_tool_definitions, - ) - return [Tool.model_validate(d) for d in get_virtual_tool_definitions()] # Get mcp_servers from context variable @@ -816,12 +875,18 @@ if MCP_AVAILABLE: } } return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta}) + except HTTPException as e: + from mcp.shared.exceptions import McpError + from mcp.types import INVALID_REQUEST, ErrorData + + raise McpError(ErrorData(code=INVALID_REQUEST, message=_http_detail_message(e.detail))) from e except Exception as e: verbose_logger.exception("Error in list_tools endpoint: %s", e) # Return empty list instead of failing completely # This prevents the HTTP stream from failing and allows the client to get a response return [] finally: + _otel_reset_mcp_request_destinations(_destinations_token) _otel_reset_mcp_transport_span(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: @@ -860,6 +925,12 @@ if MCP_AVAILABLE: verbose_logger.debug("Host progressToken captured: %s...", str(host_token)[:8]) return forward_progress + def _reject_mcp_proxy_operation() -> NoReturn: + from mcp.shared.exceptions import McpError + from mcp.types import METHOD_NOT_FOUND, ErrorData + + raise McpError(ErrorData(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy")) + async def _build_virtual_call_logging_obj( name: str, arguments: dict[str, object], @@ -911,17 +982,95 @@ if MCP_AVAILABLE: Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so the caller falls through to normal tool routing. """ + from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K from litellm.proxy._experimental.mcp_server.tool_search import ( AGENT_SEARCH_TOOL_NAME, DEFAULT_AGENT_SEARCH_TOP_K, + MCP_PROXY_CALL_TOOL_NAME, + MCP_PROXY_TOOL_NAMES, MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, VIRTUAL_TOOL_NAMES, coerce_top_k, handle_agent_search, + handle_mcp_proxy_tool, handle_mcp_tool_call, handle_mcp_tool_search, + handle_skill_search, ) + if _mcp_proxy_mode.get() and name not in MCP_PROXY_TOOL_NAMES: + return CallToolResult( + content=[ # mutable-ok: MCP result content + TextContent(type="text", text=f"Tool {name} is unavailable on /mcp/proxy") + ], + isError=True, + ) + + if _mcp_proxy_mode.get() and name in MCP_PROXY_TOOL_NAMES: + assert user_api_key_auth is not None + proxy_call_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes + proxy_logging_obj: Final = ( + await _build_virtual_call_logging_obj( + name=name, + arguments=arguments or {}, # mutable-ok: logging pipeline payload + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + if name == MCP_PROXY_CALL_TOOL_NAME + else None + ) + try: + proxy_result: Final = await handle_mcp_proxy_tool( + name=name, + arguments=arguments or {}, # mutable-ok: proxy handler payload + user_api_key_dict=user_api_key_auth, + client_ip=client_ip, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + litellm_logging_obj=proxy_logging_obj, + ) + except Exception as exc: + if proxy_logging_obj is not None: + from litellm.proxy.proxy_server import proxy_logging_obj as request_logging_obj + + failure_end: Final = datetime.now() # noqa: DTZ005 # matches the logging pipeline start time + failure_traceback: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + try: + proxy_logging_obj.failure_handler(exc, failure_traceback, proxy_call_start, failure_end) + await proxy_logging_obj.async_failure_handler( + exc, failure_traceback, proxy_call_start, failure_end + ) + if not isinstance(exc, MCPUpstreamAuthError): + await request_logging_obj.post_call_failure_hook( + request_data={ # mutable-ok: failure hook mutates its request payload + "name": name, + "arguments": arguments, + "litellm_logging_obj": proxy_logging_obj, + }, + original_exception=exc, + user_api_key_dict=user_api_key_auth, + route="/mcp/call_tool", + traceback_str=failure_traceback, + ) + except Exception: # noqa: BLE001 # a failing failure hook must not mask the tool call's own error + verbose_logger.exception("Error logging failed MCP proxy tool call") + raise + if proxy_logging_obj is not None: + return await _fire_mcp_tool_call_logging( + logging_obj=proxy_logging_obj, + result=proxy_result, + start_time=proxy_call_start, + end_time=datetime.now(), # noqa: DTZ005 # matches the logging pipeline start time + user_api_key_auth=user_api_key_auth, + request_data=types.MappingProxyType({"name": name, "arguments": arguments}), + ) + return proxy_result + if name not in VIRTUAL_TOOL_NAMES: return None @@ -961,6 +1110,12 @@ if MCP_AVAILABLE: top_k=coerce_top_k(args.get("top_k", DEFAULT_AGENT_SEARCH_TOP_K), default=DEFAULT_AGENT_SEARCH_TOP_K), user_api_key_dict=user_api_key_auth, ) + if name == SKILL_SEARCH_TOOL_NAME: + return await handle_skill_search( + query=str(args.get("query", "")), + top_k=coerce_top_k(args.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K), + user_api_key_dict=user_api_key_auth, + ) virtual_logging_obj: Final = await _build_virtual_call_logging_obj( name=name, arguments=args, @@ -1006,10 +1161,12 @@ if MCP_AVAILABLE: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None _transport_token = None + _destinations_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) _transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx)) + _destinations_token = _otel_set_mcp_request_destinations(req_ctx) # Validate arguments ( user_api_key_auth, @@ -1086,6 +1243,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, + client_ip=_client_ip, host_progress_callback=host_progress_callback, **data, # for logging ) @@ -1119,7 +1277,7 @@ if MCP_AVAILABLE: except HTTPException as e: verbose_logger.error("HTTPException in MCP tool call: %s", e) return CallToolResult( - content=[TextContent(text=f"Error: {e.detail}", type="text")], + content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")], isError=True, ) except MCPUpstreamAuthError as e: @@ -1147,6 +1305,7 @@ if MCP_AVAILABLE: return response finally: + _otel_reset_mcp_request_destinations(_destinations_token) _otel_reset_mcp_transport_span(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: @@ -1157,6 +1316,8 @@ if MCP_AVAILABLE: """ List all available prompts """ + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1214,8 +1375,8 @@ if MCP_AVAILABLE: Returns: GetPromptResult: Getting prompt execution results """ - - # Validate arguments + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1252,6 +1413,8 @@ if MCP_AVAILABLE: @server.list_resources() async def list_resources() -> list[Resource]: """List all available resources.""" + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1296,6 +1459,8 @@ if MCP_AVAILABLE: @server.list_resource_templates() async def list_resource_templates() -> list[ResourceTemplate]: """List all available resource templates.""" + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1341,6 +1506,8 @@ if MCP_AVAILABLE: @server.read_resource() async def read_resource(url: AnyUrl) -> list[ReadResourceContents]: + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1383,7 +1550,7 @@ if MCP_AVAILABLE: ######################################################## async def _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers: list[str] | None, + mcp_servers: Sequence[str] | None, allowed_mcp_servers: list[MCPServer], ) -> list[MCPServer]: """ @@ -1404,13 +1571,10 @@ if MCP_AVAILABLE: server_name_matched = False for server in allowed_mcp_servers: - if server: - match_list = [s.lower() for s in iter_known_server_prefixes(server) if s] - - if server_or_group.lower() in match_list: - filtered_server[server.server_id] = server - server_name_matched = True - break + if server and _server_answers_to(server, server_or_group): + filtered_server[server.server_id] = server + server_name_matched = True + break if not server_name_matched: try: @@ -1440,6 +1604,72 @@ if MCP_AVAILABLE: return allowed_mcp_servers + def _http_detail_message(detail: object) -> str: + return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail) + + def _server_answers_to(server: MCPServer, name: str) -> bool: + requested: Final = name.lower() + return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known) + + class _McpDeniedDetail(TypedDict): + error: ReadOnly[str] + + async def raise_denied_scoped_mcp_access( + requested_names: Sequence[str], + user_api_key_auth: UserAPIKeyAuth | None, + client_ip: str | None = None, + ) -> None: + """A scoped request (``/mcp/`` path or ``x-mcp-servers`` header) resolved to zero + allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy + server with no tools. Unknown, unauthorized, and access-group names all share one generic + error so scoping cannot probe which servers exist; the agent variant fires only when the + same request resolves once the agent binding is stripped, proving the binding caused the veto.""" + agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None + if user_api_key_auth is not None and agent_id: + resolved_without_agent: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})), + mcp_servers=requested_names, + client_ip=client_ip, + ) + + def _resolved_to_server(name: str) -> bool: + return any(_server_answers_to(server, name) for server in resolved_without_agent) + + vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None) + if vetoed_server is not None: + agent_denial: Final[_McpDeniedDetail] = { + "error": ( + f"MCP server '{vetoed_server}' is not available to this key: the key is bound to " + f"agent '{agent_id}', whose MCP grants do not include this server. Add the server " + f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or " + f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." + ) + } + raise HTTPException(status_code=403, detail=agent_denial) + vetoed_group: Final = next( + ( + name + for name in requested_names + if not _resolved_to_server(name) + and any(name in (server.access_groups or ()) for server in resolved_without_agent) + ), + None, + ) + if vetoed_group is not None: + group_denial: Final[_McpDeniedDetail] = { + "error": ( + f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to " + f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the " + f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or " + f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." + ) + } + raise HTTPException(status_code=403, detail=group_denial) + generic_denial: Final[_McpDeniedDetail] = { + "error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}" + } + raise HTTPException(status_code=403, detail=generic_denial) + def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool: """ Check if a tool name matches any name in the filter list. @@ -1532,7 +1762,7 @@ if MCP_AVAILABLE: async def _get_allowed_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, - mcp_servers: list[str] | None, + mcp_servers: Sequence[str] | None, client_ip: str | None = None, ) -> list[MCPServer]: """Return allowed MCP servers for a request after applying filters. @@ -1818,18 +2048,44 @@ if MCP_AVAILABLE: return texts[0][1] return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts) + async def _raise_if_initialize_grants_no_mcp_servers( + allowed: Sequence[MCPServer], + user_api_key_auth: UserAPIKeyAuth | None, + mcp_servers: Sequence[str] | None, + client_ip: str | None, + ) -> None: + if allowed or user_api_key_auth is None or not user_api_key_auth.api_key: + return + if mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) + no_servers_denial: Final[_McpDeniedDetail] = { + "error": ( + "The key has no MCP servers granted, or none of its granted servers is loaded and allowed for " + "this client IP. Grant servers or access groups to the key, its team, or its organization " + "(object_permission.mcp_servers), check the server's allowed IPs, and reconnect." + ) + } + raise HTTPException(status_code=403, detail=no_servers_denial) + @contextlib.asynccontextmanager async def _gateway_initialize_instructions_request_scope( user_api_key_auth: UserAPIKeyAuth | None, mcp_servers: list[str] | None, client_ip: str | None, scoped_server_endpoint: bool = False, + is_initialize: bool = False, ) -> AsyncIterator[None]: allowed: Final = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip, ) + if is_initialize: + await _raise_if_initialize_grants_no_mcp_servers(allowed, user_api_key_auth, mcp_servers, client_ip) if allowed: # return_exceptions=True: a per-server probe failure (incl. CancelledError # bubbled from anyio task group teardown on connection refused) must not @@ -1876,6 +2132,7 @@ if MCP_AVAILABLE: litellm_trace_id: str | None = None, request_tags: list[str] | None = None, client_ip: str | None = None, + mcp_proxy_mode: bool = False, ) -> AggregateToolListing: """ Helper method to fetch tools from MCP servers based on server filtering criteria. @@ -1968,6 +2225,12 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) # Pre-fetch OAuth credentials only when at least one server uses OAuth2, # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. @@ -2049,9 +2312,14 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, ) - # Apply display-name/description overrides last so that - # permission filtering always works against original names. - filtered_tools = apply_tool_overrides(filtered_tools, server) + if mcp_proxy_mode: + from litellm.proxy._experimental.mcp_server.tool_search import with_mcp_proxy_identity + + filtered_tools = [ # mutable-ok: MCP tool pipeline + with_mcp_proxy_identity(tool, server.server_id) for tool in filtered_tools + ] + else: + filtered_tools = apply_tool_overrides(filtered_tools, server) verbose_logger.debug( "Successfully fetched %s tools from server %s, %s after filtering", @@ -2363,6 +2631,7 @@ if MCP_AVAILABLE: log_list_tools_to_spendlogs: bool = False, list_tools_log_source: str | None = None, client_ip: str | None = None, + mcp_proxy_mode: bool = False, ) -> AggregateToolListing: """ List all available MCP tools. @@ -2392,9 +2661,12 @@ if MCP_AVAILABLE: log_list_tools_to_spendlogs=log_list_tools_to_spendlogs, list_tools_log_source=list_tools_log_source, client_ip=client_ip, + mcp_proxy_mode=mcp_proxy_mode, ) verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools)) return listing + except HTTPException: + raise except Exception as e: verbose_logger.exception("Error getting tools from managed MCP servers: %s", e) # Continue with an empty listing instead of failing completely @@ -3067,6 +3339,43 @@ if MCP_AVAILABLE: ) return result + async def fire_mcp_tool_call_failure_logging( + logging_obj: LiteLLMLoggingObj | None, + exception: Exception, + start_time: datetime, + user_api_key_auth: UserAPIKeyAuth | None, + request_data: Mapping[str, object], + ) -> None: + """Failure logging shared by the ``/mcp`` path and the REST endpoint. Call from + inside the ``except`` block so the traceback is still available. + + The failure handlers run first because ``_ProxyDBLogger.async_post_call_failure_hook`` + builds the failure spend-log row from the ``standard_logging_object`` they produce; + both gate on ``should_run_logging``, so the ``@client`` wrapper does not log twice. + A relayed upstream 401 (``MCPUpstreamAuthError``) is an expected caller-must-reauth + signal and skips ``post_call_failure_hook``, which fires the ``llm_exceptions`` alert. + """ + from litellm.proxy.proxy_server import proxy_logging_obj + + traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + if logging_obj is not None: + end_time: Final = datetime.now() # noqa: DTZ005 # naive to match `start_time`, which it is subtracted from + logging_obj.failure_handler(exception, traceback_str, start_time, end_time) + await logging_obj.async_failure_handler(exception, traceback_str, start_time, end_time) + + if isinstance(exception, MCPUpstreamAuthError) or not proxy_logging_obj or user_api_key_auth is None: + return + sanitized_request_data: Final = { + key: value for key, value in request_data.items() if key not in _MCP_CREDENTIAL_REQUEST_FIELDS + } + await proxy_logging_obj.post_call_failure_hook( + request_data=sanitized_request_data, + original_exception=exception, + user_api_key_dict=user_api_key_auth, + route="/mcp/call_tool", + traceback_str=traceback_str, + ) + @client async def call_mcp_tool( name: str, @@ -3077,6 +3386,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, **kwargs: Any, ) -> CallToolResult: """ @@ -3107,6 +3417,12 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, allowed_mcp_servers=allowed_mcp_servers, ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) if not allowed_mcp_servers: raise HTTPException( status_code=403, @@ -3126,40 +3442,8 @@ if MCP_AVAILABLE: raw_headers=raw_headers, **kwargs, ) - except MCPUpstreamAuthError: - # A client-forwarded pass-through upstream 401 is an expected caller-must-reauth signal, so - # re-raise it without post_call_failure_hook, which fires the proxy's llm_exceptions alert. - # mcp_server_tool_call then downgrades it to an informational isError result for the - # streamable client. Note: this function is @client-decorated, so the decorator's standard - # failure logging still records the event (spend log / OTel); only the extra alert sink is - # skipped here. - raise except Exception as e: - traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) - from litellm.proxy.proxy_server import proxy_logging_obj - - # Ordering is load-bearing. ``_ProxyDBLogger.async_post_call_failure_hook``, - # reached below, writes the failure spend-log row from this logger's - # ``standard_logging_object``, which only exists once the failure handlers - # have run. Flush them first or the row lands with - # ``guardrail_information=None`` and a guardrail block is never counted. - # - # Not double-logged: both handlers gate on ``should_run_logging`` and then - # mark it, so the ``@client`` wrapper's own post-raise logging no-ops on this - # logger, same as ``_fire_mcp_tool_call_logging`` does for ``isError=True``. - if litellm_logging_obj is not None: - end_time: Final = datetime.now() # noqa: DTZ005 # naive to match `start_time`, which it is subtracted from - litellm_logging_obj.failure_handler(e, traceback_str, start_time, end_time) - await litellm_logging_obj.async_failure_handler(e, traceback_str, start_time, end_time) - - if proxy_logging_obj and user_api_key_auth: - await proxy_logging_obj.post_call_failure_hook( - request_data=kwargs, - original_exception=e, - user_api_key_dict=user_api_key_auth, - route="/mcp/call_tool", - traceback_str=traceback_str, - ) + await fire_mcp_tool_call_failure_logging(litellm_logging_obj, e, start_time, user_api_key_auth, kwargs) raise if litellm_logging_obj: @@ -3282,7 +3566,9 @@ if MCP_AVAILABLE: server_name: str | None, session_id: str | None = None, ) -> StandardLoggingMCPToolCall: - mcp_server: Final = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + mcp_server: Final = global_mcp_server_manager._get_mcp_server_from_tool_name( + add_server_prefix_to_name(name, server_name) if server_name else name + ) namespaced_tool_name: Final = f"{server_name}/{name}" if server_name else name if mcp_server: mcp_info: Final = mcp_server.mcp_info or {} @@ -3848,12 +4134,14 @@ if MCP_AVAILABLE: # then exchanges. A tool-call-time 401 would be wrapped into a JSON-RPC error and the # header lost, so the discovery flow needs this pre-emptive challenge. if server and server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers: - from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # lazy: adapter pulls MCP subgraph raise_token_exchange_challenge, ) - from litellm.proxy.utils import get_server_root_path # noqa: PLC0415 + from litellm.proxy.middleware.per_request_root_path_middleware import ( # noqa: PLC0415 # lazy: middleware imports proxy utils + get_request_root_path, + ) - raise_token_exchange_challenge(server, root_path=get_server_root_path()) + raise_token_exchange_challenge(server, root_path=get_request_root_path()) # Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run # the exchange here at the transport edge, so a rejected subject raises the RFC 9728 @@ -4219,13 +4507,15 @@ if MCP_AVAILABLE: raw_headers=raw_headers, scope=dict(scope), mcp_servers=mcp_servers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, client_ip=_client_ip, ) - if _debug_headers: - send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers) + diagnostics: Final = MCPAuthDiagnostics() if _debug_headers else None + if diagnostics is not None: + scope[MCP_AUTH_DIAGNOSTICS_SCOPE_KEY] = diagnostics + send = MCPDebug.wrap_send_with_debug_headers( + send, _debug_headers, diagnostics.headers, request_method=scope.get("method") + ) # Ensure session managers are initialized if not _SESSION_MANAGERS_INITIALIZED: @@ -4397,6 +4687,7 @@ if MCP_AVAILABLE: async def _dispatch() -> None: _otel_publish_transport_span_on_scope(scope) + _otel_publish_request_destinations_on_scope(scope) auth_user: Final = _set_or_update_auth_context( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, @@ -4423,6 +4714,7 @@ if MCP_AVAILABLE: mcp_servers, _client_ip, scoped_server_endpoint=scoped_server_endpoint, + is_initialize=is_initialize, ): await target_manager.handle_request(scope, receive, local_send) if use_stateful and session_id and scope.get("method") == "DELETE": @@ -4559,6 +4851,7 @@ if MCP_AVAILABLE: mcp_servers, _sse_client_ip, scoped_server_endpoint=scoped_server_endpoint, + is_initialize=scope.get("method") == "GET", ): await sse_session_manager.handle_request(scope, receive, send) except MCPUpstreamAuthError as e: diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index af02c11ad86..2c73f9b863b 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -1,5 +1,6 @@ from __future__ import annotations +import hashlib import json from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -11,6 +12,7 @@ from pydantic import ValidationError from typing_extensions import ReadOnly, Required, assert_never import litellm +from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K from litellm.proxy.common_utils.semantic_text_index import ( Embedder, @@ -29,8 +31,17 @@ if TYPE_CHECKING: MCP_TOOL_SEARCH_SETTINGS_KEY: Final[str] = "mcp_tool_search" MCP_TOOL_SEARCH_TOOL_NAME: Final[str] = "mcp_tool_search" MCP_TOOL_CALL_TOOL_NAME: Final[str] = "mcp_tool_call" +MCP_PROXY_SEARCH_TOOL_NAME: Final[str] = "search_tools" +MCP_PROXY_SCHEMA_TOOL_NAME: Final[str] = "get_tool_schema" +MCP_PROXY_CALL_TOOL_NAME: Final[str] = "call_tool" +MCP_PROXY_TOOL_NAMES: Final = frozenset( + (MCP_PROXY_SEARCH_TOOL_NAME, MCP_PROXY_SCHEMA_TOOL_NAME, MCP_PROXY_CALL_TOOL_NAME) +) AGENT_SEARCH_TOOL_NAME: Final[str] = "agent_search" -VIRTUAL_TOOL_NAMES: Final = frozenset((MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME)) +SKILL_SEARCH_TOOL_NAME: Final[str] = "skill_search" +VIRTUAL_TOOL_NAMES: Final = frozenset( + (MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, SKILL_SEARCH_TOOL_NAME) +) def coerce_top_k(value: Any, default: int = 5) -> int: @@ -47,6 +58,29 @@ class ToolSearchResult(TypedDict, total=False): score: ReadOnly[float] +class MCPProxySearchResult(TypedDict, total=False): + tool_id: Required[ReadOnly[str]] + name: Required[ReadOnly[str]] + description: Required[ReadOnly[str]] + score: ReadOnly[float] + + +class MCPProxySchemaResult(MCPProxySearchResult, total=False): + inputSchema: Required[ReadOnly[Mapping[str, object]]] + outputSchema: ReadOnly[Mapping[str, object]] + + +class MCPProxyToolIdentity(TypedDict): + server_id: ReadOnly[str] + tool_name: ReadOnly[str] + + +@dataclass(frozen=True, slots=True) +class MCPToolSearchHit: + tool: Tool + score: float | None = None + + @dataclass(frozen=True, slots=True) class SemanticToolRanker: embed: Embedder @@ -72,6 +106,55 @@ def _scored_result(tool: Tool, score: float) -> ToolSearchResult: return {"name": tool.name, "description": tool.description or "", "inputSchema": tool.inputSchema, "score": score} +_MCP_PROXY_IDENTITY_META_KEY: Final[str] = "litellm.ai/proxy_tool_identity" + + +def with_mcp_proxy_identity(tool: Tool, server_id: str) -> Tool: + identity: Final[MCPProxyToolIdentity] = {"server_id": server_id, "tool_name": tool.name} + return tool.model_copy( # mutable-ok: Pydantic requires mutable update and metadata mappings + update={ # mutable-ok: Pydantic update payload + "meta": {**(tool.meta or {}), _MCP_PROXY_IDENTITY_META_KEY: identity} # mutable-ok: metadata mapping + } + ) + + +def _mcp_proxy_identity(tool: Tool) -> MCPProxyToolIdentity: + identity: Final = (tool.meta or {}).get(_MCP_PROXY_IDENTITY_META_KEY) # mutable-ok: absent metadata default + if not isinstance(identity, Mapping): + raise TypeError("MCP proxy tool identity is missing") + server_id: Final = identity.get("server_id") + tool_name: Final = identity.get("tool_name") + if not isinstance(server_id, str) or not isinstance(tool_name, str): + raise TypeError("MCP proxy tool identity is invalid") + return {"server_id": server_id, "tool_name": tool_name} # mutable-ok: TypedDict identity payload + + +def mcp_proxy_tool_id(tool: Tool) -> str: + identity: Final = _mcp_proxy_identity(tool) + return hashlib.sha256(f"{identity['server_id']}\0{identity['tool_name']}".encode()).hexdigest()[:32] + + +def _proxy_search_result(hit: MCPToolSearchHit) -> MCPProxySearchResult: + base: Final[MCPProxySearchResult] = { + "tool_id": mcp_proxy_tool_id(hit.tool), + "name": hit.tool.name, + "description": hit.tool.description or "", + } + return {**base, "score": hit.score} if hit.score is not None else base # mutable-ok: wire result payload + + +def _proxy_schema_result(tool: Tool) -> MCPProxySchemaResult: + base: Final[MCPProxySchemaResult] = { + "tool_id": mcp_proxy_tool_id(tool), + "name": tool.name, + "description": tool.description or "", + "inputSchema": tool.inputSchema, + } + if tool.outputSchema is None: + return base + return {**base, "outputSchema": tool.outputSchema} # mutable-ok: wire schema payload + + def _tool_text(tool: Tool) -> str: return "\n".join(part for part in (tool.name, tool.description or "") if part) @@ -103,6 +186,38 @@ def search_tools(query: str, tools: Sequence[Tool], top_k: int = 5) -> tuple[Too return tuple(_tool_result(tool) for _, tool in _top_hits(tools, scores, minimum=1.0, limit=top_k)) +async def rank_mcp_tools( + query: str, + tools: Sequence[Tool], + top_k: int, + settings: MCPToolSearchSettings, + ranker: SemanticToolRanker | None, +) -> tuple[MCPToolSearchHit, ...] | EmbeddingFailed: + core, rest = _split_core_tools(tools, settings.core_tools) + core_hits: Final = tuple(MCPToolSearchHit(tool) for tool in core) + if not query: + return core_hits + limit: Final = min(top_k, settings.top_k) + if ranker is None: + scores: Final = tuple(_keyword_score(query, tool) for tool in rest) + return ( + *core_hits, + *(MCPToolSearchHit(tool) for _, tool in _top_hits(rest, scores, minimum=1.0, limit=limit)), + ) + semantic_scores: Final = await ranker.index.scores( + query, tuple(_tool_text(tool) for tool in rest), ranker.embed, ranker.embedding_model + ) + if isinstance(semantic_scores, EmbeddingFailed): + return semantic_scores + return ( + *core_hits, + *( + MCPToolSearchHit(tool, score) + for score, tool in _top_hits(rest, semantic_scores, settings.similarity_threshold, limit) + ), + ) + + async def search_mcp_tools( query: str, tools: Sequence[Tool], @@ -110,21 +225,12 @@ async def search_mcp_tools( settings: MCPToolSearchSettings, ranker: SemanticToolRanker | None, ) -> tuple[ToolSearchResult, ...] | EmbeddingFailed: - """Core tools the caller can access come first, then up to `top_k` ranked matches from the remaining tools.""" - core, rest = _split_core_tools(tools, settings.core_tools) - limit: Final = min(top_k, settings.top_k) - core_results: Final = tuple(_tool_result(tool) for tool in core) - if ranker is None: - return (*core_results, *search_tools(query, rest, limit)) - if not query: - return core_results - scores: Final = await ranker.index.scores( - query, tuple(_tool_text(tool) for tool in rest), ranker.embed, ranker.embedding_model + hits: Final = await rank_mcp_tools(query, tools, top_k, settings, ranker) + if isinstance(hits, EmbeddingFailed): + return hits + return tuple( + _scored_result(hit.tool, hit.score) if hit.score is not None else _tool_result(hit.tool) for hit in hits ) - if isinstance(scores, EmbeddingFailed): - return scores - hits: Final = _top_hits(rest, scores, minimum=settings.similarity_threshold, limit=limit) - return (*core_results, *(_scored_result(tool, score) for score, tool in hits)) class _ToolParamSchema(TypedDict, total=False): @@ -199,8 +305,66 @@ _AGENT_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { } +_SKILL_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { + "name": SKILL_SEARCH_TOOL_NAME, + "description": "Find registered skills by describing what you need in natural language. Returns the best " + "matching skills you can access, ranked by semantic similarity, each with its skill_id, display_title, " + "description, and score.", + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "What you need the skill to do, in natural language."}, + "top_k": { + "type": "integer", + "description": "Maximum number of skills to return.", + "default": DEFAULT_SKILL_SEARCH_TOP_K, + }, + }, + "required": _json_array("query"), + }, +} + + +_MCP_PROXY_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { + "name": MCP_PROXY_SEARCH_TOOL_NAME, + "description": "Search accessible MCP tools by describing what you need. Returns opaque tool IDs.", + "inputSchema": { + "type": "object", + "properties": {"query": {"type": "string", "description": "What the tool should do."}}, + "required": _json_array("query"), + }, +} + +_MCP_PROXY_SCHEMA_DEFINITION: Final[VirtualToolDefinition] = { + "name": MCP_PROXY_SCHEMA_TOOL_NAME, + "description": "Return the complete schema for an accessible MCP tool ID.", + "inputSchema": { + "type": "object", + "properties": {"tool_id": {"type": "string", "description": "Opaque ID from search_tools."}}, + "required": _json_array("tool_id"), + }, +} + +_MCP_PROXY_CALL_DEFINITION: Final[VirtualToolDefinition] = { + "name": MCP_PROXY_CALL_TOOL_NAME, + "description": "Call an accessible MCP tool by opaque ID with schema-valid arguments.", + "inputSchema": { + "type": "object", + "properties": { + "tool_id": {"type": "string", "description": "Opaque ID from search_tools."}, + "arguments": {"type": "object", "description": "Arguments validated against the selected tool schema."}, + }, + "required": _json_array("tool_id"), + }, +} + + def get_virtual_tool_definitions() -> tuple[VirtualToolDefinition, ...]: - return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION) + return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION, _SKILL_SEARCH_DEFINITION) + + +def get_mcp_proxy_tool_definitions() -> tuple[VirtualToolDefinition, ...]: + return (_MCP_PROXY_SEARCH_DEFINITION, _MCP_PROXY_SCHEMA_DEFINITION, _MCP_PROXY_CALL_DEFINITION) def _text_tool_result(text: str, is_error: bool) -> CallToolResult: @@ -223,7 +387,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI ) from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj await check_feature_access_for_user(user_api_key_dict, "agents") outcome: Final = await search_agents( @@ -234,6 +398,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) match outcome: case AgentSearchHits(hits): @@ -245,6 +410,39 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI assert_never(outcome) +async def handle_skill_search(query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth) -> CallToolResult: + from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler + from litellm.llms.litellm_proxy.skills.skill_search import ( + MAX_SKILL_SEARCH_TOP_K, + SkillSearchEmbeddingFailed, + SkillSearchHits, + SkillSearchNotConfigured, + global_skill_search_index, + search_skills, + skill_search_result, + ) + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj + + outcome: Final = await search_skills( + query=query, + skills=await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict), + top_k=min(max(top_k, 1), MAX_SKILL_SEARCH_TOP_K), + router=llm_router, + embedding_model=litellm.skill_search_embedding_model, + index=global_skill_search_index, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + match outcome: + case SkillSearchHits(hits): + results: Final = tuple(skill_search_result(hit).model_dump() for hit in hits) + return _text_tool_result(json.dumps(results), is_error=False) + case SkillSearchNotConfigured(reason) | SkillSearchEmbeddingFailed(reason): + return _text_tool_result(reason, is_error=True) + case _: + assert_never(outcome) + + async def handle_mcp_tool_search( query: str, top_k: int, @@ -256,8 +454,10 @@ async def handle_mcp_tool_search( oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, ) -> CallToolResult: - from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools - from litellm.proxy.proxy_server import llm_router + from litellm.proxy._experimental.mcp_server.server import ( + _list_mcp_tools, # pyright: ignore[reportPrivateUsage] # shared catalog owner + ) + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj settings: Final = mcp_tool_search_settings() if isinstance(settings, ValidationError): @@ -271,7 +471,7 @@ async def handle_mcp_tool_search( ) ranker: Final = ( SemanticToolRanker( - embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict), + embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict, proxy_logging_obj), embedding_model=settings.embedding_model, index=global_mcp_tool_search_index, ) @@ -293,6 +493,97 @@ async def handle_mcp_tool_search( return _text_tool_result(json.dumps(results), is_error=False) +async def handle_mcp_proxy_tool( + name: str, + arguments: dict[str, object], # mutable-ok: MCP dispatcher passes mutable call arguments + user_api_key_dict: UserAPIKeyAuth, + client_ip: str | None = None, + mcp_servers: list[str] | None = None, # mutable-ok: preserve MCP scope container for existing resolver + mcp_auth_header: str | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, # mutable-ok: preserve forwarded headers + oauth2_headers: dict[str, str] | None = None, # mutable-ok: preserve forwarded headers + raw_headers: dict[str, str] | None = None, # mutable-ok: preserve request headers + litellm_logging_obj: LiteLLMLoggingObj | None = None, +) -> CallToolResult: + from fastapi import HTTPException + from jsonschema import ValidationError as JsonSchemaValidationError + from jsonschema import validate + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.server import ( # pyright: ignore[reportPrivateUsage] # shared catalog owner + _list_mcp_tools, # pyright: ignore[reportPrivateUsage] # shared catalog owner + ) + + listing: Final = await _list_mcp_tools( + user_api_key_auth=user_api_key_dict, + mcp_servers=mcp_servers, + client_ip=client_ip, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + mcp_proxy_mode=True, + ) + tools_by_id: Final = {mcp_proxy_tool_id(tool): tool for tool in listing.tools} # mutable-ok: lookup index + + if name == MCP_PROXY_SEARCH_TOOL_NAME: + llm_router: Final = proxy_server.llm_router + proxy_logging_obj: Final = proxy_server.proxy_logging_obj + settings: Final = mcp_tool_search_settings() + if isinstance(settings, ValidationError): + return _text_tool_result(str(settings), is_error=True) + if settings.embedding_model is not None and llm_router is None: + return _text_tool_result( + f"litellm_settings.{MCP_TOOL_SEARCH_SETTINGS_KEY}.embedding_model needs a model_list so it can be called", + is_error=True, + ) + ranker: Final = ( + SemanticToolRanker( + embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict, proxy_logging_obj), + embedding_model=settings.embedding_model, + index=global_mcp_tool_search_index, + ) + if settings.embedding_model is not None and llm_router is not None + else None + ) + results: Final = await rank_mcp_tools(str(arguments.get("query", "")), listing.tools, 5, settings, ranker) + if isinstance(results, EmbeddingFailed): + return _text_tool_result(results.reason, is_error=True) + return _text_tool_result(json.dumps(tuple(_proxy_search_result(hit) for hit in results)), is_error=False) + + tool_id: Final = arguments.get("tool_id") + tool: Final = tools_by_id.get(tool_id) if isinstance(tool_id, str) else None + if tool is None: + return _text_tool_result("Unknown or unauthorized tool_id", is_error=True) + + if name == MCP_PROXY_SCHEMA_TOOL_NAME: + return _text_tool_result(json.dumps(_proxy_schema_result(tool)), is_error=False) + if name != MCP_PROXY_CALL_TOOL_NAME: + raise HTTPException(status_code=400, detail=f"Unknown MCP proxy tool: {name}") + + tool_arguments: Final = arguments.get("arguments", {}) # mutable-ok: JSON Schema validator consumes mapping + if not isinstance(tool_arguments, dict): + return _text_tool_result("arguments must be an object", is_error=True) + try: + validate(instance=tool_arguments, schema=tool.inputSchema) + except JsonSchemaValidationError as exc: + return _text_tool_result(f"Invalid arguments: {exc.message}", is_error=True) + + return await handle_mcp_tool_call( + tool_name=_mcp_proxy_identity(tool)["tool_name"], + arguments=tool_arguments, + user_api_key_dict=user_api_key_dict, + requested_server_id=_mcp_proxy_identity(tool)["server_id"], + client_ip=client_ip, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + litellm_logging_obj=litellm_logging_obj, + ) + + async def handle_mcp_tool_call( tool_name: str, arguments: dict[str, Any], @@ -304,10 +595,12 @@ async def handle_mcp_tool_call( oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, litellm_logging_obj: LiteLLMLoggingObj | None = None, + requested_server_id: str | None = None, ) -> CallToolResult: from litellm.proxy._experimental.mcp_server.server import ( _get_allowed_mcp_servers, execute_mcp_tool, + raise_denied_scoped_mcp_access, ) allowed_mcp_servers: Final = await _get_allowed_mcp_servers( @@ -315,6 +608,12 @@ async def handle_mcp_tool_call( mcp_servers=mcp_servers, client_ip=client_ip, ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_dict, + client_ip=client_ip, + ) # Reject before dispatch when the key has no accessible servers; otherwise an # unprefixed local tool name would fall through to the local registry in @@ -335,4 +634,5 @@ async def handle_mcp_tool_call( oauth2_headers=oauth2_headers, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, + requested_server_id=requested_server_id, ) diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index ecaaf35e817..48bad178927 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -132,10 +132,18 @@ async def update_mcp_toolset( data: UpdateMCPToolsetRequest, touched_by: str, ) -> MCPToolset | None: - data_dict: Final = data.model_dump(exclude_none=True, exclude={"toolset_id"}) - if "tools" in data_dict: - data_dict["tools"] = json.dumps(data_dict["tools"]) - data_dict["updated_by"] = touched_by + """A partial update: absent keeps, null clears. A toolset always has a name and a + tool list, so a null ``toolset_name`` or ``tools`` is a no-op rather than a clear; + emptying the tool selection is an explicit ``[]``, which cannot be mistaken for a + caller that left the field out.""" + data_dict: Final = dict( # mutable-ok: Prisma requires a plain dict for JSON query serialization + ( + (field, json.dumps(value) if field == "tools" else value) + for field, value in data.model_dump(exclude_unset=True).items() + if field != "toolset_id" and (field not in ("toolset_name", "tools") or value is not None) + ), + updated_by=touched_by, + ) try: row: Final = await _toolset_table(prisma_client).update( where={"toolset_id": data.toolset_id}, diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 3f90e6c0a7a..dd1180b30ad 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -8,14 +8,17 @@ omits each feature's routes until the feature is warmed. import asyncio import importlib -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass, field +from types import MappingProxyType from typing import TYPE_CHECKING, Final +from starlette.routing import BaseRoute, Match from starlette.types import Receive, Scope, Send from litellm._logging import verbose_proxy_logger +from litellm.proxy.route_priority import hot_routes_first if TYPE_CHECKING: from fastapi import APIRouter, FastAPI @@ -185,6 +188,31 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( module_path="litellm.proxy.management_endpoints.config_override_endpoints", path_prefixes=("/config_overrides",), ), + LazyFeature( + name="llm_passthrough", + module_path="litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints", + path_prefixes=( + "/anthropic/", + "/assemblyai/", + "/azure/", + "/azure_ai/", + "/bedrock/", + "/cohere/", + "/comprehendmedical", + "/cursor/", + "/eu.assemblyai/", + "/gemini/", + "/gigachat/", + "/milvus/", + "/mistral/", + "/openai/", + "/openai_passthrough/", + "/vertex-ai/", + "/vertex_ai/", + "/vllm/", + "/watsonx/", + ), + ), LazyFeature( name="realtime", module_path="litellm.proxy.realtime_endpoints.endpoints", @@ -295,24 +323,86 @@ class LazyFeatureMiddleware: # Short-circuit once every feature has loaded. if scope["type"] in ("http", "websocket") and len(self._loaded) < len(self._features): path = scope.get("path", "") - # Strip SERVER_ROOT_PATH so prefix matching works under a server - # root path. Without this, requests like /api/v1/policies/... never - # match the registered prefixes (/policies/...) and lazy features - # stay unloaded — every endpoint under them returns 404. The + # Strip the request's root_path so prefix matching works under a + # server root path. Without this, requests like /api/v1/policies/... + # never match the registered prefixes (/policies/...) and lazy + # features stay unloaded — every endpoint under them returns 404. + # scope["root_path"] wins over the cached env scalar: FastAPI + # stamps SERVER_ROOT_PATH there, and PerRequestRootPathMiddleware + # resolves SERVER_ROOT_PATHS prefixes there per request. The # `+ "/"` boundary prevents false-positive matches (e.g. /apiv2 - # against root /api). If the path doesn't start with the prefix - # (e.g. a reverse proxy already stripped it), we leave it alone. - if self._root_path and path.startswith(self._root_path + "/"): - path = path[len(self._root_path) :] + # against root /api); a pre-stripped path is left alone. + root_path: Final = str(scope.get("root_path", "")).rstrip("/") or self._root_path + if root_path and path.startswith(root_path + "/"): + path = path[len(root_path) :] # rebind-ok: local strip after the boundary check above for feat in self._features: - if feat.module_path in self._loaded: + if feat.module_path in self._loaded or not feat.matches(path): continue - if feat.matches(path): - await _force_load(self._fastapi_app, feat) + if _eager_route_wins(self._fastapi_app, feat, scope): + continue + await _force_load(self._fastapi_app, feat, self._features) await self.app(scope, receive, send) -async def _force_load(app: "FastAPI", feat: LazyFeature) -> bool: +def _lazy_slots(app: "FastAPI") -> Mapping[str, BaseRoute | None]: + return app.state.lazy_slots if hasattr(app.state, "lazy_slots") else MappingProxyType({}) + + +def reserve_lazy_slot(app: "FastAPI", name: str, features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> None: + """Record the route the feature's router used to be included after, so its routes + are spliced back in there once it loads and keep the same precedence. Anchoring on + the route rather than its index survives later reordering of the table.""" + feat: Final = next(f for f in features if f.name == name) + anchor: Final = app.router.routes[-1] if app.router.routes else None + app.state.lazy_slots = MappingProxyType({**_lazy_slots(app), feat.module_path: anchor}) + + +def _slot_index(routes: Sequence[BaseRoute], anchor: BaseRoute | None) -> int: + if anchor is None: + return 0 + return next((i + 1 for i, route in enumerate(routes) if route is anchor), len(routes)) + + +def _eager_route_wins(app: "FastAPI", feat: LazyFeature, scope: Scope) -> bool: + """Routes ahead of a feature's reserved slot beat its routes in Starlette's scan, + so a request one of them fully matches never needs the feature loaded.""" + slots: Final = _lazy_slots(app) + if feat.module_path not in slots: + return False + ahead: Final = app.router.routes[: _slot_index(app.router.routes, slots[feat.module_path])] + return any(route.matches(scope)[0] is Match.FULL for route in ahead) + + +def _in_registry_order( + routes: Sequence[BaseRoute], + lazy_routes: Mapping[str, tuple[BaseRoute, ...]], + features: tuple[LazyFeature, ...], + slots: Mapping[str, BaseRoute | None], +) -> tuple[BaseRoute, ...]: + """Lazy routers land in registry order, not first-request order, so overlapping + paths (/openai/{endpoint:path} vs /openai/v1/realtime/calls) resolve the same + way no matter which feature a deployment happens to hit first. Features with a + reserved slot go back where they were eagerly included; the rest follow every + eager route.""" + rank: Final = MappingProxyType({f.module_path: i for i, f in enumerate(features)}) + modules: Final = tuple(sorted(lazy_routes, key=lambda m: rank.get(m, len(rank)))) + lazy_ids: Final = frozenset(id(route) for module_path in modules for route in lazy_routes[module_path]) + eager: Final = tuple(route for route in routes if id(route) not in lazy_ids) + + def slot_of(module_path: str) -> int: + return _slot_index(eager, slots[module_path]) if module_path in slots else len(eager) + + return tuple( + route + for index in range(len(eager) + 1) + for route in ( + *(r for module_path in modules if slot_of(module_path) == index for r in lazy_routes[module_path]), + *eager[index : index + 1], + ) + ) + + +async def _force_load(app: "FastAPI", feat: LazyFeature, features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> bool: """Import + register a lazy feature exactly once per (app, module). Shared by the middleware and the /lazy/warm endpoint.""" if not hasattr(app.state, "lazy_loaded"): @@ -327,7 +417,18 @@ async def _force_load(app: "FastAPI", feat: LazyFeature) -> bool: # mutates app.router.routes, so it stays on the loop thread. loop: Final = asyncio.get_running_loop() module: Final = await loop.run_in_executor(None, importlib.import_module, feat.module_path) + before: Final = len(app.router.routes) feat.register_fn(app, module) + previous: Final[Mapping[str, tuple[BaseRoute, ...]]] = ( + app.state.lazy_routes if hasattr(app.state, "lazy_routes") else MappingProxyType({}) + ) + lazy_routes: Final[Mapping[str, tuple[BaseRoute, ...]]] = MappingProxyType( + {**previous, feat.module_path: tuple(app.router.routes[before:])} + ) + app.state.lazy_routes = lazy_routes # rebind-ok: the app owns the record of which routes each feature added + app.router.routes[:] = hot_routes_first( # rebind-ok: the app owns its route table + _in_registry_order(app.router.routes, lazy_routes, features, _lazy_slots(app)) + ) app.state.lazy_loaded.add(feat.module_path) app.openapi_schema = None verbose_proxy_logger.info( diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 67a0b2a2d15..53af85baac6 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -2634,6 +2634,20 @@ ], "title": "Mcp Tool Permissions" }, + "mcp_toolsets": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Mcp Toolsets" + }, "models": { "anyOf": [ { @@ -4321,7 +4335,7 @@ "/anthropic/{endpoint}": { "delete": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete_2", "parameters": [ { "in": "path", @@ -4365,7 +4379,7 @@ }, "get": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__get", + "operationId": "anthropic_proxy_route_anthropic__endpoint__get_2", "parameters": [ { "in": "path", @@ -4409,7 +4423,7 @@ }, "patch": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__patch", + "operationId": "anthropic_proxy_route_anthropic__endpoint__patch_2", "parameters": [ { "in": "path", @@ -4453,7 +4467,7 @@ }, "post": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__post", + "operationId": "anthropic_proxy_route_anthropic__endpoint__post_2", "parameters": [ { "in": "path", @@ -4497,7 +4511,7 @@ }, "put": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__put_2", "parameters": [ { "in": "path", @@ -4687,6 +4701,17 @@ "title": "Created At", "type": "string" }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, "display_title": { "anyOf": [ { @@ -4713,6 +4738,17 @@ ], "title": "Latest Version" }, + "search_score": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Search Score" + }, "source": { "title": "Source", "type": "string" @@ -4781,7 +4817,7 @@ "paths": { "/v1/skills": { "get": { - "description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\" -H \"x-litellm-model: claude-account-1\"\n```\n\nReturns: ListSkillsResponse with list of skills", + "description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\" -H \"x-litellm-model: claude-account-1\"\n```\n\nPass `?custom_llm_provider=litellm_proxy&query=` to rank the LiteLLM-hosted skills you can\naccess by semantic similarity instead of paging through the whole registry:\n```bash\ncurl \"http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5\" -H \"Authorization: Bearer your-key\"\n```\n\nReturns: ListSkillsResponse with list of skills", "operationId": "list_skills_v1_skills_get", "parameters": [ { @@ -4849,6 +4885,39 @@ "default": "anthropic", "title": "Custom Llm Provider" } + }, + { + "description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.", + "in": "query", + "name": "query", + "required": false, + "schema": { + "anyOf": [ + { + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.", + "title": "Query" + } + }, + { + "description": "With query: the maximum number of ranked skills to return.", + "in": "query", + "name": "top_k", + "required": false, + "schema": { + "default": 5, + "description": "With query: the maximum number of ranked skills to return.", + "maximum": 100, + "minimum": 1, + "title": "Top K", + "type": "integer" + } } ], "responses": { @@ -5315,7 +5384,7 @@ "additionalProperties": { "type": "string" }, - "description": "Git source reference", + "description": "Plugin source reference", "title": "Source", "type": "object" }, @@ -5342,7 +5411,7 @@ "type": "object" }, "RegisterPluginRequest": { - "description": "Request body for registering a plugin in the marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket and referenced by their git source.", + "description": "Request body for registering a plugin in the marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket or as a zip archive on any https host and referenced by their source.", "properties": { "author": { "anyOf": [ @@ -5440,7 +5509,7 @@ "additionalProperties": { "type": "string" }, - "description": "Git source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}", + "description": "Plugin source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}\n- Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''}", "title": "Source", "type": "object" }, @@ -5584,7 +5653,7 @@ "additionalProperties": { "type": "string" }, - "description": "Git source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}", + "description": "Plugin source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}\n- Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''}", "title": "Source", "type": "object" }, @@ -5652,8 +5721,26 @@ "paths": { "/claude-code/marketplace.json": { "get": { - "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add \n- claude plugin install @\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin install my-plugin@litellm\n ```", + "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add \n- claude plugin install @\n\nWithout `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`\nthe key is authenticated and the catalog also holds the disabled plugins granted\nto it through `object_permission.skills` on the key or its team.\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin marketplace add \"http://localhost:4000/claude-code/marketplace.json?key=sk-...\"\n claude plugin install my-plugin@litellm\n ```", "operationId": "get_marketplace_claude_code_marketplace_json_get", + "parameters": [ + { + "in": "query", + "name": "key", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key" + } + } + ], "responses": { "200": { "content": { @@ -5662,6 +5749,16 @@ } }, "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" } }, "summary": "Get Marketplace", @@ -5719,7 +5816,7 @@ ] }, "post": { - "description": "Register a new plugin in the LiteLLM marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket. Claude Code will clone from the git source\nwhen users install.\n\nThis endpoint is create-only and never overwrites. If a plugin with\nthe same name already exists it returns 409 Conflict; use\nPUT /claude-code/plugins/{plugin_name} to update an existing plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - name: Plugin name (kebab-case)\n - source: Git source reference (github, url, or git-subdir format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n\nReturns:\n Registration status (action is always \"created\") and plugin information.\n\nExample:\n ```bash\n curl -X POST http://localhost:4000/claude-code/plugins \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"name\": \"my-plugin\",\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"1.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", + "description": "Register a new plugin in the LiteLLM marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket or as a zip archive on any https host (e.g. S3).\nClaude Code clones the git source or downloads the archive when users install.\n\nThis endpoint is create-only and never overwrites. If a plugin with\nthe same name already exists it returns 409 Conflict; use\nPUT /claude-code/plugins/{plugin_name} to update an existing plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - name: Plugin name (kebab-case)\n - source: Plugin source reference (github, url, git-subdir, or archive format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n\nReturns:\n Registration status (action is always \"created\") and plugin information.\n\nExample:\n ```bash\n curl -X POST http://localhost:4000/claude-code/plugins \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"name\": \"my-plugin\",\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"1.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", "operationId": "register_plugin_claude_code_plugins_post", "requestBody": { "content": { @@ -5854,7 +5951,7 @@ ] }, "put": { - "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default (e.g. an omitted version is cleared, not kept).\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: Name of the plugin to update (path parameter)\n - source: Git source reference (github, url, or git-subdir format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n\nReturns:\n Update status (action is always \"updated\") and plugin information.\n\nExample:\n ```bash\n curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"2.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", + "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default (e.g. an omitted version is cleared, not kept).\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: Name of the plugin to update (path parameter)\n - source: Plugin source reference (github, url, git-subdir, or archive format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n\nReturns:\n Update status (action is always \"updated\") and plugin information.\n\nExample:\n ```bash\n curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"2.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", "operationId": "update_plugin_claude_code_plugins__plugin_name__put", "parameters": [ { @@ -9597,7 +9694,9 @@ { "enum": [ "warn", - "end_session" + "end_session", + "block", + "alert" ], "type": "string" }, @@ -9605,7 +9704,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -11273,6 +11372,12 @@ ], "description": "Threshold configuration for Lakera guardrail categories" }, + "ccr_retrieval": { + "default": true, + "description": "Inject the Headroom retrieval tool for hashes declared by the compression service.", + "title": "Ccr Retrieval", + "type": "boolean" + }, "checks": { "anyOf": [ { @@ -11908,7 +12013,9 @@ { "enum": [ "warn", - "end_session" + "end_session", + "block", + "alert" ], "type": "string" }, @@ -11916,7 +12023,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -15729,6 +15836,4976 @@ } } }, + "llm_passthrough": { + "components": { + "schemas": { + "Body_image_edit_api_openai_deployments__model__images_edits_post": { + "properties": { + "image": { + "anyOf": [ + { + "items": { + "contentMediaType": "application/octet-stream", + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Image" + }, + "image[]": { + "anyOf": [ + { + "items": { + "contentMediaType": "application/octet-stream", + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Image[]" + }, + "mask": { + "anyOf": [ + { + "items": { + "contentMediaType": "application/octet-stream", + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Mask" + }, + "mask[]": { + "anyOf": [ + { + "items": { + "contentMediaType": "application/octet-stream", + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Mask[]" + } + }, + "title": "Body_image_edit_api_openai_deployments__model__images_edits_post", + "type": "object" + }, + "ErrorResponse": { + "properties": { + "detail": { + "additionalProperties": true, + "example": { + "error": { + "code": "error_code", + "message": "Error message", + "param": "error_param", + "type": "error_type" + } + }, + "title": "Detail", + "type": "object" + } + }, + "required": [ + "detail" + ], + "title": "ErrorResponse", + "type": "object" + }, + "HTTPValidationError": { + "properties": { + "detail": { + "items": { + "$ref": "#/components/schemas/ValidationError" + }, + "title": "Detail", + "type": "array" + } + }, + "title": "HTTPValidationError", + "type": "object" + }, + "RealtimeClientSecretResponse": { + "description": "Response from POST /v1/realtime/client_secrets.\n\nBoth the top-level `value` and `session.client_secret.value`\nwill contain the encrypted token instead of the raw ephemeral key.\nThe `session` field is kept as a raw dict so unknown fields pass through.", + "properties": { + "expires_at": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Expires At" + }, + "session": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Session" + }, + "value": { + "title": "Value", + "type": "string" + } + }, + "required": [ + "value" + ], + "title": "RealtimeClientSecretResponse", + "type": "object" + }, + "RealtimeTranscriptionSessionResponse": { + "additionalProperties": true, + "description": "Response from POST /v1/realtime/transcription_sessions.\n\n`client_secret.value` contains the encrypted token instead of the raw\nephemeral key. Unknown fields pass through unchanged.", + "properties": { + "client_secret": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Client Secret" + } + }, + "title": "RealtimeTranscriptionSessionResponse", + "type": "object" + }, + "ValidationError": { + "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, + "loc": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "title": "Location", + "type": "array" + }, + "msg": { + "title": "Message", + "type": "string" + }, + "type": { + "title": "Error Type", + "type": "string" + } + }, + "required": [ + "loc", + "msg", + "type" + ], + "title": "ValidationError", + "type": "object" + } + } + }, + "paths": { + "/anthropic/{endpoint}": { + "delete": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Anthropic Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", + "operationId": "anthropic_proxy_route_anthropic__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Anthropic Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", + "operationId": "anthropic_proxy_route_anthropic__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Anthropic Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", + "operationId": "anthropic_proxy_route_anthropic__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Anthropic Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", + "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Anthropic Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/assemblyai/{endpoint}": { + "delete": { + "operationId": "assemblyai_proxy_route_assemblyai__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "operationId": "assemblyai_proxy_route_assemblyai__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "operationId": "assemblyai_proxy_route_assemblyai__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "operationId": "assemblyai_proxy_route_assemblyai__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "operationId": "assemblyai_proxy_route_assemblyai__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/azure/{endpoint}": { + "delete": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/azure_ai/{endpoint}": { + "delete": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure_ai__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure_ai__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure_ai__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure_ai__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Call any azure endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/azure/{endpoint:path}`\n\nChecks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.", + "operationId": "azure_proxy_route_azure_ai__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Azure Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/bedrock/{endpoint}": { + "delete": { + "description": "This is the v1 passthrough for Bedrock.\nV2 is handled by the `/bedrock/v2` endpoint.\n[Docs](https://docs.litellm.ai/docs/pass_through/bedrock)", + "operationId": "bedrock_proxy_route_bedrock__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Bedrock Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "This is the v1 passthrough for Bedrock.\nV2 is handled by the `/bedrock/v2` endpoint.\n[Docs](https://docs.litellm.ai/docs/pass_through/bedrock)", + "operationId": "bedrock_proxy_route_bedrock__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Bedrock Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "This is the v1 passthrough for Bedrock.\nV2 is handled by the `/bedrock/v2` endpoint.\n[Docs](https://docs.litellm.ai/docs/pass_through/bedrock)", + "operationId": "bedrock_proxy_route_bedrock__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Bedrock Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "This is the v1 passthrough for Bedrock.\nV2 is handled by the `/bedrock/v2` endpoint.\n[Docs](https://docs.litellm.ai/docs/pass_through/bedrock)", + "operationId": "bedrock_proxy_route_bedrock__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Bedrock Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "This is the v1 passthrough for Bedrock.\nV2 is handled by the `/bedrock/v2` endpoint.\n[Docs](https://docs.litellm.ai/docs/pass_through/bedrock)", + "operationId": "bedrock_proxy_route_bedrock__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Bedrock Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/cohere/{endpoint}": { + "delete": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/cohere)", + "operationId": "cohere_proxy_route_cohere__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cohere Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/cohere)", + "operationId": "cohere_proxy_route_cohere__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cohere Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/cohere)", + "operationId": "cohere_proxy_route_cohere__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cohere Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/cohere)", + "operationId": "cohere_proxy_route_cohere__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cohere Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/cohere)", + "operationId": "cohere_proxy_route_cohere__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cohere Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/comprehendmedical": { + "post": { + "description": "AWS-SDK-shaped pass-through for Amazon Comprehend Medical: point the SDK's\n`endpoint_url` at `/comprehendmedical` and the operation is read from the\n`X-Amz-Target` header, per the AWS JSON 1.1 protocol.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/comprehend_medical)", + "operationId": "comprehend_medical_sdk_proxy_route_comprehendmedical_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Comprehend Medical Sdk Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/comprehendmedical/{operation}": { + "post": { + "description": "Pass-through for Amazon Comprehend Medical, e.g. `POST /comprehendmedical/DetectEntitiesV2`.\n\nThe request body is forwarded as-is to the AWS JSON 1.1 API and signed with SigV4\nusing the proxy's AWS credentials.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/comprehend_medical)", + "operationId": "comprehend_medical_proxy_route_comprehendmedical__operation__post", + "parameters": [ + { + "in": "path", + "name": "operation", + "required": true, + "schema": { + "title": "Operation", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Comprehend Medical Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/cursor/chat/completions": { + "post": { + "description": "Cursor BYOK endpoint. Accepts both request shapes Cursor sends to its OpenAI-compatible\nbase URL and always answers in chat completions format.\n\nCursor agent mode sends Responses API format bodies (`input`, flat tool defs, `reasoning`,\ncustom tools) to the chat/completions path while expecting chat completions responses;\nthose are routed through the Responses API pipeline and converted back. Genuine chat\ncompletions bodies (`messages` present) are routed through the standard chat completions\npipeline, after normalizing each level of the `tools` array and `tool_choice` to the chat\ncompletions shapes OpenAI requires. Cursor mixes Responses API shapes into chat bodies\nper level, independently: a flat tool def (`{\"type\": \"custom\", \"name\": \"ApplyPatch\", ...}`)\ngets nested under `custom`, and a flat grammar format\n(`{\"type\": \"grammar\", \"definition\", \"syntax\"}`) gets wrapped as\n`{\"type\": \"grammar\", \"grammar\": {...}}` wherever it appears, including inside tool defs\nCursor already sent pre-nested.\n\n```bash\ncurl -X POST http://localhost:4000/cursor/chat/completions -H \"Content-Type: application/json\" -H \"Authorization: Bearer sk-1234\" -d '{\n \"model\": \"gpt-4o\",\n \"input\": [{\"role\": \"user\", \"content\": \"Hello\"}]\n}'\nResponds back in chat completions format.\n```", + "operationId": "cursor_chat_completions_cursor_chat_completions_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cursor Chat Completions", + "tags": [ + "llm_passthrough" + ] + } + }, + "/cursor/models": { + "get": { + "description": "OpenAI-compatible model listing for the Cursor BYOK base URL.\n\nClients pointed at `/cursor` as an OpenAI-compatible base URL resolve and\nverify models via `GET {base}/models` (the OpenAI SDK contract). Without this\nroute those requests fall through to the Cursor Cloud Agents passthrough, which\ndemands a Cursor API key and 401s, so key verification silently fails before any\nchat request is ever sent. Delegates to the standard `/v1/models` handler.", + "operationId": "cursor_model_list_cursor_models_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cursor Model List", + "tags": [ + "llm_passthrough" + ] + } + }, + "/cursor/v1/models": { + "get": { + "description": "OpenAI-compatible model listing for the Cursor BYOK base URL.\n\nClients pointed at `/cursor` as an OpenAI-compatible base URL resolve and\nverify models via `GET {base}/models` (the OpenAI SDK contract). Without this\nroute those requests fall through to the Cursor Cloud Agents passthrough, which\ndemands a Cursor API key and 401s, so key verification silently fails before any\nchat request is ever sent. Delegates to the standard `/v1/models` handler.", + "operationId": "cursor_model_list_cursor_v1_models_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cursor Model List", + "tags": [ + "llm_passthrough" + ] + } + }, + "/cursor/{endpoint}": { + "delete": { + "description": "Pass-through endpoint for the Cursor Cloud Agents API.\n\nSupports all Cursor Cloud Agents endpoints:\n- GET /v0/agents \u2014 List agents\n- POST /v0/agents \u2014 Launch an agent\n- GET /v0/agents/{id} \u2014 Agent status\n- GET /v0/agents/{id}/conversation \u2014 Agent conversation\n- POST /v0/agents/{id}/followup \u2014 Add follow-up\n- POST /v0/agents/{id}/stop \u2014 Stop an agent\n- DELETE /v0/agents/{id} \u2014 Delete an agent\n- GET /v0/me \u2014 API key info\n- GET /v0/models \u2014 List models\n- GET /v0/repositories \u2014 List GitHub repositories\n\nUses Basic Authentication (base64-encoded `API_KEY:`).\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. litellm.credential_list (credentials added via UI)\n3. CURSOR_API_KEY environment variable", + "operationId": "cursor_proxy_route_cursor__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cursor Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Pass-through endpoint for the Cursor Cloud Agents API.\n\nSupports all Cursor Cloud Agents endpoints:\n- GET /v0/agents \u2014 List agents\n- POST /v0/agents \u2014 Launch an agent\n- GET /v0/agents/{id} \u2014 Agent status\n- GET /v0/agents/{id}/conversation \u2014 Agent conversation\n- POST /v0/agents/{id}/followup \u2014 Add follow-up\n- POST /v0/agents/{id}/stop \u2014 Stop an agent\n- DELETE /v0/agents/{id} \u2014 Delete an agent\n- GET /v0/me \u2014 API key info\n- GET /v0/models \u2014 List models\n- GET /v0/repositories \u2014 List GitHub repositories\n\nUses Basic Authentication (base64-encoded `API_KEY:`).\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. litellm.credential_list (credentials added via UI)\n3. CURSOR_API_KEY environment variable", + "operationId": "cursor_proxy_route_cursor__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cursor Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Pass-through endpoint for the Cursor Cloud Agents API.\n\nSupports all Cursor Cloud Agents endpoints:\n- GET /v0/agents \u2014 List agents\n- POST /v0/agents \u2014 Launch an agent\n- GET /v0/agents/{id} \u2014 Agent status\n- GET /v0/agents/{id}/conversation \u2014 Agent conversation\n- POST /v0/agents/{id}/followup \u2014 Add follow-up\n- POST /v0/agents/{id}/stop \u2014 Stop an agent\n- DELETE /v0/agents/{id} \u2014 Delete an agent\n- GET /v0/me \u2014 API key info\n- GET /v0/models \u2014 List models\n- GET /v0/repositories \u2014 List GitHub repositories\n\nUses Basic Authentication (base64-encoded `API_KEY:`).\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. litellm.credential_list (credentials added via UI)\n3. CURSOR_API_KEY environment variable", + "operationId": "cursor_proxy_route_cursor__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cursor Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Pass-through endpoint for the Cursor Cloud Agents API.\n\nSupports all Cursor Cloud Agents endpoints:\n- GET /v0/agents \u2014 List agents\n- POST /v0/agents \u2014 Launch an agent\n- GET /v0/agents/{id} \u2014 Agent status\n- GET /v0/agents/{id}/conversation \u2014 Agent conversation\n- POST /v0/agents/{id}/followup \u2014 Add follow-up\n- POST /v0/agents/{id}/stop \u2014 Stop an agent\n- DELETE /v0/agents/{id} \u2014 Delete an agent\n- GET /v0/me \u2014 API key info\n- GET /v0/models \u2014 List models\n- GET /v0/repositories \u2014 List GitHub repositories\n\nUses Basic Authentication (base64-encoded `API_KEY:`).\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. litellm.credential_list (credentials added via UI)\n3. CURSOR_API_KEY environment variable", + "operationId": "cursor_proxy_route_cursor__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cursor Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Pass-through endpoint for the Cursor Cloud Agents API.\n\nSupports all Cursor Cloud Agents endpoints:\n- GET /v0/agents \u2014 List agents\n- POST /v0/agents \u2014 Launch an agent\n- GET /v0/agents/{id} \u2014 Agent status\n- GET /v0/agents/{id}/conversation \u2014 Agent conversation\n- POST /v0/agents/{id}/followup \u2014 Add follow-up\n- POST /v0/agents/{id}/stop \u2014 Stop an agent\n- DELETE /v0/agents/{id} \u2014 Delete an agent\n- GET /v0/me \u2014 API key info\n- GET /v0/models \u2014 List models\n- GET /v0/repositories \u2014 List GitHub repositories\n\nUses Basic Authentication (base64-encoded `API_KEY:`).\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. litellm.credential_list (credentials added via UI)\n3. CURSOR_API_KEY environment variable", + "operationId": "cursor_proxy_route_cursor__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cursor Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/eu.assemblyai/{endpoint}": { + "delete": { + "operationId": "assemblyai_proxy_route_eu_assemblyai__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "operationId": "assemblyai_proxy_route_eu_assemblyai__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "operationId": "assemblyai_proxy_route_eu_assemblyai__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "operationId": "assemblyai_proxy_route_eu_assemblyai__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "operationId": "assemblyai_proxy_route_eu_assemblyai__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Assemblyai Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/gemini/{endpoint}": { + "delete": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio)", + "operationId": "gemini_proxy_route_gemini__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Gemini Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio)", + "operationId": "gemini_proxy_route_gemini__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Gemini Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio)", + "operationId": "gemini_proxy_route_gemini__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Gemini Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio)", + "operationId": "gemini_proxy_route_gemini__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Gemini Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio)", + "operationId": "gemini_proxy_route_gemini__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Gemini Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/gigachat/{endpoint}": { + "delete": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/gigachat)", + "operationId": "gigachat_proxy_route_gigachat__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Gigachat Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/gigachat)", + "operationId": "gigachat_proxy_route_gigachat__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Gigachat Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/gigachat)", + "operationId": "gigachat_proxy_route_gigachat__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Gigachat Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/gigachat)", + "operationId": "gigachat_proxy_route_gigachat__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Gigachat Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/gigachat)", + "operationId": "gigachat_proxy_route_gigachat__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Gigachat Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/milvus/{endpoint}": { + "delete": { + "description": "Enable using Milvus `/vectors` endpoint as a pass-through endpoint.", + "operationId": "milvus_proxy_route_milvus__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Milvus Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Enable using Milvus `/vectors` endpoint as a pass-through endpoint.", + "operationId": "milvus_proxy_route_milvus__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Milvus Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Enable using Milvus `/vectors` endpoint as a pass-through endpoint.", + "operationId": "milvus_proxy_route_milvus__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Milvus Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Enable using Milvus `/vectors` endpoint as a pass-through endpoint.", + "operationId": "milvus_proxy_route_milvus__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Milvus Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Enable using Milvus `/vectors` endpoint as a pass-through endpoint.", + "operationId": "milvus_proxy_route_milvus__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Milvus Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/mistral/{endpoint}": { + "delete": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/mistral)", + "operationId": "mistral_proxy_route_mistral__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Mistral Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/mistral)", + "operationId": "mistral_proxy_route_mistral__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Mistral Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/mistral)", + "operationId": "mistral_proxy_route_mistral__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Mistral Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/mistral)", + "operationId": "mistral_proxy_route_mistral__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Mistral Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/mistral)", + "operationId": "mistral_proxy_route_mistral__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Mistral Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/deployments/{model}/chat/completions": { + "post": { + "description": "Follows the exact same API spec as `OpenAI's Chat API https://platform.openai.com/docs/api-reference/chat`\n\n```bash\ncurl -X POST http://localhost:4000/v1/chat/completions \n-H \"Content-Type: application/json\" \n-H \"Authorization: Bearer sk-1234\" \n-d '{\n \"model\": \"gpt-4o\",\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": \"Hello!\"\n }\n ]\n}'\n```", + "operationId": "chat_completion_openai_deployments__model__chat_completions_post", + "parameters": [ + { + "in": "path", + "name": "model", + "required": true, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful response" + }, + "400": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "ContentPolicyViolationError" + }, + "401": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "AuthenticationError" + }, + "403": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "PermissionDeniedError" + }, + "404": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "NotFoundError" + }, + "408": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "Timeout" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "UnprocessableEntityError" + }, + "429": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + }, + "500": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "JSONSchemaValidationError" + }, + "503": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponse" + } + } + }, + "description": "APIConnectionError" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Chat Completion", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/deployments/{model}/completions": { + "post": { + "description": "Follows the exact same API spec as `OpenAI's Completions API https://platform.openai.com/docs/api-reference/completions`\n\n```bash\ncurl -X POST http://localhost:4000/v1/completions \n-H \"Content-Type: application/json\" \n-H \"Authorization: Bearer sk-1234\" \n-d '{\n \"model\": \"gpt-3.5-turbo-instruct\",\n \"prompt\": \"Once upon a time\",\n \"max_tokens\": 50,\n \"temperature\": 0.7\n}'\n```", + "operationId": "completion_openai_deployments__model__completions_post", + "parameters": [ + { + "in": "path", + "name": "model", + "required": true, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Completion", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/deployments/{model}/embeddings": { + "post": { + "description": "Follows the exact same API spec as `OpenAI's Embeddings API https://platform.openai.com/docs/api-reference/embeddings`\n\n```bash\ncurl -X POST http://localhost:4000/v1/embeddings \n-H \"Content-Type: application/json\" \n-H \"Authorization: Bearer sk-1234\" \n-d '{\n \"model\": \"text-embedding-ada-002\",\n \"input\": \"The quick brown fox jumps over the lazy dog\"\n}'\n```", + "operationId": "embeddings_openai_deployments__model__embeddings_post", + "parameters": [ + { + "in": "path", + "name": "model", + "required": true, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Embeddings", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/deployments/{model}/images/edits": { + "post": { + "description": "Follows the OpenAI Images API spec: https://platform.openai.com/docs/api-reference/images/create\n\n```bash\ncurl -s -D >(grep -i x-request-id >&2) -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) -X POST \"http://localhost:4000/v1/images/edits\" -H \"Authorization: Bearer sk-1234\" -F \"model=gpt-image-1\" -F \"image[]=@soap.png\" -F 'prompt=Create a studio ghibli image of this'\n```", + "operationId": "image_edit_api_openai_deployments__model__images_edits_post", + "parameters": [ + { + "in": "path", + "name": "model", + "required": true, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + } + } + ], + "requestBody": { + "content": { + "multipart/form-data": { + "schema": { + "$ref": "#/components/schemas/Body_image_edit_api_openai_deployments__model__images_edits_post" + } + } + } + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Image Edit Api", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/deployments/{model}/images/generations": { + "post": { + "operationId": "image_generation_openai_deployments__model__images_generations_post", + "parameters": [ + { + "in": "path", + "name": "model", + "required": true, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Image Generation", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/v1/realtime/calls": { + "post": { + "operationId": "proxy_realtime_calls_openai_v1_realtime_calls_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Realtime Calls", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/v1/realtime/client_secrets": { + "post": { + "operationId": "create_realtime_client_secret_openai_v1_realtime_client_secrets_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RealtimeClientSecretResponse" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Create Realtime Client Secret", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/v1/realtime/transcription_sessions": { + "post": { + "description": "Create an ephemeral Realtime transcription session\n(POST /v1/realtime/transcription_sessions) for the WebRTC/WebSocket flow.\n\nMirrors the client_secrets route but targets the transcription_sessions\nendpoint and encrypts the ephemeral key returned under `client_secret.value`.", + "operationId": "create_realtime_transcription_session_openai_v1_realtime_transcription_sessions_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RealtimeTranscriptionSessionResponse" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Create Realtime Transcription Session", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/v1/responses": { + "post": { + "description": "Follows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses\n\nSupports background mode with polling_via_cache for partial response retrieval.\nWhen background=true and polling_via_cache is enabled, returns a polling_id immediately\nand streams the response in the background, updating Redis cache.\n\n```bash\n# Normal request\ncurl -X POST http://localhost:4000/v1/responses -H \"Content-Type: application/json\" -H \"Authorization: Bearer sk-1234\" -d '{\n \"model\": \"gpt-4o\",\n \"input\": \"Tell me about AI\"\n}'\n\n# Background request with polling\ncurl -X POST http://localhost:4000/v1/responses -H \"Content-Type: application/json\" -H \"Authorization: Bearer sk-1234\" -d '{\n \"model\": \"gpt-4o\",\n \"input\": \"Tell me about AI\",\n \"background\": true\n}'\n```", + "operationId": "responses_api_openai_v1_responses_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Responses Api", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/v1/responses/compact": { + "post": { + "description": "Compact a response by running a compaction pass over a conversation.\n\nReturns encrypted, opaque items that can be used to reduce context size.\n\nFollows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/compact\n\n```bash\ncurl -X POST http://localhost:4000/v1/responses/compact -H \"Content-Type: application/json\" -H \"Authorization: Bearer sk-1234\" -d '{\n \"model\": \"gpt-4o\",\n \"input\": [{\"role\": \"user\", \"content\": \"Hello\"}]\n}'\n```", + "operationId": "compact_response_openai_v1_responses_compact_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Compact Response", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/v1/responses/input_tokens": { + "post": { + "description": "Count the input tokens of a Responses API request without calling the model.\n\nFollows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/input-tokens\n\n```bash\ncurl -X POST http://localhost:4000/v1/responses/input_tokens -H \"Content-Type: application/json\" -H \"Authorization: Bearer sk-1234\" -d '{\n \"model\": \"gpt-4o\",\n \"input\": \"Hello, how are you?\"\n}'\n```\n\nReturns: `{\"object\": \"response.input_tokens\", \"input_tokens\": }`", + "operationId": "responses_input_tokens_openai_v1_responses_input_tokens_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Responses Input Tokens", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/v1/responses/{response_id}": { + "delete": { + "description": "Delete a response by ID.\n\nSupports both:\n- Polling IDs (litellm_poll_*): Deletes from Redis cache\n- Provider response IDs: Passes through to provider API\n\nFollows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/delete\n\n```bash\ncurl -X DELETE http://localhost:4000/v1/responses/resp_abc123 -H \"Authorization: Bearer sk-1234\"\n```", + "operationId": "delete_response_openai_v1_responses__response_id__delete", + "parameters": [ + { + "in": "path", + "name": "response_id", + "required": true, + "schema": { + "title": "Response Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Delete Response", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Get a response by ID.\n\nSupports both:\n- Polling IDs (litellm_poll_*): Returns cumulative cached content from background responses\n- Provider response IDs: Passes through to provider API\n\nFollows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/get\n\n```bash\n# Get polling response\ncurl -X GET http://localhost:4000/v1/responses/litellm_poll_abc123 -H \"Authorization: Bearer sk-1234\"\n\n# Get provider response\ncurl -X GET http://localhost:4000/v1/responses/resp_abc123 -H \"Authorization: Bearer sk-1234\"\n```", + "operationId": "get_response_openai_v1_responses__response_id__get", + "parameters": [ + { + "in": "path", + "name": "response_id", + "required": true, + "schema": { + "title": "Response Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Response", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/v1/responses/{response_id}/cancel": { + "post": { + "description": "Cancel a response by ID.\n\nSupports both:\n- Polling IDs (litellm_poll_*): Cancels background response and updates status in Redis\n- Provider response IDs: Passes through to provider API\n\nFollows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/cancel\n\n```bash\n# Cancel polling response\ncurl -X POST http://localhost:4000/v1/responses/litellm_poll_abc123/cancel -H \"Authorization: Bearer sk-1234\"\n\n# Cancel provider response\ncurl -X POST http://localhost:4000/v1/responses/resp_abc123/cancel -H \"Authorization: Bearer sk-1234\"\n```", + "operationId": "cancel_response_openai_v1_responses__response_id__cancel_post", + "parameters": [ + { + "in": "path", + "name": "response_id", + "required": true, + "schema": { + "title": "Response Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Cancel Response", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/v1/responses/{response_id}/input_items": { + "get": { + "description": "List input items for a response.", + "operationId": "get_response_input_items_openai_v1_responses__response_id__input_items_get", + "parameters": [ + { + "in": "path", + "name": "response_id", + "required": true, + "schema": { + "title": "Response Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Response Input Items", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai/{endpoint}": { + "delete": { + "description": "Pass-through endpoint for OpenAI API calls.\n\nAvailable on both routes:\n- /openai/{endpoint:path} - Standard OpenAI passthrough route\n- /openai_passthrough/{endpoint:path} - Dedicated passthrough route (recommended for Responses API)\n\nUse /openai_passthrough/* when you need guaranteed passthrough to OpenAI without conflicts\nwith LiteLLM's native implementations (e.g., for the Responses API at /v1/responses).\n\nExamples:\n Standard route:\n - /openai/v1/chat/completions\n - /openai/v1/assistants\n - /openai/v1/threads\n\n Dedicated passthrough (for Responses API):\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_proxy_route_openai__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Pass-through endpoint for OpenAI API calls.\n\nAvailable on both routes:\n- /openai/{endpoint:path} - Standard OpenAI passthrough route\n- /openai_passthrough/{endpoint:path} - Dedicated passthrough route (recommended for Responses API)\n\nUse /openai_passthrough/* when you need guaranteed passthrough to OpenAI without conflicts\nwith LiteLLM's native implementations (e.g., for the Responses API at /v1/responses).\n\nExamples:\n Standard route:\n - /openai/v1/chat/completions\n - /openai/v1/assistants\n - /openai/v1/threads\n\n Dedicated passthrough (for Responses API):\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_proxy_route_openai__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Pass-through endpoint for OpenAI API calls.\n\nAvailable on both routes:\n- /openai/{endpoint:path} - Standard OpenAI passthrough route\n- /openai_passthrough/{endpoint:path} - Dedicated passthrough route (recommended for Responses API)\n\nUse /openai_passthrough/* when you need guaranteed passthrough to OpenAI without conflicts\nwith LiteLLM's native implementations (e.g., for the Responses API at /v1/responses).\n\nExamples:\n Standard route:\n - /openai/v1/chat/completions\n - /openai/v1/assistants\n - /openai/v1/threads\n\n Dedicated passthrough (for Responses API):\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_proxy_route_openai__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Pass-through endpoint for OpenAI API calls.\n\nAvailable on both routes:\n- /openai/{endpoint:path} - Standard OpenAI passthrough route\n- /openai_passthrough/{endpoint:path} - Dedicated passthrough route (recommended for Responses API)\n\nUse /openai_passthrough/* when you need guaranteed passthrough to OpenAI without conflicts\nwith LiteLLM's native implementations (e.g., for the Responses API at /v1/responses).\n\nExamples:\n Standard route:\n - /openai/v1/chat/completions\n - /openai/v1/assistants\n - /openai/v1/threads\n\n Dedicated passthrough (for Responses API):\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_proxy_route_openai__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Pass-through endpoint for OpenAI API calls.\n\nAvailable on both routes:\n- /openai/{endpoint:path} - Standard OpenAI passthrough route\n- /openai_passthrough/{endpoint:path} - Dedicated passthrough route (recommended for Responses API)\n\nUse /openai_passthrough/* when you need guaranteed passthrough to OpenAI without conflicts\nwith LiteLLM's native implementations (e.g., for the Responses API at /v1/responses).\n\nExamples:\n Standard route:\n - /openai/v1/chat/completions\n - /openai/v1/assistants\n - /openai/v1/threads\n\n Dedicated passthrough (for Responses API):\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_proxy_route_openai__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/openai_passthrough/{endpoint}": { + "delete": { + "description": "Dedicated pass-through to the OpenAI API with no overlap with LiteLLM's native\nimplementations (e.g. the Responses API at /v1/responses).\n\nExamples:\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_passthrough_route_openai_passthrough__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Passthrough Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Dedicated pass-through to the OpenAI API with no overlap with LiteLLM's native\nimplementations (e.g. the Responses API at /v1/responses).\n\nExamples:\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_passthrough_route_openai_passthrough__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Passthrough Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Dedicated pass-through to the OpenAI API with no overlap with LiteLLM's native\nimplementations (e.g. the Responses API at /v1/responses).\n\nExamples:\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_passthrough_route_openai_passthrough__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Passthrough Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Dedicated pass-through to the OpenAI API with no overlap with LiteLLM's native\nimplementations (e.g. the Responses API at /v1/responses).\n\nExamples:\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_passthrough_route_openai_passthrough__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Passthrough Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Dedicated pass-through to the OpenAI API with no overlap with LiteLLM's native\nimplementations (e.g. the Responses API at /v1/responses).\n\nExamples:\n - /openai_passthrough/v1/responses\n - /openai_passthrough/v1/responses/{response_id}\n - /openai_passthrough/v1/responses/{response_id}/input_items\n\n[Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough)", + "operationId": "openai_passthrough_route_openai_passthrough__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openai Passthrough Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/vertex_ai/discovery/{endpoint}": { + "delete": { + "description": "Call any vertex discovery endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/vertex_ai/discovery/{endpoint:path}`\n\nTarget url: `https://discoveryengine.googleapis.com`", + "operationId": "vertex_discovery_proxy_route_vertex_ai_discovery__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Vertex Discovery Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Call any vertex discovery endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/vertex_ai/discovery/{endpoint:path}`\n\nTarget url: `https://discoveryengine.googleapis.com`", + "operationId": "vertex_discovery_proxy_route_vertex_ai_discovery__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Vertex Discovery Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Call any vertex discovery endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/vertex_ai/discovery/{endpoint:path}`\n\nTarget url: `https://discoveryengine.googleapis.com`", + "operationId": "vertex_discovery_proxy_route_vertex_ai_discovery__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Vertex Discovery Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Call any vertex discovery endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/vertex_ai/discovery/{endpoint:path}`\n\nTarget url: `https://discoveryengine.googleapis.com`", + "operationId": "vertex_discovery_proxy_route_vertex_ai_discovery__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Vertex Discovery Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Call any vertex discovery endpoint using the proxy.\n\nJust use `{PROXY_BASE_URL}/vertex_ai/discovery/{endpoint:path}`\n\nTarget url: `https://discoveryengine.googleapis.com`", + "operationId": "vertex_discovery_proxy_route_vertex_ai_discovery__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Vertex Discovery Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/vertex_ai/{endpoint}": { + "delete": { + "description": "Call LiteLLM proxy via Vertex AI SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai)", + "operationId": "vertex_proxy_route_vertex_ai__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vertex Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Call LiteLLM proxy via Vertex AI SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai)", + "operationId": "vertex_proxy_route_vertex_ai__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vertex Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Call LiteLLM proxy via Vertex AI SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai)", + "operationId": "vertex_proxy_route_vertex_ai__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vertex Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Call LiteLLM proxy via Vertex AI SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai)", + "operationId": "vertex_proxy_route_vertex_ai__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vertex Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Call LiteLLM proxy via Vertex AI SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai)", + "operationId": "vertex_proxy_route_vertex_ai__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vertex Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/vllm/{endpoint}": { + "delete": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/vllm)", + "operationId": "vllm_proxy_route_vllm__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vllm Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/vllm)", + "operationId": "vllm_proxy_route_vllm__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vllm Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/vllm)", + "operationId": "vllm_proxy_route_vllm__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vllm Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/vllm)", + "operationId": "vllm_proxy_route_vllm__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vllm Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "[Docs](https://docs.litellm.ai/docs/pass_through/vllm)", + "operationId": "vllm_proxy_route_vllm__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Vllm Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/watsonx/{endpoint}": { + "delete": { + "description": "Watsonx pass-through endpoint.\nAllows using Watsonx APIs with automatic IAM token management and version parameter injection.\n\nExample:\n POST /watsonx/ml/v1/text/tokenization\n POST /watsonx/ml/v1/text/generation", + "operationId": "watsonx_proxy_route_watsonx__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Watsonx Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Watsonx pass-through endpoint.\nAllows using Watsonx APIs with automatic IAM token management and version parameter injection.\n\nExample:\n POST /watsonx/ml/v1/text/tokenization\n POST /watsonx/ml/v1/text/generation", + "operationId": "watsonx_proxy_route_watsonx__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Watsonx Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Watsonx pass-through endpoint.\nAllows using Watsonx APIs with automatic IAM token management and version parameter injection.\n\nExample:\n POST /watsonx/ml/v1/text/tokenization\n POST /watsonx/ml/v1/text/generation", + "operationId": "watsonx_proxy_route_watsonx__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Watsonx Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Watsonx pass-through endpoint.\nAllows using Watsonx APIs with automatic IAM token management and version parameter injection.\n\nExample:\n POST /watsonx/ml/v1/text/tokenization\n POST /watsonx/ml/v1/text/generation", + "operationId": "watsonx_proxy_route_watsonx__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Watsonx Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Watsonx pass-through endpoint.\nAllows using Watsonx APIs with automatic IAM token management and version parameter injection.\n\nExample:\n POST /watsonx/ml/v1/text/tokenization\n POST /watsonx/ml/v1/text/generation", + "operationId": "watsonx_proxy_route_watsonx__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Watsonx Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + } + } + }, "mcp_app": { "components": { "schemas": { @@ -16349,6 +21426,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -16943,6 +22025,134 @@ "mcp_app" ] } + }, + "/mcp/proxy": { + "delete": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_delete", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "get": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "head": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_head", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "options": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_options", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "patch": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_patch", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "post": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "put": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_put", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + } } } }, @@ -17972,6 +23182,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -18853,6 +24068,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -19792,6 +25012,46 @@ ] } }, + "/authorize/flow": { + "get": { + "operationId": "authorize_flow_authorize_flow_get", + "parameters": [ + { + "in": "query", + "name": "flow", + "required": true, + "schema": { + "title": "Flow", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Authorize Flow", + "tags": [ + "mcp_discoverable" + ] + } + }, "/callback": { "get": { "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.", @@ -20862,6 +26122,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -22336,6 +27601,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -22856,6 +28126,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -25365,6 +30640,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -31655,7 +36935,7 @@ "paths": { "/openai/v1/realtime/calls": { "post": { - "operationId": "proxy_realtime_calls_openai_v1_realtime_calls_post", + "operationId": "proxy_realtime_calls_openai_v1_realtime_calls_post_2", "responses": { "200": { "content": { @@ -31674,7 +36954,7 @@ }, "/openai/v1/realtime/client_secrets": { "post": { - "operationId": "create_realtime_client_secret_openai_v1_realtime_client_secrets_post", + "operationId": "create_realtime_client_secret_openai_v1_realtime_client_secrets_post_2", "responses": { "200": { "content": { @@ -31701,7 +36981,7 @@ "/openai/v1/realtime/transcription_sessions": { "post": { "description": "Create an ephemeral Realtime transcription session\n(POST /v1/realtime/transcription_sessions) for the WebRTC/WebSocket flow.\n\nMirrors the client_secrets route but targets the transcription_sessions\nendpoint and encrypts the ephemeral key returned under `client_secret.value`.", - "operationId": "create_realtime_transcription_session_openai_v1_realtime_transcription_sessions_post", + "operationId": "create_realtime_transcription_session_openai_v1_realtime_transcription_sessions_post_2", "responses": { "200": { "content": { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f83011835fd..7b0e92d81ae 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3,6 +3,7 @@ import json import os from collections.abc import Callable, Mapping from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple import httpx @@ -496,12 +497,16 @@ class LiteLLMRoutes(enum.Enum): "/v1/messages/count_tokens", "/v1/skills", "/v1/skills/{skill_id}", + "/claude-code/marketplace.json", + "/claude-code/plugins", + "/claude-code/plugins/{plugin_name}", ] # MCP tool-call / passthrough routes — data-plane. Gated by DISABLE_LLM_API_ENDPOINTS. mcp_inference_routes = [ "/mcp", "/mcp/", + "/mcp/proxy", "/mcp/{subpath}", "/mcp/tools", "/mcp/tools/list", @@ -698,6 +703,7 @@ class LiteLLMRoutes(enum.Enum): "/spend/logs", "/spend/logs/v2", "/spend/logs/ui", + "/spend/logs/ui/{request_id}", "/spend/logs/session/ui", "/key/spend/report", "/user/spend/report", @@ -835,6 +841,14 @@ class LiteLLMRoutes(enum.Enum): "/team/daily/activity/aggregated", "/team/spend/by_user", "/team/{team_id}/members/me", + # POST/GET the team's logging callbacks, and DELETE one of them. Every + # handler calls _verify_team_access, which admits only a proxy admin, an + # org admin for the team, or an admin of this team. + # + # team_id is a free-form string, so it spells these with the same path + # converter the router uses; the gate matches that converter. + "/team/{team_id:path}/callback", + "/team/{team_id:path}/callback/{callback_name}", "/model/new", "/model/update", "/model/delete", @@ -848,6 +862,7 @@ class LiteLLMRoutes(enum.Enum): "/model/{model_id}/update", "/prompt/list", "/prompt/info", + "/vector_store/info", # Project read routes - endpoint scopes results to caller's teams (non-admin) "/project/list", "/project/info", @@ -865,6 +880,8 @@ class LiteLLMRoutes(enum.Enum): # proxy admin, or team admin naming their own team via team_id "/auto_router/test_routing", "/auto_router/validate_complexity_router_config", + # Per-session auto-router read - the endpoint scopes the row to the caller's own key hash + "/auto_router/session", # Agent registry - reads are role-scoped and writes are proxy-admin-gated # inside agent_endpoints/endpoints.py *agent_management_routes, @@ -918,10 +935,10 @@ class LiteLLMRoutes(enum.Enum): # PROXY_ADMIN_VIEW_ONLY — the route gate must match). "/customer/list", "/customer/info", - # UI Logs page detail drawer (single + session) and the filter facets. - # The list endpoint `/spend/logs/ui` is covered via - # spend_tracking_routes below. - "/spend/logs/ui/{logId}", + # UI Logs page session detail drawer and the end-user filter facet. + # The list endpoint `/spend/logs/ui` and the single-log detail route + # `/spend/logs/ui/{request_id}` are covered via spend_tracking_routes + # below. "/spend/logs/session/ui", "/management/v1/spend_logs/end_users", "/management/v1/spend_logs/users", @@ -1113,6 +1130,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): models: list[str] | None = None search_tools: list[str] | None = None mcp_tool_search_enabled: bool | None = None + skills: list[str] | None = None from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 @@ -1276,6 +1294,7 @@ class UpdateKeyRequest(KeyRequestBase): # else they will get overwritten duration: str | None = None spend: float | None = None + soft_budget: float | None = None metadata: dict | None = None temp_budget_increase: float | None = None temp_budget_expiry: datetime | None = None @@ -1283,6 +1302,13 @@ class UpdateKeyRequest(KeyRequestBase): rotation_interval: str | None = None organization_id: str | None = None + @model_validator(mode="before") + @classmethod + def drop_blank_team_id(cls, values: object) -> object: + if isinstance(values, Mapping) and values.get("team_id") == "": + return MappingProxyType({k: v for k, v in values.items() if k != "team_id"}) + return values + @field_validator("organization_id", mode="before") @classmethod def treat_cleared_organization_id_as_unset(cls, v: object) -> object: @@ -1379,6 +1405,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 +1475,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 +1540,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 +1588,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 +1633,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, @@ -2365,6 +2434,13 @@ class CoordinationRedisParams(LiteLLMPydanticObjectBase): ) sentinel_password: str | None = Field(None, description="password for the sentinel nodes") service_name: str | None = Field(None, description="sentinel service name") + aws_iam_auth: bool | str | None = Field(None, description="enable AWS ElastiCache IAM authentication") + aws_iam_user_name: str | None = Field(None, description="AWS ElastiCache IAM user name") + aws_iam_cache_name: str | None = Field(None, description="AWS ElastiCache cache name") + aws_iam_region: str | None = Field(None, description="AWS region for ElastiCache IAM authentication") + aws_iam_serverless: bool | str | None = Field( + None, description="the ElastiCache cache is serverless rather than a self-designed cluster" + ) def has_connection_target(self) -> bool: return any(value is not None for value in (self.host, self.url, self.startup_nodes, self.sentinel_nodes)) @@ -2528,6 +2604,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): global_max_parallel_requests: int | None = Field( None, description="global max parallel requests to allow for a proxy instance." ) + user_api_key_cache_max_size: int | None = Field( + None, + gt=0, + description=( + "max number of entries (virtual keys, teams, users, end users, memberships, ...) each worker keeps in " + "its in-memory auth cache. Defaults to 200. Raise this if you have more active keys than that or auth " + "lookups keep hitting the DB" + ), + ) max_request_size_mb: int | None = Field( None, description="max request size in MB, if a request is larger than this size it will be rejected", @@ -2774,6 +2859,37 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "UI username/password login. Default is False." ), ) + disable_responses_id_security: bool | None = Field( + None, + description=( + "If True, disables ownership enforcement on Responses API ids. " + "Keys may then retrieve, cancel, delete, and chain from any response id, " + "including ids belonging to another user or team and ids this proxy never issued. " + "WARNING: this removes tenant isolation on /v1/responses" + ), + ) + allow_unmanaged_response_ids: bool | None = Field( + None, + description=( + "If True, lets keys address Responses API ids that this proxy did not issue " + "(raw provider ids, or ids issued before response-id encryption was configured). " + "Such an id carries no owner, so no ownership check can run on it; ids this proxy " + "did issue keep full ownership enforcement. Off by default, in which case an " + "unrecognized response id is rejected with 403" + ), + ) + disable_env_credential_login: bool | None = Field( + None, + description=( + "If True, disables signing in to the Admin UI with the environment credentials: " + "UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback " + "means env-credential login is always live by default). Database users with passwords " + "are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password " + "before enabling, or nobody can sign in to the UI. A locked-out admin can still " + "administer the proxy over the API with the master key, and can unset this setting " + "and restart the proxy to restore env-credential login. Default is False." + ), + ) disable_budget_reservation: bool | None = Field( None, description=( @@ -2788,7 +2904,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "Enable only if your deployment is experiencing phantom " "BudgetExceededError responses caused by leaked reservations " "(see GitHub issue #27639). " - "A proxy-level WARNING is logged on every request while this flag " + "An INFO notice is logged once per worker at config load while this flag " "is active as a reminder that hard enforcement is relaxed." ), ) @@ -3539,7 +3655,9 @@ class AllCallbacks(LiteLLMPydanticObjectBase): ui_callback_name="OpenTelemetry", litellm_callback_params=[ "OTEL_EXPORTER", + "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_ENDPOINT", + "OTEL_TRACES_ENDPOINT", "OTEL_HEADERS", ], ) @@ -3650,6 +3768,15 @@ class AllCallbacks(LiteLLMPydanticObjectBase): ], ) + pointfive: CallbackOnUI = CallbackOnUI( + litellm_callback_name="pointfive", + ui_callback_name="PointFive", + litellm_callback_params=[ # mutable-ok: the registry field is typed list + "POINTFIVE_API_KEY", + "POINTFIVE_API_URL", + ], + ) + class SpendLogsRouterMetadata(TypedDict): """ @@ -5092,9 +5219,26 @@ class CostEstimateRequest(LiteLLMPydanticObjectBase): model: str = Field(description="Model name (from /model_group/info)") input_tokens: int = Field(description="Expected input tokens per request", ge=0) output_tokens: int = Field(description="Expected output tokens per request", ge=0) + cache_read_input_tokens: int = Field( + default=0, description="Input tokens read from the prompt cache; counted within input_tokens", ge=0 + ) + cache_creation_input_tokens: int = Field( + default=0, description="Input tokens written to the prompt cache; counted within input_tokens", ge=0 + ) + reasoning_tokens: int = Field( + default=0, description="Reasoning tokens the model emits; counted within output_tokens", ge=0 + ) num_requests_per_day: int | None = Field(default=None, description="Number of requests per day", ge=0) num_requests_per_month: int | None = Field(default=None, description="Number of requests per month", ge=0) + @model_validator(mode="after") + def validate_token_subsets(self) -> "CostEstimateRequest": + if self.cache_read_input_tokens + self.cache_creation_input_tokens > self.input_tokens: + raise ValueError("cache_read_input_tokens plus cache_creation_input_tokens cannot exceed input_tokens") + if self.reasoning_tokens > self.output_tokens: + raise ValueError("reasoning_tokens cannot exceed output_tokens") + return self + class CostEstimateResponse(LiteLLMPydanticObjectBase): """Response body for /cost/estimate endpoint.""" @@ -5102,6 +5246,9 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase): model: str input_tokens: int output_tokens: int + cache_read_input_tokens: int = 0 + cache_creation_input_tokens: int = 0 + reasoning_tokens: int = 0 num_requests_per_day: int | None = None num_requests_per_month: int | None = None # Per-request costs @@ -5109,17 +5256,33 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase): input_cost_per_request: float = Field(description="Input token cost per request (before margin)") output_cost_per_request: float = Field(description="Output token cost per request (before margin)") margin_cost_per_request: float = Field(default=0.0, description="Margin/fee added per request") + cache_read_cost_per_request: float = Field(default=0.0, description="Cache-read share of input_cost_per_request") + cache_creation_cost_per_request: float = Field( + default=0.0, description="Cache-write share of input_cost_per_request" + ) + reasoning_cost_per_request: float = Field(default=0.0, description="Reasoning share of output_cost_per_request") # Daily costs (if num_requests_per_day provided) daily_cost: float | None = Field(default=None, description="Total daily cost (includes margin)") daily_input_cost: float | None = Field(default=None, description="Daily input token cost") daily_output_cost: float | None = Field(default=None, description="Daily output token cost") daily_margin_cost: float | None = Field(default=None, description="Daily margin/fee") + daily_cache_read_cost: float | None = Field(default=None, description="Cache-read share of daily_input_cost") + daily_cache_creation_cost: float | None = Field(default=None, description="Cache-write share of daily_input_cost") + daily_reasoning_cost: float | None = Field(default=None, description="Reasoning share of daily_output_cost") # Monthly costs (if num_requests_per_month provided) monthly_cost: float | None = Field(default=None, description="Total monthly cost (includes margin)") monthly_input_cost: float | None = Field(default=None, description="Monthly input token cost") monthly_output_cost: float | None = Field(default=None, description="Monthly output token cost") monthly_margin_cost: float | None = Field(default=None, description="Monthly margin/fee") - # Pricing info - input_cost_per_token: float | None = None - output_cost_per_token: float | None = None + monthly_cache_read_cost: float | None = Field(default=None, description="Cache-read share of monthly_input_cost") + monthly_cache_creation_cost: float | None = Field( + default=None, description="Cache-write share of monthly_input_cost" + ) + monthly_reasoning_cost: float | None = Field(default=None, description="Reasoning share of monthly_output_cost") + # Pricing info: the rates this request's usage bills at, after token tiers and regional multipliers + input_cost_per_token: float | None = Field(default=None, description="Rate billed per input token") + output_cost_per_token: float | None = Field(default=None, description="Rate billed per output token") + cache_read_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-read token") + cache_creation_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-write token") + output_cost_per_reasoning_token: float | None = Field(default=None, description="Rate billed per reasoning token") provider: str | None = None diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 28882484db4..95c34f70d7b 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -144,31 +144,32 @@ def _validate_push_notification_url(url: str) -> None: raise HTTPException(status_code=400, detail=str(e)) from e -def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> dict[str, str]: - headers: Final[dict[str, str]] = {} - if user_api_key_dict.user_id: - headers["X-LiteLLM-User-Id"] = user_api_key_dict.user_id - if user_api_key_dict.team_id: - headers["X-LiteLLM-Team-Id"] = user_api_key_dict.team_id - return headers +def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, str]: + return MappingProxyType( + { + name: value + for name, value in ( + ("X-LiteLLM-User-Id", user_api_key_dict.user_id), + ("X-LiteLLM-Team-Id", user_api_key_dict.team_id), + ) + if value + } + ) def _forwarding_headers( - user_api_key_dict: UserAPIKeyAuth, + caller_identity: Mapping[str, str], request_data: Mapping[str, object], agent_extra_headers: Mapping[str, str] | None, -) -> Mapping[str, str] | None: - sanitized: Final = ( - {k: v for k, v in agent_extra_headers.items() if not k.lower().startswith("x-litellm-")} - if agent_extra_headers - else None +) -> dict[str, str] | None: + passthrough: Final = tuple( + (name, value) + for name, value in (agent_extra_headers.items() if agent_extra_headers else ()) + if not name.lower().startswith("x-litellm-") ) - merged: Final = merge_agent_headers(dynamic_headers=sanitized, static_headers=None) or {} - identity: Final = _caller_identity_headers(user_api_key_dict) trace_id: Final = request_data.get("litellm_trace_id") - if trace_id: - identity["X-LiteLLM-Trace-Id"] = str(trace_id) - merged.update(identity) + trace: Final = (("X-LiteLLM-Trace-Id", str(trace_id)),) if trace_id else () + merged: Final = dict((*passthrough, *caller_identity.items(), *trace)) return merged or None @@ -755,6 +756,7 @@ async def invoke_agent_a2a( ProxyBaseLLMRequestProcessing, ) + caller_identity: Final = _caller_identity_headers(user_api_key_dict) processor: Final = ProxyBaseLLMRequestProcessing(data=body) data, logging_obj = await processor.common_processing_pre_call_logic( request=request, @@ -793,9 +795,13 @@ async def invoke_agent_a2a( if header_name: dynamic_headers[header_name] = val - agent_extra_headers = merge_agent_headers( - dynamic_headers=dynamic_headers or None, - static_headers=static_headers or None, + agent_extra_headers = _forwarding_headers( + caller_identity=caller_identity, + request_data=data, + agent_extra_headers=merge_agent_headers( + dynamic_headers=dynamic_headers or None, + static_headers=static_headers or None, + ), ) # Databricks App endpoints require a short-lived OAuth M2M token rather @@ -942,12 +948,7 @@ async def invoke_agent_a2a( "method": method, "params": params, } - caller_headers: Final = _forwarding_headers( - user_api_key_dict=user_api_key_dict, - request_data=data, - agent_extra_headers=agent_extra_headers, - ) - result = await _forward_jsonrpc(agent_url, forward_body, extra_headers=caller_headers) + result = await _forward_jsonrpc(agent_url, forward_body, extra_headers=agent_extra_headers) if method == "agent/getAuthenticatedExtendedCard": card: Final = result.get("result") if isinstance(card, dict): @@ -988,16 +989,11 @@ async def invoke_agent_a2a( "method": method, "params": params, } - sse_caller_headers: Final = _forwarding_headers( - user_api_key_dict=user_api_key_dict, - request_data=data, - agent_extra_headers=agent_extra_headers, - ) return await _forward_jsonrpc_sse( agent_url, forward_body, request_id=request_id, - extra_headers=sse_caller_headers, + extra_headers=agent_extra_headers, proxy_logging_obj=proxy_logging_obj, user_api_key_dict=user_api_key_dict, request_data=data, diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py index 76e3fe6c5ad..65a89bb2c7a 100644 --- a/litellm/proxy/agent_endpoints/agent_search.py +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -18,6 +18,7 @@ from litellm.types.agents import AgentResponse if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging from litellm.router import Router DEFAULT_AGENT_SEARCH_TOP_K: Final = 5 @@ -132,6 +133,7 @@ async def search_agents( embedding_model: str | None, index: AgentSearchIndex, user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, ) -> AgentSearchOutcome: if embedding_model is None: return AgentSearchNotConfigured( @@ -139,6 +141,5 @@ async def search_agents( ) if router is None: return AgentSearchNotConfigured(reason="agent search needs a model_list so the embedding model can be called") - return await index.search( - query, agents, top_k, router_embedder(router, embedding_model, user_api_key_dict), embedding_model - ) + embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj) + return await index.search(query, agents, top_k, embed, embedding_model) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index cc17672553b..aa8979a73c6 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -249,7 +249,7 @@ def _agent_search_error(status_code: int, error: str, message: str) -> HTTPExcep async def _rank_agents_by_query( query: str, agents: Sequence[AgentResponse], top_k: int, user_api_key_dict: UserAPIKeyAuth ) -> tuple[AgentResponse, ...]: - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj outcome: Final = await search_agents( query=query, @@ -259,6 +259,7 @@ async def _rank_agents_by_query( embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) match outcome: case AgentSearchHits(hits): diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 65bc46edfaf..7c6a4571948 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -2,14 +2,15 @@ CLAUDE CODE MARKETPLACE Provides a registry/discovery layer for Claude Code plugins. -Plugins are stored as metadata + git source references in LiteLLM database. -Actual plugin files are hosted on GitHub/GitLab/Bitbucket. +Plugins are stored as metadata + source references in LiteLLM database. +Actual plugin files are hosted on GitHub/GitLab/Bitbucket or as a zip archive on +any HTTPS host (S3, Artifactory, a static file server). Endpoints: -/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated) +/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated; `?key=` adds the key's granted skills) /claude-code/plugins - POST - Register a new plugin (create-only, proxy admin only) -/claude-code/plugins - GET - List plugins (any authenticated key) -/claude-code/plugins/{name} - GET - Get plugin details (any authenticated key) +/claude-code/plugins - GET - List plugins visible to the key (enabled, plus granted disabled ones) +/claude-code/plugins/{name} - GET - Get plugin details (403 on a disabled plugin the key is not granted) /claude-code/plugins/{name} - PUT - Update an existing plugin (proxy admin only) /claude-code/plugins/{name}/enable - POST - Enable a plugin (proxy admin only) /claude-code/plugins/{name}/disable - POST - Disable a plugin (proxy admin only) @@ -21,12 +22,17 @@ import re from collections.abc import Mapping, Sequence from datetime import datetime, timezone from typing import Annotated, Final, Protocol, TypedDict +from urllib.parse import urlsplit -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import JSONResponse from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, ProxyException, UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( + SkillVisibility, + skill_visibility, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.repositories.table_repositories import ClaudeCodePluginRepository @@ -82,7 +88,7 @@ async def _get_prisma_client() -> object: "/claude-code/marketplace.json", tags=["Claude Code Marketplace"], ) -async def get_marketplace(): +async def get_marketplace(request: Request, key: str | None = None): """ Serve marketplace.json for Claude Code plugin discovery. @@ -90,24 +96,35 @@ async def get_marketplace(): - claude plugin marketplace add - claude plugin install @ + Without `key` the catalog holds the enabled (public) plugins. With `?key=sk-...` + the key is authenticated and the catalog also holds the disabled plugins granted + to it through `object_permission.skills` on the key or its team. + Returns: Marketplace catalog with list of available plugins and their git sources. Example: ```bash claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json + claude plugin marketplace add "http://localhost:4000/claude-code/marketplace.json?key=sk-..." claude plugin install my-plugin@litellm ``` """ try: prisma_client: Final = await _get_prisma_client() + caller: Final[UserAPIKeyAuth | None] = ( + await user_api_key_auth(request=request, api_key=f"Bearer {key}") if key else None + ) + visibility: Final[SkillVisibility] = skill_visibility(caller) plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many( - where={"enabled": True} + where=visibility.where() ) plugin_list: Final = [] for plugin in plugins: + if not visibility.allows(plugin): + continue try: manifest: Mapping[str, object] = json.loads(plugin.manifest_json or "{}") except json.JSONDecodeError: @@ -147,7 +164,7 @@ async def get_marketplace(): return JSONResponse(content=marketplace) - except HTTPException: + except (HTTPException, ProxyException): raise except Exception as e: verbose_proxy_logger.exception("Error generating marketplace: %s", e) @@ -162,6 +179,15 @@ async def get_marketplace(): # alphanumeric characters, dots, hyphens, and underscores. # This implicitly blocks '..', leading '/', backslashes, and percent-encoded sequences. _VALID_GIT_SUBDIR_PATH_RE: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a-zA-Z0-9][a-zA-Z0-9._-]*)*$") +_VALID_SHA256_RE: Final = re.compile(r"^[0-9a-fA-F]{64}$") + + +def _is_https_url_with_host(url: str) -> bool: + try: + parts: Final = urlsplit(url) + except ValueError: + return False + return parts.scheme == "https" and bool(parts.hostname) def _validate_plugin_source(source: Mapping[str, str]) -> None: @@ -199,10 +225,24 @@ def _validate_plugin_source(source: Mapping[str, str]) -> None: "error": "git-subdir 'path' must be a relative path of the form 'segment/segment' (alphanumeric, dots, hyphens, underscores only)" }, ) + elif source_type == "archive": + if not _is_https_url_with_host(source.get("url", "")): + raise HTTPException( + status_code=400, + detail={ + "error": "archive source must include an https 'url' field " + "(e.g., 'https://bucket.s3.amazonaws.com/plugins/plugin-name.zip')" + }, + ) + if "sha256" in source and not _VALID_SHA256_RE.match(source["sha256"]): + raise HTTPException( + status_code=400, + detail={"error": "archive 'sha256' must be a 64-character hex digest"}, + ) else: raise HTTPException( status_code=400, - detail={"error": "source.source must be 'github', 'url', or 'git-subdir'"}, + detail={"error": "source.source must be 'github', 'url', 'git-subdir', or 'archive'"}, ) @@ -248,8 +288,8 @@ async def register_plugin( Register a new plugin in the LiteLLM marketplace. LiteLLM acts as a registry/discovery layer. Plugins are hosted on - GitHub/GitLab/Bitbucket. Claude Code will clone from the git source - when users install. + GitHub/GitLab/Bitbucket or as a zip archive on any https host (e.g. S3). + Claude Code clones the git source or downloads the archive when users install. This endpoint is create-only and never overwrites. If a plugin with the same name already exists it returns 409 Conflict; use @@ -259,7 +299,7 @@ async def register_plugin( Parameters: - name: Plugin name (kebab-case) - - source: Git source reference (github, url, or git-subdir format) + - source: Plugin source reference (github, url, git-subdir, or archive format) - version: Semantic version (optional) - description: Plugin description (optional) - author: Author information (optional) @@ -370,13 +410,15 @@ async def list_plugins( try: prisma_client: Final = await _get_prisma_client() - where: Final = {"enabled": True} if enabled_only else {} + visibility: Final[SkillVisibility] = skill_visibility(user_api_key_dict) plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many( - where=where + where={"enabled": True} if enabled_only else visibility.where() ) plugin_list: Final = [] for p in plugins: + if not visibility.allows(p): + continue # Parse manifest to get additional fields manifest = json.loads(p.manifest_json) if p.manifest_json else {} @@ -448,6 +490,12 @@ async def get_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) + if not skill_visibility(user_api_key_dict).allows(plugin): + raise HTTPException( + status_code=403, + detail={"error": f"Plugin '{plugin_name}' is not granted to this key"}, + ) + manifest: Final[Mapping[str, object]] = json.loads(plugin.manifest_json or "{}") if plugin.manifest_json else {} return { @@ -503,7 +551,7 @@ async def update_plugin( Parameters: - plugin_name: Name of the plugin to update (path parameter) - - source: Git source reference (github, url, or git-subdir format) + - source: Plugin source reference (github, url, git-subdir, or archive format) - version: Semantic version (optional) - description: Plugin description (optional) - author: Author information (optional) diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py new file mode 100644 index 00000000000..e1e1ee6f160 --- /dev/null +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py @@ -0,0 +1,67 @@ +""" +Claude Code marketplace visibility: enabled plugins are public, disabled plugins +are private and resolve only for proxy admins or keys granted them via +``object_permission.skills``. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Protocol + +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin + +if TYPE_CHECKING: + from prisma.types import LiteLLM_ClaudeCodePluginTableWhereInput + + +class _SkillRecord(Protocol): + name: str + enabled: bool + + +def _skills_of(permission: LiteLLM_ObjectPermissionTable | None) -> frozenset[str]: + return frozenset(permission.skills or ()) if permission is not None else frozenset() + + +def granted_skills(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: + """Key grant intersected with the team grant when both are non-empty; either alone applies as is. + + An empty list is the Prisma column default for every object-permission row, so it means + "no private grants configured here" and defers to the other scope, same as the agents check. + """ + key_skills: Final = _skills_of(user_api_key_dict.object_permission) + team_skills: Final = _skills_of(user_api_key_dict.team_object_permission) + match (bool(key_skills), bool(team_skills)): + case (True, True): + return key_skills & team_skills + case (True, False): + return key_skills + case _: + return team_skills + + +@dataclass(frozen=True, slots=True) +class SkillVisibility: + granted: frozenset[str] + sees_private: bool + + def allows(self, skill: _SkillRecord) -> bool: + return skill.enabled or self.sees_private or skill.name in self.granted + + def where(self) -> "LiteLLM_ClaudeCodePluginTableWhereInput": + if self.sees_private: + return {} + if not self.granted: + return {"enabled": True} + return {"OR": [{"enabled": True}, {"name": {"in": sorted(self.granted)}}]} + + +PUBLIC_ONLY: Final = SkillVisibility(granted=frozenset(), sees_private=False) + + +def skill_visibility(user_api_key_dict: UserAPIKeyAuth | None) -> SkillVisibility: + if user_api_key_dict is None: + return PUBLIC_ONLY + if is_proxy_admin(user_api_key_dict): + return SkillVisibility(granted=frozenset(), sees_private=True) + return SkillVisibility(granted=granted_skills(user_api_key_dict), sees_private=False) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index b243b737b0a..d4cb3b84ee4 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -25,6 +25,11 @@ from litellm.proxy.common_request_processing import ( proxy_exception_from_http_exception, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.types.utils import TokenCountResponse router: Final = APIRouter() @@ -243,9 +248,9 @@ async def anthropic_response( return _anthropic_error_json_response( ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), headers=headers, ), request, diff --git a/litellm/proxy/anthropic_endpoints/skills_endpoints.py b/litellm/proxy/anthropic_endpoints/skills_endpoints.py index 9390bf4c537..4426c0b547a 100644 --- a/litellm/proxy/anthropic_endpoints/skills_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/skills_endpoints.py @@ -2,11 +2,23 @@ Anthropic Skills API endpoints - /v1/skills """ -from typing import Final +from types import MappingProxyType +from typing import Annotated, Final import orjson -from fastapi import APIRouter, Depends, Request, Response +from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response +from typing_extensions import ReadOnly, TypedDict, assert_never +import litellm +from litellm.llms.litellm_proxy.skills.skill_search import ( + DEFAULT_SKILL_SEARCH_TOP_K, + SkillSearchEmbeddingFailed, + SkillSearchHits, + SkillSearchNotConfigured, + SkillSearchUnsupportedProvider, + global_skill_search_index, + search_hosted_skills, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -23,6 +35,51 @@ from litellm.types.llms.anthropic_skills import ( router: Final = APIRouter() +class _SkillSearchErrorDetail(TypedDict): + error: ReadOnly[str] + message: ReadOnly[str] + + +def _skill_search_error(status_code: int, error: str, message: str) -> HTTPException: + detail: Final[_SkillSearchErrorDetail] = {"error": error, "message": message} + return HTTPException(status_code=status_code, detail=detail) + + +async def _search_skills( + custom_llm_provider: str | None, query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth +) -> ListSkillsResponse: + from litellm.llms.litellm_proxy.skills.transformation import ( + LiteLLMSkillsTransformationHandler, + ) + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj + + outcome: Final = await search_hosted_skills( + custom_llm_provider=custom_llm_provider, + query=query, + top_k=top_k, + router=llm_router, + embedding_model=litellm.skill_search_embedding_model, + index=global_skill_search_index, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + to_response: Final = LiteLLMSkillsTransformationHandler().db_skill_to_response + match outcome: + case SkillSearchHits(hits): + skills: Final = [ # mutable-ok: ListSkillsResponse.data requires list[Skill]; never mutated after + to_response(hit.skill).model_copy(update=MappingProxyType({"search_score": hit.score})) for hit in hits + ] + return ListSkillsResponse(data=skills, has_more=False, next_page=None) + case SkillSearchUnsupportedProvider(reason): + raise _skill_search_error(400, "skill_search_unsupported_provider", reason) + case SkillSearchNotConfigured(reason): + raise _skill_search_error(400, "skill_search_not_configured", reason) + case SkillSearchEmbeddingFailed(reason): + raise _skill_search_error(503, "skill_search_unavailable", reason) + case _: + assert_never(outcome) + + @router.post( "/v1/skills", tags=["[beta] Anthropic Skills API"], @@ -134,32 +191,58 @@ async def list_skills( after_id: str | None = None, before_id: str | None = None, custom_llm_provider: str | None = "anthropic", + query: Annotated[ + str | None, + Query( + min_length=1, + description="Describe what you need in natural language to rank the skills you can access by " + "semantic similarity over their title and description. Each result carries a search_score. " + "Only supported for custom_llm_provider=litellm_proxy. Requires " + "litellm_settings.skill_search_embedding_model.", + ), + ] = None, + top_k: Annotated[ + int, + Query(ge=1, le=100, description="With query: the maximum number of ranked skills to return."), + ] = DEFAULT_SKILL_SEARCH_TOP_K, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ List skills on Anthropic. - + Requires `?beta=true` query parameter. - + Model-based routing (for multi-account support): - Pass model via header: `x-litellm-model: claude-account-1` - Pass model via query: `?model=claude-account-1` - Pass model via body: `{"model": "claude-account-1"}` - + Example usage: ```bash # Basic usage curl "http://localhost:4000/v1/skills?beta=true&limit=10" \ -H "Authorization: Bearer your-key" - + # With model-based routing curl "http://localhost:4000/v1/skills?beta=true&limit=10" \ -H "Authorization: Bearer your-key" \ -H "x-litellm-model: claude-account-1" ``` - + + Pass `?custom_llm_provider=litellm_proxy&query=` to rank the LiteLLM-hosted skills you can + access by semantic similarity instead of paging through the whole registry: + ```bash + curl "http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5" \ + -H "Authorization: Bearer your-key" + ``` + Returns: ListSkillsResponse with list of skills """ + if query is not None: + return await _search_skills( + custom_llm_provider=custom_llm_provider, query=query, top_k=top_k, user_api_key_dict=user_api_key_dict + ) + from litellm.proxy.proxy_server import ( general_settings, llm_router, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dc693317de0..1efc9611fe6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -94,6 +94,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_auth_cache_key, team_membership_reservation_cache_key, ) +from litellm.proxy.db.db_lookup_gate import db_lookup_gate from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( TOOL_CAPABLE_CALL_TYPES, @@ -475,6 +476,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) +_TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True}) def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool: @@ -2858,7 +2860,9 @@ class TeamNotFoundError(HTTPException): async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None ) -> "_PrismaTeamRow | None": - response = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) + response = await _team_table(TeamRepository(prisma_client)).find_unique( + where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS + ) if response is None and team_id_upsert: from litellm.proxy.management_endpoints.team_endpoints import new_team @@ -3158,7 +3162,9 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(where={"team_alias": team_alias}) + teams: Final = await _team_table(TeamRepository(prisma_client)).find_many( + where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS + ) if not teams: raise HTTPException( @@ -3445,36 +3451,37 @@ async def _fetch_key_object_from_db_with_reconnect( """ Fetch key object from DB and retry once if a DB connection error can be healed. """ - try: - return await prisma_client.get_data( - token=hashed_token, - table_name="combined_view", - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: - if PrismaDBExceptionHandler.is_database_transport_error(e): - did_reconnect = False - if hasattr(prisma_client, "attempt_db_reconnect"): - auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0) - if not isinstance(auth_reconnect_timeout, (int, float)): - auth_reconnect_timeout = 2.0 - auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1) - if not isinstance(auth_reconnect_lock_timeout, (int, float)): - auth_reconnect_lock_timeout = 0.1 - did_reconnect = await prisma_client.attempt_db_reconnect( - reason="auth_get_key_object_lookup_failure", - timeout_seconds=auth_reconnect_timeout, - lock_timeout_seconds=auth_reconnect_lock_timeout, - ) - if did_reconnect: - return await prisma_client.get_data( - token=hashed_token, - table_name="combined_view", - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - raise + async with db_lookup_gate.current(): + try: + return await prisma_client.get_data( + token=hashed_token, + table_name="combined_view", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + if PrismaDBExceptionHandler.is_database_transport_error(e): + did_reconnect = False + if hasattr(prisma_client, "attempt_db_reconnect"): + auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0) + if not isinstance(auth_reconnect_timeout, (int, float)): + auth_reconnect_timeout = 2.0 + auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1) + if not isinstance(auth_reconnect_lock_timeout, (int, float)): + auth_reconnect_lock_timeout = 0.1 + did_reconnect = await prisma_client.attempt_db_reconnect( + reason="auth_get_key_object_lookup_failure", + timeout_seconds=auth_reconnect_timeout, + lock_timeout_seconds=auth_reconnect_lock_timeout, + ) + if did_reconnect: + return await prisma_client.get_data( + token=hashed_token, + table_name="combined_view", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + raise def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d6007a2d56e..be65c3b39ec 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -11,7 +11,7 @@ from fastapi import HTTPException, Request, status from pydantic import PositiveInt, TypeAdapter, ValidationError import litellm -from litellm import Router, provider_list +from litellm import Router, constants, provider_list from litellm._logging import verbose_proxy_logger from litellm.constants import ( BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.url_utils import ( provider_url_destination_candidates, validate_url, ) +from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint from litellm.proxy._types import * from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -316,6 +317,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( "aws_profile_name", "aws_session_name", "aws_external_id", + "aws_session_tags", "vertex_credentials", # Azure managed-identity / federated-auth token. The Azure provider # transformer reads ``azure_ad_token`` (top-level or via @@ -1390,6 +1392,24 @@ def warn_once_if_custom_auth_skips_common_checks( _custom_auth_common_checks_warning_emitted = True +def log_once_if_budget_reservation_disabled( + *, + disabled: bool, + logger: Logger = verbose_proxy_logger, +) -> None: + if constants.budget_reservation_disabled_info_emitted or not disabled: + return + logger.info( + "disable_budget_reservation is enabled: skipping optimistic budget " + "reservation. Budget enforcement is read-time only. Concurrent " + "requests can each pass the spend check before their cost is recorded, " + "so a configured budget may be briefly exceeded under high concurrency. " + "Set disable_budget_reservation to False or remove it to restore " + "hard per-request budget enforcement." + ) + constants.budget_reservation_disabled_info_emitted = True # rebind-ok: process-wide one-shot sentinel + + def is_pass_through_provider_route(route: str) -> bool: PROVIDER_SPECIFIC_PASS_THROUGH_ROUTES: Final = [ "vertex-ai", @@ -1981,9 +2001,39 @@ def get_model_from_request( if vertex_match: model = vertex_match.group(1) + if route.lower().startswith("/bedrock"): + bedrock_model: Final = _model_from_bedrock_route(route) + return model if bedrock_model is None else bedrock_model + + if route.lower().startswith(("/azure/", "/azure_ai/")): + azure_model: Final = _router_model_from_azure_route(route, llm_router) + return model if azure_model is None else azure_model + return model +def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None: + if llm_router is None: + return None + endpoint: Final = re.sub(r"^/azure(?:_ai)?/", "", route, flags=re.IGNORECASE) + return azure_router_model_in_endpoint(endpoint, frozenset(llm_router.get_model_names())) + + +def _model_from_bedrock_route(route: str) -> str | None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _extract_model_from_bedrock_endpoint, + is_bedrock_count_tokens_endpoint, + ) + + bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) + if is_bedrock_count_tokens_endpoint(bedrock_endpoint): + return None + try: + return _extract_model_from_bedrock_endpoint(bedrock_endpoint) + except ValueError: + return None + + def abbreviate_api_key(api_key: str) -> str: if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH: return "sk-..." diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 0795cee7409..69091ee8344 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -53,6 +53,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.team_grants import team_model_aliases from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, get_management_object_ttl, @@ -1595,7 +1596,7 @@ class JWTAuthManager: model=requested_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=None, + team_model_aliases=team_model_aliases(team_object), ) ): is_allowed = allowed_routes_check( @@ -2132,7 +2133,7 @@ class JWTAuthManager: model=requested_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=None, + team_model_aliases=team_model_aliases(team_object), ) except ProxyException: continue diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 8d4f6f81363..c0a76a4fc20 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -85,6 +85,29 @@ def get_ui_credentials(master_key: str | None) -> tuple[str, str]: return ui_username, ui_password +def _matches_env_credentials(username: str, password: str, master_key: str | None) -> bool: + ui_username, ui_password = get_ui_credentials(master_key) + return secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest( + password.encode("utf-8"), ui_password.encode("utf-8") + ) + + +def is_env_credential_login_enabled(general_settings: Mapping[str, object]) -> bool: + """Whether a login with UI_USERNAME/UI_PASSWORD (or the master-key fallback) can succeed. + + Two settings can turn it off: `disable_env_credential_login` unconditionally, and + `disable_password_login_when_sso_enabled` as a side effect, since its gate rejects + every username/password login before the env comparison runs. Feeds both the + `authenticate_user` gate and the Admin UI warning banner, so the banner never nags + about a login path that is already unreachable. + """ + if general_settings.get("disable_env_credential_login") is True: + return False + if general_settings.get("disable_password_login_when_sso_enabled") is True and is_sso_provider_fully_configured(): + return False + return True + + class LoginResult: """Result object containing authentication data from login.""" @@ -129,7 +152,8 @@ async def authenticate_user( master_key: Master key for the proxy (required) prisma_client: Prisma database client (optional) general_settings: Proxy general_settings, checked for - `disable_password_login_when_sso_enabled` + `disable_password_login_when_sso_enabled` and + `disable_env_credential_login` Returns: LoginResult: Object containing authentication data @@ -170,8 +194,6 @@ async def authenticate_user( code=500, ) - ui_username, ui_password = get_ui_credentials(master_key) - # Check if we can find the `username` in the db. On the UI, users can enter username=their email _user_row: LiteLLM_UserTable | None = None user_role: ( @@ -197,8 +219,8 @@ async def authenticate_user( - Login with UI_USERNAME and UI_PASSWORD - Login with Invite Link `user_email` and `password` combination """ - if secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest( - password.encode("utf-8"), ui_password.encode("utf-8") + if general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials( + username, password, master_key ): # Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin user_role = LitellmUserRoles.PROXY_ADMIN @@ -340,8 +362,13 @@ async def authenticate_user( code=401, ) else: + env_credentials_hint: Final = ( + "\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file" + if is_env_credential_login_enabled(general_settings) + else "" + ) raise ProxyException( - message="Invalid credentials used to access UI.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file", + message=f"Invalid credentials used to access UI.{env_credentials_hint}", type=ProxyErrorTypes.auth_error, param="invalid_credentials", code=401, diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 4dba2497bb9..953e3cf3e88 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -497,10 +497,22 @@ class RouteChecks: def _placeholder_to_regex(match: re.Match) -> str: placeholder: Final = match.group(0).strip("{}") - if placeholder.endswith(":path"): - # allow "/" in the placeholder value, but don't eat the route suffix after ":" - return r"[^:]+" - return r"[^/]+" + if not placeholder.endswith(":path"): + return r"[^/]+" + # A ":path" placeholder takes whatever the router's own path + # converter takes, slashes and colons alike, so an id spelled with + # either (or both) still matches the template it was mounted under. + # + # Unless the template puts a ":" literal of its own after the + # placeholder: the Google routes end in ":generateContent" and + # friends, and there the value has to stop before that suffix + # rather than swallow it and match a different verb. + # + # "[\s\S]" rather than ".", because "." stops at a newline and the + # path converter does not: a %0A anywhere in the value would leave + # the route unmatched here while still reaching the handler, which + # turns this gate into a bypass for the lists built on it. + return r"[^:]+" if ":" in match.string[match.end() :] else r"[\s\S]+" pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) # Anchor the pattern to match the entire string diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py new file mode 100644 index 00000000000..1196011dcdd --- /dev/null +++ b/litellm/proxy/auth/team_grants.py @@ -0,0 +1,122 @@ +"""Project a team row (plus the caller's membership in it) onto the ``team_*`` fields of ``UserAPIKeyAuth``. + +The virtual-key path gets these fields for free from the combined-view SQL join. Every other auth path +starts from a ``LiteLLM_TeamTable`` object instead and has to copy them over by hand, which is how JWT +callers kept losing grants (aliases, permissions, limits) one field at a time. Build the badge through +``team_grants`` and the two paths cannot drift. +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Annotated, Final + +from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError +from pydantic.main import IncEx +from typing_extensions import ReadOnly, TypedDict + +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + Member, +) + +_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str]) +_JSON_COLUMNS: Final[Mapping[str, IncEx | bool]] = MappingProxyType( + {"metadata": True, "litellm_model_table": MappingProxyType({"model_aliases": True})} +) + + +def _decode_model_aliases(value: object) -> object: + """``LiteLLM_ModelTable.model_aliases`` is typed ``str | dict``; writers hand Prisma ``json.dumps(...)``, so take both.""" + if not isinstance(value, str): + return value + try: + return _MODEL_ALIASES_ADAPTER.validate_json(value) + except ValidationError: + return None + + +class TeamModelAliasTable(BaseModel): + model_config = ConfigDict(protected_namespaces=()) + + model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None + + +class _TeamJsonColumns(BaseModel): + """The two loosely typed columns on ``LiteLLM_TeamTable``, re-read with the shape the badge needs.""" + + metadata: Mapping[str, object] | None = None + litellm_model_table: TeamModelAliasTable | None = None + + +class TeamGrants(TypedDict, total=False): + """Keyword arguments for ``UserAPIKeyAuth``. Empty when the caller has no team, so the model's own defaults apply.""" + + team_alias: ReadOnly[str | None] + team_tpm_limit: ReadOnly[int | None] + team_rpm_limit: ReadOnly[int | None] + team_max_budget: ReadOnly[float | None] + team_soft_budget: ReadOnly[float | None] + team_spend: ReadOnly[float | None] + team_models: ReadOnly[Sequence[str]] + team_blocked: ReadOnly[bool] + team_metadata: ReadOnly[Mapping[str, object] | None] + team_model_aliases: ReadOnly[Mapping[str, str] | None] + team_object_permission_id: ReadOnly[str | None] + team_object_permission: ReadOnly[LiteLLM_ObjectPermissionTable | None] + team_member: ReadOnly[Member | None] + team_member_spend: ReadOnly[float | None] + team_member_tpm_limit: ReadOnly[int | None] + team_member_rpm_limit: ReadOnly[int | None] + + +def _json_columns(team_object: LiteLLM_TeamTable) -> _TeamJsonColumns: + try: + return _TeamJsonColumns.model_validate(team_object.model_dump(include=_JSON_COLUMNS)) + except ValidationError: + return _TeamJsonColumns() + + +def team_model_aliases(team_object: LiteLLM_TeamTable | None) -> Mapping[str, str] | None: + if team_object is None: + return None + alias_table: Final = _json_columns(team_object).litellm_model_table + return alias_table.model_aliases if alias_table is not None else None + + +def team_grants( + team_object: LiteLLM_TeamTable | None, + team_membership: LiteLLM_TeamMembership | None, + user_id: str | None, +) -> TeamGrants: + if team_object is None: + return TeamGrants() + json_columns: Final = _json_columns(team_object) + return TeamGrants( + team_alias=team_object.team_alias, + team_tpm_limit=team_object.tpm_limit, + team_rpm_limit=team_object.rpm_limit, + team_max_budget=team_object.max_budget, + team_soft_budget=team_object.soft_budget, + team_spend=team_object.spend, + team_models=tuple(team_object.models), + team_blocked=team_object.blocked, + team_metadata=json_columns.metadata, + team_model_aliases=( + json_columns.litellm_model_table.model_aliases if json_columns.litellm_model_table is not None else None + ), + team_object_permission_id=team_object.object_permission_id, + team_object_permission=team_object.object_permission, + team_member=next( + (m for m in team_object.members_with_roles if user_id is not None and m.user_id == user_id), + None, + ), + team_member_spend=team_membership.spend if team_membership is not None else None, + team_member_tpm_limit=( + team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None + ), + team_member_rpm_limit=( + team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None + ), + ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 93293db24c6..20ab9904f46 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -82,6 +82,7 @@ from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.resolvers import CredentialRef, Principal from litellm.proxy.auth.resolvers.store import IdentityStore from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.team_grants import team_grants from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( @@ -90,7 +91,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_query_params, _safe_set_request_parsed_body, populate_request_with_path_params, + read_raw_json_body, ) +from litellm.proxy.common_utils.model_listing_utils import claude_code_requested_group from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, @@ -182,6 +185,44 @@ def _get_model_from_request_context( ) +_CLAUDE_MODEL_ROUTES: Final = frozenset( + f"/{prefix}{endpoint}" for prefix in ("", "v1/") for endpoint in ("messages", "chat/completions", "responses") +) +_CLAUDE_MODEL_NORMALIZED: Final = "litellm.claude_model_normalized" + + +async def _normalize_claude_model( + request_data: dict, valid_token: UserAPIKeyAuth, request: Request | None, route: str +) -> None: + from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj + + if route not in _CLAUDE_MODEL_ROUTES or llm_router is None: + return + if request is not None and request.scope.get(_CLAUDE_MODEL_NORMALIZED) is True: + return + requested: Final = _get_model_from_request_context(request_data, route, request, llm_router) + if not isinstance(requested, str) or requested != request_data.get("model"): + return + if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"): + return + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + aliases: Final = settings.get("model_group_alias") if isinstance(settings, Mapping) else None + source: Final = claude_code_requested_group( + requested, llm_router, valid_token.team_id, (valid_token.aliases, valid_token.team_model_aliases, aliases) + ) + if request is not None: + request.scope[_CLAUDE_MODEL_NORMALIZED] = True + if source is None: + return + request_data["model"] = source + _safe_set_request_parsed_body(request=request, parsed_body=request_data) + if request is not None: + request._json = request_data + request._body = orjson.dumps(request_data) + + def _get_model_names_for_budget_checks( model: str | list[str] | None, ) -> list[str]: @@ -1476,24 +1517,16 @@ async def _user_api_key_auth_builder( user_id=user_id, user_email=user_email, team_id=team_id, - team_alias=(team_object.team_alias if team_object is not None else None), - team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), - team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), - team_models=(team_object.models if team_object is not None else []), - team_metadata=(team_object.metadata if team_object is not None else None), org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, jwt_claims=jwt_claims, + **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id), ) valid_token = UserAPIKeyAuth( api_key=None, team_id=team_id, - team_alias=(team_object.team_alias if team_object is not None else None), - team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), - team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), - team_models=(team_object.models if team_object is not None else []), user_role=( LitellmUserRoles(user_object.user_role) if user_object is not None and user_object.user_role is not None @@ -1507,17 +1540,8 @@ async def _user_api_key_auth_builder( user_tpm_limit=(user_object.tpm_limit if user_object is not None else None), user_rpm_limit=(user_object.rpm_limit if user_object is not None else None), user_model_max_budget=(user_object.model_max_budget if user_object is not None else None), - team_member_rpm_limit=( - team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None - ), - team_member_tpm_limit=( - team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None - ), - team_metadata=(team_object.metadata if team_object is not None else None), jwt_claims=jwt_claims, - ) - valid_token.team_object_permission = ( - team_object.object_permission if team_object is not None else None + **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id), ) # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. @@ -2038,7 +2062,7 @@ async def _user_api_key_auth_builder( fallback_spend=team_member_spend, max_budget=team_member_budget, ) - if team_member_spend > team_member_budget: + if team_member_spend >= team_member_budget: _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" raise litellm.BudgetExceededError( current_cost=team_member_spend, @@ -2666,6 +2690,7 @@ async def _run_centralized_common_checks( await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, + request=request, request_data=request_data, route=route, llm_router=llm_router, @@ -2701,19 +2726,12 @@ async def _reserve_budget_after_common_checks( general_settings: dict, end_user_id: str | None = None, end_user_object: LiteLLM_EndUserTable | None = None, + request: Request | None = None, ) -> None: user_api_key_auth_obj.budget_reservation = None if skip_budget_checks: return if general_settings.get("disable_budget_reservation") is True: - verbose_proxy_logger.warning( - "disable_budget_reservation is enabled: skipping optimistic budget " - "reservation. Budget enforcement is read-time only — concurrent " - "requests can each pass the spend check before their cost is recorded, " - "so a configured budget may be briefly exceeded under high concurrency. " - "Set disable_budget_reservation to False or remove it to restore " - "hard per-request budget enforcement." - ) return from litellm.proxy.spend_tracking.budget_reservation import ( @@ -2734,6 +2752,7 @@ async def _reserve_budget_after_common_checks( end_user_object=end_user_object, apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, + raw_body=await read_raw_json_body(request=request), ) @@ -2792,6 +2811,7 @@ async def _authorize_authenticated_request( """ ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) + await _normalize_claude_model(request_data, user_api_key_auth_obj, request, route) # Single authorization point. Builder paths MUST NOT call common_checks. # Route through the same exception handler the builder uses so @@ -2844,6 +2864,43 @@ async def _authorize_authenticated_request( return None +def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> None: + """Anchor the OTLP destinations this key or team overrides its traces to. + + Called inside the ``auth`` phase span so that span reaches the tenant's account + as well, and on the request task so the ``ContextVar`` is inherited by the logging + tasks that close the LLM span. Best-effort: trace routing must never fail auth. + + ``request`` carries the headers, so a backend this request disabled with + ``x-litellm-disable-callbacks`` resolves to no destination. + + Only destinations the published fan-out can build are anchored. Anchoring one is + what tells the operator's exporter to hold that backend's spans back under + ``override``, so an unbuildable one would leave the span with nowhere to go. + + The ``postgres`` spans under ``auth`` close before this runs, because they are the + reads that resolve the identity being read here. They never reach the tenant's + account, and they are never withheld from the operator's backend, whichever mode + is set. + """ + try: + from litellm.integrations.otel.logger import fan_out_provider + from litellm.integrations.otel.plumbing.context import set_request_destinations + from litellm.integrations.otel.plumbing.providers import deliverable_destinations + from litellm.proxy.litellm_pre_call_utils import ( + resolve_tenant_otel_destinations, + ) + + set_request_destinations( + deliverable_destinations( + resolve_tenant_otel_destinations(user_api_key_dict, _safe_get_request_headers(request)), + fan_out_provider(), + ) + ) + except Exception as exc: # noqa: BLE001 # telemetry routing is best-effort and must never break authentication + verbose_proxy_logger.debug("OTel V2: tenant destination resolution failed: %s", exc) + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -2891,6 +2948,7 @@ async def user_api_key_auth( raise body_parse_exception raise user_api_key_auth_obj.budget_reservation = None + _seed_request_destinations(user_api_key_auth_obj, request) # A body that never parsed is authenticated (so the trace carries identity # and this ``auth`` span) but not authorized: there is no model to check it @@ -3117,6 +3175,7 @@ async def _enforce_key_and_fallback_model_access( Key-level model allowlist and client fallbacks (same as standard auth). Not included in common_checks — common_checks enforces team/user/project model access only. """ + await _normalize_claude_model(request_data, valid_token, request, route) config: Final = valid_token.config if config != {}: diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index f7d9eb7da9a..2071576a943 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -490,7 +490,7 @@ lite codex exec "summarize the repo" Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. -The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, and older versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. @@ -508,9 +508,9 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_ ### Route Every Claude Code Session Through the Proxy -`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, writes the key it resolved (your fresh `lite login`, or an explicit `--api-key`) into `env.ANTHROPIC_AUTH_TOKEN` as a static token, drops any stray `ANTHROPIC_API_KEY` or `apiKeyHelper` so nothing fights that token, and leaves every other setting in the file untouched. It backs up the original file before patching it. Nothing here writes an `apiKeyHelper`: Claude Code would spawn `lite` (and its keychain check) on every credential refresh, so the key is copied in instead and `lite up` restores the file when it stops. -Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. +Two things need to already be true: you've run `lite login` (or passed a key), and the proxy is already reachable, since `lite up` does not start one for you. ```bash lite login @@ -526,17 +526,45 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi #### Making It Permanent at Login -`lite up` holds the patch only for as long as it runs. To wire Claude Code up once and leave it that way, pass `--config-claude` to `lite login`: +`lite up` holds its patch only for as long as it runs. To wire Claude Code up at login and leave it that way, pass `--config-claude` to `lite login`: ```bash lite --base-url https://your-proxy.example.com login --config-claude ``` -It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. +It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and the key this login minted as `env.ANTHROPIC_AUTH_TOKEN`, but persistently: no foreground process to keep alive, and `lite unconfigure claude` restores what it changed (see below). Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag -Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`. +The key in the file is the login's own, so it expires with it (24h by default): run `lite login --config-claude` again after that, which rewrites the key in place. Earlier versions wrote an `apiKeyHelper` that ran `lite auth print-token` instead, so a later login refreshed Claude Code by itself; that meant Claude Code spawning a full `lite` start, keychain check included, on every credential refresh, so the helper is no longer written and a stale one is stripped by the next `--config-claude` or `configure claude`. Like `lite up`, the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first -Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops. +#### Configuring Claude Code Once, With a Virtual Key + +`lite configure claude` wires Claude Code up persistently with a long-lived virtual key, a pinned model and an undo, and `lite unconfigure claude` puts things back: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh +lite --base-url https://your-proxy.example.com configure claude --api-key sk-... --model claude-auto +claude +``` + +The key comes from `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) and is written into `env.ANTHROPIC_AUTH_TOKEN`; without one the command refuses, since a `lite login` credential expires within a day and keeping it fresh would mean Claude Code running `lite` through `apiKeyHelper` on every credential refresh. The command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (under `claude-router-` for a group whose id contains neither `claude` nor `anthropic`, since Claude Code lists only those) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key and as `env.ANTHROPIC_MODEL`, both of which have to be on `/v1/models` for the key. The second one matters for `claude -c` and `claude --resume`: a resumed session otherwise re-sends the model its transcript recorded, which behind an auto-router with `return_raw_model_name: true` is the tier model that answered, and a key scoped to the router alias gets a 403 for it; `ANTHROPIC_MODEL` outranks the transcript on resume. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window (the proxy appends `[1m]` for a group whose configured or known input window reaches 1M) and sends no thinking parameters for it, so name the group like a Claude model id to change that. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control + +Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt + +What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any request + +#### Routed model and savings in the status line + +`lite configure claude`, `lite login --config-claude`, `lite up` and `lite autoroute up` also install a status line (`~/.litellm/statusline.py`, registered as `statusLine` in `~/.claude/settings.json` unless you already run one) that shows which model the auto-router actually served the last turn and, once the proxy has recorded the session, what the session cost against the router's savings baseline: + +``` +claude-auto · Routed to: claude-haiku-4-5 -63% vs Claude Opus 5 +LiteLLM ████████░░░░░░░░░░░░░░░░ $0.14 +Claude Opus 5 ████████████████████████ $0.38 +``` + +The routed model comes from Claude Code's own transcript, so it only names the tier model when the auto-router deployment sets `return_raw_model_name: true` (the `lite autoroute` wizard does); otherwise it shows the alias you requested. The cost lines come from `GET /auto_router/session?session_id=...`, which any virtual key may call for its own sessions, and are cached for five seconds under a per-user `$TMPDIR/litellm-statusline-` directory. The baseline is the priciest model in the router's hardest tier, the same counterfactual the auto-router's savings reports use. `lite unconfigure claude` removes the `statusLine` entry only while it still points at that script. + +`lite codex` registers the same script as a Codex `Stop` hook for the launch, so after each turn Codex prints the same block as a system message. Codex asks once to trust the hook; the answer is remembered for later launches. ### QA Complexity-Based Auto-Routing Against Your Real Proxy @@ -584,7 +612,7 @@ An interactive wizard. It runs the same model-group discovery as above, splits t The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key. -You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.) +You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets the top-level `model` and `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, `ANTHROPIC_DEFAULT_OPUS_MODEL` and `ANTHROPIC_DEFAULT_FABLE_MODEL` to `autorouter` in `~/.claude/settings.json` (and `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when missing, like every other wiring), so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.) You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first. diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 7a0ae9dc955..ea1eed65505 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -1,9 +1,12 @@ +import json import os +import re import shutil import subprocess import sys from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from pathlib import Path from types import MappingProxyType from typing import Final, TypeAlias @@ -12,10 +15,12 @@ import requests from pydantic import BaseModel, TypeAdapter, ValidationError from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login +from .claude_settings import ClaudeSettingsError, install_statusline_script from .cmd_quoting import quote_for_cmd from .pi import ( LITELLM_PROXY_API_KEY_ENV, PI_PROVIDER_NAME, + ListingFailure, PiSyncError, fetch_model_ids, fetch_model_limits, @@ -165,7 +170,9 @@ def prepare_pi( """ ids: Final = fetch_model_ids(base_url, api_key, get=get) if isinstance(ids, PiSyncError): - raise AgentRunError(ids.message) + raise AgentRunError( + f"{ids.message} pi would have nothing to run." if ids.kind is ListingFailure.EMPTY else ids.message + ) limits: Final = fetch_model_limits(base_url, api_key, get=get) path: Final = models_json_path(base_env) error: Final = sync_models_json(path, base_url, ids, limits) @@ -175,10 +182,52 @@ def prepare_pi( return ("--model", f"{PI_PROVIDER_NAME}/{ids[0]}") +def _warn(message: str) -> None: + click.echo(message, err=True) + + +_CODEX_STOP_HOOKS_DECLARED: Final = re.compile( + r"^\s*(\[\[\s*\"?hooks\"?\s*\.\s*\"?Stop\"?\s*\]\]|\"?hooks\"?(?:\s*\.\s*\"?Stop\"?)?\s*=|\[\s*\"?hooks\"?\s*\])", + re.MULTILINE, +) + + +def codex_config_path(base_env: Mapping[str, str]) -> Path: + return Path(base_env.get("CODEX_HOME") or Path.home() / ".codex") / "config.toml" + + +def codex_declares_stop_hooks(config_path: Path) -> bool: + """A config that cannot be read or decoded declares nothing we can see; Codex reports its own + TOML failure at launch, so the pre-check must not be the thing that stops `lite codex`.""" + try: + return _CODEX_STOP_HOOKS_DECLARED.search(config_path.read_text(encoding="utf-8")) is not None + except (OSError, UnicodeDecodeError): + return False + + +def prepare_codex( + base_url: str, + api_key: str, + base_env: Mapping[str, str], + *, + install: Callable[[], str] = install_statusline_script, + warn: Callable[[str], None] = _warn, +) -> tuple[str, ...]: + """A `-c hooks.Stop=` session flag replaces the user's whole Stop list, so their own hooks win over ours.""" + if codex_declares_stop_hooks(codex_config_path(base_env)): + warn("litellm: your Codex config already declares hooks; not adding the routed-model Stop hook") + return () + try: + command: Final = install() + except ClaudeSettingsError as e: + raise AgentRunError(str(e)) from e + return ("-c", f'hooks.Stop=[{{hooks=[{{type="command",command={json.dumps(command)}}}]}}]') + + _Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]] _PREPARERS: Final[Mapping[str, _Preparer]] = MappingProxyType( - {"pi": prepare_pi} # mutable-ok: MappingProxyType freezes the provider registry + {"pi": prepare_pi, "codex": prepare_codex} # mutable-ok: MappingProxyType freezes the provider registry ) @@ -440,10 +489,6 @@ def _restore_controlling_terminal() -> None: os.close(fd) -def _warn(message: str) -> None: - click.echo(message, err=True) - - def run_agent( base_url: str, api_key: str, @@ -535,7 +580,7 @@ def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify started_interactive: Final = _is_interactive() api_key: Final = resolve_api_key(ctx) - display_name, _ = agent_profile(binary) + display_name, _profiles = agent_profile(binary) click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") try: diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 12a288202b6..98af32fa7aa 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -1,3 +1,4 @@ +import os import sys import time import webbrowser @@ -40,10 +41,15 @@ from litellm.litellm_core_utils.cli_token_utils import ( ) from .claude_settings import ( - CLAUDE_SETTINGS_PATH, - SETTINGS_FILE_OWNERS, + STARTING_MODEL_ROLE, ClaudeSettingsError, - write_claude_settings, + KeepModel, + StaticToken, + claude_settings_path, + configure_claude_settings, + configure_state_path, + refuse_while_owned, + settings_file_owners, ) from .pkce_login import ( Http, @@ -777,27 +783,47 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None: return None -def _configure_claude_code(base_url: str) -> None: - """Point Claude Code at base_url by patching ~/.claude/settings.json.""" +def _configure_claude_code(base_url: str, api_key: str) -> None: + """Write the key this login just minted into Claude Code's settings.json as a static token, undoable with + `lite unconfigure claude`. The key expires with the login, so the flag is the re-wire step of each login + rather than a one-time setup: no apiKeyHelper is written, since Claude Code would spawn `lite` (and its + keychain probe) on every credential refresh to keep one fresh.""" + settings_path: Final = claude_settings_path(os.environ) try: - write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS) + configure_claude_settings( + base_url, + StaticToken(api_key), + KeepModel(), + settings_path, + configure_state_path(settings_path), + settings_file_owners(settings_path), + ) except ClaudeSettingsError as e: raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}") - click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.") - click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.") + click.echo(f"\nConfigured Claude Code: {settings_path} now routes through {base_url.rstrip('/')}.") + click.echo( + "This login's key is stored in the file, so run `lite login --config-claude` again after it expires. " + "Your other Claude Code settings were left untouched. Restart Claude Code to pick this up. " + f"Undo with `lite unconfigure claude`; `lite configure claude --model` sets {STARTING_MODEL_ROLE}." + ) def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: SecretSave) -> None: + """Claude Code is configured from the key in hand, so it does not wait on the CLI's own store: a login whose + token file or keychain refused it still has a usable key, and `--config-claude` asked for exactly that + key to be written into settings.json.""" from litellm.proxy.client.cli.interface import show_commands click.echo("\nLogin successful!") click.echo(f"JWT Token: {api_key[:20]}...") click.echo(storage_notice(stored)) + if config_claude: + _configure_claude_code(base_url, api_key) if isinstance(stored, (CredentialNotSaved, CredentialNotRecorded)): + if config_claude: + click.echo("Claude Code was configured with this key even though the CLI itself could not keep it.") return click.echo("You can now use the CLI without specifying --api-key") - if config_claude: - _configure_claude_code(base_url) click.echo("\n" + "=" * 60) show_commands() @@ -832,8 +858,8 @@ def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault) -> None: is_flag=True, default=False, help=( - "After logging in, update ~/.claude/settings.json so Claude Code routes through this proxy. " - "Unrelated settings are preserved." + "After logging in, write this login's key into ~/.claude/settings.json so Claude Code routes through " + "this proxy; run it again after the key expires. Unrelated settings are preserved." ), ) @click.option( @@ -853,6 +879,12 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] + if config_claude: + settings_path: Final = claude_settings_path(os.environ) + try: + refuse_while_owned(settings_path, settings_file_owners(settings_path)) + except ClaudeSettingsError as e: + raise click.ClickException(f"Cannot configure Claude Code, so not logging in: {e}") try: if pkce: diff --git a/litellm/proxy/client/cli/commands/autoroute/commands.py b/litellm/proxy/client/cli/commands/autoroute/commands.py index 26d45138a27..5d91fc81350 100644 --- a/litellm/proxy/client/cli/commands/autoroute/commands.py +++ b/litellm/proxy/client/cli/commands/autoroute/commands.py @@ -1,5 +1,4 @@ import atexit -import json import secrets import signal import threading @@ -14,11 +13,15 @@ from ..claude_settings import ( AUTOROUTE_BACKUP_PATH, CLAUDE_SETTINGS_PATH, ClaudeSettingsError, + StaticToken, + install_statusline_script, load_json_or_empty, + merge_claude_settings, + write_claude_settings, ) from ..up import BackupRecord as ClaudeBackupRecord from ..up import restore_claude_settings, write_backup -from .config import master_key_from_config +from .config import AUTOROUTER_MODEL_NAME, master_key_from_config from .process import ( CONFIG_PATH, DEFAULT_AUTOROUTE_PORT, @@ -37,7 +40,6 @@ from .process import ( terminate, write_pid_record, ) -from .settings import merge_claude_settings_static_token from .wizard import run_configure_wizard _GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) @@ -150,16 +152,23 @@ def up(port: int) -> None: raise click.ClickException(str(e)) try: + status_line: Final = install_statusline_script() original_existed: Final = CLAUDE_SETTINGS_PATH.exists() original_settings: Final = load_json_or_empty(CLAUDE_SETTINGS_PATH) write_backup( ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None), AUTOROUTE_BACKUP_PATH, ) - merged: Final = merge_claude_settings_static_token(original_settings, base_url, master_key) + merged: Final = merge_claude_settings( + original_settings, + base_url, + StaticToken(master_key), + AUTOROUTER_MODEL_NAME, + AUTOROUTER_MODEL_NAME, + status_line=status_line, + ) CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True) - with secure_create(CLAUDE_SETTINGS_PATH) as f: - json.dump(merged, f, indent=2) + write_claude_settings(CLAUDE_SETTINGS_PATH, merged) except ClaudeSettingsError as e: terminate(process.pid) clear_pid_record() diff --git a/litellm/proxy/client/cli/commands/autoroute/config.py b/litellm/proxy/client/cli/commands/autoroute/config.py index 9dfc4ad079b..1f3ad34e3d9 100644 --- a/litellm/proxy/client/cli/commands/autoroute/config.py +++ b/litellm/proxy/client/cli/commands/autoroute/config.py @@ -162,6 +162,7 @@ def build_generated_model_list(config: AutorouteConfig) -> list[JsonValue]: complexity_router_config: Final[dict[str, JsonValue]] = { "tiers": {tier: list(models) for tier, models in config.tiers.items()}, "default_model": config.default_model, + "return_raw_model_name": True, } if isinstance(config.classifier, LLMClassifier): complexity_router_config["classifier_type"] = "llm" diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py deleted file mode 100644 index 60729b5410d..00000000000 --- a/litellm/proxy/client/cli/commands/autoroute/settings.py +++ /dev/null @@ -1,51 +0,0 @@ -from typing import Final - -from pydantic import JsonValue - -from .config import AUTOROUTER_MODEL_NAME - -ENV_KEY: Final = "env" -API_KEY_HELPER_KEY: Final = "apiKeyHelper" -ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" -ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN" -ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" -ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" -ENABLE_TOOL_SEARCH_VALUE: Final = "true" -# Force every one of Claude Code's own model tiers to request the auto-router by name. -# Router's auto-router registry is keyed by the literal requested model string -# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*" -# model_name can never work as a catch-all -- these overrides are what actually makes -# Claude Code send "autorouter" regardless of /model or its own version-specific defaults. -ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = ( - "ANTHROPIC_DEFAULT_SONNET_MODEL", - "ANTHROPIC_DEFAULT_HAIKU_MODEL", - "ANTHROPIC_DEFAULT_OPUS_MODEL", -) - - -def merge_claude_settings_static_token( - settings: dict[str, JsonValue], base_url: str, auth_token: str -) -> dict[str, JsonValue]: - """Return a new settings dict wired to a local ephemeral proxy with a static token. - - Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real - remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key is the - locally persisted autoroute master key, so a plain env var is simpler and correct. Any - existing apiKeyHelper is cleared so it can't fight with the static token. - """ - raw_env: Final = settings.get(ENV_KEY, {}) - base_env: Final = raw_env if isinstance(raw_env, dict) else {} - env: Final[dict[str, JsonValue]] = { - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, - **base_env, - ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), - ANTHROPIC_AUTH_TOKEN_KEY: auth_token, - **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS}, - } - env.pop(ANTHROPIC_API_KEY_KEY, None) - merged: Final[dict[str, JsonValue]] = {**settings, ENV_KEY: env} - merged.pop(API_KEY_HELPER_KEY, None) - return merged - - -__all__ = ["merge_claude_settings_static_token"] diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 5e3ce95f088..e6231f3cac9 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -1,37 +1,80 @@ """Shared handling of Claude Code's ~/.claude/settings.json. -`lite up` patches this file temporarily and restores it on exit; `lite login ---config-claude` patches it persistently. Both need the same merge and the same -apiKeyHelper command, and `up` already imports from `auth`, so the shared parts -live here rather than in either command module. +`lite up` and `lite autoroute up` patch this file temporarily and restore it on +exit; `lite configure claude` patches it persistently and records how to undo it. +All of them need the same merge, and `up` already imports from `auth`, so the +shared parts live here rather than in any one command module. The credential is +always a static token in `env.ANTHROPIC_AUTH_TOKEN`: Claude Code's `apiKeyHelper` +would spawn a `lite` process on every credential refresh, and that process touches +the keychain, so nothing here writes one; a helper left by an earlier version is +owned like any other key and stripped. """ +import hashlib +import json import shlex -import shutil import sys -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from functools import reduce +from itertools import chain from pathlib import Path -from typing import Final +from types import MappingProxyType +from typing import Final, TypeAlias -from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError -from litellm.litellm_core_utils.private_json import write_private_json +from litellm.litellm_core_utils.private_json import ( + commit_staged_json, + discard_staged_json, + ensure_private_dir, + stage_private_json, + write_private_bytes, +) +from . import statusline_script from .cmd_quoting import quote_for_cmd ENV_KEY: Final = "env" API_KEY_HELPER_KEY: Final = "apiKeyHelper" +MODEL_KEY: Final = "model" +STATUS_LINE_KEY: Final = "statusLine" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" +ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" ENABLE_TOOL_SEARCH_VALUE: Final = "true" ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1" +ANTHROPIC_MODEL_KEY: Final = "ANTHROPIC_MODEL" +ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = ( + "ANTHROPIC_DEFAULT_SONNET_MODEL", + "ANTHROPIC_DEFAULT_HAIKU_MODEL", + "ANTHROPIC_DEFAULT_OPUS_MODEL", + "ANTHROPIC_DEFAULT_FABLE_MODEL", +) +OWNED_ENV_KEYS: Final = ( + ENABLE_TOOL_SEARCH_KEY, + ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, + ANTHROPIC_BASE_URL_KEY, + ANTHROPIC_AUTH_TOKEN_KEY, + ANTHROPIC_API_KEY_KEY, + ANTHROPIC_MODEL_KEY, +) +OWNED_TOP_LEVEL_KEYS: Final = (API_KEY_HELPER_KEY, MODEL_KEY, STATUS_LINE_KEY) +OWNED_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in OWNED_ENV_KEYS), *OWNED_TOP_LEVEL_KEYS) +_CREDENTIAL_ENV_KEYS: Final = frozenset((ANTHROPIC_API_KEY_KEY, ANTHROPIC_AUTH_TOKEN_KEY)) +_CREDENTIAL_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in sorted(_CREDENTIAL_ENV_KEYS)), API_KEY_HELPER_KEY) +_BASE_URL_PATH: Final = f"{ENV_KEY}.{ANTHROPIC_BASE_URL_KEY}" +_MODEL_PATHS: Final = (MODEL_KEY, f"{ENV_KEY}.{ANTHROPIC_MODEL_KEY}") +STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts and resumes on" CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" +CLAUDE_CONFIG_DIR_ENV: Final = "CLAUDE_CONFIG_DIR" BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json" +CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json" +STATUSLINE_SCRIPT_PATH: Final = Path.home() / ".litellm" / "statusline.py" @dataclass(frozen=True, slots=True) @@ -55,6 +98,121 @@ class ClaudeSettingsError(Exception): """Raised for any user-actionable failure while reading or writing Claude Code settings.""" +def claude_settings_path(environ: Mapping[str, str]) -> Path: + """The settings.json Claude Code reads: under CLAUDE_CONFIG_DIR when set, else ~/.claude/settings.json.""" + config_dir: Final = environ.get(CLAUDE_CONFIG_DIR_ENV, "") + if not config_dir: + return CLAUDE_SETTINGS_PATH + return Path(config_dir).expanduser() / "settings.json" + + +def _is_default_settings_file(settings_path: Path) -> bool: + return settings_path.resolve() == CLAUDE_SETTINGS_PATH.resolve() + + +def settings_file_owners(settings_path: Path) -> tuple[SettingsFileOwner, ...]: + """The commands whose backups guard settings_path: `lite up` and `lite autoroute up` only ever manage the default file.""" + return SETTINGS_FILE_OWNERS if _is_default_settings_file(settings_path) else () + + +def configure_state_path(settings_path: Path) -> Path: + """The receipt describing settings_path: the default file keeps CONFIGURE_STATE_PATH, and any other file + (a CLAUDE_CONFIG_DIR) gets its own beside it, keyed by its resolved path, so two settings files never + share one undo record.""" + if _is_default_settings_file(settings_path): + return CONFIGURE_STATE_PATH + digest: Final = hashlib.sha256(str(settings_path.resolve()).encode()).hexdigest() + return CONFIGURE_STATE_PATH.parent / CONFIGURE_STATE_PATH.stem / f"{digest}.json" + + +@dataclass(frozen=True, slots=True) +class StaticToken: + """A long-lived virtual key, written into env.ANTHROPIC_AUTH_TOKEN.""" + + token: str + + +@dataclass(frozen=True, slots=True) +class KeepModel: + """Leave the top-level `model` as it is, the user's or an earlier configure's (a re-login).""" + + +@dataclass(frozen=True, slots=True) +class UnpinModel: + """Let go of a `model` an earlier configure pinned; one the user set themselves stays.""" + + +@dataclass(frozen=True, slots=True) +class StartOn: + """Pin `model` and `env.ANTHROPIC_MODEL`: the row Claude Code starts on, and the one a resumed session stays + on, since resume otherwise re-sends the transcript's served model, which a raw-model router made a tier + model the key may not reach.""" + + model: str + + +ModelChoice: TypeAlias = KeepModel | UnpinModel | StartOn + + +class OwnedValue(BaseModel): + """What one key held at a moment in time; `present=False` is an absent key, not a null one.""" + + model_config = ConfigDict(frozen=True) + + present: bool + value: JsonValue = None + + +class ConfigureReceipt(BaseModel): + """What `lite configure claude` found and what it owns, keyed by dotted path (`env.X` or a top-level key). + + Ownership moves only by a write: `written` fingerprints the keys some configure changed, at the + value it wrote; a repeat configure refreshes a fingerprint only for a key its merge changed and + carries the earlier one otherwise, so a key the user edited in between stops matching and is left + alone. `previous` is what each key held before configure took it over; a repeat keeps the earlier + snapshot while the key still holds our value and snapshots afresh otherwise, so whatever the + repeat displaces is what comes back. `endpoints` is the ANTHROPIC_BASE_URL each credential slot + was captured beside, so a credential is only ever put back next to the server it was issued for. + No fingerprint is a second copy of a token. + """ + + model_config = ConfigDict(frozen=True) + + file_existed: bool + env_present: bool + env_was_object: bool + previous: Mapping[str, OwnedValue] + written: Mapping[str, str] + endpoints: Mapping[str, OwnedValue] + + +@dataclass(frozen=True, slots=True) +class WithheldCredential: + """A credential left removed: captured beside `endpoint`, while the restored file points elsewhere.""" + + key: str + endpoint: str + + +@dataclass(frozen=True, slots=True) +class UnconfigureOutcome: + """Keys whose value unconfigure changed back, keys the user changed since and so were left as they + are, credentials withheld (the receipt is kept for them, so a later unconfigure can finish once the + URL points back), and whether no settings file remains.""" + + restored: tuple[str, ...] + kept: tuple[str, ...] + withheld: tuple[WithheldCredential, ...] = () + file_removed: bool = False + + +@dataclass(frozen=True, slots=True) +class _Claim: + previous: OwnedValue + written: str | None + endpoint: OwnedValue | None + + def load_json_or_empty(path: Path) -> dict[str, JsonValue]: try: content: Final = path.read_bytes() if path.exists() else b"" @@ -70,63 +228,20 @@ def load_json_or_empty(path: Path) -> dict[str, JsonValue]: ) -def merge_claude_settings( - settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str -) -> dict[str, JsonValue]: - """Return a new settings dict wired to route Claude Code through the proxy. - - Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a - stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued - token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH - defaults to true because Claude Code turns tool search off when - ANTHROPIC_BASE_URL is not a first-party Anthropic host, and - CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so the /model picker - is filled from the proxy's /v1/models; existing values of both are left - alone. Every other key is preserved untouched. - """ - raw_env: Final = settings.get(ENV_KEY, {}) - base_env: Final = raw_env if isinstance(raw_env, dict) else {} - env: Final = { - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, - ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE, - **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY}, - ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), - } - return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} - - -def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: - """Build the shell command Claude Code should run for its apiKeyHelper. - - Claude Code hands the string to the system shell, `sh` on POSIX and cmd.exe - on Windows, so every token is quoted for the shell that will read it. - - Resolves `lite` to an absolute path so the helper works regardless of the - PATH visible to whatever subprocess Claude Code spawns it from. Passing - --base-url explicitly (rather than relying on the bare invocation Claude - Code would otherwise use) makes `print-token` enforce that the cached - token was actually issued for this proxy -- without it, a token minted - for a different, previously-logged-into proxy would be handed to - whichever server the settings currently point at. - - --base-url belongs to the top-level `lite` group, so it has to precede the - subcommand; click rejects it outright after `print-token`. - """ - lite_path: Final = shutil.which("lite") - if lite_path is None: +def _env_object(settings: Mapping[str, JsonValue], path: Path) -> Mapping[str, JsonValue]: + raw_env: Final = settings.get(ENV_KEY) + if raw_env is None: + return MappingProxyType({}) + if not isinstance(raw_env, dict): raise ClaudeSettingsError( - "Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs an absolute path to it." + f'{path} has a non-object "{ENV_KEY}" value, which this would discard. Fix or remove it, then retry.' ) - quote: Final = quote_for_cmd if platform.startswith("win") else shlex.quote - return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token")) + return raw_env -def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: - """Persistently point Claude Code at base_url, preserving every unrelated setting. - - Refuses while any owner holds a backup: each restores its backup when it - stops, which would silently undo this write. - """ +def refuse_while_owned(settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: + """Refuse while `lite up` or `lite autoroute up` holds a backup it will restore over any write; a + purely local check, so commands run it before any login prompt or request.""" for owner in owners: if owner.backup_path.exists(): raise ClaudeSettingsError( @@ -134,43 +249,392 @@ def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[S f"{owner.backup_path}) and will restore it when it stops. " f"Run `{owner.stop_command}` first, then retry." ) - normalized_base_url: Final = base_url.rstrip("/") - api_key_helper: Final = resolve_api_key_helper(normalized_base_url) - existing: Final = load_json_or_empty(settings_path) - raw_env: Final = existing.get(ENV_KEY) - if raw_env is not None and not isinstance(raw_env, dict): - raise ClaudeSettingsError( - f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. ' - "Fix or remove it, then retry." - ) - merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper) - # os.replace() swaps the symlink itself for a regular file, silently detaching a - # settings.json that is symlinked into a dotfiles repo. There is no backup to undo - # that here, unlike `lite up`, so write through to the link's target instead. - target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path + + +def _write_target(settings_path: Path) -> Path: + """Write through a symlinked settings.json rather than replacing the link, which would silently + detach a file symlinked into a dotfiles repo.""" try: - write_private_json(str(target), merged) + return settings_path.resolve() if settings_path.is_symlink() else settings_path except OSError as e: - raise ClaudeSettingsError(f"Could not write {target}: {e}") from e + raise ClaudeSettingsError(f"Could not resolve {settings_path}: {e}") from e + + +def write_claude_settings(settings_path: Path, settings: Mapping[str, JsonValue]) -> None: + """The one way a settings document lands on disk: staged owner-only beside the target and renamed into + place, through a symlink rather than over it. Every writer (`configure`, `up`, `autoroute up` and the + restores) may be carrying the credential, so none creates the file under the umask or truncates it.""" + target: Final = _write_target(settings_path) + try: + commit_staged_json(stage_private_json(str(target), settings), str(target)) + except OSError as e: + raise ClaudeSettingsError(f"Could not write {settings_path}: {e}") from e + + +def _stage(path: Path, document: Mapping[str, object]) -> str: + try: + return stage_private_json(str(path), document) + except OSError as e: + raise ClaudeSettingsError(f"Could not write {path}: {e}") from e + + +def _land( + path: Path, + staged: str | None, + also_discard: Sequence[str | None] = (), + commit: Callable[[str, str], None] = commit_staged_json, +) -> None: + """Commit a staged file to `path`, or remove `path` when nothing is staged for it. The one place a + filesystem error becomes a ClaudeSettingsError; on failure the operation's other staged files are + discarded, so no temp file holding a token is left behind.""" + try: + if staged is None: + path.unlink(missing_ok=True) + else: + commit(staged, str(path)) + except OSError as e: + for other in also_discard: + if other is not None: + discard_staged_json(other) + raise ClaudeSettingsError(f"Could not {'remove' if staged is None else 'write'} {path}: {e}") from e + + +def statusline_command(script_path: Path, platform: str = sys.platform) -> str: + """This interpreter, not a bare `python3`: it is the one the apiKeyHelper already depends on.""" + quote: Final = quote_for_cmd if platform.startswith("win") else shlex.quote + return " ".join(quote(token) for token in (sys.executable, str(script_path))) + + +def install_statusline_script(script_path: Path | None = None) -> str: + target: Final = script_path or STATUSLINE_SCRIPT_PATH + try: + ensure_private_dir(target.parent) + write_private_bytes(str(target), Path(statusline_script.__file__).read_bytes()) + except OSError as e: + raise ClaudeSettingsError(f"Could not install the status line script at {target}: {e}") from e + return statusline_command(target) + + +def with_status_line(settings: Mapping[str, JsonValue], command: str) -> Mapping[str, JsonValue]: + """Ours is recognised by the script it runs, so a re-install under another interpreter is still ours.""" + existing: Final = settings.get(STATUS_LINE_KEY) + existing_command: Final = existing.get("command") if isinstance(existing, dict) else None + ours: Final = existing is None or (isinstance(existing_command, str) and command.split()[-1] in existing_command) + if not ours: + return settings + entry: Final = dict((("type", "command"), ("command", command))) # mutable-ok: JSON document + return dict(chain(settings.items(), ((STATUS_LINE_KEY, entry),))) # mutable-ok: JSON document + + +def merge_claude_settings( + settings: Mapping[str, JsonValue], + base_url: str, + credential: StaticToken, + default_model: str | None = None, + tier_model: str | None = None, + *, + status_line: str | None = None, +) -> Mapping[str, JsonValue]: + """Return a new settings mapping wired to route Claude Code through the proxy. + + The token lands in env.ANTHROPIC_AUTH_TOKEN; the other credential slots (a stray ANTHROPIC_API_KEY, + an apiKeyHelper) are removed, since Claude Code given two credentials may send the wrong one. + ENABLE_TOOL_SEARCH and CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY get their defaults only when + missing. `default_model` is the top-level `model` and env.ANTHROPIC_MODEL (see StartOn); + `tier_model` is `lite autoroute up`'s knob that points every ANTHROPIC_DEFAULT_*_MODEL at one + group. Apart from those tier keys, exactly OWNED_PATHS are touched. + """ + raw_env: Final = settings.get(ENV_KEY, {}) + current_env: Final = raw_env if isinstance(raw_env, dict) else {} + env: Final = dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping + chain( + ( + (ENABLE_TOOL_SEARCH_KEY, ENABLE_TOOL_SEARCH_VALUE), + (ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE), + ), + ((key, value) for key, value in current_env.items() if key not in _CREDENTIAL_ENV_KEYS), + ((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")), (ANTHROPIC_AUTH_TOKEN_KEY, credential.token)), + ((ANTHROPIC_MODEL_KEY, default_model),) if default_model is not None else (), + ((key, tier_model) for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS if tier_model is not None), + ) + ) + return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping + chain( + ( + (key, value) + for key, value in (with_status_line(settings, status_line) if status_line else settings).items() + if key not in (API_KEY_HELPER_KEY, ENV_KEY) + ), + ((ENV_KEY, env),), + ((MODEL_KEY, default_model),) if default_model is not None else (), + ) + ) + + +def _owned(container: Mapping[str, JsonValue], key: str) -> OwnedValue: + return OwnedValue(present=key in container, value=container.get(key)) + + +def _fingerprint(owned: OwnedValue) -> str: + return hashlib.sha256(json.dumps(owned.model_dump(mode="json"), sort_keys=True).encode()).hexdigest() + + +def _env(settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + raw_env: Final = settings.get(ENV_KEY) + return raw_env if isinstance(raw_env, dict) else MappingProxyType({}) + + +def _lookup(settings: Mapping[str, JsonValue], path: str) -> OwnedValue: + section, _, key = path.rpartition(".") + return _owned(_env(settings) if section else settings, key) + + +def _with_key(container: Mapping[str, JsonValue], key: str, owned: OwnedValue) -> Mapping[str, JsonValue]: + return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping + chain(((k, v) for k, v in container.items() if k != key), ((key, owned.value),) if owned.present else ()) + ) + + +def _with(settings: Mapping[str, JsonValue], path: str, owned: OwnedValue) -> Mapping[str, JsonValue]: + """`settings` with the key at `path` set (or removed when `owned` is absent); nothing else changes.""" + section, _, key = path.rpartition(".") + if not section: + return _with_key(settings, key, owned) + return _with_key(settings, section, OwnedValue(present=True, value=_with_key(_env(settings), key, owned))) + + +def _with_all(settings: Mapping[str, JsonValue], updates: Mapping[str, OwnedValue]) -> Mapping[str, JsonValue]: + return reduce(lambda acc, item: _with(acc, *item), updates.items(), settings) + + +def _ours(settings: Mapping[str, JsonValue], path: str, receipt: ConfigureReceipt) -> bool: + """Whether the key still holds what a configure wrote (a key no configure ever changed is never ours).""" + return receipt.written.get(path) == _fingerprint(_lookup(settings, path)) + + +def _claim( + path: str, + current: Mapping[str, JsonValue], + merged: Mapping[str, JsonValue], + earlier: ConfigureReceipt | None, + url_now: OwnedValue, +) -> _Claim: + """What this configure records for one key; see ConfigureReceipt for the rules.""" + before, after = _lookup(current, path), _lookup(merged, path) + carried: Final = earlier if earlier is not None and _ours(current, path, earlier) else None + return _Claim( + previous=before if carried is None else carried.previous.get(path, before), + written=_fingerprint(after) if before != after else (None if earlier is None else earlier.written.get(path)), + endpoint=None + if path not in _CREDENTIAL_PATHS + else (url_now if carried is None else carried.endpoints.get(path, url_now)), + ) + + +def _receipt( + current: Mapping[str, JsonValue], + merged: Mapping[str, JsonValue], + earlier: ConfigureReceipt | None, + file_exists: bool, +) -> ConfigureReceipt: + url_now: Final = _lookup(current, _BASE_URL_PATH) + claims: Final = MappingProxyType({path: _claim(path, current, merged, earlier, url_now) for path in OWNED_PATHS}) + return ConfigureReceipt( + file_existed=file_exists if earlier is None else earlier.file_existed, + env_present=ENV_KEY in current if earlier is None else earlier.env_present, + env_was_object=isinstance(current.get(ENV_KEY), dict) if earlier is None else earlier.env_was_object, + previous=MappingProxyType({path: claim.previous for path, claim in claims.items()}), + written=MappingProxyType({path: claim.written for path, claim in claims.items() if claim.written is not None}), + endpoints=MappingProxyType( + {path: claim.endpoint for path, claim in claims.items() if claim.endpoint is not None} + ), + ) + + +def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None: + if not state_path.exists(): + return None + try: + return ConfigureReceipt.model_validate_json(state_path.read_bytes()) + except (OSError, ValidationError) as e: + raise ClaudeSettingsError( + f"{state_path} is not a readable `lite configure claude` receipt ({e}). " + "Remove it and edit Claude Code's settings by hand if they still point at the proxy." + ) from e + + +def configure_claude_settings( + base_url: str, + credential: StaticToken, + model: ModelChoice, + settings_path: Path, + state_path: Path, + owners: Sequence[SettingsFileOwner], + commit: Callable[[str, str], None] = commit_staged_json, + script_path: Path | None = None, +) -> None: + """Persistently route Claude Code through base_url, recording how to undo it. + + Both files are staged before either is committed, so a full disk or a read-only directory fails + before anything changes. The two commits are still two renames: a receipt rename that fails + discards the staged settings, and a settings rename that fails after the receipt landed puts the + earlier receipt back (or removes the new one), so the receipt on disk never describes settings + that were not written. `model`: StartOn pins the starting model, UnpinModel lets go of a pin an + earlier configure made (never of the user's own), KeepModel leaves it alone (a re-login). The + status line script is installed and registered under `statusLine` unless the user runs their own; + the receipt owns that key like any other, so unconfigure removes only ours. + """ + refuse_while_owned(settings_path, owners) + current: Final = load_json_or_empty(settings_path) + _env_object(current, settings_path) + earlier: Final = read_configure_receipt(state_path) + unpinned: Final = MappingProxyType( + { + path: earlier.previous[path] + for path in _MODEL_PATHS + if isinstance(model, UnpinModel) and earlier is not None and _ours(current, path, earlier) + } + ) + existing: Final = _with_all(current, unpinned) + merged: Final = merge_claude_settings( + existing, + base_url, + credential, + model.model if isinstance(model, StartOn) else None, + status_line=install_statusline_script(script_path), + ) + receipt: Final = _receipt(current, merged, earlier, settings_path.exists()) + target: Final = _write_target(settings_path) + try: + ensure_private_dir(state_path.parent) + except OSError as e: + raise ClaudeSettingsError(f"Could not write {state_path}: {e}") from e + staged_receipt: Final = _stage(state_path, receipt.model_dump(mode="json")) + try: + staged_settings: Final = _stage(target, merged) + except ClaudeSettingsError: + discard_staged_json(staged_receipt) + raise + _land(state_path, staged_receipt, (staged_settings,), commit) + try: + _land(target, staged_settings, commit=commit) + except ClaudeSettingsError as settings_error: + try: + _land(state_path, None if earlier is None else _stage(state_path, earlier.model_dump(mode="json"))) + except ClaudeSettingsError as receipt_error: + raise ClaudeSettingsError( + f"{settings_error} The receipt at {state_path} now describes settings that were not written and " + f"could not be put back either ({receipt_error}); remove it before retrying." + ) from settings_error + raise + + +def _endpoint_text(endpoint: OwnedValue) -> str: + if not endpoint.present: + return f"no {ANTHROPIC_BASE_URL_KEY} (Anthropic's default endpoint)" + return endpoint.value if isinstance(endpoint.value, str) else json.dumps(endpoint.value) + + +def unconfigure_claude_settings( + settings_path: Path, state_path: Path, owners: Sequence[SettingsFileOwner] +) -> UnconfigureOutcome: + """Undo `lite configure claude`: put back every key still holding what configure wrote, leave the + rest alone, and withhold a credential the restored file would send to a different server than it + was issued for (the receipt stays, owning only those slots, so a later unconfigure can finish).""" + refuse_while_owned(settings_path, owners) + receipt: Final = read_configure_receipt(state_path) + if receipt is None: + raise ClaudeSettingsError( + f"Claude Code is not configured by `lite configure claude` (no receipt at {state_path}); nothing to undo." + ) + current: Final = load_json_or_empty(settings_path) + _env_object(current, settings_path) + ours: Final = tuple(path for path in receipt.written if _ours(current, path, receipt)) + kept: Final = tuple(path for path in receipt.written if path not in ours and _lookup(current, path).present) + put_back: Final = _with_all(current, MappingProxyType({path: receipt.previous[path] for path in ours})) + url_after: Final = _lookup(put_back, _BASE_URL_PATH) + withheld: Final = tuple( + WithheldCredential(path, _endpoint_text(receipt.endpoints[path])) + for path in _CREDENTIAL_PATHS + if path in ours and receipt.previous[path].present and receipt.endpoints[path] != url_after + ) + absent: Final = OwnedValue(present=False) + trimmed: Final = _with_all(put_back, MappingProxyType({item.key: absent for item in withheld})) + settings: Final = ( + trimmed + if _env(trimmed) or receipt.env_was_object + else _with_key(trimmed, ENV_KEY, OwnedValue(present=receipt.env_present, value=None)) + ) + target: Final = _write_target(settings_path) + file_removed: Final = not settings and not (receipt.file_existed and target.exists()) + kept_receipt: Final = ( # mutable-ok: pydantic serializes the update as given and rejects a mappingproxy + receipt.model_copy(update={"written": {item.key: _fingerprint(absent) for item in withheld}}) + if withheld + else None + ) + staged_settings: Final = None if file_removed else _stage(target, settings) + try: + staged_receipt: Final = ( + None if kept_receipt is None else _stage(state_path, kept_receipt.model_dump(mode="json")) + ) + except ClaudeSettingsError: + if staged_settings is not None: + discard_staged_json(staged_settings) + raise + _land(target, staged_settings, (staged_receipt,)) + _land(state_path, staged_receipt) + return UnconfigureOutcome( + restored=tuple(path for path in ours if _lookup(current, path) != _lookup(settings, path)), + kept=kept, + withheld=withheld, + file_removed=file_removed, + ) __all__ = ( "ANTHROPIC_API_KEY_KEY", + "ANTHROPIC_AUTH_TOKEN_KEY", "ANTHROPIC_BASE_URL_KEY", + "ANTHROPIC_DEFAULT_MODEL_ENV_KEYS", + "ANTHROPIC_MODEL_KEY", "API_KEY_HELPER_KEY", "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", + "CLAUDE_CONFIG_DIR_ENV", "CLAUDE_SETTINGS_PATH", + "CONFIGURE_STATE_PATH", "ENABLE_GATEWAY_MODEL_DISCOVERY_KEY", "ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE", "ENABLE_TOOL_SEARCH_KEY", "ENABLE_TOOL_SEARCH_VALUE", "ENV_KEY", + "MODEL_KEY", + "OWNED_ENV_KEYS", + "OWNED_PATHS", + "OWNED_TOP_LEVEL_KEYS", "SETTINGS_FILE_OWNERS", + "STARTING_MODEL_ROLE", + "STATUSLINE_SCRIPT_PATH", + "STATUS_LINE_KEY", "ClaudeSettingsError", + "ConfigureReceipt", + "KeepModel", + "ModelChoice", + "OwnedValue", "SettingsFileOwner", + "StartOn", + "StaticToken", + "UnconfigureOutcome", + "UnpinModel", + "WithheldCredential", + "claude_settings_path", + "configure_claude_settings", + "configure_state_path", "load_json_or_empty", "merge_claude_settings", - "resolve_api_key_helper", + "read_configure_receipt", + "refuse_while_owned", + "settings_file_owners", + "unconfigure_claude_settings", "write_claude_settings", ) diff --git a/litellm/proxy/client/cli/commands/configure.py b/litellm/proxy/client/cli/commands/configure.py new file mode 100644 index 00000000000..4acf94e16f9 --- /dev/null +++ b/litellm/proxy/client/cli/commands/configure.py @@ -0,0 +1,282 @@ +"""`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable.""" + +import os +import sys +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import click +from InquirerPy import inquirer +from InquirerPy.base.control import Choice + +from litellm.proxy.common_utils.model_listing_utils import ( + CLAUDE_CODE_CLIENT, + CLAUDE_CODE_PICKER_PATTERN, + GATEWAY_CLIENT_HEADER, +) + +from .auth import CliContextObj +from .claude_settings import ( + STARTING_MODEL_ROLE, + ClaudeSettingsError, + ModelChoice, + StartOn, + StaticToken, + UnconfigureOutcome, + UnpinModel, + claude_settings_path, + configure_claude_settings, + configure_state_path, + refuse_while_owned, + settings_file_owners, + unconfigure_claude_settings, +) +from .pi import ListedModel, ListingFailure, PiSyncError, fetch_model_listing + +_LISTED_MODELS_SHOWN: Final = 20 +_CLAUDE_TARGET: Final = "claude" +_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),) +_KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default" +_CLAUDE_CODE_VIEW: Final = MappingProxyType( + {"anthropic-version": "2023-06-01", GATEWAY_CLIENT_HEADER: CLAUDE_CODE_CLIENT} +) +_MODEL_OPTION_HELP: Final = ( + f"Proxy model to set as {STARTING_MODEL_ROLE}. Must be listed on /v1/models for the key; without it, " + "Claude Code keeps its own default and a pin an earlier configure made is let go of. Nothing pins Claude " + "Code's sub-agent or background tiers; `lite autoroute up` is the mode that does." +) + + +def resolve_credential(ctx: click.Context, api_key: str | None) -> StaticToken: + """The long-lived key written into settings.json: --api-key, `lite --api-key` or LITELLM_PROXY_API_KEY. + + A `lite login` credential is never written: it expires within a day, and keeping it fresh would mean + Claude Code running `lite` through `apiKeyHelper` on every credential refresh. + """ + ctx_obj: Final[CliContextObj] = ctx.obj + explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key")) + if not explicit: + raise ClaudeSettingsError( + "`lite configure claude` needs a long-lived virtual key: pass --api-key, `lite --api-key`, or set " + "LITELLM_PROXY_API_KEY. Your `lite login` credential expires within a day, so it is not written " + "into Claude Code's settings." + ) + return StaticToken(explicit) + + +@dataclass(frozen=True, slots=True) +class _Listing: + models: tuple[ListedModel, ...] + + @property + def ids(self) -> tuple[str, ...]: + return tuple(model.id for model in self.models) + + +def _start(ctx: click.Context, api_key: str | None) -> tuple[StaticToken, _Listing]: + """Every configure path begins the same way: the local ownership check first, so a `lite up` + session is refused before any request, then the credential, then the listing.""" + settings_path: Final = claude_settings_path(os.environ) + try: + refuse_while_owned(settings_path, settings_file_owners(settings_path)) + credential: Final = resolve_credential(ctx, api_key) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) + return credential, _listed_models(ctx.obj["base_url"], credential.token) + + +def _listing_error(base_url: str, error: PiSyncError) -> str: + """The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question.""" + if error.kind is ListingFailure.REJECTED: + return f"LiteLLM rejected your key (HTTP {error.status}). Pass a valid --api-key." + if error.kind is ListingFailure.UNREACHABLE: + return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?" + if error.kind is ListingFailure.EMPTY: + return f"{error.message} Claude Code would have nothing to run; give the key access to at least one model." + return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy." + + +def _listed_models(base_url: str, key: str) -> _Listing: + listed: Final = fetch_model_listing(base_url, key, headers=_CLAUDE_CODE_VIEW) + if isinstance(listed, PiSyncError): + raise click.ClickException(_listing_error(base_url, listed)) + return _Listing(listed) + + +def _starting_model(model: str, listing: _Listing) -> str | None: + source: Final = next((listed.id for listed in listing.models if listed.source_model == model), None) + return source or next((listed.id for listed in listing.models if listed.id == model), None) + + +def _model_choice(model: str | None) -> ModelChoice: + return StartOn(model) if model is not None else UnpinModel() + + +def _apply_claude(ctx: click.Context, credential: StaticToken, listing: _Listing, model: str | None) -> None: + ctx_obj: Final[CliContextObj] = ctx.obj + base_url: Final = ctx_obj["base_url"] + listed: Final = listing.ids + starting: Final = _starting_model(model, listing) if model is not None else None + if model is not None and starting is None: + shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN]) + more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else "" + raise click.ClickException( + f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}." + ) + settings_path: Final = claude_settings_path(os.environ) + try: + configure_claude_settings( + base_url, + credential, + _model_choice(starting), + settings_path, + configure_state_path(settings_path), + settings_file_owners(settings_path), + ) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) + in_picker: Final = sum(1 for listed_model in listed if CLAUDE_CODE_PICKER_PATTERN.search(listed_model)) + click.echo(f"Configured Claude Code: {settings_path} now routes through {base_url}.") + + click.echo("Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN.") + click.echo( + f"Starting model: {starting} ({STARTING_MODEL_ROLE}); switch any time with /model." + if starting is not None + else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or " + "pass --model to start on a proxy model. Without a pin, a resumed session re-sends the model its transcript " + "recorded, which behind a raw-model auto-router is the tier model." + ) + click.echo( + f"/model will list all {len(listed)} of the proxy's models." + if in_picker == len(listed) + else f"/model will list {in_picker} of the proxy's {len(listed)} models: Claude Code shows only ids containing " + "'claude' or 'anthropic', and this proxy does not list the rest under such names." + ) + click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.") + if settings_path.is_symlink(): + click.echo( + f"Note: {settings_path} is a symlink to {settings_path.resolve()}, so your key now lives in " + "that file; keep it out of version control.", + err=True, + ) + + +def _pick_targets() -> tuple[str, ...]: + picked: Final = inquirer.checkbox( + message="Which agents should route through LiteLLM?", + choices=[Choice(value, name=label, enabled=True) for value, label in _TARGETS], + validate=lambda chosen: len(chosen) > 0, + invalid_message="Pick at least one.", + ).execute() + return tuple(str(value) for value in picked) + + +def _pick_model(listed: Sequence[str]) -> str | None: + picked: Final = inquirer.fuzzy( + message="Model Claude Code starts on (type to filter; /model switches any time):", + choices=[_KEEP_DEFAULT_MODEL, *listed], + ).execute() + return None if picked == _KEEP_DEFAULT_MODEL else str(picked) + + +def interactive_configure( + ctx: click.Context, + pick_targets: Callable[[], tuple[str, ...]] = _pick_targets, + pick_model: Callable[[Sequence[str]], str | None] = _pick_model, +) -> None: + """`lite configure` with no agent named: ask which agents to wire and which model to pin.""" + targets: Final = pick_targets() + if _CLAUDE_TARGET not in targets: + return + credential, listing = _start(ctx, None) + _apply_claude( + ctx, credential, listing, pick_model(tuple(model.source_model or model.id for model in listing.models)) + ) + + +@click.group(name="configure", invoke_without_command=True) +@click.pass_context +def configure_group(ctx: click.Context) -> None: + """Persistently route a coding agent through your LiteLLM proxy. + + With no agent named, asks which agents to wire and which proxy model to pin. + """ + if ctx.invoked_subcommand is not None: + return + if not sys.stdin.isatty(): + raise click.ClickException( + "`lite configure` asks questions, so it needs a terminal. Non-interactively, run " + "`lite configure claude --api-key --model `." + ) + interactive_configure(ctx) + + +@click.group(name="unconfigure") +def unconfigure_group() -> None: + """Undo `lite configure` for a coding agent.""" + + +@configure_group.command(name="claude") +@click.option( + "--api-key", + "api_key", + default=None, + help="Long-lived LiteLLM virtual key written into Claude Code's settings. Defaults to the `lite --api-key` / " + "LITELLM_PROXY_API_KEY value; required, since a `lite login` credential expires within a day.", +) +@click.option("--model", default=None, help=_MODEL_OPTION_HELP) +@click.pass_context +def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None: + """Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`. + + Patches ~/.claude/settings.json in place: the proxy URL, your virtual key as a static token, + and gateway model discovery so /model lists the proxy's models; --model picks the one Claude + Code starts on and resumes with. Every other + setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back. + Assumes the proxy is already running. + """ + credential, listing = _start(ctx, api_key) + _apply_claude(ctx, credential, listing, model) + + +@unconfigure_group.command(name="claude") +def unconfigure_claude() -> None: + """Return Claude Code's settings to what they were before `lite configure claude`. + + Also undoes `lite login --config-claude`. Only keys still holding what configure wrote are + put back; anything you changed since is left as it is and named in the output. + """ + settings_path: Final = claude_settings_path(os.environ) + state_path: Final = configure_state_path(settings_path) + try: + outcome: Final = unconfigure_claude_settings(settings_path, state_path, settings_file_owners(settings_path)) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) + _report_unconfigure(settings_path, state_path, outcome) + + +def _report_unconfigure(settings_path: Path, state_path: Path, outcome: UnconfigureOutcome) -> None: + """Say what unconfigure did, naming only keys whose value it changed.""" + if outcome.file_removed: + click.echo( + f"No settings file remains at {settings_path}; it held nothing but `lite configure claude`'s own keys." + ) + elif outcome.restored: + click.echo(f"Restored in {settings_path}: {', '.join(outcome.restored)}.") + else: + click.echo(f"Nothing in {settings_path} was still ours to restore.") + if outcome.kept: + click.echo(f"Left as you changed them since: {', '.join(outcome.kept)}.") + if outcome.withheld: + click.echo( + "Left removed, since the file now points at a different server than they were issued for: " + + "; ".join(f"{item.key} (captured with {item.endpoint})" for item in outcome.withheld) + + f". They stay in {state_path}: point env.ANTHROPIC_BASE_URL back and run `lite unconfigure claude` " + "again to put them back, or delete that file to drop them." + ) + + +__all__ = ("configure_group", "interactive_configure", "resolve_credential", "unconfigure_group") diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 7b0c1970c4e..9810e81ae36 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -10,21 +10,40 @@ import os import tempfile from collections.abc import Callable, Mapping from dataclasses import dataclass +from enum import StrEnum from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Annotated, Final import requests -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator +from pydantic.types import StringConstraints PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR" PI_PROVIDER_NAME: Final = "litellm" LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY" +_REJECTED_STATUSES: Final = frozenset((401, 403)) + + +class ListingFailure(StrEnum): + """Why a proxy could not be listed, decided once where the HTTP outcome is classified. + + `unreachable` means no response at all; the other kinds prove the proxy answered, so callers + must not suggest checking whether it is running. + """ + + UNREACHABLE = "unreachable" + REJECTED = "rejected" + BAD_BODY = "bad_body" + EMPTY = "empty" + OTHER = "other" @dataclass(frozen=True, slots=True) class PiSyncError: message: str + status: int | None = None + kind: ListingFailure | None = None @dataclass(frozen=True, slots=True) @@ -33,12 +52,25 @@ class ModelLimits: max_tokens: int | None -class _Model(BaseModel): - id: str +_NonEmptyString = Annotated[str, StringConstraints(min_length=1)] + + +class ListedModel(BaseModel): + model_config = ConfigDict(frozen=True) + + id: _NonEmptyString + source_model: _NonEmptyString | None = None class _ModelList(BaseModel): - data: tuple[_Model, ...] + data: tuple[ListedModel, ...] + + @model_validator(mode="after") + def unique_id_mappings(self) -> "_ModelList": + mappings: Final = frozenset((model.id, model.source_model or model.id) for model in self.data) + if len(frozenset(model.id for model in self.data)) != len(mappings): + raise ValueError("model ids must not map to multiple source models") + return self class _ModelGroup(BaseModel): @@ -51,31 +83,47 @@ class _ModelGroupList(BaseModel): data: tuple[_ModelGroup, ...] +def fetch_model_listing( + base_url: str, + api_key: str, + *, + get: Callable[..., requests.Response] = requests.get, + headers: Mapping[str, str] = MappingProxyType({}), +) -> tuple[ListedModel, ...] | PiSyncError: + url: Final = base_url.rstrip("/") + "/v1/models" + try: + resp: Final = get( + url, + headers={"Authorization": f"Bearer {api_key}", **headers}, # mutable-ok: requests headers require a dict + timeout=10, + ) + except requests.RequestException as e: + return PiSyncError(f"Could not list models from the proxy: {e}", kind=ListingFailure.UNREACHABLE) + if resp.status_code != 200: + return PiSyncError( + f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot list models.", + resp.status_code, + ListingFailure.REJECTED if resp.status_code in _REJECTED_STATUSES else ListingFailure.OTHER, + ) + try: + listing: Final = _ModelList.model_validate(resp.json()) + except (ValueError, ValidationError) as e: + return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}", kind=ListingFailure.BAD_BODY) + models: Final = tuple(dict.fromkeys(listing.data)) + if not models: + return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY) + return models + + def fetch_model_ids( base_url: str, api_key: str, *, get: Callable[..., requests.Response] = requests.get, + headers: Mapping[str, str] = MappingProxyType({}), ) -> tuple[str, ...] | PiSyncError: - url: Final = base_url.rstrip("/") + "/v1/models" - try: - resp: Final = get( - url, - headers={"Authorization": f"Bearer {api_key}"}, # mutable-ok: requests headers require a dict - timeout=10, - ) - except requests.RequestException as e: - return PiSyncError(f"Could not list models from the proxy: {e}") - if resp.status_code != 200: - return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.") - try: - listing: Final = _ModelList.model_validate(resp.json()) - except (ValueError, ValidationError) as e: - return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}") - ids: Final = tuple(dict.fromkeys(model.id for model in listing.data)) - if not ids: - return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.") - return ids + listed: Final = fetch_model_listing(base_url, api_key, get=get, headers=headers) + return listed if isinstance(listed, PiSyncError) else tuple(dict.fromkeys(model.id for model in listed)) _NO_LIMITS: Final[Mapping[str, ModelLimits]] = MappingProxyType({}) @@ -200,10 +248,13 @@ __all__ = ( "LITELLM_PROXY_API_KEY_ENV", "PI_CONFIG_DIR_ENV", "PI_PROVIDER_NAME", + "ListedModel", + "ListingFailure", "ModelLimits", "PiSyncError", "fetch_model_ids", "fetch_model_limits", + "fetch_model_listing", "models_json_path", "provider_block", "sync_models_json", diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py new file mode 100644 index 00000000000..a8abeb68978 --- /dev/null +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -0,0 +1,392 @@ +"""Claude Code status line and Codex Stop hook for auto-routed sessions. + +`lite` copies this file verbatim to ~/.litellm/statusline.py and registers it as Claude +Code's `statusLine` command and as Codex's `[[hooks.Stop]]` command, so it must stay +standard-library only and must never import litellm. Claude Code re-runs it on every +status refresh (about every 300ms while typing), so the proxy is asked at most once per +TTL per session and every other refresh is served from a small on-disk cache that holds +only the proxy's answer, never the key. + +Claude Code pipes a JSON payload on stdin (session_id, transcript_path, model); the routed +model is the `message.model` of the latest foreground assistant line in the transcript, +which is the proxy's response `model` field. That only names the tier model when the +auto-router deployment sets `return_raw_model_name: true`; otherwise it is the alias the +client requested. Codex pipes its Stop event instead (hook_event_name, session_id) and has +no transcript to read, so the routed model comes from the proxy's session record and the +result is printed as a `systemMessage` for the transcript. The proxy key is read from the +agent's own environment (the static token `lite configure claude` writes); nothing here +spawns a credential helper. + +Cost figures come from GET /auto_router/session on the proxy, which reads the per-session +rollup written by the spend flush. That flush is asynchronous, so a turn's cost lands a +second or two after the turn; the cache TTL absorbs it. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import sys +import tempfile +import time +import urllib.error +import urllib.request +from collections.abc import Callable, Mapping +from pathlib import Path +from types import MappingProxyType +from typing import IO, Final, NamedTuple, Protocol +from urllib.parse import urlencode + +SESSION_ENDPOINT: Final = "/auto_router/session" +CACHE_TTL_SECONDS: Final = 5.0 +FETCH_TIMEOUT_SECONDS: Final = 3 +BAR_WIDTH: Final = 24 +BAR_FULL: Final = "\u2588" +BAR_EMPTY: Final = "\u2591" +SEPARATOR: Final = " \u00b7 " +TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024 +CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",) +CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY") +CODEX_BASE_URL_ENV_KEYS: Final = ("OPENAI_BASE_URL",) +CODEX_API_KEY_ENV_KEYS: Final = ("OPENAI_API_KEY",) +CODEX_STOP_EVENT: Final = "Stop" +SYNTHETIC_MODEL: Final = "" +LITELLM_LABEL: Final = "LiteLLM" +RESET: Final = "\033[0m" +BOLD: Final = "\033[1m" +DIM: Final = "\033[90m" +LITELLM_COLOR: Final = "\033[38;2;79;70;229m" +BASELINE_COLOR: Final = "\033[38;2;217;119;87m" +EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) +EMPTY_ENV: Final[Mapping[str, str]] = MappingProxyType({}) + + +class Session(NamedTuple): + router_name: str + last_model: str + spend: float + baseline_spend: float + baseline_model: str | None + + +class Credentials(NamedTuple): + base_url: str + api_key: str + + @property + def usable(self) -> bool: + return bool(self.base_url and self.api_key) + + +class Fetched(NamedTuple): + session: Session | None + definitive: bool + + +class Fetch(Protocol): + def __call__(self, credentials: Credentials, session_id: str) -> Fetched: ... + + +def as_mapping(value: object) -> Mapping[str, object]: + return value if isinstance(value, dict) else EMPTY + + +def as_str(value: object) -> str: + return value if isinstance(value, str) else "" + + +def printable(value: object) -> str: + """Labels come from the transcript, the proxy, and Claude Code's model cache, none of which this script + controls, and every one is written to a terminal: a control character (ESC, BEL, C1) in a model name + could redraw the screen or set the clipboard, so only printable text survives.""" + return "".join(character for character in as_str(value) if character.isprintable()) + + +def load_json(raw: bytes | str) -> object: + try: + return json.loads(raw) + except ValueError: + return None + + +def resolve_base_url(env: Mapping[str, str], keys: tuple[str, ...]) -> str: + raw: Final = next((env[key] for key in keys if env.get(key)), "").strip().rstrip("/") + return raw.removesuffix("/v1") + + +def resolve_api_key(env: Mapping[str, str], keys: tuple[str, ...]) -> str: + return next((env[key] for key in keys if env.get(key)), "").strip() + + +def claude_credentials(env: Mapping[str, str]) -> Credentials: + """Claude Code's own resolution order, so the key-scoped lookup runs as the principal that wrote the rows: + ANTHROPIC_AUTH_TOKEN, then ANTHROPIC_API_KEY. A `lite` variable such as LITELLM_PROXY_API_KEY is not a + key Claude Code ever sends, so honoring it would ask as someone else. An apiKeyHelper is never run: a + status line refreshes every few hundred milliseconds, and spawning a credential helper that often is + how a keychain prompt ends up on screen a hundred times.""" + return Credentials(resolve_base_url(env, CLAUDE_BASE_URL_ENV_KEYS), resolve_api_key(env, CLAUDE_API_KEY_ENV_KEYS)) + + +def codex_credentials(env: Mapping[str, str]) -> Credentials: + return Credentials(resolve_base_url(env, CODEX_BASE_URL_ENV_KEYS), resolve_api_key(env, CODEX_API_KEY_ENV_KEYS)) + + +def _transcript_line_model(line: bytes) -> str: + """A `` model is Claude Code's own marker for a locally produced message (an API error, a + resume note), not a served model, so it is skipped like a sidechain line.""" + item: Final = as_mapping(load_json(line)) + if item.get("type") != "assistant" or item.get("isSidechain") is True or item.get("agentId"): + return "" + model: Final = printable(as_mapping(item.get("message")).get("model")) + return "" if model == SYNTHETIC_MODEL else model + + +def latest_transcript_model(transcript_path: str) -> str: + if not transcript_path: + return "" + try: + with Path(transcript_path).open("rb") as transcript: + size: Final = transcript.seek(0, os.SEEK_END) + transcript.seek(max(0, size - TRANSCRIPT_SCAN_LIMIT_BYTES)) + tail: Final = transcript.read() + except OSError: + return "" + return next((model for line in reversed(tail.split(b"\n")) if (model := _transcript_line_model(line))), "") + + +def model_label(model: str, config_dir: Path) -> str: + bare: Final = model.rsplit("/", 1)[-1] + try: + raw: Final = (config_dir / "cache" / "gateway-models.json").read_bytes() + except OSError: + return bare + listed: Final = as_mapping(load_json(raw)).get("models") + if not isinstance(listed, list): + return bare + entries: Final = tuple(as_mapping(entry) for entry in listed) + return next( + ( + printable(entry.get("display_name")) + for entry in entries + if entry.get("id") in (model, bare) and printable(entry.get("display_name")) + ), + bare, + ) + + +def baseline_label(model: str, config_dir: Path) -> str: + labelled: Final = model_label(model, config_dir) + if labelled != model.rsplit("/", 1)[-1]: + return labelled + return " ".join(word.capitalize() for word in labelled.replace("-", " ").split()) + + +def fetch_session(credentials: Credentials, session_id: str) -> Fetched: + """Any 4xx is this credential's definite answer (no row, no access, expired login) and is cached for the + TTL; a 5xx or transport failure is not, so the next refresh tries again.""" + query: Final = urlencode((("session_id", session_id),)) + request: Final = urllib.request.Request( + f"{credentials.base_url}{SESSION_ENDPOINT}?{query}", + headers={ # mutable-ok: urllib.request.Request takes a dict + "Authorization": f"Bearer {credentials.api_key}", + "Accept": "application/json", + }, + ) + try: + with urllib.request.urlopen(request, timeout=FETCH_TIMEOUT_SECONDS) as response: + raw: Final[bytes] = response.read() + except urllib.error.HTTPError as error: + return Fetched(session=None, definitive=400 <= error.code < 500) + except (urllib.error.URLError, OSError): + return Fetched(session=None, definitive=False) + session: Final = _session_from_payload(as_mapping(load_json(raw))) + return Fetched(session=session, definitive=session is not None) + + +def _session_from_payload(payload: Mapping[str, object]) -> Session | None: + router_name: Final = printable(payload.get("router_name")) + last_model: Final = printable(payload.get("last_model")) + spend: Final = payload.get("spend") + baseline_spend: Final = payload.get("baseline_spend") + if not router_name or not last_model: + return None + if not isinstance(spend, (int, float)) or not isinstance(baseline_spend, (int, float)): + return None + return Session( + router_name=router_name, + last_model=last_model, + spend=float(spend), + baseline_spend=float(baseline_spend), + baseline_model=printable(payload.get("baseline_model")) or None, + ) + + +def cache_path(cache_dir: Path, credentials: Credentials, session_id: str) -> Path: + identity: Final = "\n".join((credentials.base_url, credentials.api_key, session_id)) + return cache_dir / hashlib.sha256(identity.encode()).hexdigest() + + +def load_session( + credentials: Credentials, + session_id: str, + cache_dir: Path, + fetch: Fetch = fetch_session, + now: Callable[[], float] = time.time, +) -> Session | None: + path: Final = cache_path(cache_dir, credentials, session_id) + cached: Final = _read_cache(path) + fetched_at: Final = cached.get("fetched_at") + if isinstance(fetched_at, (int, float)) and now() - fetched_at < CACHE_TTL_SECONDS: + return _session_from_payload(as_mapping(cached.get("session"))) + fetched: Final = fetch(credentials, session_id) + if fetched.definitive: + _write_cache(path, fetched.session, now()) + return fetched.session + + +NOFOLLOW: Final = getattr(os, "O_NOFOLLOW", 0) + + +def cache_dir_name() -> str: + return f"litellm-statusline-{os.getuid()}" if hasattr(os, "getuid") else "litellm-statusline" + + +def _own_private_dir(directory: Path) -> bool: + """A shared temp root lets another local user pre-create the directory, so it must be ours and private + before anything is read or written under it. Windows has no uids or POSIX mode bits and a per-user temp + directory already, so there it only has to exist and not be a link.""" + try: + directory.mkdir(mode=0o700, parents=True, exist_ok=True) + status: Final = directory.lstat() + except OSError: + return False + if not os.path.isdir(directory) or os.path.islink(directory): + return False + if not hasattr(os, "getuid"): + return True + return status.st_uid == os.getuid() and not status.st_mode & 0o077 + + +def _read_cache(path: Path) -> Mapping[str, object]: + if not _own_private_dir(path.parent): + return EMPTY + try: + descriptor: Final = os.open(path, os.O_RDONLY | NOFOLLOW) + with os.fdopen(descriptor, "rb") as handle: + return as_mapping(load_json(handle.read())) + except OSError: + return EMPTY + + +def _write_cache(path: Path, session: Session | None, fetched_at: float) -> None: + """Staged beside the entry and renamed into place, so a refresh reading the entry never sees a torn write.""" + entry: Final = session._asdict() if session else None + body: Final = json.dumps({"fetched_at": fetched_at, "session": entry}) # mutable-ok: json.dumps takes a dict + if not _own_private_dir(path.parent): + return + try: + descriptor, staged = tempfile.mkstemp(dir=path.parent, prefix=".tmp-") + except OSError: + return + try: + with os.fdopen(descriptor, "w") as handle: + handle.write(body) + os.replace(staged, path) + except OSError: + Path(staged).unlink(missing_ok=True) + + +def _bar(fraction: float, color: str, width: int, use_color: bool) -> str: + filled: Final = round(max(0.0, min(1.0, fraction)) * width) + if not use_color: + return BAR_FULL * filled + BAR_EMPTY * (width - filled) + return f"{color}{BAR_FULL * filled}{DIM}{BAR_EMPTY * (width - filled)}{RESET}" + + +def render(model: str, session: Session | None, config_dir: Path, use_color: bool, bar_width: int = BAR_WIDTH) -> str: + def paint(code: str, text: str) -> str: + return f"{code}{text}{RESET}" if use_color else text + + routed: Final = paint(BOLD, f"Routed to: {model}") + if session is None: + return routed + header: Final = f"{session.router_name}{SEPARATOR}{routed}" + if session.baseline_model is None or session.baseline_spend <= 0: + return header + reference: Final = baseline_label(session.baseline_model, config_dir) + pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100 + delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}") + peak: Final = max(session.spend, session.baseline_spend) + label_width: Final = max(len(LITELLM_LABEL), len(reference)) + rows: Final = ( + (LITELLM_LABEL, session.spend, LITELLM_COLOR), + (reference, session.baseline_spend, BASELINE_COLOR), + ) + lines: Final = ( + f"{paint(DIM, label.ljust(label_width))} {_bar(amount / peak, color, bar_width, use_color)} " + f"{paint(DIM, f'${amount:.2f}')}" + for label, amount, color in rows + ) + return "\n".join((f"{header} {delta}", *lines)) + + +def color_enabled(env: Mapping[str, str]) -> bool: + return env.get("NO_COLOR") is None and env.get("TERM", "") not in ("", "dumb") + + +def status_line( + payload: Mapping[str, object], env: Mapping[str, str], config_dir: Path, cache_dir: Path, fetch: Fetch +) -> str: + fallback: Final = printable(as_mapping(payload.get("model")).get("display_name")) + served: Final = latest_transcript_model(as_str(payload.get("transcript_path"))) + if not served: + return fallback or "claude" + label: Final = model_label(served, config_dir) + session_id: Final = as_str(payload.get("session_id")) + credentials: Final = claude_credentials(env) + if not session_id or not credentials.usable: + return render(label, None, config_dir, color_enabled(env)) + session: Final = load_session(credentials, session_id, cache_dir, fetch) + return render(label, session, config_dir, color_enabled(env)) + + +def codex_stop_message( + payload: Mapping[str, object], env: Mapping[str, str], config_dir: Path, cache_dir: Path, fetch: Fetch +) -> str: + """No cache here: the Stop hook runs once per turn, and a first turn's cached absence would hide the record + the next turn finds.""" + session_id: Final = as_str(payload.get("session_id")) + credentials: Final = codex_credentials(env) + if not session_id or not credentials.usable: + return "" + session: Final = fetch(credentials, session_id).session + if session is None: + return "" + text: Final = render(model_label(session.last_model, config_dir), session, config_dir, use_color=False) + return json.dumps({"systemMessage": f"\n{text}"}) # mutable-ok: json.dumps takes a dict + + +def run(stdin: IO[str], stdout: IO[str], env: Mapping[str, str], fetch: Fetch = fetch_session) -> None: + """A failure renders each mode's own quiet fallback: Claude Code gets the label it already knows, Codex gets + nothing at all rather than a bare string it would reject as hook JSON.""" + body: Final = as_mapping(load_json(stdin.read())) + codex: Final = body.get("hook_event_name") == CODEX_STOP_EVENT + config_dir: Final = Path(env.get("CLAUDE_CONFIG_DIR") or Path.home() / ".claude") + cache_dir: Final = ( + Path(env.get("TMPDIR") or env.get("TEMP") or env.get("TMP") or tempfile.gettempdir()) / cache_dir_name() + ) + try: + text: Final = ( + codex_stop_message(body, env, config_dir, cache_dir, fetch) + if codex + else status_line(body, env, config_dir, cache_dir, fetch) + ) + except Exception: # noqa: BLE001 # a status line must never break the agent session + stdout.write("" if codex else printable(as_mapping(body.get("model")).get("display_name")) or "claude") + return + stdout.write(text) + + +if __name__ == "__main__": + run(sys.stdin, sys.stdout, os.environ) diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index b7c02866d6f..f2624797a5f 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -24,9 +24,11 @@ from .claude_settings import ( BACKUP_PATH, CLAUDE_SETTINGS_PATH, ClaudeSettingsError, + StaticToken, + install_statusline_script, load_json_or_empty, merge_claude_settings, - resolve_api_key_helper, + write_claude_settings, ) @@ -97,8 +99,7 @@ def restore_claude_settings(settings_path: Path | None = None, backup_path: Path return None if record.existed and record.content is not None: resolved_settings_path.parent.mkdir(parents=True, exist_ok=True) - with open(resolved_settings_path, "w") as f: - json.dump(record.content, f, indent=2) + write_claude_settings(resolved_settings_path, record.content) elif resolved_settings_path.exists(): resolved_settings_path.unlink() resolved_backup_path.unlink() @@ -123,7 +124,7 @@ def _stored_login_is_pkce(vault: SecretVault) -> bool: return token_data is not None and token_data.get("refresh_token") is not None -def _ensure_fresh_login(ctx: click.Context) -> None: +def ensure_fresh_login(ctx: click.Context) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"].rstrip("/") vault: Final = context_secret_vault(ctx) @@ -133,15 +134,12 @@ def _ensure_fresh_login(ctx: click.Context) -> None: pkce: Final = _stored_login_is_pkce(vault) login_command: Final = "lite login --pkce" if pkce else "lite login" if not sys.stdin.isatty(): - raise UpError( - f"No fresh LiteLLM login found for this proxy. Run `{login_command}` first (apiKeyHelper " - "reads this token on every Claude Code request)." - ) + raise UpError(f"No fresh LiteLLM login found for this proxy. Run `{login_command}` first.") click.echo("No fresh LiteLLM login found for this proxy; starting login...") - ctx.invoke(login, pkce=pkce) + ctx.invoke(login, config_claude=False, pkce=pkce) if not _usable_login(get_stored_api_key(expected_base_url=base_url, vault=vault), vault): - raise UpError("Login did not produce a usable token; cannot start `lite up`.") + raise UpError("Login did not produce a usable token.") def _restore_and_report() -> None: @@ -161,7 +159,9 @@ def up(ctx: click.Context) -> None: """Route every Claude Code session through your LiteLLM proxy until stopped. Patches ~/.claude/settings.json so Claude Code picks up the proxy on its own - next startup, from any terminal -- no need to launch it through `lite`. + next startup, from any terminal -- no need to launch it through `lite`. The + key written is the one this command resolved (your fresh `lite login`, or an + explicit --api-key), copied in as a static token for as long as `up` runs. Press Ctrl-C to stop and restore your original settings. Assumes the proxy is already running (this does not start one for you). Cursor is not supported: it has no equivalent file-based config to patch. @@ -169,7 +169,7 @@ def up(ctx: click.Context) -> None: base_url: Final = ctx.obj["base_url"] try: - _ensure_fresh_login(ctx) + ensure_fresh_login(ctx) api_key: Final = resolve_api_key(ctx) verify_proxy_key(base_url, api_key) @@ -179,7 +179,7 @@ def up(ctx: click.Context) -> None: "running (or crashed without cleanup). Run `lite down` first." ) - api_key_helper: Final = resolve_api_key_helper(base_url) + status_line: Final = install_statusline_script() original_existed: Final = CLAUDE_SETTINGS_PATH.exists() original_settings: Final = load_json_or_empty(CLAUDE_SETTINGS_PATH) write_backup( @@ -190,9 +190,10 @@ def up(ctx: click.Context) -> None: ) CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True) - merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper) - with open(CLAUDE_SETTINGS_PATH, "w") as f: - json.dump(merged, f, indent=2) + merged: Final = merge_claude_settings( + original_settings, base_url, StaticToken(api_key), status_line=status_line + ) + write_claude_settings(CLAUDE_SETTINGS_PATH, merged) except (AgentRunError, ClaudeSettingsError) as e: raise click.ClickException(str(e)) @@ -247,7 +248,6 @@ __all__ = [ "load_json_or_empty", "merge_claude_settings", "read_backup", - "resolve_api_key_helper", "restore_claude_settings", "up", "write_backup", diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index eae1b0f5bc9..b0e81a222c0 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -13,6 +13,7 @@ from .commands.auth import auth_group, context_secret_vault, get_stored_api_key, from .commands.autoroute.commands import autoroute_group from .commands.chat import chat from .commands.config import config_commands, get_config_value, hidden_command_names +from .commands.configure import configure_group, unconfigure_group from .commands.credentials import credentials from .commands.debug import debug from .commands.encryption import encryption @@ -162,6 +163,9 @@ cli.add_command(model_groups) # Add the autoroute command group (QA auto-routing against your real proxy) cli.add_command(autoroute_group, name="autoroute") cli.add_command(config_commands) +# Add configure/unconfigure (persistently wire a coding agent to the proxy with a virtual key) +cli.add_command(configure_group) +cli.add_command(unconfigure_group) if __name__ == "__main__": diff --git a/litellm/proxy/collector.py b/litellm/proxy/collector.py new file mode 100644 index 00000000000..3ff2a83860b --- /dev/null +++ b/litellm/proxy/collector.py @@ -0,0 +1,220 @@ +"""Collector sidecar: consume spend events from the pod's inference workers and run the cost pipeline. + +Runs the proxy startup lifespan (config, Prisma, Redis transaction buffer, scheduled spend flushes) +without serving HTTP, then listens on ``LITELLM_COLLECTOR_ADDRESS`` for newline-delimited spend +events. Each event goes through the unchanged ``_ProxyDBLogger._PROXY_track_cost_callback``, so +spend logs, spend counters, budget reservation reconciliation and cache updates happen exactly as +they would in-process, just in this container. Events are handled in order per producer connection +(one per uvicorn worker); a slow pipeline fills the socket buffer and the producer's bounded queue, +which is the backpressure that triggers its fallback or drop policy. ``SIGTERM`` stops accepting +connections, half-closes every producer connection so the producers switch to their unavailable +policy, finishes the events already sent, then runs the proxy shutdown (which flushes the buffered +spend transactions). + +``DATABASE_URL`` is assembled from the same ``DATABASE_*`` inputs as the proxy container, and when +``LITELLM_PGBOUNCER_ENABLED`` is set it points at the PgBouncer that container already runs on the +pod's loopback, so the sidecar must see the same env as the proxy. Under ``IAM_TOKEN_DB_AUTH`` or +``AZURE_POSTGRESQL_AUTH`` that PgBouncer only accepts the token the proxy container minted, so the +sidecar goes to Postgres directly and mints its own. Works from any image that has ``litellm`` +installed: + + python -m litellm.proxy.collector [--address unix:///path.sock] +""" + +import asyncio +import logging +import os +import signal +import sys +from collections.abc import Awaitable, Callable, Mapping, Sequence +from pathlib import Path +from typing import Final + +from litellm._logging import verbose_logger, verbose_proxy_logger, verbose_router_logger +from litellm.proxy.db.db_url_settings import DatabaseURLSettings +from litellm.proxy.db.pgbouncer import ( + PgBouncerError, + PgBouncerSettings, + export_pooled_database_url, + pooled_database_url, +) +from litellm.proxy.spend_tracking.spend_event_producer import ( + COLLECTOR_JOB_ROLE, + AddressError, + CollectorAddress, + CollectorSettings, + TcpAddress, + UnixAddress, + parse_collector_address, +) + +MAX_EVENT_BYTES: Final = 64 * 1024 * 1024 + + +class SpendEventConsumer: + """Accepts producer connections and runs ``handler`` on every line each one sends, in order.""" + + def __init__(self, handler: Callable[[bytes], Awaitable[None]]) -> None: + self._handler = handler + self._open_connections: set[asyncio.StreamWriter] = set() # mutable-ok: live producer connections + self._idle = asyncio.Event() + self._idle.set() + self._received = 0 + self._handled = 0 + self._failed = 0 + + @property + def received(self) -> int: + return self._received + + @property + def handled(self) -> int: + return self._handled + + @property + def failed(self) -> int: + return self._failed + + async def serve(self, address: CollectorAddress) -> asyncio.Server: + match address: + case UnixAddress(path=path): + socket_path: Final = Path(path) + socket_path.parent.mkdir(parents=True, exist_ok=True) + socket_path.unlink(missing_ok=True) + return await asyncio.start_unix_server(self._on_connection, path=path, limit=MAX_EVENT_BYTES) + case TcpAddress(host=host, port=port): + return await asyncio.start_server(self._on_connection, host=host, port=port, limit=MAX_EVENT_BYTES) + + async def _on_connection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + self._open_connections.add(writer) + self._idle.clear() + try: + while line := await reader.readline(): + if not line.endswith(b"\n"): + verbose_proxy_logger.error("collector: discarding truncated spend event (%d bytes)", len(line)) + break + self._received += 1 + await self._handle(line) + except (ConnectionError, asyncio.IncompleteReadError, asyncio.LimitOverrunError) as error: + verbose_proxy_logger.warning("collector: producer connection ended abnormally: %s", error) + finally: + writer.close() + self._open_connections.discard(writer) + if not self._open_connections: + self._idle.set() + + async def _handle(self, line: bytes) -> None: + try: + await self._handler(line) + self._handled += 1 + except Exception: # noqa: BLE001 # the cost pipeline raises anything; one bad event must not stop the sidecar + self._failed += 1 + verbose_proxy_logger.exception("collector: spend event failed") + + async def drain(self, timeout: float) -> int: + """Half-close every producer connection, then keep reading until each producer hangs up or ``timeout``. + + Returns how many producer connections were still open when the timeout hit. + """ + for writer in tuple(self._open_connections): + if writer.is_closing() or not writer.can_write_eof(): + continue + try: + writer.write_eof() + except (OSError, RuntimeError) as error: + verbose_proxy_logger.debug("collector: producer already gone before half-close: %s", error) + try: + await asyncio.wait_for(self._idle.wait(), timeout) + except TimeoutError: + pass + return len(self._open_connections) + + +def _install_stop_signals(loop: asyncio.AbstractEventLoop, stop: asyncio.Event) -> None: + for signum in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(signum, stop.set) + + +async def run_collector(address: CollectorAddress, drain_timeout: float) -> None: + from fastapi import FastAPI + + from litellm.proxy.hooks.proxy_track_cost_callback import run_spend_event + from litellm.proxy.proxy_server import proxy_startup_event + + stop: Final = asyncio.Event() + _install_stop_signals(asyncio.get_running_loop(), stop) + consumer: Final = SpendEventConsumer(handler=run_spend_event) + async with proxy_startup_event(FastAPI()): + server: Final = await consumer.serve(address) + verbose_proxy_logger.info("collector: listening on %s", address) + await stop.wait() + server.close() + still_open: Final = await consumer.drain(drain_timeout) + verbose_proxy_logger.info( + "collector: stopping. received=%d handled=%d failed=%d connections_cut=%d", + consumer.received, + consumer.handled, + consumer.failed, + still_open, + ) + + +def address_argument(argv: Sequence[str], default: str) -> str | AddressError: + match tuple(argv): + case (): + return default + case ("--address", value): + return value + case _: + return AddressError(f"usage: python -m litellm.proxy.collector [--address ADDRESS], got {tuple(argv)}") + + +def apply_log_level(litellm_log: str | None) -> None: + """Mirror the proxy's ``LITELLM_LOG`` handling: the sidecar has no CLI flags to turn logging on.""" + level: Final = logging.getLevelNamesMapping().get((litellm_log or "").upper()) + if level is None: + return + for logger in (verbose_logger, verbose_router_logger, verbose_proxy_logger): + logger.setLevel(level) + + +def pod_pgbouncer_database_url( + pgbouncer: PgBouncerSettings, environ: Mapping[str, str], *, token_auth: bool +) -> str | PgBouncerError | None: + """The proxy container's PgBouncer URL for ``environ["DATABASE_URL"]``, or None to connect to Postgres directly. + + Direct is the answer when PgBouncer is off, and also under token auth: that PgBouncer's auth file + only holds the token its own container minted, which this container cannot present. + """ + if not pgbouncer.enabled or token_auth: + return None + upstream_url: Final = environ.get("DATABASE_URL") + if upstream_url is None: + return PgBouncerError("LITELLM_PGBOUNCER_ENABLED is set but no DATABASE_URL could be assembled") + return pooled_database_url(upstream_url, pgbouncer) + + +def main(argv: Sequence[str]) -> None: + os.environ.setdefault("LITELLM_JOB_ROLE", COLLECTOR_JOB_ROLE) + apply_log_level(os.environ.get("LITELLM_LOG")) + database: Final = DatabaseURLSettings.from_env() + database.apply_to_env() + pooled: Final = pod_pgbouncer_database_url( + PgBouncerSettings(), + os.environ, + token_auth=database.iam_token_db_auth or database.azure_postgresql_auth, + ) + if isinstance(pooled, PgBouncerError): + sys.exit(f"LiteLLM collector: cannot use the pod's pgbouncer: {pooled.reason}") + if pooled is not None: + export_pooled_database_url(pooled) + settings: Final = CollectorSettings() + raw_address: Final = address_argument(argv, default=settings.address) + address: Final = raw_address if isinstance(raw_address, AddressError) else parse_collector_address(raw_address) + if isinstance(address, AddressError): + sys.exit(f"LiteLLM collector: {address.reason}") + asyncio.run(run_collector(address, drain_timeout=settings.drain_timeout_seconds)) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5e6c9b34332..e3a2b892721 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -54,6 +54,12 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, ) +from litellm.proxy.common_utils.openai_error_payload import ( + attribute_of, + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, coerce_keepalive_interval, @@ -184,6 +190,7 @@ from litellm.proxy.litellm_pre_call_utils import ( refresh_proxy_server_request_body_snapshot, reject_url_valued_destination, ) +from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -464,46 +471,6 @@ def _stream_usage_tracking_updates( } -def _getattr_object(value: object, name: str, default: object = None) -> object: - return getattr(value, name, default) - - -_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( - { - status.HTTP_401_UNAUTHORIZED: "authentication_error", - status.HTTP_403_FORBIDDEN: "permission_error", - status.HTTP_429_TOO_MANY_REQUESTS: "rate_limit_error", - } -) - - -def _error_status_code(exc: object, default: int) -> int: - """The HTTP status an exception carries, or ``default`` when it carries none.""" - carried: Final = _getattr_object(exc, "status_code") - return carried if isinstance(carried, int) and not isinstance(carried, bool) else default - - -def _openai_error_type(exc: object, status_code: int) -> str: - """OpenAI types ``error.type`` as a required string, so an exception carrying none - falls back to the type its status code stands for.""" - carried: Final = _getattr_object(exc, "type") - if isinstance(carried, str): - return carried - mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) - if mapped is not None: - return mapped - if status_code < status.HTTP_500_INTERNAL_SERVER_ERROR: - return "invalid_request_error" - return "internal_server_error" - - -def _openai_error_param(exc: object) -> str | None: - """OpenAI types ``error.param`` as nullable, so an exception carrying none - serializes as JSON ``null``.""" - carried: Final = _getattr_object(exc, "param") - return carried if isinstance(carried, str) else None - - class _UpstreamHttpResponse(Protocol): @property def status_code(self) -> int: ... @@ -573,15 +540,15 @@ def serialize_http_exception_detail( def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException: - raw_detail: Final = _getattr_object(exc, "detail", str(exc)) + raw_detail: Final = attribute_of(exc, "detail", str(exc)) message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) - error_status: Final = _error_status_code(exc, status.HTTP_400_BAD_REQUEST) + error_status: Final = error_status_code(exc, status.HTTP_400_BAD_REQUEST) return ProxyException( message=message, - type=_openai_error_type(exc, error_status), - param=_openai_error_param(exc), + type=openai_error_type(exc, error_status), + param=openai_error_param(exc), code=error_status, provider_specific_fields=merged_fields, headers=headers, @@ -756,7 +723,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse): content: AsyncGenerator[str, None], *, media_type: str | None = None, - headers: dict | None = None, + headers: Mapping[str, str] | None = None, status_code: int = status.HTTP_200_OK, upstream_generator: AsyncGenerator[str, None] | None = None, ) -> None: @@ -865,8 +832,8 @@ def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: are byte-identical. """ # Preserve status code from HTTPException (e.g. guardrail blocks) - error_status: Final = _error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR) - raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start") + error_status: Final = error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR) + raw_detail: Final = attribute_of(exc, "detail", "Error processing stream start") message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} @@ -874,8 +841,8 @@ def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: error_obj: Final = { "message": message, - "type": _openai_error_type(exc, error_status), - "param": _openai_error_param(exc), + "type": openai_error_type(exc, error_status), + "param": openai_error_param(exc), "code": str(error_status), } if not merged_fields: @@ -888,25 +855,39 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]: return f"data: {json.dumps({'error': error_obj})}\n\n", "data: [DONE]\n\n" +def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + """`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384).""" + return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS}) + + +async def _resolve_stream_headers( + headers: Mapping[str, str], refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None +) -> Mapping[str, str]: + if refresh_headers is None: + return headers + try: + return await refresh_headers() + except Exception as e: # noqa: BLE001 # a stream whose first chunk is already paid for must not fail over its headers + verbose_proxy_logger.exception("Error refreshing streaming response headers: %s", e) + return headers + + async def create_response( generator: AsyncGenerator[str, None], media_type: str, - headers: dict, + headers: Mapping[str, str], default_status_code: int = status.HTTP_200_OK, request: Request | None = None, + refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None = None, ) -> StreamingResponse | JSONResponse: """ Create streaming response, checking if the first chunk is an error. If the first chunk is an error, return a standard JSON error response. Otherwise, return StreamingResponse and stream all content. + + ``refresh_headers`` is consulted once the first chunk has been buffered, for + callers whose headers can only be known then. """ - # Tell buffering reverse proxies (nginx, ingress-nginx, Envoy) to flush SSE - # immediately instead of releasing the whole stream in one batch (issue #28384). - streaming_headers: Final = { - **headers, - "Cache-Control": "no-cache", - "X-Accel-Buffering": "no", - } first_chunk_value: str | None = None final_status_code = default_status_code @@ -917,6 +898,7 @@ async def create_response( # Now get the first chunk from the actual generator first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request) + resolved_headers: Final = await _resolve_stream_headers(headers, refresh_headers) if first_chunk_value is not None: try: @@ -943,7 +925,7 @@ async def create_response( return JSONResponse( status_code=final_status_code, content={"error": error_dict}, - headers=headers, + headers=resolved_headers, ) except Exception as e: verbose_proxy_logger.debug("Error parsing first chunk value: %s", e) @@ -972,7 +954,7 @@ async def create_response( return StreamingResponse( empty_gen(), media_type=media_type, - headers=streaming_headers, + headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)), status_code=default_status_code, ) except Exception as e: @@ -988,7 +970,7 @@ async def create_response( return StreamingResponse( error_gen_message(), media_type=media_type, - headers=streaming_headers, + headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)), status_code=error_status, ) @@ -1010,7 +992,7 @@ async def create_response( return _UpstreamClosingStreamingResponse( combined_generator(), media_type=media_type, - headers=streaming_headers, + headers=_sse_stream_headers(resolved_headers), status_code=final_status_code, upstream_generator=generator, ) @@ -1535,7 +1517,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _merge_passthrough_streaming_headers( response_headers: httpx.Headers | dict | None, - custom_headers: dict, + custom_headers: Mapping[str, str], ) -> dict: """ Merge upstream passthrough headers with proxy/custom headers. @@ -1868,7 +1850,6 @@ class ProxyBaseLLMRequestProcessing: data=self.data, user_api_key_dict=user_api_key_dict, ) - # Calculate request queue time after add_litellm_data_to_request # which sets arrival_time in proxy_server_request. Ends at start_time # (not a freshly captured time.time() here) so this window is exactly @@ -2016,6 +1997,12 @@ class ProxyBaseLLMRequestProcessing: data=self.data, call_type=route_type, ) + if route_type == "aget_responses": + attach_post_call_pipelines_to_retrieval( + data=self.data, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + ) # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in @@ -2143,14 +2130,45 @@ class ProxyBaseLLMRequestProcessing: return fallback_model_group @staticmethod - def _get_model_id_from_response(hidden_params: dict, data: dict) -> str: + def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str: """Extract model_id from hidden_params with fallback to litellm_metadata.""" model_id = hidden_params.get("model_id", None) or "" if not model_id: - litellm_metadata: Final = data.get("litellm_metadata", {}) or {} - model_info: Final = litellm_metadata.get("model_info", {}) or {} - model_id = model_info.get("id", "") or "" - return model_id + litellm_metadata: Final = data.get("litellm_metadata") + model_info: Final = litellm_metadata.get("model_info") if isinstance(litellm_metadata, Mapping) else None + model_id = (model_info.get("id") or "") if isinstance(model_info, Mapping) else "" + return str(model_id) if model_id else "" + + def _stream_response_headers( + self, + *, + hidden_params: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + logging_obj: LiteLLMLoggingObj, + version: str | None, + callback_headers: Mapping[str, str], + ) -> Mapping[str, str]: + """The streaming response headers describing `hidden_params`' deployment.""" + return MappingProxyType( + { + **ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=logging_obj.litellm_call_id, + model_id=self._get_model_id_from_response(hidden_params, self.data), + cache_key=hidden_params.get("cache_key") or "", + api_base=hidden_params.get("api_base") or "", + version=version, + response_cost=hidden_params.get("response_cost") or "", + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"), + request_data=self.data, + hidden_params=hidden_params, + litellm_logging_obj=logging_obj, + **(hidden_params.get("additional_headers") or MappingProxyType({})), + ), + **callback_headers, + } + ) @staticmethod def _get_deployment_model_name( @@ -2419,31 +2437,32 @@ class ProxyBaseLLMRequestProcessing: if self._is_streaming_request( data=self.data, is_streaming_request=is_streaming_request ) or self._is_streaming_response(response): # use generate_responses to stream responses - custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - call_id=logging_obj.litellm_call_id, - model_id=model_id, - cache_key=cache_key, - api_base=api_base, - version=version, - response_cost=response_cost, - model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - fastest_response_batch_completion=fastest_response_batch_completion, - request_data=self.data, - hidden_params=hidden_params, - litellm_logging_obj=logging_obj, - **additional_headers, - ) - # Call response headers hook for streaming success - callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + stream_callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook( data=self.data, user_api_key_dict=user_api_key_dict, response=response, request_headers=dict(request.headers), ) - if callback_headers: - custom_headers.update(callback_headers) + custom_headers: Final = self._stream_response_headers( + hidden_params=hidden_params, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + callback_headers=stream_callback_headers or MappingProxyType({}), + ) + + async def refresh_stream_headers() -> Mapping[str, str]: + """`custom_headers` rebuilt for whichever deployment served the stream.""" + if not getattr(response, "fallback_headers_adopted", False): + return custom_headers + return self._stream_response_headers( + hidden_params=get_hidden_params_dict(response), + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + callback_headers=stream_callback_headers or MappingProxyType({}), + ) # Preserve the original client-requested model (pre-alias mapping) for downstream # streaming generators. Pre-call processing can rewrite `self.data["model"]` for @@ -2581,6 +2600,7 @@ class ProxyBaseLLMRequestProcessing: media_type="text/event-stream", headers=custom_headers, request=request, + refresh_headers=refresh_stream_headers, ) ### CALL HOOKS ### - modify outgoing data @@ -2767,10 +2787,10 @@ class ProxyBaseLLMRequestProcessing: ``ResponsesAPIResponse`` directly. Handle both shapes so the container-ownership recording path can walk ``.output`` either way. """ - completed: Final = _getattr_object(stream_response, "completed_response") + completed: Final = attribute_of(stream_response, "completed_response") if completed is None: return None - response_obj: Final = _getattr_object(completed, "response") + response_obj: Final = attribute_of(completed, "response") if response_obj is not None: return response_obj return completed @@ -3032,7 +3052,7 @@ class ProxyBaseLLMRequestProcessing: response: Any, proxy_logging_obj: "ProxyLogging", user_api_key_dict: "UserAPIKeyAuth", - custom_headers: dict, + custom_headers: Mapping[str, str], request_headers: dict[str, str], ) -> Response | None: if not self._has_post_call_guardrails_for_passthrough(): @@ -3280,9 +3300,10 @@ class ProxyBaseLLMRequestProcessing: has completed. Guardrails routed through unified_guardrail are skipped, since they already ran - via its streaming iterator. Guardrails that override - async_post_call_success_hook directly run here, including those that implement - apply_guardrail but keep their native lifecycle hooks. + via its streaming iterator, and so are guardrails a post_call policy pipeline + manages, since the pipeline ran them against the buffered stream. Guardrails + that override async_post_call_success_hook directly run here, including those + that implement apply_guardrail but keep their native lifecycle hooks. This is audit-only — content has already been delivered to the client. @@ -3292,12 +3313,18 @@ class ProxyBaseLLMRequestProcessing: _response = assembled_response try: from litellm.proxy.proxy_server import llm_router as _global_llm_router - from litellm.proxy.utils import _check_and_merge_model_level_guardrails + from litellm.proxy.utils import ( + _check_and_merge_model_level_guardrails, + stream_gated_guardrail_names, + ) guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router) + stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict) for cb in litellm.callbacks: if not isinstance(cb, CustomGuardrail): continue + if cb.guardrail_name in stream_gated: + continue if not cb.should_run_guardrail( data=guardrail_data, event_type=GuardrailEventHooks.post_call, @@ -3420,7 +3447,7 @@ class ProxyBaseLLMRequestProcessing: headers = getattr(e, "headers", None) or {} if not headers: # Try to get headers from e.response.headers (httpx.Response) - _response: Final = _getattr_object(e, "response") + _response: Final = attribute_of(e, "response") if _response is not None: _response_headers: Final = getattr(_response, "headers", None) if _response_headers: @@ -3460,9 +3487,13 @@ class ProxyBaseLLMRequestProcessing: error_body: Final = await http_status_error.response.aread() error_text: Final = error_body.decode("utf-8") + error_headers: Final = { # mutable-ok: HTTPException takes a plain header dict + k: v if isinstance(v, str) else str(v) for k, v in safe_headers.items() + } raise HTTPException( status_code=http_status_error.response.status_code, detail={"error": error_text}, + headers=error_headers, ) error_msg: Final = f"{e}" # Check for AttributeError in the exception chain. @@ -3491,8 +3522,8 @@ class ProxyBaseLLMRequestProcessing: _code = status.HTTP_500_INTERNAL_SERVER_ERROR raise ProxyException( message=redact_internal_details_from_client_message(getattr(e, "message", error_msg)), - type=_openai_error_type(e, _code), - param=_openai_error_param(e), + type=openai_error_type(e, _code), + param=openai_error_param(e), openai_code=getattr(e, "code", None), code=_code, provider_specific_fields=getattr(e, "provider_specific_fields", None), @@ -3702,11 +3733,11 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raise e - stream_error_status: Final = _error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR) + stream_error_status: Final = error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR) proxy_exception: Final = ProxyException( message=redact_internal_details_from_client_message(getattr(e, "message", str(e))), - type=_openai_error_type(e, stream_error_status), - param=_openai_error_param(e), + type=openai_error_type(e, stream_error_status), + param=openai_error_param(e), code=stream_error_status, ) stream_completed = True diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index fb2ca6372c0..dbe11882b3c 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -196,9 +196,7 @@ class AuthCacheInvalidationSubscriber: for additional_cache in self._additional_in_memory_caches: additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl) return - in_memory_cache: Final = self._user_api_key_cache.in_memory_cache - if in_memory_cache is not None: - in_memory_cache.delete_cache(parsed.cache_key) + self._user_api_key_cache.in_memory_cache_for(parsed.cache_key).delete_cache(parsed.cache_key) for additional_cache in self._additional_in_memory_caches: additional_cache.delete_cache(parsed.cache_key) diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py index 7ee3bd8d829..c9d97068313 100644 --- a/litellm/proxy/common_utils/callback_config_validation.py +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -44,6 +44,91 @@ def _langfuse_environment_error(callback_vars: Mapping[str, str]) -> str | None: return None +# Which credential family a dynamic variable belongs to. The families are the +# integrations that share one account: every langfuse_* variable configures the +# same Langfuse project whether it rides the classic callback or the OTel one, +# and every dd_* variable configures the same Datadog account. +_VAR_FAMILIES: Final[Mapping[str, str]] = MappingProxyType( + { + "arize_": "Arize", + "dd_": "Datadog", + "gcs_": "GCS", + "humanloop_": "Humanloop", + "langfuse_": "Langfuse", + "langsmith_": "LangSmith", + "newrelic_": "New Relic", + "posthog_": "PostHog", + "wandb_": "Weights & Biases", + "weave_": "Weights & Biases", + } +) + + +def _family_of(var: str) -> str | None: + """The credential family ``var`` configures, or ``None`` if it configures none. + + ``turn_off_message_logging`` and friends belong to no backend, so they carry + no credentials anyone could redirect. + """ + return next((family for prefix, family in _VAR_FAMILIES.items() if var.startswith(prefix)), None) + + +def cross_entry_family_error( + callback_vars: Mapping[str, str] | None, + stored_vars_by_entry: Sequence[Mapping[str, str]], +) -> str | None: + """Reject an entry that changes what a family another entry holds resolves to. + + Every stored entry's variables are flattened into one dict before a request + reads them, and the flattened dict is what the exporter authenticates and + addresses with. So an entry naming only a destination is enough to redirect + credentials that were written somewhere else: a host on a second entry pairs + with the key from the first, and the request carries that key to the new + host. + + Two rules together keep the flattened dict out of the caller's hands. A + variable the family already configures has to keep the value it has, so + nothing already in use can be moved. A variable the family does not yet + configure may only carry a value the family already holds, which is what lets + the same credential go in under its other spelling (``langfuse_secret`` and + ``langfuse_secret_key`` are one key) without anything here having to list the + spellings. Between them, no value the caller chose can enter the family, and + repeating the family as it stands is still allowed -- that is how one + integration gets registered for both the success and the failure event. + + A team admin who does want to move a family deletes the entry holding it + first, which reveals nothing. + + Only the writers this endpoint newly admits are held to this, because a proxy + admin already holds every credential the proxy has. + + ``stored_vars_by_entry`` has to arrive decrypted; the credential values are + encrypted at rest and ciphertext never equals the plaintext coming in. + """ + if not callback_vars: + return None + stored_by_var: Final = { + var: value for entry in stored_vars_by_entry for var, value in entry.items() if _family_of(var) is not None + } + family_values: Final = frozenset( + (family, value) + for entry in stored_vars_by_entry + for var, value in entry.items() + if (family := _family_of(var)) is not None + ) + held_families: Final = frozenset(family for family, _ in family_values) + return next( + ( + f"{family} is already configured by another callback entry on this team. " + f"Remove that entry before setting {var} here." + for var, value, family in ((v, callback_vars[v], _family_of(v)) for v in callback_vars) + if family in held_families + and (stored_by_var[var] != value if var in stored_by_var else (family, value) not in family_values) + ), + None, + ) + + def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str | None: """Validate every ``logging`` entry of a team/key metadata payload.""" if not metadata: diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 39e74d2c8bd..561a53409f4 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,17 +1,22 @@ import copy +import json import os from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass +from itertools import accumulate from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias -from typing_extensions import assert_never +from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, + MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + ROUTING_REQUEST_TAGS_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger @@ -26,6 +31,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.types_utils.utils import get_instance_fn +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( StandardLoggingGuardrailInformation, StandardLoggingPayload, @@ -50,6 +56,15 @@ reset_color_code: Final = "\033[0m" TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY: Final = "_pillar_response_headers_trusted" GUARDRAIL_SCAN_IDS_METADATA_KEY: Final = "guardrail_scan_ids" +GUARDRAIL_SCAN_METADATA_METADATA_KEY: Final = "guardrail_scan_metadata" + + +class GuardrailScanMetadata(TypedDict): + guardrail: ReadOnly[str | None] + stage: ReadOnly[str] + provider: ReadOnly[str] + scan_id: ReadOnly[str] + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -448,6 +463,16 @@ def get_remaining_tokens_and_requests_from_request_data(data: dict) -> dict[str, return headers +def _serialize_scan_metadata_header(entries: Iterable[object], *, max_length: int) -> str | None: + """Compact JSON list of scan metadata entries, dropping trailing entries so the header fits in max_length.""" + encoded: Final = tuple(json.dumps(entry, separators=(",", ":")) for entry in entries) + lengths: Final = tuple(accumulate(len(item) + 1 for item in encoded)) + kept: Final = sum(1 for length in lengths if length + 1 <= max_length) + if kept == 0: + return None + return f"[{','.join(encoded[:kept])}]" + + def get_logging_caching_headers(request_data: dict) -> dict | None: _metadata: Final[dict] = {} metadata_bucket: Final = request_data.get("metadata") @@ -466,6 +491,15 @@ def get_logging_caching_headers(request_data: dict) -> dict | None: if scan_ids: headers["x-litellm-guardrail-scan-id"] = ",".join(scan_ids) + scan_metadata: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY) + scan_metadata_header: Final = ( + _serialize_scan_metadata_header(scan_metadata, max_length=MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH) + if isinstance(scan_metadata, (list, tuple)) + else None + ) + if scan_metadata_header: + headers["x-litellm-guardrail-scan-metadata"] = scan_metadata_header + if "applied_policies" in _metadata: headers["x-litellm-applied-policies"] = ",".join(_metadata["applied_policies"]) @@ -499,6 +533,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( "applied_policies", "applied_guardrails", GUARDRAIL_SCAN_IDS_METADATA_KEY, + GUARDRAIL_SCAN_METADATA_METADATA_KEY, "policy_sources", "guardrails", "guardrail_config", @@ -507,6 +542,8 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, + CLIENT_OUTPUT_CEILING_METADATA_KEY, + ROUTING_REQUEST_TAGS_METADATA_KEY, "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", @@ -561,21 +598,40 @@ def add_guardrail_to_applied_guardrails_header(request_data: dict, guardrail_nam _metadata["applied_guardrails"] = [guardrail_name] -def add_guardrail_scan_id(request_data: dict, scan_id: str | None) -> None: +def add_guardrail_scan_id( + request_data: dict[str, object], + scan_id: str | None, + *, + guardrail_name: str | None, + provider: str, + stage: GuardrailEventHooks, +) -> None: """ - Record a provider scan id so it can be surfaced to the caller. + Record a provider scan id, keyed to the guardrail execution that produced it, so it can be surfaced to the caller. Guardrails only return scan details to the client when they block, so allowed requests carry no - audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header. + audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header, and the + (guardrail, stage, provider, scan_id) entries become the x-litellm-guardrail-scan-metadata header. """ if not scan_id: return _, _metadata = get_or_create_metadata_bucket(request_data) existing: Final = _metadata.get(GUARDRAIL_SCAN_IDS_METADATA_KEY) - scan_ids: Final = tuple(existing) if isinstance(existing, (list, tuple)) else () + scan_ids: Final[tuple[object, ...]] = tuple(existing) if isinstance(existing, (list, tuple)) else () if scan_id not in scan_ids: _metadata[GUARDRAIL_SCAN_IDS_METADATA_KEY] = (*scan_ids, scan_id) + entry: Final[GuardrailScanMetadata] = { + "guardrail": guardrail_name, + "stage": stage.value, + "provider": provider, + "scan_id": scan_id, + } + existing_entries: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY) + entries: Final[tuple[object, ...]] = tuple(existing_entries) if isinstance(existing_entries, (list, tuple)) else () + if entry not in entries: + _metadata[GUARDRAIL_SCAN_METADATA_METADATA_KEY] = (*entries, entry) + def add_policy_to_applied_policies_header(request_data: dict, policy_name: str | None): """ diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 554a6ae8d1a..1d4024bb84e 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -147,8 +147,11 @@ async def memory_usage_in_mem_cache( llm_router.cache.in_memory_cache.ttl_dict ) - num_items_in_user_api_key_cache: Final = len(user_api_key_cache.in_memory_cache.cache_dict) + len( - user_api_key_cache.in_memory_cache.ttl_dict + num_items_in_user_api_key_cache: Final = ( + len(user_api_key_cache.in_memory_cache.cache_dict) + + len(user_api_key_cache.in_memory_cache.ttl_dict) + + len(user_api_key_cache.key_object_cache.in_memory_cache.cache_dict) + + len(user_api_key_cache.key_object_cache.in_memory_cache.ttl_dict) ) num_items_in_proxy_logging_obj_cache: Final = len( @@ -189,6 +192,8 @@ async def memory_usage_in_mem_cache_items( return { "user_api_key_cache": user_api_key_cache.in_memory_cache.cache_dict, "user_api_key_ttl": user_api_key_cache.in_memory_cache.ttl_dict, + "user_key_object_cache": user_api_key_cache.key_object_cache.in_memory_cache.cache_dict, + "user_key_object_ttl": user_api_key_cache.key_object_cache.in_memory_cache.ttl_dict, "llm_router_cache": llm_router_in_memory_cache_dict, "llm_router_ttl": llm_router_in_memory_ttl_dict, "proxy_logging_obj_cache": proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict, @@ -294,7 +299,9 @@ async def get_memory_summary( try: # User API key cache - user_cache_items: Final = len(user_api_key_cache.in_memory_cache.cache_dict) + user_cache_items: Final = len(user_api_key_cache.in_memory_cache.cache_dict) + len( + user_api_key_cache.key_object_cache.in_memory_cache.cache_dict + ) total_cache_items += user_cache_items caches["user_api_keys"] = { "count": user_cache_items, @@ -429,10 +436,16 @@ def _get_cache_memory_stats( cache_stats: Final[dict[str, object]] = {} try: # User API key cache - user_cache_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.cache_dict) - user_ttl_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.ttl_dict) + key_object_in_memory_cache: Final = user_api_key_cache.key_object_cache.in_memory_cache + user_cache_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.cache_dict) + sys.getsizeof( + key_object_in_memory_cache.cache_dict + ) + user_ttl_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.ttl_dict) + sys.getsizeof( + key_object_in_memory_cache.ttl_dict + ) cache_stats["user_api_key_cache"] = { - "num_items": len(user_api_key_cache.in_memory_cache.cache_dict), + "num_items": len(user_api_key_cache.in_memory_cache.cache_dict) + + len(key_object_in_memory_cache.cache_dict), "cache_dict_size_bytes": user_cache_size, "ttl_dict_size_bytes": user_ttl_size, "total_size_mb": round((user_cache_size + user_ttl_size) / (1024 * 1024), 2), diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 836a4a778bb..fd9b3beee46 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -1,7 +1,10 @@ import base64 import os +from collections.abc import Mapping from typing import Final, Literal, cast +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_proxy_logger # Versioned ciphertext marker for AES-256-GCM values. @@ -203,3 +206,40 @@ def decrypt_value(value: bytes, signing_key: str) -> str: return plaintext except Exception as e: raise e + + +class SecretMapDecodeError(RuntimeError): + pass + + +_SECRET_MAP: Final = TypeAdapter(Mapping[str, str]) +_STORED_SECRET_MAP: Final = TypeAdapter(Mapping[str, str] | str) +_SECRET_STRING: Final = TypeAdapter(str) + + +def encrypt_secret_map(value: Mapping[str, str], new_encryption_key: str | None = None) -> str: + if not value: + return "{}" + ciphertext: Final = _SECRET_STRING.validate_python( + encrypt_value_helper(_SECRET_MAP.dump_json(value).decode(), new_encryption_key=new_encryption_key), strict=True + ) + return _SECRET_STRING.dump_json(ciphertext).decode() + + +def decode_secret_map(value: object, *, key: str) -> Mapping[str, str] | None: + if value is None: + return None + try: + stored: Final = ( + _STORED_SECRET_MAP.validate_json(value, strict=True) + if isinstance(value, str) and value.lstrip().startswith(("{", '"')) + else _STORED_SECRET_MAP.validate_python(value, strict=True) + ) + if not isinstance(stored, str): + return stored + decrypted: Final = decrypt_value_helper( + value=stored, key=key, exception_type="debug", return_original_value=False + ) + return _SECRET_MAP.validate_json(decrypted, strict=True) + except ValidationError: + raise SecretMapDecodeError(f"Cannot decode encrypted MCP {key}; check LITELLM_SALT_KEY") from None diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 54a0f18fd63..845589aee7a 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -39,7 +39,7 @@ def _is_form_content_type(content_type: str) -> bool: return _normalize_media_type(content_type) in _FORM_CONTENT_TYPES -def _is_json_content_type(content_type: str) -> bool: +def is_json_content_type(content_type: str) -> bool: """True iff the body should be parsed as JSON.""" return _normalize_media_type(content_type) == "application/json" @@ -213,6 +213,18 @@ async def _read_request_body(request: Request | None) -> dict: return {} +async def read_raw_json_body(request: Request | None) -> bytes | None: + if request is None or _safe_get_request_parsed_body(request=request) is None: + return None + content_type: Final = _safe_get_request_headers(request=request).get("content-type", "") + if _is_form_content_type(content_type): + return None + try: + return await request.body() + except RuntimeError: + return None + + def _safe_get_request_parsed_body(request: Request | None) -> dict | None: if request is None: return None @@ -406,7 +418,7 @@ async def get_request_body(request: Request) -> dict[str, Any]: """ if request.method == "POST": content_type: Final = request.headers.get("content-type", "") - if _is_json_content_type(content_type): + if is_json_content_type(content_type): return await _read_request_body(request) elif _is_form_content_type(content_type): return await get_form_data(request) diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 213a697b3dd..4b3ff2711b3 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -10,12 +10,24 @@ legacy internal names with `general_settings.use_team_public_model_name: false`. from __future__ import annotations -from collections.abc import Mapping, Sequence +import re +from collections.abc import Container, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +import litellm + if TYPE_CHECKING: from litellm.router import Router + from litellm.types.proxy.model_listing import ModelInfoResponse + +CLAUDE_CODE_PICKER_PATTERN: Final = re.compile(r"claude|anthropic", re.IGNORECASE) +GATEWAY_CLIENT_HEADER: Final = "x-gateway-client" +CLAUDE_CODE_CLIENT: Final = "claude-code" +_CLAUDE_CODE_ALIAS_PREFIX: Final = "claude-router-" +_ONE_MILLION_SUFFIX: Final = "[1m]" +_ONE_MILLION_TOKENS: Final = 1_000_000 def configured_display_names( @@ -40,6 +52,115 @@ def configured_display_names( ) +def _unmarked(name: str) -> str: + return name[: -len(_ONE_MILLION_SUFFIX)] if name.lower().endswith(_ONE_MILLION_SUFFIX) else name + + +def _compatibility_id(model_id: str) -> str: + return f"{_CLAUDE_CODE_ALIAS_PREFIX}{model_id.encode().hex()}" + + +def _decoded_compatibility_id(view_id: str) -> str | None: + encoded: Final = _unmarked(view_id).removeprefix(_CLAUDE_CODE_ALIAS_PREFIX) + if encoded == _unmarked(view_id): + return None + try: + model_id: Final = bytes.fromhex(encoded).decode() + except (ValueError, UnicodeDecodeError): + return None + return model_id if _compatibility_id(model_id) == _unmarked(view_id) else None + + +def claude_code_model_id( + model_id: str, + max_input_tokens: float | None, + routing_names: Container[str], +) -> str: + """The collision-free id Claude Code's picker lists a model under.""" + if "*" in model_id: + return model_id + shaped: Final = model_id if CLAUDE_CODE_PICKER_PATTERN.search(model_id) else _compatibility_id(model_id) + one_million: Final = max_input_tokens is not None and max_input_tokens >= _ONE_MILLION_TOKENS + marked: Final = ( + f"{shaped}{_ONE_MILLION_SUFFIX}" if one_million and not shaped.lower().endswith(_ONE_MILLION_SUFFIX) else shaped + ) + return next( + ( + name + for name in (marked, shaped) + if name == model_id or claude_code_group_name(name, routing_names) == model_id + ), + model_id, + ) + + +def claude_code_group_name(view_id: str, routing_names: Container[str]) -> str | None: + """Decode a canonical compatibility id only when no configured route claims it.""" + if view_id in routing_names: + return None + unmarked: Final = _unmarked(view_id) + if unmarked != view_id and unmarked in routing_names: + return unmarked + model_id: Final = _decoded_compatibility_id(view_id) + return model_id if model_id and model_id in routing_names else None + + +def is_claude_code_client(headers: Mapping[str, str]) -> bool: + """Claude Code itself, or a client asking for its view of the listing the way Ramp Router's does""" + from litellm.llms.anthropic.common_utils import is_claude_code_user_agent + + return ( + is_claude_code_user_agent(headers.get("user-agent", "")) + or headers.get(GATEWAY_CLIENT_HEADER, "").lower() == CLAUDE_CODE_CLIENT + ) + + +def claude_code_view_ids( + rows: Sequence[ModelInfoResponse], + headers: Mapping[str, str], + routing_names: Container[str], +) -> Mapping[str, str]: + """served id -> Claude Code id for the requested listing view""" + if not is_claude_code_client(headers): + return MappingProxyType({}) + return MappingProxyType( + {row["id"]: claude_code_model_id(row["id"], row.get("max_input_tokens"), routing_names) for row in rows} + ) + + +@dataclass(frozen=True, slots=True) +class ClaudeCodeRoutingNames: + """Existing routes always own their names, including aliases and wildcard routes.""" + + llm_router: Router | None + team_id: str | None = None + alias_maps: tuple[object, ...] = () + + def __contains__(self, name: object) -> bool: + if not isinstance(name, str): + return False + if name in litellm.model_alias_map or any( + isinstance(aliases, Mapping) and name in aliases for aliases in self.alias_maps + ): + return True + if self.llm_router is None: + return False + return ( + name in self.llm_router.model_group_alias + or self.llm_router.has_model_id(name) + or bool(self.llm_router.get_candidate_model_ids_for_route(name, self.team_id)) + ) + + +def claude_code_requested_group( + requested: str, + llm_router: Router, + team_id: str | None, + alias_maps: tuple[object, ...] = (), +) -> str | None: + return claude_code_group_name(requested, ClaudeCodeRoutingNames(llm_router, team_id, alias_maps)) + + class TeamModelNameTranslator: """Translates internal team routing keys to their public names for the model listing/retrieve responses. Stateless; the live router and general_settings diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py new file mode 100644 index 00000000000..89f735ee8b6 --- /dev/null +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -0,0 +1,52 @@ +"""Shapes the ``error`` object the proxy answers with so it matches OpenAI's contract: +``type`` is a required string and ``param`` is nullable, neither of which the literal +string ``"None"`` satisfies.""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from fastapi import status + +_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( + { + status.HTTP_401_UNAUTHORIZED: "authentication_error", + status.HTTP_403_FORBIDDEN: "permission_error", + status.HTTP_429_TOO_MANY_REQUESTS: "rate_limit_error", + } +) + + +def attribute_of(value: object, name: str, default: object = None) -> object: + return getattr(value, name, default) + + +def error_status_code(exc: object, default: int) -> int: + """The HTTP status an exception carries as ``status_code`` or, the way ``ProxyException`` + stores it, as a stringified ``code``; ``default`` when it carries neither.""" + carried: Final = attribute_of(exc, "status_code") + if isinstance(carried, int) and not isinstance(carried, bool): + return carried + stringified: Final = attribute_of(exc, "code") + return int(stringified) if isinstance(stringified, str) and stringified.isdecimal() else default + + +def openai_error_type(exc: object, status_code: int) -> str: + """OpenAI types ``error.type`` as a required string, so an exception carrying none + falls back to the type its status code stands for.""" + carried: Final = attribute_of(exc, "type") + if isinstance(carried, str): + return carried + mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) + if mapped is not None: + return mapped + if status_code < status.HTTP_500_INTERNAL_SERVER_ERROR: + return "invalid_request_error" + return "internal_server_error" + + +def openai_error_param(exc: object) -> str | None: + """OpenAI types ``error.param`` as nullable, so an exception carrying none + serializes as JSON ``null``.""" + carried: Final = attribute_of(exc, "param") + return carried if isinstance(carried, str) else None diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 460b348e188..63da5d15207 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -125,6 +125,7 @@ async def _resync_model_deployments(model_name: str) -> bool: ) return proxy_server.llm_router is not None async with proxy_server.MODEL_RECONCILE_LOCK: + await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client) proxy_server.proxy_config._add_deployment(db_models=rows) proxy_server.llm_model_list = router.get_model_list() return True diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index f2648c8466e..35e74418628 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -224,33 +224,31 @@ def _queue_budget_linked_resets( def _queue_enduser_resets(writes: LinkedSpendResetWrites, cascade: "_BudgetCascade") -> None: - """End users are matched by id rather than budget link: rows with no - budget_id ride the default budget tier (litellm.max_end_user_budget_id). - Zero-before-decrement ordering matters here too (see - _queue_budget_linked_resets).""" - if not cascade.rollover_caps: - if cascade.endusers: - writes.queue_spend_zero( - where={"user_id": {"in": [row.user_id for row in cascade.endusers]}} - ) # mutable-ok: prisma where filter must be a dict + """End users reset on the budget link like every other gated table, plus a + NULL-budget_id branch: rows created implicitly persist no link and ride the + default tier (litellm.max_end_user_budget_id). + + Matching on the link rather than enumerating user ids keeps a statement's + bind count proportional to the expiring tiers instead of the customer + population, which past ~32,700 dependents exceeds PostgreSQL's per-statement + bind ceiling and wedges the cascade permanently (#40564). + """ + _queue_budget_linked_resets(writes, cascade, extra=_SPENT_ROWS_WHERE) + default_budget_id: Final = litellm.max_end_user_budget_id + if default_budget_id is None or default_budget_id not in cascade.budget_ids: return - tiered: Final = tuple((row.budget_id or litellm.max_end_user_budget_id, row.user_id) for row in cascade.endusers) - for budget_id, cap in cascade.rollover_caps.items(): - if not ( - user_ids := [uid for bid, uid in tiered if bid == budget_id] - ): # mutable-ok: prisma "in" filter takes a list - continue + cap: Final = cascade.rollover_caps.get(default_budget_id) + if cap is None: writes.queue_spend_zero( - where={"user_id": {"in": user_ids}, "spend": {"lte": cap}} + where={"budget_id": None, **_SPENT_ROWS_WHERE} ) # mutable-ok: prisma where filter must be a dict - writes.queue_spend_decrement( - where={"user_id": {"in": user_ids}, "spend": {"gt": cap}}, amount=cap - ) # mutable-ok: prisma where filter must be a dict - plain: Final = [ - uid for bid, uid in tiered if bid is None or bid not in cascade.rollover_caps - ] # mutable-ok: prisma "in" filter takes a list - if plain: - writes.queue_spend_zero(where={"user_id": {"in": plain}}) # mutable-ok: prisma where filter must be a dict + return + writes.queue_spend_zero( + where={"budget_id": None, "spend": {"gt": 0, "lte": cap}} + ) # mutable-ok: prisma where filter must be a dict + writes.queue_spend_decrement( + where={"budget_id": None, "spend": {"gt": cap}}, amount=cap + ) # mutable-ok: prisma where filter must be a dict @dataclass(frozen=True, slots=True) diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py index 0820459af49..b8d3595163e 100644 --- a/litellm/proxy/common_utils/semantic_text_index.py +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -5,10 +5,11 @@ from __future__ import annotations import math from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass -from itertools import chain +from itertools import chain, islice from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias +from fastapi import HTTPException from openai import OpenAIError from pydantic import BaseModel, ConfigDict @@ -16,10 +17,14 @@ from litellm.exceptions import BudgetExceededError if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging from litellm.router import Router Vector: TypeAlias = tuple[float, ...] +DEFAULT_MAX_CACHED_VECTORS: Final = 5000 +"""Ceiling on how many (embedding model, text) vectors one index keeps; the least recently searched are evicted first.""" + class Embedder(Protocol): def __call__(self, texts: Sequence[str]) -> Awaitable[Sequence[Vector]]: ... @@ -42,6 +47,16 @@ class _EmbeddingData(BaseModel): data: tuple[_EmbeddingItem, ...] +class _EmbeddingRequest(BaseModel): + """The /embeddings-shaped request as the pre-call hooks (rate limits, budgets, guardrails) hand it back.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + model: str + input: tuple[str, ...] + metadata: dict[str, object] # mutable-ok: the router mutates the metadata dict it is handed + + def cosine_similarity(left: Vector, right: Vector) -> float: dot: Final = sum(a * b for a, b in zip(left, right, strict=True)) norms: Final = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right)) @@ -57,23 +72,40 @@ def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, obj } -def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> Embedder: +def router_embedder( + router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging +) -> Embedder: + """Embeds through the router after the same key rate-limit, budget and guardrail pre-call hooks /embeddings runs.""" + async def embed(texts: Sequence[str]) -> Sequence[Vector]: - batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input + request: Final = { # mutable-ok: pre_call_hook mutates the request dict in place + "model": embedding_model, + "input": list(texts), # mutable-ok: Router.aembedding accepts only str | list input + "metadata": embedding_spend_metadata(user_api_key_dict), + } + processed: Final = _EmbeddingRequest.model_validate( + await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=request, call_type="aembedding" + ) + ) response: Final = await router.aembedding( - model=embedding_model, input=batch, metadata=embedding_spend_metadata(user_api_key_dict) + model=processed.model, + input=list(processed.input), # mutable-ok: Router.aembedding accepts only str | list input + metadata=processed.metadata, ) return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data) return embed -_NO_VECTORS: Final[Mapping[str, Vector]] = MappingProxyType({}) +_CacheKey: TypeAlias = tuple[str, str] async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed: try: vectors: Final = tuple(await embed(texts)) + except HTTPException: + raise except (OpenAIError, ValueError, BudgetExceededError) as exc: return EmbeddingFailed(reason=f"embedding the search query failed: {exc}") if len(vectors) != len(texts): @@ -111,20 +143,34 @@ async def _embed_query_and_texts( class SemanticTextIndex: - """Caches one vector per distinct text per embedding model, so repeat searches only embed the query.""" + """Caches one vector per distinct text per embedding model, so repeat searches only embed the query. - def __init__(self) -> None: - self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({}) + Holds at most ``max_entries`` vectors across all models: once full, the texts no recent search touched go first.""" - def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]: - kept: Final = MappingProxyType( + def __init__(self, max_entries: int = DEFAULT_MAX_CACHED_VECTORS) -> None: + self._max_entries: Final = max_entries + self._vectors: Mapping[_CacheKey, Vector] = MappingProxyType({}) + + def _cached(self, embedding_model: str) -> Mapping[str, Vector]: + return MappingProxyType( + {text: vector for (model, text), vector in self._vectors.items() if model == embedding_model} + ) + + def _merged(self, embedding_model: str, embedded: _Embedded, texts: Sequence[str]) -> Mapping[_CacheKey, Vector]: + dimension: Final = len(embedded.query_vector) + touched: Final = MappingProxyType({(embedding_model, text): embedded.vectors[text] for text in texts}) + untouched: Final = MappingProxyType( { - text: vector - for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items() - if len(vector) == len(embedded.query_vector) + key: vector + for key, vector in chain( + self._vectors.items(), + (((embedding_model, text), vector) for text, vector in embedded.vectors.items()), + ) + if key not in touched and (key[0] != embedding_model or len(vector) == dimension) } ) - return MappingProxyType({**kept, **embedded.vectors}) + ordered: Final = MappingProxyType({**untouched, **touched}) + return MappingProxyType(dict(islice(ordered.items(), max(len(ordered) - self._max_entries, 0), None))) async def scores( self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str @@ -132,11 +178,10 @@ class SemanticTextIndex: """Cosine similarity of `query` to each entry of `texts`, in the same order.""" if not texts: return () - cached: Final = self._vectors.get(embedding_model, _NO_VECTORS) - embedded: Final = await _embed_query_and_texts(embed, query, texts, cached) + embedded: Final = await _embed_query_and_texts(embed, query, texts, self._cached(embedding_model)) if isinstance(embedded, EmbeddingFailed): return embedded if not _same_dimension(embedded.query_vector, embedded.vectors, texts): return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions") - self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)}) + self._vectors = self._merged(embedding_model, embedded, texts) return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 76982d30306..cb72088ee4a 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -1,11 +1,15 @@ from __future__ import annotations +import re +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec @@ -14,6 +18,13 @@ if TYPE_CHECKING: T = TypeVar("T", bound=BaseModel) +_HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") + + +def is_user_key_cache_key(key: str) -> bool: + """Only user-key objects are cached under a bare ``hash_token`` digest; every other object uses a prefixed key.""" + return _HASHED_TOKEN_CACHE_KEY.fullmatch(key) is not None + class UserApiKeyCache(DualCache): """ @@ -36,10 +47,50 @@ class UserApiKeyCache(DualCache): ``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting ``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis). + User-key objects (see ``is_user_key_cache_key``) live in their own in-memory partition, + ``key_object_cache``, so churn in the other management objects cannot evict them. Both + partitions share the same Redis backend and TTL settings. + ``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous (no other methods in between) so mypy resolves ``@overload`` + implementation correctly. """ + def __init__( + self, + in_memory_cache: InMemoryCache | None = None, + redis_cache: RedisCache | None = None, + default_in_memory_ttl: float | None = None, + default_redis_ttl: float | None = None, + key_object_in_memory_cache: InMemoryCache | None = None, + ) -> None: + super().__init__( + in_memory_cache=in_memory_cache, + redis_cache=redis_cache, + default_in_memory_ttl=default_in_memory_ttl, + default_redis_ttl=default_redis_ttl, + ) + self.key_object_cache: Final = DualCache( + in_memory_cache=key_object_in_memory_cache or InMemoryCache(), + redis_cache=redis_cache, + default_in_memory_ttl=default_in_memory_ttl, + default_redis_ttl=default_redis_ttl, + ) + + def in_memory_cache_for(self, key: str) -> InMemoryCache: + return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache + + def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None) -> None: + super().update_cache_ttl(default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl) + self.key_object_cache.update_cache_ttl( + default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl + ) + + def attach_redis_cache( + self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None + ) -> None: + super().attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) + self.key_object_cache.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) + @overload def get_cache( self, @@ -71,7 +122,11 @@ class UserApiKeyCache(DualCache): ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - cached: Final = super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) + cached: Final = ( + self.key_object_cache.get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) + if is_user_key_cache_key(key) + else super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) + ) if model_type is None: return cached if cached is None: @@ -117,8 +172,14 @@ class UserApiKeyCache(DualCache): ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - cached: Final = await super().async_get_cache( - key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + cached: Final = ( + await self.key_object_cache.async_get_cache( + key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + ) + if is_user_key_cache_key(key) + else await super().async_get_cache( + key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + ) ) if model_type is None: return cached @@ -137,20 +198,49 @@ class UserApiKeyCache(DualCache): def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + if key is not None and is_user_key_cache_key(key): + return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + if key is not None and is_user_key_cache_key(key): + return await self.key_object_cache.async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs: object) -> None: + def delete_cache(self, key: str) -> None: + if is_user_key_cache_key(key): + self.key_object_cache.delete_cache(key) + return + super().delete_cache(key) + + async def async_delete_cache(self, key: str) -> None: + if is_user_key_cache_key(key): + await self.key_object_cache.async_delete_cache(key) + return + await super().async_delete_cache(key) + + def flush_cache(self) -> None: + super().flush_cache() + self.key_object_cache.in_memory_cache.flush_cache() + + async def async_set_cache_pipeline( + self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs: object + ) -> None: """ Batch writes with the same Codec boundary as ``async_set_cache`` without ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. """ - normalized: Final = [(key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list] - return await super().async_set_cache_pipeline(cache_list=normalized, local_only=local_only, **kwargs) + normalized: Final = tuple((key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list) + key_object_entries: Final = tuple(entry for entry in normalized if is_user_key_cache_key(entry[0])) + other_entries: Final = tuple(entry for entry in normalized if not is_user_key_cache_key(entry[0])) + if key_object_entries: + await self.key_object_cache.async_set_cache_pipeline( + cache_list=key_object_entries, local_only=local_only, **kwargs + ) + if other_entries: + await super().async_set_cache_pipeline(cache_list=other_entries, local_only=local_only, **kwargs) #: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row, diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index 9c637a62dc1..3a61da164d0 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -37,7 +37,9 @@ CACHE_TTL_1H_SECONDS: Final = 3600 AUTOROUTER_BENCHMARKS_SQL: Final = """ WITH windowed AS ( SELECT * FROM "LiteLLM_AutoRouterSession" - WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp + WHERE last_turn_at >= $1::timestamp + AND first_turn_at < $2::timestamp + AND ($3::text IS NULL OR api_key = $3::text) ), tier_maps AS ( SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns @@ -73,6 +75,8 @@ SELECT COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens, COALESCE(SUM(spend), 0)::float8 AS spend, COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, + COALESCE(SUM(classifier_cost), 0)::float8 AS classifier_cost, + COALESCE(SUM(classifier_cost_recorded_turns), 0)::int AS classifier_cost_recorded_turns, COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds FROM windowed GROUP BY router_name, router_type @@ -93,11 +97,13 @@ class AutoRouterTurnTransaction: total_tokens: int spend: float saved_spend: float + classifier_cost: float covered: bool cache_hit: bool cache_ttl_seconds: int | None cache_touched: bool tier: str | None = None + baseline_model: str | None = None class TurnCacheFacts(NamedTuple): @@ -163,7 +169,7 @@ def _write_ttl_seconds(usage_object: Mapping[str, object] | None) -> int | None: SESSION_ID_MAX_CHARS: Final = 256 -def _bounded_session_id(session_id: str) -> str: +def bounded_session_id(session_id: str) -> str: """The session id as stored, bounded so a caller-chosen identifier cannot exceed Postgres's B-tree index entry limit through the composite primary key. Oversized ids map to a stable digest, so their turns still aggregate into one session.""" @@ -189,6 +195,9 @@ def build_autorouter_turn_transaction( classifier_cost folded into this turn's spend: the excluded classifier row is how it was billed, the decision is how it is attributed. Cache facts are derived from the payload's own usage record through the savings owner, never handed in beside it. + The baseline the turn's saved_spend was priced against travels with the turn, so the + row can name the counterfactual for the money it holds even after the router is + reconfigured or removed. """ if payload.get("status") != "success": return None @@ -211,18 +220,21 @@ def build_autorouter_turn_transaction( usage_object_raw: Final = metadata.get("usage_object") cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None) tier_raw: Final = routing_decision.get("tier") + baseline_raw: Final = routing_decision.get("savings_baseline_model") classifier_cost: Final = classifier_cost_from_decision(routing_decision) return AutoRouterTurnTransaction( api_key=api_key, - session_id=_bounded_session_id(session_id), + session_id=bounded_session_id(session_id), router_name=router_name, router_type=str(routing_decision.get("router_type") or "unknown"), tier=tier_raw if isinstance(tier_raw, str) and tier_raw else None, + baseline_model=baseline_raw if isinstance(baseline_raw, str) and baseline_raw else None, model=model, turn_at=turn_at, total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0), spend=float(payload.get("spend") or 0.0) + (classifier_cost or 0.0), saved_spend=saved_spend, + classifier_cost=classifier_cost or 0.0, covered=cache.covered, cache_hit=cache.read_tokens > 0, cache_ttl_seconds=cache.write_ttl_seconds, @@ -247,6 +259,10 @@ _CACHE_TTL: Final = _p("cache_ttl_seconds") _TOUCHED: Final = _p("cache_touched") _TIER: Final = f"{_p('tier')}::text" _TIER_DELTA: Final = f"(CASE WHEN {_TIER} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_TIER}, 1) END)" +_BASELINE: Final = f"{_p('baseline_model')}::text" +_BASELINE_DELTA: Final = ( + f"(CASE WHEN {_BASELINE} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_BASELINE}, 1) END)" +) _IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at" _SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}" @@ -264,7 +280,8 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t ( last_model, models, turns, unordered_turns, covered_turns, cache_hits, same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, return_turns, return_hits, return_expired_misses, return_within_ttl_misses, - ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns + ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns, + baseline_models ) VALUES ( {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, @@ -275,13 +292,15 @@ VALUES ( (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END), (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END), {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8, - {_TIER_DELTA} + {_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA} ) ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET turns = t.turns + 1, total_tokens = t.total_tokens + EXCLUDED.total_tokens, spend = t.spend + EXCLUDED.spend, saved_spend = t.saved_spend + EXCLUDED.saved_spend, + classifier_cost = t.classifier_cost + EXCLUDED.classifier_cost, + classifier_cost_recorded_turns = t.classifier_cost_recorded_turns + 1, covered_turns = t.covered_turns + EXCLUDED.covered_turns, cache_hits = t.cache_hits + EXCLUDED.cache_hits, ttl_5m_turns = t.ttl_5m_turns + EXCLUDED.ttl_5m_turns, @@ -309,6 +328,9 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET tier_turns = (CASE WHEN {_TIER} IS NOT NULL AND t.router_type = {_p("router_type")} THEN t.tier_turns || jsonb_build_object({_TIER}, COALESCE((t.tier_turns ->> {_TIER})::int, 0) + 1) ELSE t.tier_turns END), + baseline_models = (CASE WHEN {_BASELINE} IS NOT NULL + THEN t.baseline_models || jsonb_build_object({_BASELINE}, COALESCE((t.baseline_models ->> {_BASELINE})::int, 0) + 1) + ELSE t.baseline_models END), first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at), last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at) """ diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index b7a07d4eeea..6e2a06c96e1 100644 --- a/litellm/proxy/db/check_migration.py +++ b/litellm/proxy/db/check_migration.py @@ -46,17 +46,32 @@ def extract_sql_commands(diff_output: str) -> list[str]: def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: """Checks for differences between current database and Prisma schema. + + Never raises: a diff that cannot be produced, because the runner is missing, + because the command failed, or because it outlived its budget, is reported as + "no diff" so boot continues. + Returns: A tuple containing: - A boolean indicating if differences were found (True) or not (False). - - A string with the diff output or error message. - Raises: - subprocess.CalledProcessError: If the Prisma command fails. - Exception: For any other errors during execution. + - The SQL commands that would close the diff, empty when there is none. """ - verbose_logger.debug("Checking for Prisma schema diff...") try: - result: Final = subprocess.run( + from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_COMMAND_TIMEOUT_ENV_VAR, + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}" + ) + return False, [] + + verbose_logger.debug("Checking for Prisma schema diff...") + timeout: Final = prisma_command_timeout() + try: + result: Final = run_prisma( [ "prisma", "migrate", @@ -67,12 +82,10 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: "./schema.prisma", "--script", ], - capture_output=True, - text=True, - check=True, + timeout=timeout, + env=os.environ.copy(), ) - # return True, "Migration diff generated successfully." sql_commands: Final = extract_sql_commands(result.stdout) if sql_commands: @@ -83,6 +96,12 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: return True, sql_commands else: return False, [] + except subprocess.TimeoutExpired: + print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + f"Timed out after {timeout}s generating the migration diff. " + f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer." + ) + return False, [] except subprocess.CalledProcessError as e: error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}" print(error_message) # noqa: T201 diff --git a/litellm/proxy/db/db_lookup_gate.py b/litellm/proxy/db/db_lookup_gate.py new file mode 100644 index 00000000000..2fd427687bd --- /dev/null +++ b/litellm/proxy/db/db_lookup_gate.py @@ -0,0 +1,23 @@ +import asyncio +from typing import Final + +from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY + + +class LoopBoundSemaphore: + __slots__ = ("_loop", "_semaphore", "_value") + + def __init__(self, value: int) -> None: + self._value: Final = value + self._loop: asyncio.AbstractEventLoop | None = None + self._semaphore: asyncio.Semaphore | None = None + + def current(self) -> asyncio.Semaphore: + loop: Final = asyncio.get_running_loop() + if self._semaphore is None or self._loop is not loop: + self._semaphore = asyncio.Semaphore(self._value) + self._loop = loop + return self._semaphore + + +db_lookup_gate: Final = LoopBoundSemaphore(PROXY_DB_LOOKUP_MAX_CONCURRENCY) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e6880d521f1..eaa03c5d7f7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -14,6 +14,7 @@ import time import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload import litellm @@ -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, @@ -380,6 +502,7 @@ class DBSpendUpdateWriter: llm_router=get_llm_router, cost_breakdown=metadata.get("cost_breakdown"), recorded_autorouter_savings=metadata.get("autorouter_savings"), + billed_at=payload.get("endTime"), ) transaction: Final = build_autorouter_turn_transaction( payload=payload, @@ -940,7 +1063,7 @@ class DBSpendUpdateWriter: await enqueue_spend_logs(prisma_client, (payload,)) if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: - request_spend_log_flush() + request_spend_log_flush(prisma_client) else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") @@ -2066,6 +2189,7 @@ class DBSpendUpdateWriter: usage_object=usage_obj, cost_breakdown=_metadata.get("cost_breakdown"), recorded_autorouter_savings=_metadata.get("autorouter_savings"), + billed_at=payload.get("endTime"), ) daily_transaction: Final = BaseDailySpendTransaction( diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 4be1331e955..bc67617e444 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -1,10 +1,11 @@ import asyncio import json +import logging from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_cache import RedisCache, log_redis_failure from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.types.services import ServiceTypes @@ -109,7 +110,7 @@ end ) return False except Exception as e: - verbose_proxy_logger.error("Error acquiring Redis lock for %s: %s", cronjob_id, e) + log_redis_failure(verbose_proxy_logger, logging.ERROR, f"Error acquiring Redis lock for {cronjob_id}", e) return False async def release_lock( @@ -151,7 +152,7 @@ end cronjob_id, ) except Exception as e: - verbose_proxy_logger.error("Error releasing Redis lock for %s: %s", cronjob_id, e) + log_redis_failure(verbose_proxy_logger, logging.ERROR, f"Error releasing Redis lock for {cronjob_id}", e) async def _compare_and_delete_lock(self, lock_key: str) -> int: """ diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py index 221c142d9d3..b756bfeb6f6 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py +++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py @@ -16,7 +16,12 @@ keeps the batched-DELETE path, so existing deployments are untouched. import re from collections.abc import Callable from datetime import date, datetime, timedelta, timezone -from typing import TYPE_CHECKING, Final, TypeAlias +from typing import ( + TYPE_CHECKING, + Final, + TypeAlias, + cast, # noqa: TID251 # db.tx is reached through untyped __getattr__ delegation +) from litellm._logging import verbose_proxy_logger from litellm.constants import ( @@ -25,6 +30,8 @@ from litellm.constants import ( ) if TYPE_CHECKING: + from prisma.client import TransactionManager + from litellm.proxy.utils import PrismaClient SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs" @@ -116,6 +123,21 @@ def select_partitions_to_drop(partitions: list[tuple[str, datetime | None]], cut return [name for name, upper in partitions if upper is not None and upper <= cutoff] +_TX_COMMIT_SLACK: Final = timedelta(seconds=5) + + +def _bounded_tx(prisma_client: "PrismaClient", timeout_ms: int) -> "TransactionManager": + """ + Open an interactive transaction that outlives the statement bound it + carries. prisma's default 5s transaction timeout would close it mid + lock-wait, after which the engine answers the next call with a 422. + """ + return cast( # cast-ok: PrismaWrapper delegates tx via __getattr__ (untyped) + "TransactionManager", + prisma_client.db.tx(timeout=timedelta(milliseconds=timeout_ms) + _TX_COMMIT_SLACK), + ) + + class SpendLogsPartitionManager: def __init__( self, @@ -137,7 +159,7 @@ class SpendLogsPartitionManager: if budget_ms is None: return False try: - async with prisma_client.db.tx() as tx: + async with _bounded_tx(prisma_client, budget_ms) as tx: await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}") rows: Final = await tx.query_raw( """ @@ -172,7 +194,7 @@ class SpendLogsPartitionManager: wait for the lock and statement_timeout bounds the work itself, so a partition this run cannot get is simply left for the next one. """ - async with prisma_client.db.tx() as tx: + async with _bounded_tx(prisma_client, timeout_ms) as tx: await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}") await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}") await tx.execute_raw(statement) @@ -209,7 +231,7 @@ class SpendLogsPartitionManager: async def _list_partitions( self, prisma_client: "PrismaClient", timeout_ms: int ) -> list[tuple[str, datetime | None]]: - async with prisma_client.db.tx() as tx: + async with _bounded_tx(prisma_client, timeout_ms) as tx: await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}") rows: Final = await tx.query_raw( """ diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index f28e505246a..e93bc96da69 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -34,16 +34,25 @@ writer's connection params (pool size, timeouts, pgbouncer mode) for the ones the reader URL does not pin itself. """ +import _ssl +import hashlib import os +import socket +import ssl +import struct +import sys +import tempfile import urllib.parse -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence from functools import partial +from pathlib import Path from types import MappingProxyType -from typing import Annotated, Final, cast +from typing import Annotated, Final, Protocol, TypeAlias, cast from pydantic import AliasChoices, BeforeValidator, Field from pydantic_settings import BaseSettings, SettingsConfigDict +from litellm.proxy.db.pgbouncer import database_url_is_pooled from litellm.proxy.db.token_auth import ( AZURE_POSTGRESQL_AUTH_ENV_VAR, DEFAULT_POSTGRES_PORT, @@ -64,6 +73,7 @@ DISABLE_PREPARED_STATEMENTS_ENV_VAR: Final = "DATABASE_DISABLE_PREPARED_STATEMEN DisablePreparedStatementsFlag = Annotated[ bool, BeforeValidator(partial(token_auth_flag_enabled, env_var=DISABLE_PREPARED_STATEMENTS_ENV_VAR)) ] +MAX_IDLE_CONNECTION_LIFETIME_ENV_VAR: Final = "DATABASE_MAX_IDLE_CONNECTION_LIFETIME" # schema.prisma pins `provider = "postgresql"`, so these are the only schemes # Prisma can actually connect with. @@ -125,21 +135,100 @@ def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"}) +PEM_CERT_HEADER: Final = b"-----BEGIN CERTIFICATE-----" +PG_SSL_REQUEST: Final = struct.pack("!ii", 8, 80877103) +TLS_PROBE_TIMEOUT_SECONDS: Final = 10.0 + +RootCertResolver: TypeAlias = Callable[[str, str, int], str] # mutable-ok: Callable parameter syntax -def translate_libpq_ssl_params(url: str) -> str: +class _VerifiedChainSource(Protocol): + def get_verified_chain(self) -> Sequence[_ssl.Certificate] | None: ... + + +def _verified_chain_der(tls: ssl.SSLSocket) -> tuple[bytes, ...]: + if sys.version_info >= (3, 13): + return tuple(tls.get_verified_chain()) + legacy: Final = cast( # cast-ok: the stub omits _sslobj, the C object has get_verified_chain since 3.10 + "_VerifiedChainSource | None", + tls._sslobj, # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # public API only from 3.13 + ) + chain: Final = () if legacy is None else legacy.get_verified_chain() or () + return tuple(cert.public_bytes(_ssl.ENCODING_DER) for cert in chain) + + +def _server_trust_anchor(cafile: str, host: str, port: int) -> bytes | None: + try: + context: Final = ssl.create_default_context(cafile=cafile) + with socket.create_connection((host, port), timeout=TLS_PROBE_TIMEOUT_SECONDS) as raw: + raw.sendall(PG_SSL_REQUEST) + if raw.recv(1) != b"S": + return None + with context.wrap_socket(raw, server_hostname=host) as tls: + chain: Final = _verified_chain_der(tls) + except (OSError, ValueError): + return None + return chain[-1] if chain else None + + +def pin_bundle_root(cert_path: str, host: str, port: int) -> str: + """Reduce a multi-root CA bundle to the one root that verifies ``host``. + + Prisma's ``sslcert`` loads a single PEM certificate (native-tls + ``Certificate::from_pem``), so pointing it at a bundle such as the AWS RDS + global bundle trusts only the first of its 108 regional roots and the + handshake fails with "unable to get local issuer certificate" for every + other region. A single-certificate file is returned as is. For a bundle, + one verifying handshake (chain and hostname, whole bundle as trust store) + identifies the trust anchor the server actually chains to, which is + written to a single-certificate file for Prisma. If the probe fails the + bundle path is returned unchanged, so Prisma fails closed exactly as + before rather than trusting anything the bundle would not. + """ + try: + if Path(cert_path).read_bytes().count(PEM_CERT_HEADER) < 2: + return cert_path + except OSError: + return cert_path + root: Final = _server_trust_anchor(cert_path, host, port) + if root is None: + return cert_path + pinned: Final = Path(tempfile.gettempdir()) / f"litellm-sslcert-{hashlib.sha256(root).hexdigest()[:16]}.pem" + return str(pinned) if _replace_file(pinned, ssl.DER_cert_to_PEM_cert(root)) else cert_path + + +def _replace_file(target: Path, content: str) -> bool: + """Write ``content`` to a private temp file and rename it over ``target``, so + readers never see a partial file and a symlink planted at ``target`` is + replaced rather than followed.""" + try: + fd, staged = tempfile.mkstemp(dir=target.parent, prefix=f"{target.name}.") + except OSError: + return False + try: + with os.fdopen(fd, "w") as handle: + handle.write(content) + os.replace(staged, target) + except OSError: + Path(staged).unlink(missing_ok=True) + return False + return True + + +def translate_libpq_ssl_params(url: str, resolve_root_cert: RootCertResolver = pin_bundle_root) -> str: """Rewrite libpq's certificate-verification params into Prisma's dialect. Prisma's engine only knows ``sslmode=disable|prefer|require``, ``sslcert`` - (the CA bundle) and ``sslaccept=strict``. It silently discards + (a single CA certificate) and ``sslaccept=strict``. It silently discards ``sslrootcert`` and downgrades ``sslmode=verify-ca`` / ``verify-full`` to ``prefer``, so a URL copied from libpq / RDS docs connects over TLS with no certificate check at all. ``verify-ca`` and ``verify-full`` both become ``require`` (Prisma has no CA-only mode), ``sslrootcert`` becomes - ``sslcert``, and either one turns on ``sslaccept=strict`` (chain and - hostname), matching libpq where a root cert makes ``require`` verify. - Prisma params the operator pinned themselves win; anything else is left - untouched. + ``sslcert`` (run through ``resolve_root_cert``, which pins a multi-root + bundle down to the server's root), and either one turns on + ``sslaccept=strict`` (chain and hostname), matching libpq where a root + cert makes ``require`` verify. Prisma params the operator pinned + themselves win; anything else is left untouched. """ parsed: Final = urllib.parse.urlsplit(url) pairs: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True)) @@ -153,7 +242,9 @@ def translate_libpq_ssl_params(url: str) -> str: if key != "sslrootcert" ) root_cert: Final = tuple( - ("sslcert", value) for key, value in pairs if key == "sslrootcert" and "sslcert" not in keys + ("sslcert", resolve_root_cert(value, parsed.hostname or "", parsed.port or int(DEFAULT_POSTGRES_PORT))) + for key, value in pairs + if key == "sslrootcert" and "sslcert" not in keys ) strict: Final = () if "sslaccept" in keys else (("sslaccept", "strict"),) query: Final = urllib.parse.urlencode(translated + root_cert + strict) @@ -217,6 +308,9 @@ class DatabaseURLSettings(BaseSettings): disable_prepared_statements: DisablePreparedStatementsFlag = Field( default=False, validation_alias=DISABLE_PREPARED_STATEMENTS_ENV_VAR ) + max_idle_connection_lifetime: int | None = Field( + default=None, validation_alias=MAX_IDLE_CONNECTION_LIFETIME_ENV_VAR + ) # Writer database_url: str | None = Field(default=None, validation_alias="DATABASE_URL") @@ -265,8 +359,12 @@ class DatabaseURLSettings(BaseSettings): Raises ``RuntimeError`` (naming the offending vars) when token auth is enabled but a required field is missing — the proxy cannot recover from this and a clear startup error beats a Prisma connect failure. + A ``DATABASE_URL`` the supervisor pointed at the in-container PgBouncer + is kept even under token auth: the pooler renews the token upstream. """ auth: Final = self.token_auth() + if auth is not None and database_url_is_pooled(): + return None if auth is not None: missing: Final = tuple( env @@ -453,6 +551,12 @@ class DatabaseURLSettings(BaseSettings): if url: os.environ[env_var] = add_missing_query_params(url, MappingProxyType({"pgbouncer": "true"})) + lifetime_params: Final = idle_lifetime_params(self.max_idle_connection_lifetime) + for env_var in ("DATABASE_URL", "DIRECT_URL"): + url = os.environ.get(env_var) + if url: + os.environ[env_var] = add_missing_query_params(url, lifetime_params) + # The reader inherits the writer's connection params (pool size, timeouts, # pgbouncer mode). Without this the reader pool ignores the configured cap # and falls back to Prisma's `num_physical_cpus * 2 + 1` default. diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index f469587ab8e..2cee5128c66 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -221,6 +221,18 @@ class PrismaDBExceptionHandler: or "write conflict or a deadlock" in error_message ) + @staticmethod + def is_read_only_transaction_error(e: Exception) -> bool: + """True iff ``e`` is Postgres SQLSTATE 25006 surfaced through prisma: the + pooled session answers reads but rejects writes, so the connection is + poisoned until the client is recreated.""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + error_message: Final = str(e).lower() + return '"25006"' in error_message or "read-only transaction" in error_message + @staticmethod def is_prisma_engine_internal_error(e: Exception) -> bool: """True iff ``e`` is a non-``PrismaError`` exception raised from inside diff --git a/litellm/proxy/db/gateway_request_tracking.py b/litellm/proxy/db/gateway_request_tracking.py index bebd74e877c..c9ace68db33 100644 --- a/litellm/proxy/db/gateway_request_tracking.py +++ b/litellm/proxy/db/gateway_request_tracking.py @@ -10,13 +10,28 @@ strings rather than passing the raw path through. Nothing a caller sends can add a key, so the fold and the table it commits to are bounded by (days x routes) however much traffic arrives, and the response path carries no unbounded queue that would block once full. + +A flush commits its whole snapshot as one multi-row ``INSERT ... ON CONFLICT DO +UPDATE`` rather than one upsert per key, so a worker costs the primary one +statement per interval however many routes it served. With +``use_redis_transaction_buffer`` on, workers instead push their snapshot to a +Redis list and one lock-holding pod folds every entry and writes the table, so +the deployment as a whole costs the primary one statement per interval. """ -from dataclasses import asdict +import json +from collections.abc import AsyncIterator, Iterable from datetime import datetime, timezone -from typing import TYPE_CHECKING, Final +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.caching import RedisCache +from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory from litellm.types.proxy.gateway_requests import ( GatewayRequestCounts, @@ -28,6 +43,15 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient _EMPTY: Final = GatewayRequestCounts(successful_requests=0, failed_requests=0) +_TABLE: Final = '"LiteLLM_DailyGatewayRequests"' +_COLUMNS_PER_ROW: Final = 5 +_UTC_NOW: Final = "(NOW() AT TIME ZONE 'UTC')" +GATEWAY_REQUESTS_JOB_NAME: Final = "update_gateway_requests_job" + +_BufferedRows: TypeAlias = tuple[tuple[str, str, str, int, int], ...] +_BUFFERED_ROWS: Final = TypeAdapter(_BufferedRows) +_BUFFERED_ENTRIES: Final = TypeAdapter(tuple[str | bytes, ...]) +_NO_COUNTS: Final[GatewayRequestSnapshot] = MappingProxyType({}) def _utc_date() -> str: @@ -59,20 +83,54 @@ class GatewayRequestAccumulator: route) however long the database is unreachable. This buys at-least-once, not exactly-once, and the cost is worth stating. - The batch commits inside its context manager's ``__aexit__``, so a failure - raised after the transaction committed (a connection dropped while reading - the acknowledgement) restores counts that are already persisted, and the - next flush increments them a second time. Exactly-once would need a dedup - key the upserts could ignore on replay. For a traffic-volume metric a rare + The statement commits on the server before its acknowledgement is read, so + a failure raised after the commit (a connection dropped while reading the + acknowledgement) restores counts that are already persisted, and the next + flush increments them a second time. Exactly-once would need a dedup key + the upsert could ignore on replay. For a traffic-volume metric a rare overcount on a dropped acknowledgement beats losing a whole interval to every database blip, so the trade is deliberate. """ - for key, counts in snapshot.items(): - existing = self._counts.get(key, _EMPTY) - self._counts[key] = GatewayRequestCounts( - successful_requests=existing.successful_requests + counts.successful_requests, - failed_requests=existing.failed_requests + counts.failed_requests, - ) + self._counts = dict(fold_counts(chain(self._counts.items(), snapshot.items()))) # mutable-ok: fold replaced + + +def fold_counts(items: Iterable[tuple[GatewayRequestKey, GatewayRequestCounts]]) -> GatewayRequestSnapshot: + """Sum counts key-wise; the result stays bounded by (date x category x route).""" + folded: Final[dict[GatewayRequestKey, GatewayRequestCounts]] = {} # mutable-ok: local fold returned once + for key, counts in items: + existing = folded.get(key, _EMPTY) + folded[key] = GatewayRequestCounts( + successful_requests=existing.successful_requests + counts.successful_requests, + failed_requests=existing.failed_requests + counts.failed_requests, + ) + return folded + + +def build_gateway_requests_upsert(snapshot: GatewayRequestSnapshot) -> tuple[str, tuple[str | int, ...]]: + """ + One ``INSERT ... ON CONFLICT DO UPDATE`` that increments every (date, category, + route) in the snapshot. Rows are ordered by the conflict key so concurrent + writers lock rows in the same order and cannot deadlock. + """ + ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route)) + rows: Final = ", ".join( + f"(${base + 1}::text, ${base + 2}::text, ${base + 3}::text, ${base + 4}::bigint, ${base + 5}::bigint, {_UTC_NOW})" + for base in range(0, len(ordered) * _COLUMNS_PER_ROW, _COLUMNS_PER_ROW) + ) + sql: Final = ( + f'INSERT INTO {_TABLE} ("date", "category", "route", "successful_requests", "failed_requests", "updated_at")\n' + f"VALUES {rows}\n" + 'ON CONFLICT ("date", "category", "route") DO UPDATE SET\n' + f' "successful_requests" = {_TABLE}."successful_requests" + EXCLUDED."successful_requests",\n' + f' "failed_requests" = {_TABLE}."failed_requests" + EXCLUDED."failed_requests",\n' + f' "updated_at" = {_UTC_NOW}' + ) + params: Final[tuple[str | int, ...]] = tuple( + value + for key, counts in ordered + for value in (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests) + ) + return sql, params async def commit_gateway_requests_to_db( @@ -80,50 +138,130 @@ async def commit_gateway_requests_to_db( prisma_client: "PrismaClient", snapshot: GatewayRequestSnapshot, ) -> None: - """Upsert one incrementing row per (date, category, route).""" + """Increment every (date, category, route) in the snapshot with a single statement.""" if not snapshot: return - ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route)) + sql, params = build_gateway_requests_upsert(snapshot) + await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client - # pyright: ignore[reportAny] on both lines -- prisma's generated client is untyped, - # so .db and every table action off it resolve to Any at this boundary. The dict - # literals below are the shape prisma's generated inputs require. - async with prisma_client.db.batch_() as batcher: # pyright: ignore[reportAny] # untyped prisma client - for key, counts in ordered: - columns = asdict(key) - batcher.litellm_dailygatewayrequests.upsert( # pyright: ignore[reportAny] # untyped prisma client - where={"date_category_route": columns}, # mutable-ok: prisma input is dict-shaped - data={ # mutable-ok: prisma input is dict-shaped - "create": { # mutable-ok: prisma input is dict-shaped - **columns, - "successful_requests": counts.successful_requests, - "failed_requests": counts.failed_requests, - }, - "update": { # mutable-ok: prisma input is dict-shaped - "successful_requests": {"increment": counts.successful_requests}, # mutable-ok: as above - "failed_requests": {"increment": counts.failed_requests}, # mutable-ok: as above - }, - }, + verbose_proxy_logger.debug( + "Gateway request tracking - committed %d aggregated rows in one statement", len(snapshot) + ) + + +class GatewayRequestRedisBuffer: + """ + Folds every worker's snapshot through one Redis list so a single pod per + interval writes the table, mirroring the spend writer's transaction buffer. + + Each entry is one worker's snapshot as JSON rows; the lock holder pops them, + sums them, and commits one statement. A commit failure pushes the summed + rows back so the next holder retries, keeping the at-least-once guarantee. + If that push fails too, the rows go back to the holder's own accumulator so + they ride along with its next flush instead of vanishing with the pop. + """ + + def __init__(self, *, redis_cache: RedisCache, pod_lock_manager: PodLockManager) -> None: + self._redis_cache: Final = redis_cache + self._pod_lock_manager: Final = pod_lock_manager + + async def push(self, snapshot: GatewayRequestSnapshot) -> None: + if not snapshot: + return + rows: Final[_BufferedRows] = tuple( + (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests) + for key, counts in snapshot.items() + ) + await self._redis_cache.async_rpush(key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, values=(json.dumps(rows),)) + + async def _pop_batch(self) -> tuple[str | bytes, ...]: + popped: Final[object] = await self._redis_cache.async_lpop( # pyright: ignore[reportAny] # redis returns Any + key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT + ) + if not popped: + return () + return _BUFFERED_ENTRIES.validate_python(popped if isinstance(popped, list) else (popped,)) + + async def _pop_all(self) -> AsyncIterator[str | bytes]: + while True: + batch = await self._pop_batch() + for entry in batch: + yield entry + if len(batch) < MAX_REDIS_BUFFER_DEQUEUE_COUNT: + return + + async def pop(self) -> GatewayRequestSnapshot: + entries: Final = tuple([entry async for entry in self._pop_all()]) + return fold_counts( + ( + GatewayRequestKey(date=date, category=category, route=route), + GatewayRequestCounts(successful_requests=succeeded, failed_requests=failed), ) + for entry in entries + for date, category, route, succeeded, failed in _BUFFERED_ROWS.validate_json(entry) + ) - verbose_proxy_logger.debug("Gateway request tracking - committed %d aggregated rows", len(ordered)) + async def commit_if_leader(self, prisma_client: "PrismaClient") -> GatewayRequestSnapshot: + """ + Drain the list and write it as one statement, but only on the pod holding the job lock. + + The lock is a lease, never released: the holder re-enters it on every flush and + keeps committing alone until the TTL lapses, so the primary sees one statement + per flush interval deployment-wide instead of one per worker. + + Returns the popped rows that could be neither committed nor re-queued, for the + caller to keep in memory. Empty on success. + """ + if not await self._pod_lock_manager.acquire_lock(cronjob_id=GATEWAY_REQUESTS_JOB_NAME): + return _NO_COUNTS + buffered: Final = await self.pop() + try: + await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=buffered) + except Exception: # noqa: BLE001 -- a failed commit must not stop the scheduler + verbose_proxy_logger.warning( + "Gateway request tracking - failed to commit %d buffered rows, re-queuing to Redis for the next flush", + len(buffered), + exc_info=True, + ) + return await self._requeue(buffered) + return _NO_COUNTS + + async def _requeue(self, snapshot: GatewayRequestSnapshot) -> GatewayRequestSnapshot: + try: + await self.push(snapshot) + except Exception: # noqa: BLE001 -- the rows go back to the caller's accumulator instead + verbose_proxy_logger.warning( + "Gateway request tracking - Redis re-queue failed, keeping %d rows in memory for the next flush", + len(snapshot), + exc_info=True, + ) + return snapshot + return _NO_COUNTS async def flush_gateway_requests( prisma_client: "PrismaClient", accumulator: GatewayRequestAccumulator, + redis_buffer: GatewayRequestRedisBuffer | None = None, ) -> None: """ Scheduler entrypoint. Never raises: a metering failure must not kill the job. + With ``redis_buffer`` the snapshot goes to Redis and only the lease holder + writes to Postgres. Shutdown passes no buffer so a departing worker writes its + own counts directly instead of parking them behind a lease it may not hold. + ``CancelledError`` is deliberately not caught, so a flush cancelled during shutdown drops its snapshot rather than restoring counts onto an accumulator the process is about to discard. """ snapshot: Final = accumulator.drain() try: - await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot) + if redis_buffer is None: + await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot) + else: + await redis_buffer.push(snapshot) except Exception: # noqa: BLE001 -- a failed flush must not stop the scheduler accumulator.restore(snapshot) verbose_proxy_logger.warning( @@ -131,3 +269,13 @@ async def flush_gateway_requests( len(snapshot), exc_info=True, ) + return + if redis_buffer is None: + return + try: + accumulator.restore(await redis_buffer.commit_if_leader(prisma_client)) + except Exception: # noqa: BLE001 -- entries still in Redis are drained by the next flush + verbose_proxy_logger.warning( + "Gateway request tracking - leader drain failed, buffered rows stay in Redis for the next flush", + exc_info=True, + ) diff --git a/litellm/proxy/db/health_check_latest.py b/litellm/proxy/db/health_check_latest.py new file mode 100644 index 00000000000..35bc838379c --- /dev/null +++ b/litellm/proxy/db/health_check_latest.py @@ -0,0 +1,96 @@ +""" +Latest health-check row per model, deduplicated by Postgres. + +prisma-client-py's ``find_many(distinct=...)`` dedups client-side: the emitted +SQL carries no DISTINCT, so the whole append-only history table streams to the +worker on every call. ``SELECT DISTINCT ON`` keeps the transfer at one row per +(model_id, model_name) and is served by the matching descending index. +""" + +from __future__ import annotations + +import json +from collections.abc import Sequence +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, field_validator + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +LATEST_HEALTH_CHECKS_SQL: Final = """ +SELECT DISTINCT ON ("model_id", "model_name") + "health_check_id", "model_name", "model_id", "status", + "healthy_count", "unhealthy_count", "error_message", + "response_time_ms", "details", "checked_by", + "checked_at", "created_at", "updated_at" +FROM "LiteLLM_HealthCheckTable" +ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC +""" + +LATEST_HEALTH_CHECKS_FOR_MODELS_SQL: Final = """ +SELECT DISTINCT ON ("model_id", "model_name") + "health_check_id", "model_name", "model_id", "status", + "healthy_count", "unhealthy_count", "error_message", + "response_time_ms", "details", "checked_by", + "checked_at", "created_at", "updated_at" +FROM "LiteLLM_HealthCheckTable" +WHERE "model_name" = ANY($1) +ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC +""" + + +class LatestHealthCheckRow(BaseModel): + model_config = ConfigDict(frozen=True, protected_namespaces=()) + + health_check_id: str + model_name: str + model_id: str | None = None + status: str + healthy_count: int = 0 + unhealthy_count: int = 0 + error_message: str | None = None + response_time_ms: float | None = None + details: JsonValue | None = None + checked_by: str | None = None + checked_at: datetime + created_at: datetime + updated_at: datetime + + @field_validator("details", mode="before") + @classmethod + def _decode_json_text(cls, value: object) -> object: + return json.loads(value) if isinstance(value, str) else value + + @field_validator("checked_at", "created_at", "updated_at") + @classmethod + def _assume_utc(cls, value: datetime) -> datetime: + return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value + + +_ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...]) + + +async def fetch_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]: + try: + rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL) + return _ROWS_ADAPTER.validate_python(rows) + except Exception as query_err: # noqa: BLE001 # health decorates other reads; a driver error must not fail them + verbose_proxy_logger.error("Error getting all latest health checks: %s", query_err) + return () + + +async def fetch_latest_health_checks_for_models( + prisma_client: PrismaClient, model_names: Sequence[str] +) -> tuple[LatestHealthCheckRow, ...]: + if not model_names: + return () + try: + rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, list(model_names)) + return _ROWS_ADAPTER.validate_python(rows) + except Exception as query_err: # noqa: BLE001 # a paged model list must not fail on its health decoration + verbose_proxy_logger.error("Error getting latest health checks for models: %s", query_err) + return () diff --git a/litellm/proxy/db/pgbouncer.py b/litellm/proxy/db/pgbouncer.py new file mode 100644 index 00000000000..c9fbabd3585 --- /dev/null +++ b/litellm/proxy/db/pgbouncer.py @@ -0,0 +1,743 @@ +"""In-container PgBouncer shared by every proxy worker. + +Each uvicorn worker owns a Prisma query engine with its own pool of +``connection_limit`` server connections, so the connections a pod holds open +against Postgres scale as ``workers * connection_limit`` and a database with a +fixed connection ceiling runs out of room as pods and workers are added. + +When ``LITELLM_PGBOUNCER_ENABLED`` is set, the supervisor process starts one +PgBouncer next to the workers (no extra network hop: it listens on loopback +inside the pod) in transaction pooling mode, points ``DATABASE_URL`` at it +with ``pgbouncer=true`` so Prisma stops using server-side prepared statements, +and keeps it running for the life of the proxy. Every worker's pool then +becomes cheap client connections to PgBouncer while the upstream connection +count is capped at ``LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS`` per pod, no matter +how many workers run. + +Migrations and the schema diff run in the supervisor before the pooler is +started, so they always go straight to Postgres. ``DATABASE_URL_READ_REPLICA`` +is left untouched. + +The workers never hold the upstream credential: they log in to PgBouncer as +``litellm_pgbouncer`` with a random password made at startup, and PgBouncer +takes the database user's password from its auth file. Under +``IAM_TOKEN_DB_AUTH`` or ``AZURE_POSTGRESQL_AUTH`` that password is a +short-lived token, so the supervisor mints a new one before it expires, +rewrites the auth file and asks PgBouncer to reload; only new upstream +connections authenticate, so live ones are unaffected. The pooled +``DATABASE_URL`` then carries a static password, and the workers must not run +their own token refresh against it: ``LITELLM_PGBOUNCER_POOLED_DATABASE_URL`` +tells them so, while a read replica keeps refreshing its own token. +""" + +from __future__ import annotations + +import atexit +import functools +import os +import re +import secrets +import shlex +import shutil +import signal +import socket +import subprocess +import tempfile +import threading +import time +import urllib.parse +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from types import MappingProxyType +from typing import Final + +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.db.token_auth import ( + DatabaseTokenAuth, + IAMEndpoint, + mint_database_token, + parse_database_token_expiration, + parse_iam_endpoint_from_url, +) + +PGBOUNCER_ENV_PREFIX: Final = "LITELLM_PGBOUNCER_" +PGBOUNCER_POOLED_ENV_VAR: Final = "LITELLM_PGBOUNCER_POOLED_DATABASE_URL" +PGBOUNCER_LISTEN_ADDR: Final = "127.0.0.1" +PGBOUNCER_POOL_USER: Final = "litellm_pgbouncer" +PGBOUNCER_INI_NAME: Final = "pgbouncer.ini" +PGBOUNCER_USERLIST_NAME: Final = "userlist.txt" +PGBOUNCER_CA_NAME: Final = "server-ca.pem" +PGBOUNCER_RESTART_DELAY_SECONDS: Final = 1.0 +PGBOUNCER_READY_TIMEOUT_SECONDS: Final = 15.0 +PGBOUNCER_STOP_GRACE_SECONDS: Final = 10.0 +PGBOUNCER_UNPRIVILEGED_USER: Final = "nobody" +PGBOUNCER_MIN_VERSION: Final = (1, 19) +PGBOUNCER_MAX_PASSWORD_BYTES: Final = 2048 +PGBOUNCER_VERSION_PATTERN: Final = re.compile(r"PgBouncer (\d+)\.(\d+)") +PGBOUNCER_TOKEN_REFRESH_BUFFER_SECONDS: Final = 180.0 +PGBOUNCER_TOKEN_FALLBACK_REFRESH_SECONDS: Final = 600.0 +PGBOUNCER_TOKEN_RETRY_SECONDS: Final = 30.0 + +# Prisma's client-side TLS params describe the hop to Postgres, which becomes +# PgBouncer's server side. They move into ``server_tls_*`` and must not stay on +# the loopback URL: the listener speaks plain TCP and Prisma would refuse it +# under ``sslmode=require`` or ``channel_binding=require``. +PRISMA_TLS_PARAM_KEYS: Final[frozenset[str]] = frozenset( + {"sslmode", "sslcert", "sslaccept", "sslidentity", "sslpassword", "channel_binding", "gssencmode"} +) +POOLED_URL_DROPPED_KEYS: Final[frozenset[str]] = PRISMA_TLS_PARAM_KEYS | frozenset(("options", "pgbouncer")) +PGBOUNCER_SSLMODES: Final[frozenset[str]] = frozenset( + {"disable", "allow", "prefer", "require", "verify-ca", "verify-full"} +) + + +class PgBouncerSettings(BaseSettings): + """``LITELLM_PGBOUNCER_*`` env vars, read once in the supervisor.""" + + model_config = SettingsConfigDict( + env_prefix=PGBOUNCER_ENV_PREFIX, case_sensitive=False, extra="ignore", frozen=True + ) + + enabled: bool = False + port: int = Field(default=6432, ge=1, le=65535) + max_db_connections: int = Field(default=20, ge=1) + max_client_conn: int = Field(default=1000, ge=1) + binary: str = "pgbouncer" + + +@dataclass(frozen=True, slots=True) +class PgBouncerPlan: + ini: str + pooled_url: str + upstream_user: str + upstream_password: str | None + pool_password: str + ca_source: str | None = None + + def userlist(self, upstream_password: str) -> str: + return "".join( + f"{_userlist_quote(user)} {_userlist_quote(password)}\n" + for user, password in ((self.upstream_user, upstream_password), (PGBOUNCER_POOL_USER, self.pool_password)) + ) + + +@dataclass(frozen=True, slots=True) +class PgBouncerError: + reason: str + + +def _single_quoted(value: str) -> str: + """Quote for SQL and for PgBouncer's ``[databases]`` connection string: both double a literal ``'``.""" + return "'" + value.replace("'", "''") + "'" + + +def _userlist_quote(value: str) -> str: + return '"' + value.replace('"', '""') + '"' + + +def _option_settings(tokens: Sequence[str]) -> tuple[str, ...] | None: + """The ``name=value`` settings in a libpq ``options`` string, or None if it holds anything else. + + Accepts ``-c name=value``, ``-cname=value`` and ``--name=value``; a + detached ``-c`` is folded into the token that follows it first. + """ + folded: Final = tuple( + f"-c{tokens[index + 1]}" if token == "-c" and index + 1 < len(tokens) else token + for index, token in enumerate(tokens) + if index == 0 or tokens[index - 1] != "-c" + ) + settings: Final = tuple(token[2:] for token in folded if token.startswith(("-c", "--")) and "=" in token[2:]) + return settings if len(settings) == len(folded) else None + + +def _connect_query(options: str) -> str | PgBouncerError: + """Turn Prisma's ``options=-c name=value ...`` startup param into ``SET`` statements. + + PgBouncer rejects any ``-c`` setting in ``options`` that is not one of the + handful it tracks (``statement_timeout`` and ``lock_timeout`` are not), so + the settings are applied to each new server connection instead. Every + client shares them, which is what the single ``DATABASE_URL`` gave anyway. + """ + settings: Final = _option_settings(tuple(shlex.split(options))) + if settings is None: + return PgBouncerError(f"cannot translate the DATABASE_URL options {options!r} into PgBouncer settings") + return "; ".join( + f"SET {name.strip()} TO {_single_quoted(value.strip())}" + for name, value in (setting.split("=", 1) for setting in settings) + ) + + +def _server_tls_settings(sslmode: str, sslcert: str, sslaccept: str, ca_path: Path) -> tuple[str, ...] | PgBouncerError: + """``server_tls_*`` lines naming ``ca_path``, the runtime-dir copy of the bundle: the original (or the + 0600 root pinned by ``pin_bundle_root``) is often unreadable for the user PgBouncer drops to.""" + if sslmode not in PGBOUNCER_SSLMODES: + return PgBouncerError(f"unsupported sslmode {sslmode!r} on DATABASE_URL") + verify: Final = sslmode in ("verify-ca", "verify-full") or (sslmode == "require" and sslaccept == "strict") + if verify and not sslcert: + return PgBouncerError( + "DATABASE_URL asks for a verified TLS connection but names no CA bundle; " + "add sslcert= (or sslrootcert=) so the in-container PgBouncer can verify Postgres" + ) + mode: Final = "verify-full" if verify else sslmode + return (f"server_tls_sslmode = {mode}", *((f"server_tls_ca_file = {ca_path}",) if sslcert else ())) + + +def plan_pgbouncer( + upstream_url: str, + settings: PgBouncerSettings, + runtime_dir: Path, + run_as_user: str | None, +) -> PgBouncerPlan | PgBouncerError: + """Render the PgBouncer config for ``upstream_url`` and the loopback URL Prisma uses instead. + + Params describing Prisma's own pool (``connection_limit``, ``pool_timeout``, + ...) stay on the pooled URL; the TLS params and ``options`` describe the hop + to Postgres and move into the PgBouncer config. The upstream password is + left out of the config on purpose: PgBouncer then takes it from the auth + file, which can be rewritten while it runs. ``run_as_user`` is the + unprivileged user PgBouncer drops to when the proxy runs as root, which + PgBouncer itself refuses to do. + """ + parsed: Final = urllib.parse.urlsplit(upstream_url) + params: Final[Mapping[str, str]] = MappingProxyType( + dict(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True)) + ) + dbname: Final = urllib.parse.unquote(parsed.path.lstrip("/")) + username: Final = urllib.parse.unquote(parsed.username or "") + password: Final = None if parsed.password is None else urllib.parse.unquote(parsed.password) + if not parsed.hostname or not username or not dbname: + return PgBouncerError("DATABASE_URL must carry a host, user and database name for the in-container PgBouncer") + if username == PGBOUNCER_POOL_USER: + return PgBouncerError( + f"the database user cannot be named {PGBOUNCER_POOL_USER!r}: that is the user the workers log in to the " + "in-container PgBouncer as, and PgBouncer keeps one password per user" + ) + if "sslidentity" in params: + return PgBouncerError("client certificates (sslidentity) are not supported with the in-container PgBouncer") + tls: Final = _server_tls_settings( + params.get("sslmode", "prefer"), + params.get("sslcert", ""), + params.get("sslaccept", ""), + runtime_dir / PGBOUNCER_CA_NAME, + ) + if isinstance(tls, PgBouncerError): + return tls + connect_query: Final = _connect_query(params["options"]) if params.get("options") else "" + if isinstance(connect_query, PgBouncerError): + return connect_query + upstream: Final = " ".join( + ( + f"host={_single_quoted(parsed.hostname)}", + f"port={parsed.port or 5432}", + f"dbname={_single_quoted(dbname)}", + f"user={_single_quoted(username)}", + *((f"connect_query={_single_quoted(connect_query)}",) if connect_query else ()), + ) + ) + ini: Final = "\n".join( + ( + "[databases]", + f"{dbname} = {upstream}", + "", + "[pgbouncer]", + f"listen_addr = {PGBOUNCER_LISTEN_ADDR}", + f"listen_port = {settings.port}", + f"unix_socket_dir = {runtime_dir}", + f"auth_file = {runtime_dir / PGBOUNCER_USERLIST_NAME}", + "auth_type = scram-sha-256", + f"stats_users = {PGBOUNCER_POOL_USER}", + "pool_mode = transaction", + f"max_client_conn = {settings.max_client_conn}", + f"default_pool_size = {settings.max_db_connections}", + f"max_db_connections = {settings.max_db_connections}", + "ignore_startup_parameters = extra_float_digits", + *tls, + *((f"user = {run_as_user}",) if run_as_user else ()), + "", + ) + ) + pooled_query: Final = urllib.parse.urlencode( + (*((key, value) for key, value in params.items() if key not in POOLED_URL_DROPPED_KEYS), ("pgbouncer", "true")) + ) + pool_password: Final = secrets.token_urlsafe(32) + pooled_url: Final = urllib.parse.urlunsplit( + parsed._replace( + netloc=f"{PGBOUNCER_POOL_USER}:{pool_password}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}", query=pooled_query + ) + ) + return PgBouncerPlan( + ini=ini, + pooled_url=pooled_url, + upstream_user=username, + upstream_password=password, + pool_password=pool_password, + ca_source=params.get("sslcert") or None, + ) + + +def pooled_database_url(upstream_url: str, settings: PgBouncerSettings) -> str | PgBouncerError: + """The loopback URL of a PgBouncer another container in the pod already runs for ``upstream_url``. + + Only the container that started PgBouncer knows the pool user's password, so + this logs in as the upstream user, whom the auth file lists as well. + """ + plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir=Path("/nonexistent"), run_as_user=None) + if isinstance(plan, PgBouncerError): + return plan + password: Final = urllib.parse.urlsplit(upstream_url).password or "" + credentials: Final = f"{urllib.parse.quote(plan.upstream_user, safe='')}:{password}" + return urllib.parse.urlunsplit( + urllib.parse.urlsplit(plan.pooled_url)._replace(netloc=f"{credentials}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}") + ) + + +def _write_private(path: Path, content: str, run_as_user: str | None) -> None: + with open(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600), "w", encoding="utf-8") as handle: + handle.write(content) + if run_as_user is not None: + shutil.chown(path, user=run_as_user) + + +def write_pgbouncer_ini(plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None) -> Path | PgBouncerError: + """Write the ini (mode 0600) and the CA copy, and return the ini path. + + ``run_as_user`` is the user PgBouncer drops to when started as root; it has + to own the files it re-reads on reload and the socket directory. + """ + ini_path: Final = runtime_dir / PGBOUNCER_INI_NAME + ca_path: Final = runtime_dir / PGBOUNCER_CA_NAME + if plan.ca_source is not None: + try: + shutil.copyfile(plan.ca_source, ca_path) + except OSError as error: + return PgBouncerError(f"cannot read the CA bundle {plan.ca_source!r} named by sslcert: {error}") + _write_private(ini_path, plan.ini, run_as_user) + if run_as_user is not None: + runtime_dir.chmod(0o700) + for path in (runtime_dir, *((ca_path,) if plan.ca_source is not None else ())): + shutil.chown(path, user=run_as_user) + return ini_path + + +def write_userlist(userlist: str, runtime_dir: Path, run_as_user: str | None) -> Path: + """Replace the auth file in one step, so a PgBouncer starting or reloading meanwhile reads the old or the new one whole.""" + userlist_path: Final = runtime_dir / PGBOUNCER_USERLIST_NAME + staged_path: Final = runtime_dir / f".{PGBOUNCER_USERLIST_NAME}.next" + _write_private(staged_path, userlist, run_as_user) + os.replace(staged_path, userlist_path) + return userlist_path + + +def export_pooled_database_url(pooled_url: str) -> None: + os.environ["DATABASE_URL"] = pooled_url + os.environ[PGBOUNCER_POOLED_ENV_VAR] = "true" + + +def database_url_is_pooled(environ: Mapping[str, str] = os.environ) -> bool: + return environ.get(PGBOUNCER_POOLED_ENV_VAR) == "true" + + +@dataclass(frozen=True, slots=True) +class PgBouncerTokenSource: + auth: DatabaseTokenAuth + endpoint: IAMEndpoint + + def mint(self) -> str: + """The token as Postgres expects it: ``mint_database_token`` returns it percent-encoded for a URL.""" + return urllib.parse.unquote(mint_database_token(self.auth, self.endpoint)) + + def expires_at(self, token: str) -> datetime | None: + return parse_database_token_expiration(self.auth, token) + + +def _utcnow() -> datetime: + return datetime.now(timezone.utc).replace(tzinfo=None) + + +class PgBouncerTokenRefresher: + """Keeps the token in PgBouncer's auth file current from a daemon thread. + + ``install`` gets each fresh token and is expected to rewrite the auth file + and reload PgBouncer. The next refresh is due ``buffer_seconds`` before the + token expires, or ``fallback_seconds`` later when the expiry cannot be read. + A refresh that fails leaves the previous auth file in place and is retried + after ``retry_seconds``: the old token stays good until it expires, so a + transient credential-provider error costs nothing unless it persists. + """ + + def __init__( + self, + source: PgBouncerTokenSource, + install: Callable[[str], None], + *, + buffer_seconds: float = PGBOUNCER_TOKEN_REFRESH_BUFFER_SECONDS, + fallback_seconds: float = PGBOUNCER_TOKEN_FALLBACK_REFRESH_SECONDS, + retry_seconds: float = PGBOUNCER_TOKEN_RETRY_SECONDS, + now: Callable[[], datetime] = _utcnow, + ) -> None: + self._source: Final = source + self._install: Final = install + self._buffer_seconds: Final = buffer_seconds + self._fallback_seconds: Final = fallback_seconds + self._retry_seconds: Final = retry_seconds + self._now: Final = now + self._stopping: Final = threading.Event() + self._delay: float = 0.0 + self._thread: threading.Thread | None = None + + def refresh(self) -> float | PgBouncerError: + label: Final = self._source.auth.label + try: + token: Final = self._source.mint() + except Exception as mint_error: + return PgBouncerError(f"could not mint a {label} for the in-container pgbouncer: {mint_error!r}") + if len(token.encode()) >= PGBOUNCER_MAX_PASSWORD_BYTES: + return PgBouncerError( + f"the {label} is {len(token.encode())} bytes long, but PgBouncer's auth file holds passwords of at " + f"most {PGBOUNCER_MAX_PASSWORD_BYTES - 1} bytes" + ) + try: + self._install(token) + except OSError as install_error: + return PgBouncerError(f"could not install the {label} into the pgbouncer auth file: {install_error}") + expires_at: Final = self._source.expires_at(token) + if expires_at is None: + return self._fallback_seconds + return max(self._retry_seconds, (expires_at - self._now()).total_seconds() - self._buffer_seconds) + + def start(self) -> PgBouncerError | None: + primed: Final = self.refresh() + if isinstance(primed, PgBouncerError): + return primed + self._delay = primed + self._thread = threading.Thread(target=self._run, daemon=True, name="litellm-pgbouncer-token-refresh") + self._thread.start() + return None + + def _run(self) -> None: + while not self._stopping.wait(self._delay): + self._delay = self._refresh_and_report() + + def _refresh_and_report(self) -> float: + outcome: Final = self.refresh() + if isinstance(outcome, PgBouncerError): + verbose_proxy_logger.error( + "In-container pgbouncer keeps its current %s (%s); retrying in %.0fs.", + self._source.auth.label, + outcome.reason, + self._retry_seconds, + ) + return self._retry_seconds + verbose_proxy_logger.info( + "In-container pgbouncer picked up a fresh %s; the next one is due in %.0fs.", + self._source.auth.label, + outcome, + ) + return outcome + + def stop(self) -> None: + self._stopping.set() + if self._thread is not None: + self._thread.join() + + +def _port_open(port: int) -> bool: + try: + with socket.create_connection((PGBOUNCER_LISTEN_ADDR, port), timeout=0.5): + return True + except OSError: + return False + + +def _unix_socket_open(path: Path) -> bool: + with socket.socket(socket.AF_UNIX) as probe: + probe.settimeout(0.5) + try: + probe.connect(str(path)) + except OSError: + return False + return True + + +def unix_socket_path(runtime_dir: Path, port: int) -> Path: + return runtime_dir / f".s.PGSQL.{port}" + + +def pgbouncer_version(binary: str) -> tuple[int, int] | PgBouncerError: + """``(major, minor)`` from `` --version``. + + Readiness relies on PgBouncer exiting when it cannot bind its TCP port, + which it does from 1.19 on. Older releases log a warning and serve the unix + socket alone, so their socket would vouch for a port held by someone else. + """ + try: + output: Final = subprocess.run( + (binary, "--version"), capture_output=True, text=True, check=False, timeout=10 + ).stdout + except (OSError, subprocess.TimeoutExpired) as run_error: + return PgBouncerError(f"could not run {binary!r} --version: {run_error}") + found: Final = PGBOUNCER_VERSION_PATTERN.search(output) + if found is None: + return PgBouncerError(f"{binary!r} --version did not report a PgBouncer version: {output.strip()!r}") + return int(found[1]), int(found[2]) + + +def _end(process: subprocess.Popen[bytes]) -> None: + if process.poll() is not None: + return + process.terminate() + try: + process.wait(timeout=PGBOUNCER_STOP_GRACE_SECONDS) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + +class PgBouncerProcess: + """Runs ``argv`` as a foreground child and restarts it whenever it exits on its own. + + Prisma reconnects by itself after a failed query, so a PgBouncer crash + costs the requests in flight plus one failed query per idle pooled + connection the crash severed, and nothing else once the replacement is + listening again. A replacement that cannot be spawned, finds its port + taken, exits again or never starts listening is retried every + ``restart_delay_seconds`` until ``stop`` is called. + + A connect probe of ``port`` cannot tell the child from another process + that grabbed the port after the availability check, so readiness also + needs ``socket_path``: the unix socket PgBouncer creates in the private + runtime directory, which it only does once every TCP listener is bound + (PgBouncer 1.19 or newer, see ``pgbouncer_version``). + """ + + def __init__( + self, + argv: Sequence[str], + port: int, + socket_path: Path, + restart_delay_seconds: float = PGBOUNCER_RESTART_DELAY_SECONDS, + ready_timeout_seconds: float = PGBOUNCER_READY_TIMEOUT_SECONDS, + ) -> None: + self.argv: Final = tuple(argv) + self.port: Final = port + self.socket_path: Final = socket_path + self.restart_delay_seconds: Final = restart_delay_seconds + self.ready_timeout_seconds: Final = ready_timeout_seconds + self._stopping: Final = threading.Event() + self._lock: Final = threading.Lock() + self._process: subprocess.Popen[bytes] | None = None + + @property + def pid(self) -> int | None: + with self._lock: + return None if self._process is None else self._process.pid + + def _spawn(self) -> subprocess.Popen[bytes] | PgBouncerError | None: + """Start a child, or None once ``stop`` ran; both take the lock so no child can slip in after a stop. + + The port has to be free first: a listener that is already there would + pass the readiness check while the child fails to bind. + """ + with self._lock: + if self._stopping.is_set(): + return None + if _port_open(self.port): + return PgBouncerError(f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is already in use by another process") + try: + process: Final = subprocess.Popen(self.argv) + except OSError as spawn_error: + return PgBouncerError(f"could not start {self.argv[0]!r}: {spawn_error}") + self._process = process + return process + + def _wait_ready(self, process: subprocess.Popen[bytes]) -> PgBouncerError | None: + deadline: Final = time.monotonic() + self.ready_timeout_seconds + while time.monotonic() < deadline: + if process.poll() is not None: + return PgBouncerError(f"pgbouncer exited with status {process.returncode} during startup") + if _port_open(self.port) and _unix_socket_open(self.socket_path): + return None + time.sleep(0.1) + if _port_open(self.port): + return PgBouncerError( + f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is served by another process, not the pgbouncer that was started" + ) + return PgBouncerError( + f"pgbouncer did not start listening on {PGBOUNCER_LISTEN_ADDR}:{self.port} " + f"within {self.ready_timeout_seconds:.0f}s" + ) + + def start(self) -> PgBouncerError | None: + """Spawn PgBouncer, wait until it listens on port and unix socket, then supervise it from a daemon thread.""" + process: Final = self._spawn() + if process is None: + return PgBouncerError("pgbouncer was stopped before it started") + if isinstance(process, PgBouncerError): + return process + not_ready: Final = self._wait_ready(process) + if not_ready is not None: + self.stop() + return not_ready + self._watch(process) + return None + + def _watch(self, process: subprocess.Popen[bytes]) -> None: + threading.Thread( + target=self._supervise, args=(process,), daemon=True, name="litellm-pgbouncer-supervisor" + ).start() + + def _supervise(self, process: subprocess.Popen[bytes]) -> None: + status: Final = process.wait() + if self._stopping.is_set(): + return + verbose_proxy_logger.error( + "In-container pgbouncer (pid %s) exited with status %s; restarting in %.1fs.", + process.pid, + status, + self.restart_delay_seconds, + ) + self._restart_after_delay() + + def _restart_after_delay(self) -> None: + time.sleep(self.restart_delay_seconds) + process: Final = self._spawn() + if process is None: + return + if isinstance(process, PgBouncerError): + self._retry_restart(process.reason) + return + not_ready: Final = self._wait_ready(process) + if not_ready is None: + self._watch(process) + return + _end(process) + self._retry_restart(not_ready.reason) + + def _retry_restart(self, reason: str) -> None: + if self._stopping.is_set(): + return + verbose_proxy_logger.error( + "In-container pgbouncer could not be restarted (%s); retrying in %.1fs.", reason, self.restart_delay_seconds + ) + threading.Thread(target=self._restart_after_delay, daemon=True, name="litellm-pgbouncer-supervisor").start() + + def reload(self) -> None: + with self._lock: + if self._process is not None: + self._process.send_signal(signal.SIGHUP) + + def stop(self) -> None: + with self._lock: + self._stopping.set() + process: Final = self._process + if process is not None: + _end(process) + + +def install_pgbouncer_token( + plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None, pooler: PgBouncerProcess, token: str +) -> None: + write_userlist(plan.userlist(token), runtime_dir, run_as_user) + pooler.reload() + + +def _install_upstream_password( + plan: PgBouncerPlan, + runtime_dir: Path, + run_as_user: str | None, + pooler: PgBouncerProcess, + token_auth: DatabaseTokenAuth | None, + upstream_url: str, +) -> PgBouncerTokenRefresher | None | PgBouncerError: + if token_auth is None: + if plan.upstream_password is None: + return PgBouncerError( + "DATABASE_URL carries no password and neither IAM_TOKEN_DB_AUTH nor AZURE_POSTGRESQL_AUTH is on, " + "so the in-container PgBouncer has nothing to authenticate to Postgres with" + ) + write_userlist(plan.userlist(plan.upstream_password), runtime_dir, run_as_user) + return None + refresher: Final = PgBouncerTokenRefresher( + PgBouncerTokenSource(auth=token_auth, endpoint=parse_iam_endpoint_from_url(upstream_url)), + functools.partial(install_pgbouncer_token, plan, runtime_dir, run_as_user, pooler), + ) + failed: Final = refresher.start() + if failed is not None: + return failed + return refresher + + +def _only_in_this_process(action: Callable[[], None]) -> Callable[[], None]: + """An exit hook that does nothing in a forked child, which inherits the parent's ``atexit`` table.""" + owner_pid: Final = os.getpid() + + def run() -> None: + if os.getpid() == owner_pid: + action() + + return run + + +def start_in_container_pgbouncer( + settings: PgBouncerSettings, + upstream_url: str, + token_auth: DatabaseTokenAuth | None = None, + register_exit_hook: Callable[[Callable[[], None]], object] = atexit.register, +) -> str | PgBouncerError: + """Start the pooler for ``upstream_url`` and return the loopback URL the workers must use. + + The pooler lives as long as this process: it is stopped from the exit hooks + once the worker manager has returned, and only by the process that started + it (gunicorn forks its workers, so they carry the hooks too). PgBouncer + refuses to run as root, so a root proxy (the default image) has it drop to + ``nobody``. With ``token_auth`` the password on ``upstream_url`` is ignored: + the pooler mints its own tokens and renews them for as long as it runs. + """ + version: Final = pgbouncer_version(settings.binary) + if isinstance(version, PgBouncerError): + return version + if version < PGBOUNCER_MIN_VERSION: + return PgBouncerError( + f"PgBouncer {version[0]}.{version[1]} keeps running after failing to bind its TCP port, so the proxy " + f"cannot tell it apart from another listener; {PGBOUNCER_MIN_VERSION[0]}.{PGBOUNCER_MIN_VERSION[1]} " + "or newer is required" + ) + runtime_dir: Final = Path(tempfile.mkdtemp(prefix="litellm-pgbouncer-")) + register_exit_hook(_only_in_this_process(lambda: shutil.rmtree(runtime_dir, ignore_errors=True))) + run_as_user: Final = PGBOUNCER_UNPRIVILEGED_USER if os.geteuid() == 0 else None + plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir, run_as_user) + if isinstance(plan, PgBouncerError): + return plan + ini_path: Final = write_pgbouncer_ini(plan, runtime_dir, run_as_user) + if isinstance(ini_path, PgBouncerError): + return ini_path + pooler: Final = PgBouncerProcess( + argv=(settings.binary, str(ini_path)), + port=settings.port, + socket_path=unix_socket_path(runtime_dir, settings.port), + ) + refresher: Final = _install_upstream_password(plan, runtime_dir, run_as_user, pooler, token_auth, upstream_url) + if isinstance(refresher, PgBouncerError): + return refresher + failed: Final = pooler.start() + if failed is not None: + if refresher is not None: + refresher.stop() + return failed + register_exit_hook(_only_in_this_process(pooler.stop)) + if refresher is not None: + register_exit_hook(_only_in_this_process(refresher.stop)) + verbose_proxy_logger.info( + "In-container pgbouncer (pid %s) listening on %s:%s; capping this pod at %s upstream database connections%s.", + pooler.pid, + PGBOUNCER_LISTEN_ADDR, + settings.port, + settings.max_db_connections, + "" if token_auth is None else f" and renewing its {token_auth.label} before each one expires", + ) + return plan.pooled_url diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 2190ae55fd2..21b73f27a80 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -16,6 +16,7 @@ from datetime import datetime, timedelta from typing import Any, Final, Protocol from litellm._logging import verbose_proxy_logger +from litellm.proxy.db.db_url_settings import add_missing_query_params, connection_params_from_url from litellm.proxy.db.token_auth import ( DEFAULT_POSTGRES_PORT, DatabaseTokenAuth, @@ -438,7 +439,10 @@ class PrismaWrapper: return None endpoint: Final = self._iam_endpoint if self._iam_endpoint is not None else self._endpoint_from_env() - db_url: Final = endpoint.build_url(mint_database_token(auth, endpoint)) + db_url: Final = add_missing_query_params( + endpoint.build_url(mint_database_token(auth, endpoint)), + connection_params_from_url(os.environ.get(self._db_url_env_var, "")), + ) os.environ[self._db_url_env_var] = db_url return db_url @@ -937,9 +941,17 @@ class PrismaManager: use_v2_resolver=use_v2_resolver, ) else: + try: + from litellm_proxy_extras.prisma_toolchain import ( + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e) + return False + PrismaManager._raise_if_partitioned_spend_logs() - # Use prisma db push with increased timeout - subprocess.run( + run_prisma( [ "prisma", "db", @@ -947,13 +959,15 @@ class PrismaManager: "--accept-data-loss", "--skip-generate", ], - timeout=60, - check=True, + timeout=prisma_command_timeout(), + env=os.environ.copy(), + stdout=None, + stderr=None, ) PrismaManager._apply_replica_identity_full_if_requested() return True - except subprocess.TimeoutExpired: - verbose_proxy_logger.warning("Attempt %s timed out", attempt + 1) + except subprocess.TimeoutExpired as e: + verbose_proxy_logger.warning("Attempt %s timed out after %.0fs", attempt + 1, e.timeout) time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: attempts_left = 3 - attempt diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index a38b8a47dbd..e35b1c8c82b 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -23,6 +23,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.db.db_lookup_gate import db_lookup_gate from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, @@ -121,30 +122,33 @@ class SpendCounterReseed: if SpendCounterReseed._is_key_or_team_window_counter(counter_key): return None try: - if counter_key.startswith("spend:key:"): - token: Final = counter_key[len("spend:key:") :] - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) - elif counter_key.startswith("spend:team_member:"): - suffix: Final = counter_key[len("spend:team_member:") :] - if ":" not in suffix: + async with db_lookup_gate.current(): + if counter_key.startswith("spend:key:"): + token: Final = counter_key[len("spend:key:") :] + row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) + elif counter_key.startswith("spend:team_member:"): + suffix: Final = counter_key[len("spend:team_member:") :] + if ":" not in suffix: + return None + user_id, team_id = suffix.rsplit(":", 1) + row = await TeamMembershipRepository(prisma_client).table.find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} + ) + elif counter_key.startswith("spend:team:"): + team_id = counter_key[len("spend:team:") :] + row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + elif counter_key.startswith("spend:user:"): + user_id = counter_key[len("spend:user:") :] + row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"): + return None + elif counter_key.startswith("spend:org:"): + org_id: Final = counter_key[len("spend:org:") :] + row = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": org_id} + ) + else: return None - user_id, team_id = suffix.rsplit(":", 1) - row = await TeamMembershipRepository(prisma_client).table.find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} - ) - elif counter_key.startswith("spend:team:"): - team_id = counter_key[len("spend:team:") :] - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) - elif counter_key.startswith("spend:user:"): - user_id = counter_key[len("spend:user:") :] - row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"): - return None - elif counter_key.startswith("spend:org:"): - org_id: Final = counter_key[len("spend:org:") :] - row = await OrganizationRepository(prisma_client).table.find_unique(where={"organization_id": org_id}) - else: - return None except Exception: verbose_proxy_logger.exception("SpendCounterReseed.from_db: failed for %s", counter_key) return None diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 98707e7ddca..c37b9fff1f0 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -80,17 +80,19 @@ def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRout def policy_for_model( llm_router: "Router | None", model_alias: str, - team_id: str | None, + request_kwargs: Mapping[str, object], request_tags: Sequence[str], ) -> AutoRouterCompressionPolicy | None: - """The compression policy of the auto router marker `model_alias` resolves to. + """The compression policy of the auto router marker `model_alias` resolves to for this caller. - Pre-call arming and the routing hook both resolve through here, so an alias with - several tag-scoped markers cannot suppress under one and then route under another. + Pre-call arming and the routing hook both resolve through here, and here resolves through the + router's own request-scoped deployment lookup, so an alias with several tag-scoped markers + cannot suppress under one and then route under another, and a team router reached by its + public name carries its policy for every principal that can reach it. """ if llm_router is None: return None - deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or () + deployments: Final = llm_router.deployments_for_request(model_alias, request_kwargs) markers: Final = tuple( litellm_params for deployment in deployments @@ -108,17 +110,6 @@ def policy_for_model( return next((policy for policy in candidates if policy is not None), None) -def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: - """The caller's team id, from whichever metadata bucket this surface writes to.""" - for meta_key in ("metadata", "litellm_metadata"): - meta = request_kwargs.get(meta_key) - if isinstance(meta, Mapping): - team_id = meta.get("user_api_key_team_id") - if isinstance(team_id, str): - return team_id - return None - - def _compression_guardrail_classes() -> tuple[type, ...]: """The registered guardrail classes whose provider compresses prompts.""" from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry @@ -172,7 +163,7 @@ async def arm_pre_call( policy: Final = policy_for_model( llm_router=llm_router, model_alias=model_alias, - team_id=team_id_from_request(data), + request_kwargs=data, request_tags=_get_tags_from_request_kwargs(data), ) if policy is None: diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 744b2959c73..afb9997f2e6 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -1202,7 +1202,13 @@ async def patch_guardrail( litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True) litellm_params_dict.update(requested_litellm_params) merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict) - litellm_params = LitellmParams(**merged_litellm_params) + try: + litellm_params = LitellmParams(**merged_litellm_params) + except ValidationError as validation_error: + raise HTTPException( + status_code=422, + detail=f"Invalid guardrail configuration, update rejected: {validation_error}", + ) from validation_error # Update guardrail_info if provided guardrail_info: Final = ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 7204839f6d3..6a7ac4361b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -44,7 +44,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -917,7 +917,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source, ) return BedrockGuardrailResponse() - credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + credentials, aws_region_name = await run_aws_signing( + self._load_credentials, bearer_token=bedrock_bearer_token(api_key) + ) allow_chunking: Final = not self._content_uses_contextual_grounding(content) completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator @@ -1178,7 +1180,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): **base_request_data, "content": content, } # mutable-ok: outbound JSON request body - prepared_request: Final = self._prepare_request( + prepared_request: Final = await run_aws_signing( + self._prepare_request, credentials=credentials, data=bedrock_request_data, optional_params=self.optional_params, @@ -1875,10 +1878,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return BedrockGuardrailResponse() api_key: Final[str | None] = request_data.get("api_key") if request_data else None - credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + credentials, aws_region_name = await run_aws_signing( + self._load_credentials, bearer_token=bedrock_bearer_token(api_key) + ) body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} - prepared_request: Final = self._prepare_request( + prepared_request: Final = await run_aws_signing( + self._prepare_request, credentials=credentials, data=body, optional_params=self.optional_params, diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py index d569802ce89..b7e4d9275fa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py @@ -36,6 +36,7 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> default_on=litellm_params.default_on or False, unreachable_fallback=litellm_params.unreachable_fallback, timeout=litellm_params.timeout, + ccr_retrieval=litellm_params.ccr_retrieval, ) litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] _callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 685b90f1754..fa113aa4d33 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -6,6 +6,7 @@ import re import time import uuid from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard import httpx @@ -60,7 +61,7 @@ _STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset( # stalled service holds the caller's request and a pooled connection for 600s or more. _COMPRESS_TIMEOUT_SECONDS: Final = 60.0 HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" -_HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})") +_HASH_PATTERN: Final = re.compile(r"[a-f0-9]{12,24}") _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 # Narrows the base class's bare-dict ``request_data`` at the boundary so its # untranslated messages can be read with concrete types (values pass through by @@ -239,15 +240,25 @@ def _protected_indices( tool exchanges the way ``compress()`` expands it, so a protected assistant tool call cannot end up answered by a marker standing in for the result the model just asked for. + + Every assistant row is then withheld without expanding its tool exchange: + the service protects assistant text blocks but has no gate for assistant + strings, and the Anthropic adapter hands assistant blocks over as strings, + so the model's own earlier tables came back rewritten and it imitated the + shape. The tool results those turns asked for stay compressible. """ protected: Final = frozenset(get_protected_indices(messages)) | _retrieval_result_indices( messages, extra_retrieve_call_ids ) - return protected | frozenset( - index - for group in group_tool_exchanges(messages) - if any(member in protected for member in group) - for index in group + return ( + protected + | frozenset( + index + for group in group_tool_exchanges(messages) + if any(member in protected for member in group) + for index in group + ) + | frozenset(index for index, message in enumerate(messages) if message.get("role") == "assistant") ) @@ -290,19 +301,23 @@ def _build_compress_failure_detail(status_code: int, body: str) -> dict[str, obj return {"status_code": status_code, "body": body} -def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]: - hashes: Final[list[str]] = [] - for msg in messages: - content = msg.get("content") - if isinstance(content, str): - hashes.extend(_HASH_PATTERN.findall(content)) - elif isinstance(content, list): - for block in content: - if isinstance(block, dict): - text = block.get("text") - if isinstance(text, str): - hashes.extend(_HASH_PATTERN.findall(text)) - return hashes +def _read_ccr_hashes(body: Mapping[str, object]) -> frozenset[str]: + ccr_hashes: Final = body.get("ccr_hashes") + if not isinstance(ccr_hashes, list): + return frozenset() + return frozenset( + hash_value.lower() + for hash_value in ccr_hashes + if isinstance(hash_value, str) and _HASH_PATTERN.fullmatch(hash_value.lower()) + ) + + +@dataclass(frozen=True, slots=True) +class _CompressResult: + messages: list[dict[str, object]] + succeeded: bool + stats: dict[str, object] + ccr_hashes: frozenset[str] = frozenset() def _build_headroom_retrieve_tool() -> dict[str, object]: @@ -319,7 +334,7 @@ def _build_headroom_retrieve_tool() -> dict[str, object]: "properties": { "hash": { "type": "string", - "description": "The 24-character hex hash from the compression marker.", + "description": "The hex hash from the compression marker.", }, "query": { "type": "string", @@ -479,6 +494,7 @@ class HeadroomGuardrail(CustomGuardrail): default_on: bool = False, unreachable_fallback: str | None = None, timeout: float | None = None, + ccr_retrieval: bool = True, ): self.headroom_api_base = (api_base or get_secret_str("HEADROOM_API_BASE") or "").rstrip("/") if not self.headroom_api_base: @@ -492,6 +508,7 @@ class HeadroomGuardrail(CustomGuardrail): "fail_open" if unreachable_fallback == "fail_open" else "fail_closed" ) self.timeout: httpx.Timeout = self._resolve_timeout(timeout) + self.ccr_retrieval = ccr_retrieval self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, ) @@ -569,7 +586,7 @@ class HeadroomGuardrail(CustomGuardrail): self, messages: list[dict[str, object]], model: str | None, - ) -> tuple[list[dict[str, object]], bool, dict[str, object]]: + ) -> _CompressResult: payload: Final[dict[str, object]] = {"messages": messages} if model: payload["model"] = model @@ -582,7 +599,7 @@ class HeadroomGuardrail(CustomGuardrail): timeout=self.timeout, ) except httpx.HTTPStatusError as e: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned an error", @@ -592,7 +609,7 @@ class HeadroomGuardrail(CustomGuardrail): {}, ) except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service unreachable", @@ -604,7 +621,7 @@ class HeadroomGuardrail(CustomGuardrail): response: Final[HttpxResponse] = raw_response if response.status_code != 200: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned an error", @@ -617,7 +634,7 @@ class HeadroomGuardrail(CustomGuardrail): try: body: Final[object] = response.json() except ValueError: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned non-JSON response", @@ -627,7 +644,7 @@ class HeadroomGuardrail(CustomGuardrail): {}, ) if not _is_str_object_dict(body): - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned unexpected response shape", @@ -639,7 +656,7 @@ class HeadroomGuardrail(CustomGuardrail): compressed_messages: Final = body.get("messages") if not _is_object_list(compressed_messages): - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service response missing 'messages'", @@ -651,7 +668,7 @@ class HeadroomGuardrail(CustomGuardrail): filtered: Final = [item for item in compressed_messages if _is_str_object_dict(item)] if not filtered: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned empty message list", @@ -664,7 +681,7 @@ class HeadroomGuardrail(CustomGuardrail): if len(filtered) != len(messages): # Rows are matched positionally when the never-compressed messages # are put back, so a reshaped conversation cannot be applied at all. - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service changed the message count", @@ -705,7 +722,7 @@ class HeadroomGuardrail(CustomGuardrail): # tokens_saved, which the live compression service omits; derive it # so savings are counted, but let a service-sent value win. stats["tokens_saved"] = tokens_before - tokens_after - return filtered, True, stats + return _CompressResult(filtered, True, stats, _read_ccr_hashes(body)) async def _call_retrieve(self, hash_value: str, query: str | None = None) -> str: params: Final[dict[str, str]] = {} @@ -793,7 +810,7 @@ class HeadroomGuardrail(CustomGuardrail): model: Final = self.headroom_model or request_data.get("model") start_time: Final = time.time() - returned, compression_succeeded, stats = await self._call_compress( + result: Final = await self._call_compress( messages=_flatten_messages_for_compression(compressible), model=model if isinstance(model, str) else None, ) @@ -803,7 +820,7 @@ class HeadroomGuardrail(CustomGuardrail): add_guardrail_to_applied_guardrails_header, ) - if not compression_succeeded: + if not result.succeeded: self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response={"error": "headroom compression unavailable; request forwarded uncompressed"}, request_data=request_data, @@ -822,12 +839,12 @@ class HeadroomGuardrail(CustomGuardrail): compressed: Final = _restore_protected_messages( messages=messages, - compressed=_restore_content_shapes(originals=compressible, returned=returned), + compressed=_restore_content_shapes(originals=compressible, returned=result.messages), protected_indices=protected_indices, ) self.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response=stats, + guardrail_json_response=result.stats, request_data=request_data, guardrail_status="success", guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER, @@ -837,7 +854,7 @@ class HeadroomGuardrail(CustomGuardrail): ) add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) - hashes: Final = extract_hashes_from_messages(compressed) + hashes: Final = result.ccr_hashes if self.ccr_retrieval else frozenset() if not hashes: return {**inputs, "structured_messages": compressed} # pyright: ignore[reportReturnType] @@ -918,14 +935,15 @@ class HeadroomGuardrail(CustomGuardrail): retrieved: Final[list[tuple[dict[str, object], str]]] = [] for tc in tool_calls: arguments = tc.get("arguments", {}) - hash_value = arguments.get("hash", "") if isinstance(arguments, dict) else "" + raw_hash = arguments.get("hash", "") if isinstance(arguments, dict) else "" + hash_value = str(raw_hash).lower() query = arguments.get("query") if isinstance(arguments, dict) else None # A hash is only honored if it was issued by *this request's own* # Headroom /v1/compress call, scoped by litellm_call_id. Scoping by # message text alone is forgeable -- an attacker can plant a # hash-shaped string in their own prompt, and a hash issued for one # request would validate for any other request that echoes it back. - if str(hash_value) not in valid_hashes: + if hash_value not in valid_hashes: verbose_proxy_logger.warning( "Headroom CCR: rejecting hash=%s not produced by current request compression", hash_value, @@ -933,7 +951,7 @@ class HeadroomGuardrail(CustomGuardrail): content = f"[Headroom: hash={hash_value} was not produced by the current request]" else: content = await self._call_retrieve( - hash_value=str(hash_value), + hash_value=hash_value, query=str(query) if query else None, ) verbose_proxy_logger.debug("Headroom CCR: retrieved hash=%s (%d chars)", hash_value, len(content)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 722f96ef814..1e684c514de 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -9,6 +9,7 @@ import asyncio import json import os import re +import time from collections.abc import AsyncGenerator, Coroutine, Mapping, Sequence from datetime import datetime from re import Pattern @@ -1702,6 +1703,7 @@ class ContentFilterGuardrail(CustomGuardrail): start_time: datetime, masked_entity_count: dict[str, int], exception_str: str, + duration: float | None = None, ) -> None: """ Log guardrail information to request_data metadata. @@ -1713,6 +1715,7 @@ class ContentFilterGuardrail(CustomGuardrail): start_time: Start time of guardrail execution masked_entity_count: Count of masked entities by type exception_str: Exception string if guardrail failed + duration: Seconds spent inside the guardrail; defaults to the wall clock since start_time """ # Convert TypedDict detections to regular dicts for JSON serialization guardrail_json_response: Exception | str | dict | list[dict] = [dict(detection) for detection in detections] @@ -1741,7 +1744,7 @@ class ContentFilterGuardrail(CustomGuardrail): guardrail_status=status, start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), - duration=(datetime.now() - start_time).total_seconds(), + duration=(datetime.now() - start_time).total_seconds() if duration is None else duration, masked_entity_count=masked_entity_count, tracing_detail=GuardrailTracingDetail(**tracing_kw), ) @@ -1971,6 +1974,7 @@ class ContentFilterGuardrail(CustomGuardrail): buffer_size: Final = 50 # Increased buffer to catch patterns split across many chunks start_time: Final = datetime.now() + scan_seconds: float = 0.0 # rebind-ok: accumulates per-chunk scan time across the stream detections: list[ContentFilterDetection] = [] masked_entity_count: Final[dict[str, int]] = {} status: GuardrailStatus = "success" @@ -2007,6 +2011,7 @@ class ContentFilterGuardrail(CustomGuardrail): # Add a space at the end if it's the final chunk to trigger word boundaries (\b) text_to_scan = text_to_check + (" " if is_final else "") choice_detections: list[ContentFilterDetection] = [] + scan_started = time.perf_counter() try: # _filter_single_text scans the whole accumulated @@ -2024,6 +2029,8 @@ class ContentFilterGuardrail(CustomGuardrail): except Exception as e: verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e) masked_text = text_to_scan # Fallback to current text + finally: + scan_seconds += time.perf_counter() - scan_started # Determine how much can be safely yielded if is_final: @@ -2074,6 +2081,7 @@ class ContentFilterGuardrail(CustomGuardrail): start_time=start_time, masked_entity_count=masked_entity_count, exception_str=exception_str, + duration=scan_seconds, ) @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py index d53a4157e0e..1607dfff63e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Final, Literal, Optional import litellm from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( @@ -20,10 +20,7 @@ def initialize_guardrail( if not guardrail_name: raise ValueError("MCP Security: guardrail_name is required") - on_violation: Final[Literal["block", "alert"]] = cast( - Literal["block", "alert"], - getattr(litellm_params, "on_violation", "block"), - ) + on_violation: Final[Literal["block", "alert"]] = "block" if litellm_params.on_violation == "block" else "alert" mcp_security_guardrail: Final = MCPSecurityGuardrail( guardrail_name=guardrail_name, diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 10683550f85..c22d35509c1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -17,7 +17,8 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.proxy.common_utils.callback_utils import add_guardrail_scan_id +from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations from litellm.types.utils import ( GenericGuardrailAPIInputs, GuardrailStatus, @@ -218,6 +219,13 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): metadata: Final = request_data.get("metadata") or {} request_data["metadata"] = metadata metadata["_openai_moderation_response"] = moderation_response.model_dump() + add_guardrail_scan_id( + request_data=request_data, + scan_id=moderation_response.id, + guardrail_name=self.guardrail_name, + provider=SupportedGuardrailIntegrations.OPENAI_MODERATION.value, + stage=GuardrailEventHooks.post_call if input_type == "response" else GuardrailEventHooks.pre_call, + ) # Check if content is flagged and raise exception if needed self._check_moderation_result(moderation_response) diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index b73d3adb99e..3bc0dfabefc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -721,10 +721,18 @@ class PanwPrismaAirsHandler(CustomGuardrail): } } - def _record_scan_id(self, request_data: dict[str, object], scan_result: Mapping[str, object]) -> None: + def _record_scan_id( + self, request_data: dict[str, object], scan_result: Mapping[str, object], stage: GuardrailEventHooks + ) -> None: """Surface the AIRS scan id on the response, so allowed calls are auditable too.""" scan_id: Final = scan_result.get("scan_id") - add_guardrail_scan_id(request_data=request_data, scan_id=str(scan_id) if scan_id else None) + add_guardrail_scan_id( + request_data=request_data, + scan_id=str(scan_id) if scan_id else None, + guardrail_name=self.guardrail_name, + provider=self._PROVIDER_NAME, + stage=stage, + ) def _handle_api_error_with_logging( self, @@ -948,7 +956,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): event_type=GuardrailEventHooks.post_call, ) add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) - self._record_scan_id(request_data, scan_result) + self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call) def _check_and_mark_scanned(self, data: dict, scan_type: str) -> bool: """ @@ -1078,7 +1086,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.pre_call, ) - self._record_scan_id(data, scan_result) + self._record_scan_id(data, scan_result, GuardrailEventHooks.pre_call) action: Final = scan_result.get("action", "block") category: Final = scan_result.get("category", "unknown") @@ -1199,7 +1207,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.post_call, ) - self._record_scan_id(data, scan_result) + self._record_scan_id(data, scan_result, GuardrailEventHooks.post_call) action: Final = scan_result.get("action", "block") category: Final = scan_result.get("category", "unknown") @@ -1401,7 +1409,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.post_call, ) - self._record_scan_id(request_data, scan_result) + self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call) # Add guardrail to applied guardrails header for observability add_guardrail_to_applied_guardrails_header( @@ -1475,7 +1483,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) continue - self._record_scan_id(request_data, scan_result) + self._record_scan_id( + request_data, + scan_result, + GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call, + ) action = scan_result.get("action", "block") masked_args = self._masked_tool_call_arguments( @@ -1829,7 +1841,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): new_texts.append(text) continue - self._record_scan_id(request_data, scan_result) + self._record_scan_id( + request_data, + scan_result, + GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call, + ) action = scan_result.get("action", "block") masked_text = self._get_masked_text(scan_result, is_response=is_response) @@ -1901,7 +1917,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) # If we reach here, fallback_on_error="allow" else: - self._record_scan_id(request_data, mcp_scan_result) + self._record_scan_id(request_data, mcp_scan_result, GuardrailEventHooks.pre_call) action = mcp_scan_result.get("action", "block") masked_text = self._get_masked_text(mcp_scan_result, is_response=False) if action == "allow": diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py index 0aaba4016cd..88cf92a4a8c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py @@ -21,6 +21,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" event_hook=litellm_params.mode, default_on=litellm_params.default_on, file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None), + block_on_file_modify=getattr(litellm_params, "block_on_file_modify", None), ) litellm.logging_callback_manager.add_litellm_callback(_prompt_security_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 84c4f118b00..0954fe1698a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -93,6 +93,7 @@ class PromptSecurityGuardrail(CustomGuardrail): check_tool_results: bool | None = None, file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS, file_sanitization_fail_open: bool | None = None, + block_on_file_modify: bool | None = None, **kwargs, ): kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) @@ -124,6 +125,7 @@ class PromptSecurityGuardrail(CustomGuardrail): self.poll_interval = 2 # Seconds between polling attempts self.file_sanitization_timeout = file_sanitization_timeout self.file_sanitization_fail_open = file_sanitization_fail_open is not False + self.block_on_file_modify = block_on_file_modify is not False super().__init__(**kwargs) @@ -372,13 +374,7 @@ class PromptSecurityGuardrail(CustomGuardrail): result = await self.sanitize_file_content( file_data, filename, user_api_key_alias=user_api_key_alias ) - - if result.get("action") == "block": - violations = result.get("violations", []) - raise HTTPException( - status_code=400, - detail=f"Image blocked by Prompt Security. Violations: {', '.join(violations)}", - ) + self._raise_if_file_blocked(result, "Image") except HTTPException: raise except Exception as e: @@ -408,7 +404,7 @@ class PromptSecurityGuardrail(CustomGuardrail): file_data: bytes, filename: str, user_api_key_alias: str | None = None, - ) -> dict: + ) -> _SanitizeResult: """ Sanitize file content using Prompt Security API. Returns: dict with keys 'action', 'content', 'metadata' @@ -528,6 +524,17 @@ class PromptSecurityGuardrail(CustomGuardrail): raise HTTPException(status_code=408, detail="File sanitization timeout") + def _raise_if_file_blocked(self, sanitization_result: _SanitizeResult, resource_name: str) -> None: + action: Final = sanitization_result.get("action") + if action != "block" and not (action == "modify" and self.block_on_file_modify): + return + + violations: Final = sanitization_result.get("violations", ()) + raise HTTPException( + status_code=400, + detail=f"{resource_name} blocked by Prompt Security. Violations: {', '.join(violations)}", + ) + async def _process_image_url_item(self, item: dict, user_api_key_alias: str | None) -> dict: """Process and sanitize image_url items.""" image_url_data: Final = item.get("image_url", {}) @@ -547,13 +554,7 @@ class PromptSecurityGuardrail(CustomGuardrail): file_data, filename, user_api_key_alias=user_api_key_alias ) action: Final = sanitization_result.get("action") - - if action == "block": - violations: Final = sanitization_result.get("violations", []) - raise HTTPException( - status_code=400, - detail=f"File blocked by Prompt Security. Violations: {', '.join(violations)}", - ) + self._raise_if_file_blocked(sanitization_result, "File") if action == "modify": sanitized_content: Final = sanitization_result.get("content", "") @@ -615,13 +616,7 @@ class PromptSecurityGuardrail(CustomGuardrail): file_data, filename, user_api_key_alias=user_api_key_alias ) action: Final = sanitization_result.get("action") - - if action == "block": - violations: Final = sanitization_result.get("violations", []) - raise HTTPException( - status_code=400, - detail=f"Document blocked by Prompt Security. Violations: {', '.join(violations)}", - ) + self._raise_if_file_blocked(sanitization_result, "Document") if action == "modify": sanitized_content: Final = sanitization_result.get("content", "") diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index a8b33109900..e20f0b320b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -547,7 +547,7 @@ class ToolPermissionGuardrail(CustomGuardrail): for _tool_call, is_allowed, _rule_id, message in checked: if not is_allowed and message is not None: - verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=message, blocked_content=True @@ -809,7 +809,7 @@ class ToolPermissionGuardrail(CustomGuardrail): new_tools: Final = self._collect_request_tools(data) if not new_tools: - verbose_proxy_logger.warning( + verbose_proxy_logger.debug( "Tool Permission Guardrail: not running guardrail. No tools or functions in data" ) return data @@ -820,7 +820,7 @@ class ToolPermissionGuardrail(CustomGuardrail): is_allowed, _, message = self._check_tool_permission(tool_name, tool_type) if not is_allowed and message is not None: - verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": raise HTTPException( status_code=400, diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index ffa322da288..64a47f4f4ff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -19,7 +19,7 @@ from litellm.cost_calculator import _infer_call_type from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.llms import load_guardrail_translation_mappings +from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( @@ -69,6 +69,36 @@ def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTran return translation +def resolve_endpoint_translation( + user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None +) -> "tuple[str, BaseTranslation] | None": + """ + Resolve the endpoint guardrail translation for a streamed response: the + request route wins, falling back to inferring the call type from the first + response chunk (the same resolution order the streaming iterator hook uses). + Returns None when the call type is unresolvable or has no translation. + """ + route_call_types: Final = ( + get_call_types_for_route(user_api_key_dict.request_route) if user_api_key_dict.request_route else None + ) + call_type: Final = ( + route_call_types[0].value + if route_call_types + else ( + _infer_call_type(call_type=None, completion_response=first_response_item) + if first_response_item is not None + else None + ) + ) + if call_type is None: + return None + try: + handler_cls: Final = get_guardrail_translation_mapping(CallTypes(call_type)) + except ValueError: + return None + return call_type, handler_cls() + + def _chunk_choices(item: object) -> Sequence[object]: choices: Final[Sequence[object]] = getattr(item, "choices", None) or [] return choices @@ -343,7 +373,7 @@ class UnifiedLLMGuardrails(CustomLogger): return response - async def _handle_streaming_block( + async def handle_streaming_block( self, exc: "ModifyResponseException", endpoint_translation: _EndpointTranslation, @@ -399,7 +429,7 @@ class UnifiedLLMGuardrails(CustomLogger): return None return call_type - async def _emit_streaming_http_error( + async def emit_streaming_http_error( self, exc: HTTPException, call_type: str | None, @@ -592,7 +622,7 @@ class UnifiedLLMGuardrails(CustomLogger): except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), @@ -601,7 +631,7 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk raise _StreamTerminated() except HTTPException as e: - async for error_item in self._emit_streaming_http_error( + async for error_item in self.emit_streaming_http_error( e, call_type, responses_so_far, @@ -781,7 +811,7 @@ class UnifiedLLMGuardrails(CustomLogger): except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), @@ -869,6 +899,14 @@ class UnifiedLLMGuardrails(CustomLogger): choices: Final = _chunk_choices(item) return any(getattr(choice, "finish_reason", None) is not None for choice in choices) + def resolve_streaming_flag(self, guardrail_to_apply: CustomGuardrail | None, name: str, default: object) -> object: + """Streaming flag resolution order (later wins): default < guardrail + attribute < guardrail_config dict < this callback's optional_params.""" + attribute_value: Final = default if guardrail_to_apply is None else getattr(guardrail_to_apply, name, default) + config: Final = None if guardrail_to_apply is None else getattr(guardrail_to_apply, "guardrail_config", None) + config_value: Final = config.get(name, attribute_value) if isinstance(config, dict) else attribute_value + return self.optional_params.get(name, config_value) + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -897,17 +935,8 @@ class UnifiedLLMGuardrails(CustomLogger): if guardrail_to_apply is None: guardrail_to_apply = request_data.pop("guardrail_to_apply", None) - # Get streaming configuration. Resolution order (later wins): default - # < guardrail attribute < guardrail_config dict < this callback's - # optional_params. def _streaming_flag(name: str, default: object) -> Any: - value = default - if guardrail_to_apply is not None: - value = getattr(guardrail_to_apply, name, value) - config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {}) - if isinstance(config, dict): - value = config.get(name, value) - return self.optional_params.get(name, value) + return self.resolve_streaming_flag(guardrail_to_apply, name, default) sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). @@ -1091,7 +1120,7 @@ class UnifiedLLMGuardrails(CustomLogger): # The current chunk was appended to responses_so_far but not # yet yielded, so exclude it: the continuation must reflect # only what the client has actually received. - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=chunks_yielded, @@ -1101,7 +1130,7 @@ class UnifiedLLMGuardrails(CustomLogger): return except HTTPException as e: # Response already started (we already yielded chunks); cannot send 400. - async for error_item in self._emit_streaming_http_error( + async for error_item in self.emit_streaming_http_error( e, call_type, responses_so_far, @@ -1175,7 +1204,7 @@ class UnifiedLLMGuardrails(CustomLogger): # terminating SSE sequence with the block message rather than # propagating into a bare error blob that truncates the stream. # The withheld original chunks are never released. - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), @@ -1184,7 +1213,7 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk return except HTTPException as e: - async for error_item in self._emit_streaming_http_error( + async for error_item in self.emit_streaming_http_error( e, call_type, responses_so_far, diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 1785a2f0992..bf527e1e868 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -7,7 +7,7 @@ import secrets import time import traceback from collections.abc import Iterable, Mapping -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Any, Final, Literal, TypedDict, cast import fastapi @@ -41,6 +41,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.health_check_latest import LatestHealthCheckRow from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers from litellm.proxy.health_check import ( ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, @@ -209,6 +210,7 @@ services = ( "arize", "galileo", "newrelic", + "pointfive", "sqs", ] | str @@ -296,6 +298,7 @@ async def health_services_endpoint( "arize", "galileo", "newrelic", + "pointfive", "sqs", ]: raise HTTPException( @@ -320,7 +323,7 @@ async def health_services_endpoint( service == "openmeter" or service == "braintrust" or service == "generic_api" - or (service_in_success_callbacks and service != "langfuse") + or (service_in_success_callbacks and service not in ("langfuse", "pointfive")) ): _ = await litellm.acompletion( model="openai/litellm-mock-response-model", @@ -412,6 +415,27 @@ async def health_services_endpoint( ), } + elif service == "pointfive": + if not _is_proxy_admin(user_api_key_dict): + non_admin_detail: Final[_ServiceTestErrorDetail] = { + "error": "Only proxy admins can trigger the PointFive liveness ping." + } + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=non_admin_detail) + from litellm.integrations.pointfive import PointFiveLogger + + try: + pointfive_logger: Final = PointFiveLogger(start_periodic_flush=False) + except ValueError as missing_key: + # No key configured is the answer the operator asked for, not a server error. + no_key: Final[_ServiceTestSuccessResponse] = {"status": "unhealthy", "message": str(missing_key)} + return no_key + response = await pointfive_logger.async_health_check() + pointfive_health: Final[_ServiceTestSuccessResponse] = { + "status": response["status"], + "message": (response["error_message"] if response["status"] == "unhealthy" else "PointFive is healthy") + or "PointFive is healthy", + } + return pointfive_health if service == "webhook": user_info: Final = CallInfo( token=user_api_key_dict.token or "", @@ -724,13 +748,42 @@ def _aggregate_health_check_results( return model_results +class _AggregatedHealthResult(TypedDict): + """One entry of ``_aggregate_health_check_results``: a model's counts for this cycle.""" + + model_name: ReadOnly[str] + model_id: ReadOnly[str | None] + healthy_count: ReadOnly[int] + unhealthy_count: ReadOnly[int] + error_message: ReadOnly[str | None] + + +def _new_health_status(result: _AggregatedHealthResult) -> str: + return "healthy" if result["healthy_count"] > 0 else "unhealthy" + + +def _should_persist_health_check_result( + result: _AggregatedHealthResult, latest_checks_map: Mapping[str, LatestHealthCheckRow] +) -> bool: + """ + True when this result has to be written: no previous row, the status changed, or the + previous row is older than one hour (periodic refresh while the status is stable). + """ + lookup_key: Final = result["model_id"] if result["model_id"] else result["model_name"] + last_check: Final = latest_checks_map.get(lookup_key) + if last_check is None or last_check.status != _new_health_status(result): + return True + time_since_last_check: Final = (datetime.now(timezone.utc) - last_check.checked_at).total_seconds() + return time_since_last_check >= 3600 # 1 hour threshold + + async def _save_health_check_results_if_changed( prisma_client, model_results: dict, latest_checks_map: dict, start_time: float, checked_by: str | None = None, -): +) -> bool: """ Save health check results to database, but only if status changed or >1 hour since last save. @@ -741,47 +794,39 @@ async def _save_health_check_results_if_changed( - Status changes: Immediate write (no delay) - Result: ~92% reduction in DB writes for stable systems, while maintaining real-time updates on changes + The writes are awaited rather than detached so the caller learns whether this cycle's + persistence completed. + Args: prisma_client: Database client model_results: Dictionary of aggregated health check results per model latest_checks_map: Dictionary mapping model_id/model_name to latest health check start_time: Start time of health check for calculating response time checked_by: Identifier for who/what performed the check + + Returns: + True when every row that needed writing was written (including when nothing needed + writing); False when any write failed. """ - for result in model_results.values(): - new_status = "healthy" if result["healthy_count"] > 0 else "unhealthy" - - # Check if we should save this result - should_save = True - lookup_key = result["model_id"] if result["model_id"] else result["model_name"] - if lookup_key in latest_checks_map: - last_check = latest_checks_map[lookup_key] - # Only save if status changed or if it's been a while since last check - if last_check.status == new_status: - # Check if last check was recent (within 1 hour) - if last_check.checked_at: - from datetime import datetime, timezone - - time_since_last_check = (datetime.now(timezone.utc) - last_check.checked_at).total_seconds() - # Only skip if status unchanged AND checked recently (within 1 hour) - # This ensures we still get periodic updates even if status is stable - if time_since_last_check < 3600: # 1 hour threshold - should_save = False - - if should_save: - asyncio.create_task( - prisma_client.save_health_check_result( - model_name=result["model_name"], - model_id=result["model_id"], - status=new_status, - healthy_count=result["healthy_count"], - unhealthy_count=result["unhealthy_count"], - error_message=result["error_message"], - response_time_ms=(time.time() - start_time) * 1000, - details=None, - checked_by=checked_by, - ) - ) + to_write: Final = tuple( + result for result in model_results.values() if _should_persist_health_check_result(result, latest_checks_map) + ) + writes: Final = tuple( + prisma_client.save_health_check_result( + model_name=result["model_name"], + model_id=result["model_id"], + status=_new_health_status(result), + healthy_count=result["healthy_count"], + unhealthy_count=result["unhealthy_count"], + error_message=result["error_message"], + response_time_ms=(time.time() - start_time) * 1000, + details=None, + checked_by=checked_by, + ) + for result in to_write + ) + rows: Final = await asyncio.gather(*writes) + return all(row is not None for row in rows) async def _save_background_health_checks_to_db( @@ -791,7 +836,7 @@ async def _save_background_health_checks_to_db( unhealthy_endpoints: list, start_time: float, checked_by: str | None = None, -): +) -> bool: """ Save background health check results to database for each model. @@ -800,9 +845,13 @@ async def _save_background_health_checks_to_db( OPTIMIZATION: Only saves to database if the status has changed from the last saved check. This dramatically reduces database writes when health status remains stable. + + Returns: + True when this cycle's persistence completed; False when it was skipped or any step + failed. Never raises: a database failure must not break the health check loop. """ if prisma_client is None: - return + return False try: # Step 1: Build mapping from model parameter to model info @@ -825,7 +874,7 @@ async def _save_background_health_checks_to_db( latest_checks_map[key] = check # Step 4: Save aggregated results, but only if status changed - await _save_health_check_results_if_changed( + return await _save_health_check_results_if_changed( prisma_client, model_results, latest_checks_map, @@ -835,6 +884,7 @@ async def _save_background_health_checks_to_db( except Exception as db_error: verbose_proxy_logger.warning("Failed to save background health checks to database: %s", db_error) # Continue execution - don't let database save failure break health checks + return False _PROXY_ADMIN_ROLES: Final = frozenset( @@ -1582,6 +1632,13 @@ async def _show_no_redis_warning() -> bool: return await count_live_proxy_workers(prisma_client) != 1 +def _show_env_credential_login_warning() -> bool: + from litellm.proxy.auth.login_utils import is_env_credential_login_enabled + from litellm.proxy.proxy_server import general_settings + + return is_env_credential_login_enabled(general_settings) + + async def _get_health_readiness_details( response: Response | None = None, ) -> dict[str, Any]: @@ -1623,6 +1680,7 @@ async def _get_health_readiness_details( log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel()) is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG) show_no_redis_warning: Final = await _show_no_redis_warning() + show_env_credential_login_warning: Final = _show_env_credential_login_warning() # check DB if prisma_client is not None: # if db passed in, check if it's connected @@ -1650,6 +1708,7 @@ async def _get_health_readiness_details( "log_level": log_level_name, "is_detailed_debug": is_detailed_debug, "show_no_redis_warning": show_no_redis_warning, + "show_env_credential_login_warning": show_env_credential_login_warning, } else: return { @@ -1662,6 +1721,7 @@ async def _get_health_readiness_details( "log_level": log_level_name, "is_detailed_debug": is_detailed_debug, "show_no_redis_warning": show_no_redis_warning, + "show_env_credential_login_warning": show_env_credential_login_warning, } except Exception as e: raise HTTPException(status_code=503, detail=f"Service Unhealthy ({e})") diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index 32bccca2ab0..1a410593854 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -9,6 +9,7 @@ the reservation is refunded when the batch reaches a terminal state """ import asyncio +import logging import math import time import uuid @@ -19,6 +20,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.caching.redis_cache import log_redis_failure from litellm.constants import BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, BATCH_ENQUEUED_TOKEN_TTL_SECONDS from litellm.proxy._types import UserAPIKeyAuth @@ -233,8 +235,11 @@ class BatchEnqueuedTokenStore: try: return await self._reserve_via_redis(reserve_script, refund_script, tokens=tokens, scopes=scopes) except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters - verbose_proxy_logger.warning( - "Redis enqueued-token reserve failed, falling back to in-memory: %s", str(e) + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "Redis enqueued-token reserve failed, falling back to in-memory", + e, ) return await self._reserve_in_memory(tokens=tokens, scopes=scopes, span=litellm_parent_otel_span) @@ -374,8 +379,11 @@ class BatchEnqueuedTokenStore: (serialized, ttl), ) except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record - verbose_proxy_logger.warning( - "Redis enqueued-token reservation save failed, falling back to in-memory: %s", str(e) + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "Redis enqueued-token reservation save failed, falling back to in-memory", + e, ) else: return @@ -421,8 +429,11 @@ class BatchEnqueuedTokenStore: await pop_script((self._record_key(batch_id),), (BATCH_ENQUEUED_TOKEN_TTL_SECONDS,)) ) except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record - verbose_proxy_logger.warning( - "Redis enqueued-token reservation pop failed, falling back to in-memory: %s", str(e) + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "Redis enqueued-token reservation pop failed, falling back to in-memory", + e, ) return None diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index de8834449de..0339cf4dfea 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -32,6 +32,10 @@ from litellm.proxy.hooks.rate_limiter_utils import ( resolve_llm_provider_for_rate_limit, ) from litellm.proxy.utils import InternalUsageCache +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral @@ -55,6 +59,10 @@ def _get_priority_settings() -> "PriorityReservationSettings": return settings +def _is_latin1_encodable(value: object) -> bool: + return all(ord(char) < 256 for char in str(value)) + + class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): """ Saturation-aware priority-based rate limiter using v3 infrastructure. @@ -659,22 +667,18 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): data=data, user_api_key_dict=user_api_key_dict, response=response ) - # Add additional priority-specific headers - if isinstance(response, ModelResponse): + if response_has_hidden_params(response): priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) - - # Get existing additional headers - additional_headers: Final = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {} - - # Add priority information - additional_headers["x-litellm-priority"] = priority or "default" + additional_headers: Final = ensure_response_additional_headers(response) + priority_header: Final = priority or "default" + if _is_latin1_encodable(priority_header): + additional_headers["x-litellm-priority"] = priority_header + else: + verbose_proxy_logger.debug( + "Skipping x-litellm-priority header: priority %r is not Latin-1 encodable", priority + ) additional_headers["x-litellm-rate-limiter-version"] = "v3" - # Update response - if not hasattr(response, "_hidden_params"): - response._hidden_params = {} - response._hidden_params["additional_headers"] = additional_headers - return response except Exception as e: diff --git a/litellm/proxy/hooks/litellm_skills/__init__.py b/litellm/proxy/hooks/litellm_skills/__init__.py index 751122ac51c..d24cdd37161 100644 --- a/litellm/proxy/hooks/litellm_skills/__init__.py +++ b/litellm/proxy/hooks/litellm_skills/__init__.py @@ -21,10 +21,7 @@ from litellm.llms.litellm_proxy.skills import ( code_execution_handler, get_litellm_code_execution_tool, ) -from litellm.proxy.hooks.litellm_skills.main import ( - SkillsInjectionHook, - skills_injection_hook, -) +from litellm.proxy.hooks.litellm_skills.main import SkillsInjectionHook __all__ = [ "LITELLM_CODE_EXECUTION_TOOL", @@ -35,5 +32,4 @@ __all__ = [ "SkillsSandboxExecutor", "code_execution_handler", "get_litellm_code_execution_tool", - "skills_injection_hook", ] diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 9edbc6dbf1c..8848878d15b 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -29,6 +29,7 @@ import json from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol +import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger @@ -493,7 +494,6 @@ class SkillsInjectionHook(CustomLogger): Returns the final response with generated files inline. """ - import litellm from litellm.llms.litellm_proxy.skills.code_execution import ( LiteLLMInternalTools, ) @@ -723,7 +723,6 @@ print('No executable skill module found') Returns the final response with generated files inline. """ - import litellm from litellm.llms.litellm_proxy.skills.code_execution import ( LiteLLMInternalTools, ) @@ -913,11 +912,3 @@ print('No executable skill module found') verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to response", len(generated_files)) return response - - -# Global instance for registration -skills_injection_hook: Final = SkillsInjectionHook() - -import litellm - -litellm.logging_callback_manager.add_litellm_callback(skills_injection_hook) diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 0b8e4e65258..e07b96e5773 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -14,11 +14,13 @@ Works across multiple proxy instances via DualCache (in-memory + Redis). Follows the same pattern as max_iterations_limiter.py. """ +import logging import os from typing import TYPE_CHECKING, Any, Final from litellm import DualCache from litellm._logging import verbose_proxy_logger +from litellm.caching.redis_cache import log_redis_failure from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -215,9 +217,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return float(result) return 0.0 except Exception as e: - verbose_proxy_logger.warning( - "MaxBudgetPerSessionHandler: Redis GET failed, falling back to in-memory: %s", - str(e), + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "MaxBudgetPerSessionHandler: Redis GET failed, falling back to in-memory", + e, ) result = await self.internal_usage_cache.async_get_cache( @@ -239,9 +243,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): ) return float(result) except Exception as e: - verbose_proxy_logger.warning( - "MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, falling back to in-memory: %s", - str(e), + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, falling back to in-memory", + e, ) return await self._in_memory_increment_spend(cache_key, amount) diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 8b82842353c..bec12ba8201 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,7 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference. Reduces context window size and improves tool selection accuracy. """ -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException @@ -104,7 +104,7 @@ class SemanticToolFilterHook(CustomLogger): ) # Parse to separate MCP tools from other tools - mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + mcp_tools, _ = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_tools: return [] @@ -173,7 +173,11 @@ class SemanticToolFilterHook(CustomLogger): return [name for name in names if name] @staticmethod - def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]: + async def _narrow_mcp_references( + tools: Sequence[Mapping[str, object]], + selected_tool_names: list[str], + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] | None = None, + ) -> list[object]: """ Restrict each litellm_proxy MCP reference to the semantically selected tools. @@ -192,13 +196,14 @@ class SemanticToolFilterHook(CustomLogger): LiteLLM_Proxy_MCP_Handler, ) + via_gateway: Final = await ( + LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools, served_names) + if served_names is not None + else LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools) + ) return [ - ( - {**tool, "allowed_tools": selected_tool_names} - if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool]) - else tool - ) - for tool in tools + {**tool, "allowed_tools": selected_tool_names} if isinstance(tool, dict) and routed else tool + for tool, routed in zip(tools, via_gateway, strict=True) ] def _is_mcp_tool(self, tool: object) -> bool: @@ -325,7 +330,7 @@ class SemanticToolFilterHook(CustomLogger): filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools) selected_tool_names: Final = self._selected_tool_names(filtered_expanded_tools) - narrowed_tools: Final = self._narrow_mcp_references(tools, selected_tool_names) + narrowed_tools: Final = await self._narrow_mcp_references(tools, selected_tool_names) data["tools"] = narrowed_tools self._emit_filter_metadata_safe( data=data, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 31437af7770..c398abff099 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -6,6 +6,7 @@ This is currently in development and not yet ready for production. import asyncio import binascii +import logging import os import uuid from collections.abc import Awaitable, Callable, Mapping, Sequence, Set @@ -26,11 +27,13 @@ from typing_extensions import NotRequired, ReadOnly from litellm import DualCache from litellm._logging import verbose_proxy_logger +from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( ESTIMATED_OUTPUT_TOKENS_FIELD, @@ -52,6 +55,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.utils import ( @@ -1223,7 +1230,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) all_cache_values.extend(group_cache_values) except Exception as e: - verbose_proxy_logger.warning("Redis Lua script failed for hash tag %s: %s", hash_tag, e) + log_redis_failure( + verbose_proxy_logger, logging.WARNING, f"Redis Lua script failed for hash tag {hash_tag}", e + ) # Fallback to in-memory cache for this group group_cache_values = await self.in_memory_cache_sliding_window( keys=group_keys, @@ -1470,7 +1479,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) counts = [max(0, int(value)) for value in raw_counts] except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror, never a 500 - verbose_proxy_logger.warning("parallel_count_script failed, using local mirror: %s", e) + log_redis_failure( + verbose_proxy_logger, logging.WARNING, "parallel_count_script failed, using local mirror", e + ) counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) else: counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) @@ -1500,7 +1511,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ], ) except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to in-memory enforcement, never a 500 - verbose_proxy_logger.warning("parallel_acquire_script failed, falling back to in-memory gauge: %s", e) + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "parallel_acquire_script failed, falling back to in-memory gauge", + e, + ) async with self._check_and_increment_lock: return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span) if int(raw[0]) == 1: @@ -1626,7 +1642,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the in-memory release, never a 500 - verbose_proxy_logger.warning("parallel_release_script failed, falling back to in-memory release: %s", e) + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "parallel_release_script failed, falling back to in-memory release", + e, + ) async with self._check_and_increment_lock: for counter_key in counter_keys: @@ -1809,12 +1830,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # state ambiguous. Refund any prior groups so Redis returns # to its pre-call state, then fall back to in-memory for the # whole call (counters there are independent of Redis). - verbose_proxy_logger.error( - "atomic_check_and_increment_by_n: Redis Lua execution failed (%s: %s). Refunding %s prior descriptors and falling back to in-memory enforcement — counters will diverge from Redis until window expires (window_size=%ss).", - type(e).__name__, + log_redis_failure( + verbose_proxy_logger, + logging.ERROR, + f"atomic_check_and_increment_by_n: Redis Lua execution failed ({type(e).__name__}). Refunding " + f"{len(applied)} prior descriptors and falling back to in-memory enforcement, counters will " + f"diverge from Redis until window expires (window_size={self.window_size}s)", e, - len(applied), - self.window_size, ) await self._refund_applied_descriptor_groups(applied) flat_meta: list[AtomicCounterMeta] = [m for _k, _a, group_meta in descriptor_groups for m in group_meta] @@ -1861,8 +1883,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): value=-entry["increment"], ) except Exception as e: - verbose_proxy_logger.warning( - "Failed to refund %s on cross-descriptor rollback: %s", entry["counter_key"], e + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + f"Failed to refund {entry['counter_key']} on cross-descriptor rollback", + e, ) def _build_atomic_response( @@ -3303,7 +3328,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, ) - raw_estimated_input_tokens: Final = self._estimate_precise_input_tokens( + raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)( data=data, model=requested_model, call_type=call_type ) estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1) @@ -3851,7 +3876,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) except Exception as e: - verbose_proxy_logger.warning("TTL preservation failed, falling back to regular pipeline: %s", e) + log_redis_failure( + verbose_proxy_logger, logging.WARNING, "TTL preservation failed, falling back to regular pipeline", e + ) # Fallback to regular pipeline on error await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( increment_list=pipeline_operations, @@ -3917,9 +3944,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) continue except Exception as e: # noqa: BLE001 # Redis failures use the plain increment fallback - verbose_proxy_logger.warning( - "Window-guarded token adjustment failed for %s: %s", - operation["key"], + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + f"Window-guarded token adjustment failed for {operation['key']}", e, ) if operation["increment_value"] > 0: @@ -4677,34 +4705,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Post-call hook to update rate limit headers in the response. """ try: - from pydantic import BaseModel - stash: Final = get_request_stash() litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None - if litellm_proxy_rate_limit_response is not None: - # Update response headers - if hasattr(response, "_hidden_params"): - _hidden_params = getattr(response, "_hidden_params") - else: - _hidden_params = None - - if _hidden_params is not None and ( - isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict) - ): - if isinstance(_hidden_params, BaseModel): - _hidden_params = _hidden_params.model_dump() - - _additional_headers: Final = self._merge_ratelimit_statuses_into_additional_headers( - additional_headers=_hidden_params.get("additional_headers", {}) or {}, + if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response): + additional_headers: Final = ensure_response_additional_headers(response) + additional_headers.update( + self._merge_ratelimit_statuses_into_additional_headers( + additional_headers={}, statuses=litellm_proxy_rate_limit_response["statuses"], ) - - setattr( - response, - "_hidden_params", - {**_hidden_params, "additional_headers": _additional_headers}, - ) + ) except Exception as e: verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index d6276c03155..1e946cc2e23 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,11 +1,12 @@ import asyncio import traceback -from collections.abc import Sequence +from collections.abc import Callable, Sequence from datetime import datetime 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 ( @@ -22,10 +23,21 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.db_spend_update_writer import ( + DBSpendUpdateWriter, debitable_model_access_groups, get_llm_router, ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.spend_event import ( + ObjectMapping, + SpendEventBuildError, + SpendEventDecodeError, + build_spend_event, + decode_spend_event, + is_offloadable_success, + spend_event_callback_args, +) +from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer from litellm.proxy.spend_tracking.spend_log_error_logger import ( should_suppress_spend_log_tracebacks, spend_log_error, @@ -33,10 +45,12 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, get_request_model_access_groups, + should_store_prompts_and_responses_in_spend_logs, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( CallTypes, + LiteLLMBatch, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) @@ -68,9 +82,52 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) +def _proxy_spend_writer() -> DBSpendUpdateWriter: + from litellm.proxy.proxy_server import proxy_logging_obj + + return proxy_logging_obj.db_spend_update_writer + + class _ProxyDBLogger(CustomLogger): - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) + def __init__( + self, + spend_event_producer: SpendEventProducer | None = None, + *, + turn_off_message_logging: bool = False, + message_logging: bool = True, + spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer, + ) -> None: + super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging) + self.spend_event_producer = spend_event_producer + self._spend_writer: Final = spend_writer + + async def async_log_success_event( + self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if self.spend_event_producer is None or not is_offloadable_success(response_obj): + await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) + return + event: Final = build_spend_event( + kwargs, + response_obj, + start_time, + end_time, + store_bodies=should_store_prompts_and_responses_in_spend_logs(), + ) + if isinstance(event, SpendEventBuildError): + verbose_proxy_logger.warning("collector: tracking cost in-process, event not buildable: %s", event.reason) + await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) + return + await self.spend_event_producer.publish(event) + + async def run_spend_event(self, line: bytes) -> None: + """Run the unchanged cost pipeline on a serialized spend event (sidecar consumer and in-process fallback).""" + event: Final = decode_spend_event(line) + if isinstance(event, SpendEventDecodeError): + verbose_proxy_logger.error("collector: discarding undecodable spend event: %s", event.reason) + return + args: Final = spend_event_callback_args(event) + await self._PROXY_track_cost_callback(args.kwargs, args.response_obj, args.start_time, args.end_time) async def async_post_call_failure_hook( self, @@ -102,8 +159,6 @@ class _ProxyDBLogger(CustomLogger): ): return - from litellm.proxy.proxy_server import proxy_logging_obj - _metadata = dict( LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) ) @@ -179,13 +234,12 @@ class _ProxyDBLogger(CustomLogger): if request_data.get("litellm_trace_id") is None: request_data["litellm_trace_id"] = getattr(_litellm_logging_obj, "litellm_trace_id", None) - # Use the actual request start time from the logging object so that - # failed requests record the real duration instead of 0. - actual_start_time = datetime.now() - if _litellm_logging_obj is not None: - obj_start: Final = getattr(_litellm_logging_obj, "start_time", None) - if obj_start is not None: - actual_start_time = obj_start + lifted_start_time: Final = request_data.get("start_time") + actual_start_time: Final = ( + lifted_start_time + if isinstance(lifted_start_time, datetime) + else getattr(_litellm_logging_obj, "start_time", None) or datetime.now() + ) # A stream that broke mid-flight still billed the provider for the # chunks already delivered. ``post_call_failure_hook`` lifts that @@ -201,7 +255,7 @@ class _ProxyDBLogger(CustomLogger): existing_metadata.get("standard_logging_guardrail_information") ) - await proxy_logging_obj.db_spend_update_writer.update_database( + await self._spend_writer().update_database( token=user_api_key_dict.api_key, response_cost=recovered_response_cost, user_id=user_api_key_dict.user_id, @@ -248,6 +302,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 +351,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 +368,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.) @@ -487,6 +555,10 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value +async def run_spend_event(line: bytes) -> None: + await _ProxyDBLogger().run_spend_event(line) + + def _is_unbilled_interaction_response(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import missing_usage_is_expected from litellm.types.interactions import InteractionsAPIResponse @@ -578,9 +650,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 +677,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 +705,7 @@ async def _update_database_and_spend_counters( finally: budget_reservation["finalized"] = True raise + return True async def _release_budget_reservation(budget_reservation: dict | None) -> None: diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index c4c15c40d1e..7e7f70d6f7e 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -32,6 +32,29 @@ if TYPE_CHECKING: _RESPONSES_API_PROVIDER_PREFIX: Final = "/openai" _RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) +_ADDRESSED_RESPONSE_ID_KEY: Final = "_litellm_addressed_response_id" +_UNMANAGED_RESPONSE_ID_DETAIL: Final = ( + "Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. " + "To let keys address responses this proxy did not issue, set " + "general_settings::allow_unmanaged_response_ids to True in the config.yaml file." +) +_PROXY_ADMIN_ROLES: Final = frozenset({LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN.value}) + + +def _proxy_general_settings() -> Mapping[str, Any]: + from litellm.proxy.proxy_server import general_settings + + return general_settings + + +def _proxy_signing_key() -> str | None: + import os + + from litellm.proxy.proxy_server import master_key + + salt_key: Final = os.getenv("LITELLM_SALT_KEY", None) + return master_key if salt_key is None else salt_key + _RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -83,8 +106,13 @@ def _is_responses_api_create_route(request_route: str | None) -> bool: class ResponsesIDSecurity(CustomLogger): - def __init__(self): - pass + def __init__( + self, + general_settings_reader: Callable[[], Mapping[str, Any]] = _proxy_general_settings, + signing_key_reader: Callable[[], str | None] = _proxy_signing_key, + ) -> None: + self._general_settings_reader: Final = general_settings_reader + self._signing_key_reader: Final = signing_key_reader async def async_pre_call_hook( self, @@ -103,30 +131,51 @@ class ResponsesIDSecurity(CustomLogger): } if call_type not in responses_api_call_types: return None - if call_type == "aresponses": - # check 'previous_response_id' if present in the data - previous_response_id: Final = data.get("previous_response_id") - if previous_response_id and self._is_encrypted_response_id(previous_response_id): - original_response_id, user_id, team_id = self._decrypt_response_id(previous_response_id) - self.check_user_access_to_response_id(user_id, team_id, user_api_key_dict) - data["previous_response_id"] = original_response_id - elif call_type in {"aget_responses", "adelete_responses", "acancel_responses", "alist_input_items"}: - response_id: Final = data.get("response_id") - - if response_id and self._is_encrypted_response_id(response_id): - original_response_id, user_id, team_id = self._decrypt_response_id(response_id) - - self.check_user_access_to_response_id(user_id, team_id, user_api_key_dict) - data["response_id"] = original_response_id + addressed_id_field: Final = "previous_response_id" if call_type == "aresponses" else "response_id" + retained_id: Final = data.get(_ADDRESSED_RESPONSE_ID_KEY) + addressed_id: Final = ( + retained_id if isinstance(retained_id, str) and retained_id else data.get(addressed_id_field) + ) + if not isinstance(addressed_id, str) or not addressed_id: + return data + authorized_id: Final = self._authorize_response_id(addressed_id, user_api_key_dict) + data[addressed_id_field] = authorized_id + data[_ADDRESSED_RESPONSE_ID_KEY] = addressed_id return data + def _authorize_response_id( + self, + response_id: str, + user_api_key_dict: "UserAPIKeyAuth", + ) -> str: + if self._is_encrypted_response_id(response_id): + original_response_id, user_id, team_id = self._decrypt_response_id(response_id) + self.check_user_access_to_response_id(user_id, team_id, user_api_key_dict) + return original_response_id + + if self._unmanaged_response_ids_allowed(user_api_key_dict): + return response_id + + raise HTTPException(status_code=403, detail=_UNMANAGED_RESPONSE_ID_DETAIL) + + def _unmanaged_response_ids_allowed(self, user_api_key_dict: "UserAPIKeyAuth") -> bool: + general_settings: Final = self._general_settings_reader() + + if general_settings.get("disable_responses_id_security", False): + return True + if general_settings.get("allow_unmanaged_response_ids", False): + return True + if self._get_signing_key() is None: + return True + return user_api_key_dict.user_role in _PROXY_ADMIN_ROLES + def check_user_access_to_response_id( self, response_id_user_id: str | None, response_id_team_id: str | None, user_api_key_dict: "UserAPIKeyAuth", ) -> bool: - from litellm.proxy.proxy_server import general_settings + general_settings: Final = self._general_settings_reader() if ( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value @@ -219,15 +268,7 @@ class ResponsesIDSecurity(CustomLogger): return response_id, None, None def _get_signing_key(self) -> str | None: - """Get the signing key for encryption/decryption.""" - import os - - from litellm.proxy.proxy_server import master_key - - salt_key = os.getenv("LITELLM_SALT_KEY", None) - if salt_key is None: - salt_key = master_key - return salt_key + return self._signing_key_reader() def _encrypt_response_id( self, @@ -274,7 +315,7 @@ class ResponsesIDSecurity(CustomLogger): This method adds response IDs to an in-memory queue, which are then batch-processed by the DBSpendUpdateWriter during regular database update cycles. """ - from litellm.proxy.proxy_server import general_settings + general_settings: Final = self._general_settings_reader() if general_settings.get("disable_responses_id_security", False): return response @@ -288,7 +329,7 @@ class ResponsesIDSecurity(CustomLogger): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict ) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]: - from litellm.proxy.proxy_server import general_settings + general_settings: Final = self._general_settings_reader() # Create a request-scoped cache for consistent encryption across streaming chunks. request_encryption_cache: Final[dict[str, str]] = {} diff --git a/litellm/proxy/hooks/sensitive_data_routing.py b/litellm/proxy/hooks/sensitive_data_routing.py index 4d846744b55..bc89dec7a11 100644 --- a/litellm/proxy/hooks/sensitive_data_routing.py +++ b/litellm/proxy/hooks/sensitive_data_routing.py @@ -10,11 +10,13 @@ this hook manages: Works across multiple proxy instances via DualCache (in-memory + Redis). """ +import logging import os from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_guardrail import get_session_id_from_request_data from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -96,9 +98,11 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger): ) return routed_model except Exception as e: - verbose_proxy_logger.warning( - "SensitiveDataRoutingHandler: Redis GET failed, falling back to in-memory: %s", - str(e), + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "SensitiveDataRoutingHandler: Redis GET failed, falling back to in-memory", + e, ) result = await self.internal_usage_cache.async_get_cache( @@ -142,9 +146,11 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger): ttl=self.ttl, ) except Exception as e: - verbose_proxy_logger.warning( - "SensitiveDataRoutingHandler: Redis SET failed, falling back to in-memory: %s", - str(e), + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "SensitiveDataRoutingHandler: Redis SET failed, falling back to in-memory", + e, ) await self.internal_usage_cache.async_set_cache( diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 06f99e4ae9c..3f044855ce8 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -20,6 +20,11 @@ from litellm.proxy.common_utils.http_parsing_utils import ( coerce_numeric_form_fields, numeric_form_fields, ) +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.proxy.route_llm_request import route_request from litellm.types.images.main import ImageEditRequestParams from litellm.types.llms.openai import ChatCompletionUserMessage @@ -200,18 +205,18 @@ async def image_generation( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), openai_code=getattr(e, "code", None), - code=getattr(e, "status_code", 500), + code=error_status_code(e, 500), ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d026c5510e6..8f7b515c22a 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -10,6 +10,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast from fastapi import HTTPException, Request +from pydantic import TypeAdapter from pydantic import ValidationError as PydanticValidationError from starlette.datastructures import Headers @@ -18,15 +19,19 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm._uuid import uuid from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + ROUTING_REQUEST_TAGS_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, + X_LITELLM_DISABLE_CALLBACKS, ) +from litellm.litellm_core_utils.core_helpers import is_codex_user_agent from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, @@ -79,10 +84,6 @@ _EXPLICIT_SESSION_HEADERS: Final = frozenset({"x-litellm-trace-id", "x-litellm-s # ``session-id``/``thread-id``; builds before the codex-api split sent # ``session_id``/``conversation_id``. Ordered session before thread. _CODEX_SESSION_ID_HEADERS: Final = ("session-id", "session_id", "thread-id", "conversation_id") -# Matches every first-party Codex originator: codex-tui, codex_cli_rs, codex_exec, -# codex_vscode, "Codex ...". A separator is required so an unrelated "codexfoo" client -# does not read as Codex. -_CODEX_CLIENT_PREFIX_RE: Final = re.compile(r"^codex[-_ /]", re.IGNORECASE) # Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores # (covers UUIDs and most common session-id formats). _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") @@ -157,6 +158,7 @@ from litellm.types.utils import ( CustomPricingLiteLLMParams, LlmProviders, ProviderSpecificHeader, + StandardCallbackDynamicParams, StandardLoggingUserAPIKeyMetadata, SupportedCacheControls, ) @@ -170,6 +172,8 @@ _ENABLE_TEAM_STALE_ALIAS_BYPASS: bool | None = None if TYPE_CHECKING: + from litellm.integrations.otel.model.destination import OtelDestination + from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext @@ -229,6 +233,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "applied_policies", "policy_sources", "guardrail_scan_ids", + "guardrail_scan_metadata", "routing_decision", GATEWAY_INJECTED_CACHE_METADATA_KEY, "pillar_response_headers", @@ -285,10 +290,12 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( "applied_policies", "policy_sources", "guardrail_scan_ids", + "guardrail_scan_metadata", "routing_decision", GATEWAY_INJECTED_CACHE_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, + ROUTING_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, "standard_logging_object", "proxy_server_request", @@ -325,7 +332,9 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg # ``attempted_fallbacks`` and ``original_model_group`` are written by the router # and read by spend logs as fact; a client value has no legitimate meaning and no # key or team setting keeps it, so the strip is never gated. -_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset({"attempted_fallbacks", "original_model_group"}) +_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( + {"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY} +) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" # Request fields whose value, when URL-valued, becomes the outbound destination @@ -762,6 +771,16 @@ def apply_missing_session_id_policy( return if policy == "omit": metadata[SESSION_ID_OMITTED_METADATA_KEY] = True + requester_metadata: Final = data.get("metadata") + requester_session_id: Final = ( + requester_metadata.get("session_id") if isinstance(requester_metadata, dict) else None + ) + if ( + (body_session_id := data.get("litellm_session_id")) + and not metadata.get("session_id") + and not requester_session_id + ): + metadata["session_id"] = body_session_id return if data.get("litellm_session_id") or metadata.get("session_id"): return @@ -789,28 +808,14 @@ def apply_missing_session_id_policy( ) -def is_claude_code_user_agent(user_agent: str) -> bool: - """Claude Code identifies itself as ``claude-cli/ ...``; the IDE - extensions and the Agent SDK run through the same CLI and share that prefix.""" - return user_agent.startswith("claude-cli/") - - -def is_codex_user_agent(user_agent: str) -> bool: - """Codex builds its user agent as ``/ ...`` and ships - several first-party originators: ``codex-tui``, ``codex_cli_rs``, - ``codex_exec`` (exec mode), ``codex_vscode`` (IDE extension) and ``Codex ...`` - (see ``is_first_party_originator`` in codex-rs). They agree only on the - ``codex`` stem, and the TUI sends a bare ``codex-tui`` with no version at all, - so match the stem plus a separator rather than any one spelling.""" - return bool(_CODEX_CLIENT_PREFIX_RE.match(user_agent)) - - def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_config: ProxyConfig) -> bool: """drop_params defaults to on for agentic CLIs so their client-specific params (e.g. Claude Code's thinking, Codex's service_tier) don't fail requests routed to providers that reject them. An explicit drop_params from the caller or in the operator's ``litellm_settings`` always wins over this default.""" + from litellm.llms.anthropic.common_utils import is_claude_code_user_agent + if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)): return False if "drop_params" in data: @@ -974,6 +979,145 @@ def _get_dynamic_logging_metadata( return callback_settings_obj +_TENANT_OTEL_PARAMS: Final = TypeAdapter(StandardCallbackDynamicParams) + + +def _tenant_otel_params(callback_vars: Mapping[str, str]) -> StandardCallbackDynamicParams: + try: + return _TENANT_OTEL_PARAMS.validate_python(callback_vars) + except PydanticValidationError: + return StandardCallbackDynamicParams() + + +_NO_REQUEST_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) + + +def _dynamically_disabled_backends( + user_api_key_dict: UserAPIKeyAuth, + request_headers: Mapping[str, str] | None, +) -> frozenset[str]: + """The callbacks this request turned off, read the way dispatch reads them. + + Same sources, precedence, and premium gate ``EnterpriseCallbackControls`` applies + before it skips a callback: the ``x-litellm-disable-callbacks`` header wins over the + key's stored list, team settings are not a source, and a non-premium proxy honours + neither. A destination has to agree with that decision, or a backend the key turned + off would still be exported to, now through the fan-out instead of the callback. + """ + from litellm.proxy.proxy_server import premium_user + + if litellm.allow_dynamic_callback_disabling is not True or not premium_user: + return frozenset() + header: Final = (request_headers if request_headers is not None else _NO_REQUEST_HEADERS).get( + X_LITELLM_DISABLE_CALLBACKS + ) + if header is not None: + return frozenset(name.strip().lower() for name in header.split(",")) + metadata: Final = user_api_key_dict.metadata + disabled: Final = metadata.get("litellm_disabled_callbacks") if metadata else None + if not isinstance(disabled, list): + return frozenset() + return frozenset(name.lower() for name in disabled if isinstance(name, str)) + + +def resolve_tenant_otel_destinations( + user_api_key_dict: UserAPIKeyAuth, + request_headers: Mapping[str, str] | None = None, +) -> "tuple[OtelDestination, ...]": + """The OTLP destinations this request's key or team config overrides its traces to. + + Key settings win over team settings outright, the same precedence + ``_get_dynamic_logging_metadata`` applies, so one caller never exports the same + backend to two accounts. An empty key-level list counts as configured, since that + is what disabling a key's callbacks writes. Returns empty when OTEL V2 is off, when + neither level named a destination-capable backend, or when the config is + incomplete, and the request then keeps the operator's own exporters. + + Two entries naming the same backend merge their ``callback_vars`` last-wins, the + way ``convert_key_logging_metadata_to_callback`` merges them, so the destination + and the per-request tracer routing cannot read one config two ways. + + A ``failure``-only entry is skipped: a destination is resolved during auth, before + the request has an outcome, so honouring the filter would mean holding every span + back until the call finishes. Those entries keep today's behaviour instead, where + the tenant's credentials reach the backend through per-request tracer routing and + the operator's exporter is left alone. Its ``callback_vars`` still take part in the + merge for a backend another entry made eligible, so the destination carries the + same credentials the runtime parser resolves for that request. + + A backend the request disabled dynamically, through the key's + ``litellm_disabled_callbacks`` or the ``x-litellm-disable-callbacks`` header in + ``request_headers``, resolves to no destination, so the fan-out never carries the + request tree to that account and the operator's exporter is never suppressed for + it. That leaves the request exactly where it stood before destinations existed: + the OTel V2 logger itself is not on the disable list's class registry, so its own + span still routes to the tenant's credentials the way it did then. + """ + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.integrations.otel.presets.destinations import destination_for + + if not is_otel_v2_enabled(): + return () + key_entries: Final = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) + entries: Final = ( + key_entries + if key_entries is not None + else KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) + ) + if not entries: + return () + disabled: Final = _dynamically_disabled_backends(user_api_key_dict, request_headers) + callbacks: Final = tuple( + callback + for item in entries + if (callback := _get_validated_callback_metadata(item=item, source="otel-destination")) is not None + if callback.callback_name.lower() not in disabled + ) + return tuple( + destination + for name in dict.fromkeys( + callback.callback_name for callback in callbacks if callback.callback_type != "failure" + ) + if ( + destination := destination_for( + name, + _tenant_otel_params( + MappingProxyType( + { + var: value + for callback in callbacks + if callback.callback_name == name + for var, value in callback.callback_vars.items() + } + ) + ), + _tenant_service_name(user_api_key_dict), + ) + ) + is not None + ) + + +def _tenant_service_name(user_api_key_dict: UserAPIKeyAuth) -> str | None: + """The ``service.name`` this key or team configured, the key winning over its team. + + Same fields and same precedence the request-metadata build applies, read straight + off the auth object because destinations resolve during auth, before that metadata + is assembled. + """ + sources: Final = (user_api_key_dict.metadata, user_api_key_dict.team_metadata) + return next( + ( + stripped + for source in sources + if source + for field in OTEL_SERVICE_NAME_METADATA_KEYS + if isinstance(value := source.get(field), str) and (stripped := value.strip()) + ), + None, + ) + + def clean_headers( headers: Headers, litellm_key_header_name: str | None = None, @@ -1604,7 +1748,9 @@ class LiteLLMProxyRequestSetup: callback_vars_dict.pop("success_callback", None) callback_vars_dict.pop("failure_callback", None) callback_vars_dict = { - key: (litellm.utils.get_secret(value, default_value=value) or value if isinstance(value, str) else value) + key: ( + litellm.utils.get_secret(value, default_value=value) or value if isinstance(value, str) else str(value) + ) for key, value in callback_vars_dict.items() } @@ -1870,6 +2016,7 @@ async def add_litellm_data_to_request( "method": request.method, "headers": _logging_safe_headers, "body": None, # filled in post-strip; see below + "credential_fields": tuple(sorted(name for name in _TRANSPORT_ONLY_CREDENTIAL_KEYS if name in data)), "arrival_time": arrival_time, # Track when request arrived at proxy } @@ -2882,6 +3029,28 @@ def _add_guardrails_from_policies_in_metadata( ) +def add_guardrails_from_auth_metadata( + user_api_key_dict: UserAPIKeyAuth, + data: dict, # mutable-ok: writes guardrails into the live request dict, same contract as the helpers it wraps + metadata_variable_name: str, +) -> None: + """Resolve key, team, and project guardrails, direct and via policies, onto the request metadata.""" + _add_guardrails_from_key_or_team_metadata( + key_metadata=user_api_key_dict.metadata, + team_metadata=user_api_key_dict.team_metadata, + project_metadata=user_api_key_dict.project_metadata, + data=data, + metadata_variable_name=metadata_variable_name, + ) + _add_guardrails_from_policies_in_metadata( + key_metadata=user_api_key_dict.metadata, + team_metadata=user_api_key_dict.team_metadata, + project_metadata=user_api_key_dict.project_metadata, + data=data, + metadata_variable_name=metadata_variable_name, + ) + + async def move_guardrails_to_metadata( data: dict, _metadata_variable_name: str, @@ -2914,22 +3083,8 @@ async def move_guardrails_to_metadata( data.pop("policies", None) return - # Check key/team/project-level guardrails - _add_guardrails_from_key_or_team_metadata( - key_metadata=user_api_key_dict.metadata, - team_metadata=user_api_key_dict.team_metadata, - project_metadata=project_metadata, - data=data, - metadata_variable_name=_metadata_variable_name, - ) - - ######################################################################################### - # Add guardrails from policies attached to key/team/project metadata - ######################################################################################### - _add_guardrails_from_policies_in_metadata( - key_metadata=user_api_key_dict.metadata, - team_metadata=user_api_key_dict.team_metadata, - project_metadata=project_metadata, + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_dict, data=data, metadata_variable_name=_metadata_variable_name, ) @@ -2988,6 +3143,7 @@ def _match_and_track_policies( context: "PolicyMatchContext", request_body_policies: Sequence[str], policies_override: dict[str, "Policy"] | None = None, + attachment_registry_override: "AttachmentRegistry | None" = None, ) -> tuple[list[str], dict[str, str]]: """ Match policies via attachments and request body, track them in metadata. @@ -3004,7 +3160,9 @@ def _match_and_track_policies( from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher # Get matching policies via attachments (with match reasons for attribution) - attachment_registry: Final = get_attachment_registry() + attachment_registry: Final = ( + attachment_registry_override if attachment_registry_override is not None else get_attachment_registry() + ) matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(context) matching_policy_names: Final = [m["policy_name"] for m in matches_with_reasons] policy_reasons: Final = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons} @@ -3012,9 +3170,11 @@ def _match_and_track_policies( verbose_proxy_logger.debug("Policy engine: matched policies via attachments: %s", matching_policy_names) # Combine attachment-based policies with dynamic request body policies - all_policy_names: Final = set(matching_policy_names) - if request_body_policies and isinstance(request_body_policies, list): - all_policy_names.update(request_body_policies) + request_body_policies_list: Final = ( + tuple(request_body_policies) if request_body_policies and isinstance(request_body_policies, list) else () + ) + all_policy_names: Final = tuple(dict.fromkeys((*matching_policy_names, *request_body_policies_list))) + if request_body_policies_list: verbose_proxy_logger.debug("Policy engine: added dynamic policies from request body: %s", request_body_policies) if not all_policy_names: @@ -3071,10 +3231,9 @@ def _apply_resolved_guardrails_to_metadata( if metadata_variable_name not in data: data[metadata_variable_name] = {} - # Track pipeline-managed guardrails to exclude from independent execution - pipeline_managed_guardrails: set = set() + # Record the pipelines and the guardrails they step; the hook loops skip those per pipeline mode if pipelines: - pipeline_managed_guardrails = PolicyResolver.get_pipeline_managed_guardrails(pipelines) + pipeline_managed_guardrails: Final = PolicyResolver.get_pipeline_managed_guardrails(pipelines) data[metadata_variable_name]["_guardrail_pipelines"] = pipelines data[metadata_variable_name]["_pipeline_managed_guardrails"] = pipeline_managed_guardrails verbose_proxy_logger.debug( @@ -3086,18 +3245,14 @@ def _apply_resolved_guardrails_to_metadata( if not resolved_guardrails and not pipelines: return - existing_guardrails = data[metadata_variable_name].get("guardrails", []) - if not isinstance(existing_guardrails, list): - existing_guardrails = [] + existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) + existing_guardrails_list: Final = existing_guardrails if isinstance(existing_guardrails, list) else [] # Combine existing guardrails with policy-resolved guardrails (no duplicates) - # Exclude pipeline-managed guardrails from the flat list - combined = set(existing_guardrails) - combined.update(resolved_guardrails) - combined -= pipeline_managed_guardrails - data[metadata_variable_name]["guardrails"] = list(combined) + combined: Final = list(dict.fromkeys((*existing_guardrails_list, *resolved_guardrails))) + data[metadata_variable_name]["guardrails"] = combined - verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", list(combined)) + verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) async def add_guardrails_from_policy_engine( diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 2a7813bc140..50716e5d474 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -32,11 +32,15 @@ from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.db.autorouter_session_rollup import AUTOROUTER_BENCHMARKS_SQL +from litellm.proxy.db.autorouter_session_rollup import ( + AUTOROUTER_BENCHMARKS_SQL, + bounded_session_id, +) from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, refresh_proxy_server_request_body_snapshot, ) +from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.team_repository import TeamRepository from litellm.router_strategy.complexity_router import ComplexityRouter @@ -54,6 +58,7 @@ from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterCacheStats, AutoRouterRoutingTestRequest, AutoRouterRoutingTestResponse, + AutoRouterSessionResponse, ComplexityRouterConfigValidationRequest, ComplexityRouterConfigValidationResponse, RequestComplexityRouterConfig, @@ -484,6 +489,8 @@ class _SessionAggRow(BaseModel): total_tokens: int spend: float saved_spend: float + classifier_cost: float + classifier_cost_recorded_turns: int session_seconds: float @@ -520,6 +527,7 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals: avg_tokens_per_session=row.total_tokens / sessions if sessions else 0.0, spend=row.spend, saved_spend=row.saved_spend, + classifier_cost=row.classifier_cost if row.classifier_cost_recorded_turns == row.turns else None, baseline_spend=baseline_spend, saved_pct=_pct(row.saved_spend, baseline_spend), saved_per_session=row.saved_spend / sessions if sessions else 0.0, @@ -552,6 +560,7 @@ def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup: avg_tokens_per_session=totals.avg_tokens_per_session, spend=totals.spend, saved_spend=totals.saved_spend, + classifier_cost=totals.classifier_cost, baseline_spend=totals.baseline_spend, saved_pct=totals.saved_pct, saved_per_session=totals.saved_per_session, @@ -582,6 +591,8 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow: total_tokens=sum(row.total_tokens for row in rows), spend=sum(row.spend for row in rows), saved_spend=sum(row.saved_spend for row in rows), + classifier_cost=sum(row.classifier_cost for row in rows), + classifier_cost_recorded_turns=sum(row.classifier_cost_recorded_turns for row in rows), session_seconds=sum(row.session_seconds for row in rows), ) @@ -645,6 +656,7 @@ async def get_auto_router_benchmarks( str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to 30 days before end_date)") ] = None, end_date: Annotated[str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to today)")] = None, + api_key: Annotated[str | None, Query(description="Filter to one virtual key token hash")] = None, ) -> AutoRouterBenchmarksResponse: """ Benchmarks for the auto-router dashboard: session shape, savings against the configured @@ -681,6 +693,7 @@ async def get_auto_router_benchmarks( AUTOROUTER_BENCHMARKS_SQL, start_day.isoformat(), (end_day + timedelta(days=1)).isoformat(), + api_key, ) rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ()) groups: Final = ( @@ -696,6 +709,51 @@ async def get_auto_router_benchmarks( ) +@router.get( + "/auto_router/session", + tags=("auto router",), + dependencies=(Depends(user_api_key_auth),), + response_model=AutoRouterSessionResponse, +) +async def get_auto_router_session( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + session_id: Annotated[ + str, Query(description="The client session id (x-*-session-id header) the turns were sent under") + ], +) -> AutoRouterSessionResponse: + """ + One auto-routed session, for the key that ran it: the model its last turn was routed to and the + session's spend against the router's savings baseline. Built for a coding agent's status line + or stop hook, so any virtual key may call it and only ever sees rows written under its own + key hash. Reads the LiteLLM_AutoRouterSession rollup, which the asynchronous spend flush + fills a moment after each turn; a session with no flushed auto-routed turn yet is a 404. The + id is bounded the way the writer bounded it, so an oversized client id still finds its row. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + row: Final = await AutoRouterSessionRepository(prisma_client).find_latest_for_key( + user_api_key_dict.api_key, bounded_session_id(session_id) + ) + if row is None: + raise HTTPException( + status_code=404, detail=f"No auto-routed turns recorded for session {session_id!r} under this key" + ) + return AutoRouterSessionResponse( + session_id=session_id, + router_name=row.router_name, + router_type=row.router_type, + turns=row.turns, + last_model=row.last_model, + spend=row.spend, + saved_spend=row.saved_spend, + baseline_spend=row.spend + row.saved_spend, + baseline_model=row.baseline_model, + baseline_models=row.baseline_models, + ) + + # --------------------------------------------------------------------------- # Shadow eval: pre-adoption evaluation of an auto-router against live traffic. # The job row is immutable config plus stopped_at; status, counts, spend, and errors diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index ce6a97708ab..a8aef30107c 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -14,6 +14,7 @@ from litellm.proxy._types import CommonProxyErrors from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_emails, recover_double_hashed_key_metadata, + recover_key_metadata_from_spend_logs, ) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import PrismaClient @@ -433,9 +434,29 @@ def update_breakdown_metrics( return breakdown +def _spend_logs_window(dates: AbstractSet[str | None]) -> tuple[datetime, datetime] | None: + parsed: Final = sorted(day for day in (_parse_spend_date(raw) for raw in dates) if day is not None) + if not parsed: + return None + return (parsed[0] - timedelta(days=1), parsed[-1] + timedelta(days=2)) + + +def _parse_spend_date(raw: str | None) -> datetime | None: + if not isinstance(raw, str): + return None + try: + return datetime.fromisoformat(raw) + except ValueError: + return None + + +_EMPTY_KEY_METADATA: Final[Mapping[str, _KeyMetadataDict]] = MappingProxyType({}) + + async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: AbstractSet[str], + spend_logs_window: tuple[datetime, datetime] | None = None, ) -> Mapping[str, _KeyMetadataDict]: """Get api key metadata, falling back to deleted keys table for keys not found in active table. @@ -481,11 +502,17 @@ async def get_api_key_metadata( ) still_missing: Final = api_keys - frozenset(result) - combined: Final = ( - result - if not still_missing - else MappingProxyType({**result, **(await recover_double_hashed_key_metadata(prisma_client, still_missing))}) + from_reverse_hash: Final = ( + await recover_double_hashed_key_metadata(prisma_client, still_missing) if still_missing else _EMPTY_KEY_METADATA ) + after_token_recovery: Final = MappingProxyType({**result, **from_reverse_hash}) + unresolved: Final = api_keys - frozenset(after_token_recovery) + from_spend_logs: Final = ( + await recover_key_metadata_from_spend_logs(prisma_client, unresolved, spend_logs_window) + if unresolved and spend_logs_window is not None + else _EMPTY_KEY_METADATA + ) + combined: Final = MappingProxyType({**after_token_recovery, **from_spend_logs}) return await attach_user_emails(prisma_client, combined) @@ -898,7 +925,9 @@ async def _aggregate_spend_records( api_key_metadata: dict[str, _KeyMetadataDict] = {} if api_keys: - api_key_metadata = await get_api_key_metadata(prisma_client, api_keys) + api_key_metadata = await get_api_key_metadata( + prisma_client, api_keys, _spend_logs_window(frozenset(record.date for record in records)) + ) return await asyncio.to_thread( _aggregate_spend_records_sync, @@ -1094,7 +1123,9 @@ async def _aggregate_grouping_sets_records( api_key_metadata: dict[str, _KeyMetadataDict] = {} if api_keys: - api_key_metadata = await get_api_key_metadata(prisma_client, api_keys) + api_key_metadata = await get_api_key_metadata( + prisma_client, api_keys, _spend_logs_window(frozenset(r.date for r in records)) + ) return await asyncio.to_thread( _aggregate_grouping_sets_records_sync, @@ -1357,7 +1388,9 @@ async def get_daily_activity_aggregated( r.api_key for r in entity_records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY ) entity_key_metadata: Final = ( - await get_api_key_metadata(prisma_client, entity_api_keys) + await get_api_key_metadata( + prisma_client, entity_api_keys, _spend_logs_window(frozenset(r.date for r in entity_records)) + ) if entity_api_keys else {} # mutable-ok: matches the helper's dict return ) diff --git a/litellm/proxy/management_endpoints/coordination_redis_endpoints.py b/litellm/proxy/management_endpoints/coordination_redis_endpoints.py index 86ce336c7a3..88dc09ab001 100644 --- a/litellm/proxy/management_endpoints/coordination_redis_endpoints.py +++ b/litellm/proxy/management_endpoints/coordination_redis_endpoints.py @@ -140,7 +140,7 @@ def _merge_over_saved( def _validated_params(settings: Mapping[str, object]) -> CoordinationRedisParams: """Validate settings the way startup does: resolve env refs, then require a connection target.""" try: - params: Final = CoordinationRedisParams(**_resolve_env_refs(settings)) + params: Final = CoordinationRedisParams.model_validate(_resolve_env_refs(settings)) except ValidationError as e: invalid_fields: Final = sorted({str(error["loc"][0]) for error in e.errors() if error["loc"]}) raise HTTPException( diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 204051c3715..dc0da63555f 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -18,6 +18,7 @@ from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel import litellm +from litellm._internal_context import current_billing_time, pinned_billing_time from litellm._logging import verbose_proxy_logger from litellm.cost_calculator import completion_cost from litellm.proxy._types import ( @@ -27,7 +28,15 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.utils import CostPerToken, LlmProvidersSet, ModelInfo +from litellm.types.utils import ( + CostBreakdown, + CostPerToken, + LlmProvidersSet, + ModelInfo, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) router: Final = APIRouter() @@ -46,13 +55,15 @@ def _configured_price(key: str, sources: tuple[Mapping[str, object], ...]) -> fl def _extract_custom_pricing( - litellm_params: Mapping[str, object], model_info: Mapping[str, object] + litellm_params: Mapping[str, object], model_info: Mapping[str, object], builtin: ModelInfo | None ) -> CostPerToken | None: """ Pull per-token pricing configured on a deployment so on-prem / self-hosted models (absent from the public cost map) still estimate a real cost. Pricing may live on ``litellm_params`` or ``model_info``; ``litellm_params`` - wins, matching the router's cost-map registration precedence. + wins, matching the router's cost-map registration precedence. Cache rates the + deployment leaves unset come from the backend model's built-in entry, then its + own input rate, again matching what the router registers for live billing. """ sources: Final = (litellm_params, model_info) input_price: Final = _configured_price("input_cost_per_token", sources) @@ -61,15 +72,21 @@ def _extract_custom_pricing( if input_price is None and output_price is None: return None + input_rate: Final = input_price or 0.0 + cache_sources: Final = sources if builtin is None else (*sources, builtin) + cache_read_price: Final = _configured_price("cache_read_input_token_cost", cache_sources) + cache_creation_price: Final = _configured_price("cache_creation_input_token_cost", cache_sources) return CostPerToken( - input_cost_per_token=input_price or 0.0, + input_cost_per_token=input_rate, output_cost_per_token=output_price or 0.0, + cache_read_input_token_cost=input_rate if cache_read_price is None else cache_read_price, + cache_creation_input_token_cost=input_rate if cache_creation_price is None else cache_creation_price, ) -def _lookup_model_info(model: str) -> ModelInfo | None: +def _lookup_model_info(model: str, custom_llm_provider: str | None = None) -> ModelInfo | None: try: - return litellm.get_model_info(model=model) + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: return None @@ -98,17 +115,14 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel: model_info: Final = first_deployment.get("model_info", {}) custom_llm_provider: Final = litellm_params.get("custom_llm_provider") provider: Final = str(custom_llm_provider) if custom_llm_provider is not None else None - custom_cost_per_token: Final = _extract_custom_pricing(litellm_params, model_info) - - # Check base_model first (needed for Azure custom deployment names) + # base_model wins (needed for Azure custom deployment names) base_model: Final = model_info.get("base_model") or litellm_params.get("base_model") - if base_model: - verbose_proxy_logger.debug("Resolved model '%s' to base_model '%s' from router", model, base_model) - return ResolvedCostModel(str(base_model), provider, custom_cost_per_token) - - resolved_model: Final = litellm_params.get("model") + resolved_model: Final = base_model or litellm_params.get("model") if resolved_model: verbose_proxy_logger.debug("Resolved model '%s' to '%s' from router", model, resolved_model) + custom_cost_per_token: Final = _extract_custom_pricing( + litellm_params, model_info, _lookup_model_info(str(resolved_model), provider) + ) return ResolvedCostModel(str(resolved_model), provider, custom_cost_per_token) except Exception as e: verbose_proxy_logger.debug("Could not resolve model '%s' from router: %s", model, e) @@ -117,19 +131,59 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel: return ResolvedCostModel(model, None, None) -def _calculate_period_costs(num_requests, cost_per_request, input_cost, output_cost, margin_cost): - """ - Calculate costs for a given number of requests. +@dataclass(frozen=True, slots=True) +class CostLines: + """Cost of one request split the way the spend logs split it: the cache lines are + shares of input_cost and the reasoning line is a share of output_cost.""" - Returns tuple of (total_cost, input_cost, output_cost, margin_cost) or all None if num_requests is None/0. - """ - if not num_requests: - return None, None, None, None - return ( - cost_per_request * num_requests, - input_cost * num_requests, - output_cost * num_requests, - margin_cost * num_requests, + total_cost: float + input_cost: float + output_cost: float + margin_cost: float + cache_read_cost: float + cache_creation_cost: float + reasoning_cost: float + + def times(self, num_requests: int | None) -> "CostLines | None": + if not num_requests: + return None + return CostLines( + total_cost=self.total_cost * num_requests, + input_cost=self.input_cost * num_requests, + output_cost=self.output_cost * num_requests, + margin_cost=self.margin_cost * num_requests, + cache_read_cost=self.cache_read_cost * num_requests, + cache_creation_cost=self.cache_creation_cost * num_requests, + reasoning_cost=self.reasoning_cost * num_requests, + ) + + +def _cost_lines(cost_per_request: float, cost_breakdown: CostBreakdown | None) -> CostLines: + breakdown: Final = cost_breakdown if cost_breakdown is not None else CostBreakdown() + return CostLines( + total_cost=cost_per_request, + input_cost=breakdown.get("input_cost", 0.0), + output_cost=breakdown.get("output_cost", 0.0), + margin_cost=breakdown.get("margin_total_amount", 0.0), + cache_read_cost=breakdown.get("cache_read_cost", 0.0), + cache_creation_cost=breakdown.get("cache_creation_cost", 0.0), + reasoning_cost=breakdown.get("reasoning_cost", 0.0), + ) + + +def _usage_for_estimate(request: CostEstimateRequest) -> Usage: + cache_tokens: Final = request.cache_read_input_tokens + request.cache_creation_input_tokens + return Usage( + prompt_tokens=request.input_tokens, + completion_tokens=request.output_tokens, + total_tokens=request.input_tokens + request.output_tokens, + reasoning_tokens=request.reasoning_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=request.cache_read_input_tokens, + cache_creation_tokens=request.cache_creation_input_tokens, + ) + if cache_tokens + else None, ) @@ -530,11 +584,14 @@ async def estimate_cost( - model: Model name (e.g., "gpt-4", "claude-3-opus") - input_tokens: Expected input tokens per request - output_tokens: Expected output tokens per request + - cache_read_input_tokens: Cache-read tokens per request, counted within input_tokens (optional) + - cache_creation_input_tokens: Cache-write tokens per request, counted within input_tokens (optional) + - reasoning_tokens: Reasoning tokens per request, counted within output_tokens (optional) - num_requests_per_day: Number of requests per day (optional) - num_requests_per_month: Number of requests per month (optional) Returns cost breakdown including: - - Per-request costs (input, output, margin) + - Per-request costs (input, output, margin, plus the cache-read, cache-write and reasoning shares) - Daily costs (if num_requests_per_day provided) - Monthly costs (if num_requests_per_month provided) @@ -543,14 +600,15 @@ async def estimate_cost( { "model": "gpt-4", "input_tokens": 1000, + "cache_read_input_tokens": 800, "output_tokens": 500, + "reasoning_tokens": 200, "num_requests_per_day": 100, "num_requests_per_month": 3000 } ``` """ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.utils import ModelResponse, Usage # Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4') resolved: Final = _resolve_model_for_cost_lookup(request.model) @@ -559,15 +617,8 @@ async def estimate_cost( verbose_proxy_logger.debug("Cost estimate: request.model='%s' resolved to '%s'", request.model, resolved_model) - # Create a mock response with usage for completion_cost - mock_response: Final = ModelResponse( - model=resolved_model, - usage=Usage( - prompt_tokens=request.input_tokens, - completion_tokens=request.output_tokens, - total_tokens=request.input_tokens + request.output_tokens, - ), - ) + usage: Final = _usage_for_estimate(request) + mock_response: Final = ModelResponse(model=resolved_model, usage=usage) # Create a logging object to capture cost breakdown litellm_logging_obj: Final = LiteLLMLoggingObj( @@ -580,92 +631,73 @@ async def estimate_cost( function_id="cost-estimate", ) - # Use completion_cost which handles all the logic including margins/discounts - try: - cost_per_request: Final = completion_cost( - completion_response=mock_response, - model=resolved_model, - custom_llm_provider=resolved_provider, - custom_cost_per_token=resolved.custom_cost_per_token, - litellm_logging_obj=litellm_logging_obj, - ) - except Exception as e: - raise HTTPException( - status_code=404, - detail={ - "error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}" - }, - ) + # Pinning one moment keeps an off-peak window that opens mid-quote from pricing the totals on + # one side of it and the reported rates on the other. + with pinned_billing_time(current_billing_time()): + # Use completion_cost which handles all the logic including margins/discounts + try: + cost_per_request: Final = completion_cost( + completion_response=mock_response, + model=resolved_model, + custom_llm_provider=resolved_provider, + custom_cost_per_token=resolved.custom_cost_per_token, + litellm_logging_obj=litellm_logging_obj, + ) + except Exception as e: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriceable model + raise HTTPException( + status_code=404, + detail={ + "error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}" + }, + ) - # Get cost breakdown from the logging object - cost_breakdown: Final = litellm_logging_obj.cost_breakdown + # The rates come back from the pricing call itself rather than a second lookup, so they are the + # ones the cost lines above billed at even when completion_cost infers a provider this endpoint + # never resolved (an unrouted "xai/grok-4" prices on xai's inclusive tier thresholds; a lookup + # here without that provider would report the sub-200k rate for a line billed above it). + rates: Final = litellm_logging_obj.billed_token_rates + per_request: Final = _cost_lines(cost_per_request, litellm_logging_obj.cost_breakdown) + daily: Final = per_request.times(request.num_requests_per_day) + monthly: Final = per_request.times(request.num_requests_per_month) - input_cost: Final = cost_breakdown.get("input_cost", 0.0) if cost_breakdown else 0.0 - output_cost: Final = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0 - margin_cost: Final = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0 - - model_info: Final = _lookup_model_info(resolved_model) - mapped_input_price: Final = model_info.get("input_cost_per_token") if model_info is not None else None - mapped_output_price: Final = model_info.get("output_cost_per_token") if model_info is not None else None + model_info: Final = _lookup_model_info(resolved_model, resolved_provider) mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None - - input_cost_per_token: Final = ( - resolved.custom_cost_per_token["input_cost_per_token"] - if resolved.custom_cost_per_token is not None - else mapped_input_price - ) - output_cost_per_token: Final = ( - resolved.custom_cost_per_token["output_cost_per_token"] - if resolved.custom_cost_per_token is not None - else mapped_output_price - ) custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider - # Calculate daily and monthly costs - ( - daily_cost, - daily_input_cost, - daily_output_cost, - daily_margin_cost, - ) = _calculate_period_costs( - num_requests=request.num_requests_per_day, - cost_per_request=cost_per_request, - input_cost=input_cost, - output_cost=output_cost, - margin_cost=margin_cost, - ) - ( - monthly_cost, - monthly_input_cost, - monthly_output_cost, - monthly_margin_cost, - ) = _calculate_period_costs( - num_requests=request.num_requests_per_month, - cost_per_request=cost_per_request, - input_cost=input_cost, - output_cost=output_cost, - margin_cost=margin_cost, - ) - return CostEstimateResponse( model=request.model, input_tokens=request.input_tokens, output_tokens=request.output_tokens, + cache_read_input_tokens=request.cache_read_input_tokens, + cache_creation_input_tokens=request.cache_creation_input_tokens, + reasoning_tokens=request.reasoning_tokens, num_requests_per_day=request.num_requests_per_day, num_requests_per_month=request.num_requests_per_month, - cost_per_request=cost_per_request, - input_cost_per_request=input_cost, - output_cost_per_request=output_cost, - margin_cost_per_request=margin_cost, - daily_cost=daily_cost, - daily_input_cost=daily_input_cost, - daily_output_cost=daily_output_cost, - daily_margin_cost=daily_margin_cost, - monthly_cost=monthly_cost, - monthly_input_cost=monthly_input_cost, - monthly_output_cost=monthly_output_cost, - monthly_margin_cost=monthly_margin_cost, - input_cost_per_token=input_cost_per_token, - output_cost_per_token=output_cost_per_token, + cost_per_request=per_request.total_cost, + input_cost_per_request=per_request.input_cost, + output_cost_per_request=per_request.output_cost, + margin_cost_per_request=per_request.margin_cost, + cache_read_cost_per_request=per_request.cache_read_cost, + cache_creation_cost_per_request=per_request.cache_creation_cost, + reasoning_cost_per_request=per_request.reasoning_cost, + daily_cost=daily.total_cost if daily is not None else None, + daily_input_cost=daily.input_cost if daily is not None else None, + daily_output_cost=daily.output_cost if daily is not None else None, + daily_margin_cost=daily.margin_cost if daily is not None else None, + daily_cache_read_cost=daily.cache_read_cost if daily is not None else None, + daily_cache_creation_cost=daily.cache_creation_cost if daily is not None else None, + daily_reasoning_cost=daily.reasoning_cost if daily is not None else None, + monthly_cost=monthly.total_cost if monthly is not None else None, + monthly_input_cost=monthly.input_cost if monthly is not None else None, + monthly_output_cost=monthly.output_cost if monthly is not None else None, + monthly_margin_cost=monthly.margin_cost if monthly is not None else None, + monthly_cache_read_cost=monthly.cache_read_cost if monthly is not None else None, + monthly_cache_creation_cost=monthly.cache_creation_cost if monthly is not None else None, + monthly_reasoning_cost=monthly.reasoning_cost if monthly is not None else None, + input_cost_per_token=rates.input_cost_per_token if rates is not None else None, + output_cost_per_token=rates.output_cost_per_token if rates is not None else None, + cache_read_input_token_cost=rates.cache_read_input_token_cost if rates is not None else None, + cache_creation_input_token_cost=rates.cache_creation_input_token_cost if rates is not None else None, + output_cost_per_reasoning_token=rates.output_cost_per_reasoning_token if rates is not None else None, provider=custom_llm_provider, ) diff --git a/litellm/proxy/management_endpoints/credential_migration.py b/litellm/proxy/management_endpoints/credential_migration.py index e0725119576..915cce87dbd 100644 --- a/litellm/proxy/management_endpoints/credential_migration.py +++ b/litellm/proxy/management_endpoints/credential_migration.py @@ -43,7 +43,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( _ALGO_AES_GCM, _ENCRYPTION_ALGORITHM_SETTING, _V2_GCM_PREFIX, + SecretMapDecodeError, _get_salt_key, + decode_secret_map, decrypt_value_helper, encrypt_value_helper, ) @@ -65,6 +67,20 @@ class LocationReport: # Used by --check (read-only classification): legacy: int = 0 # nacl ciphertext still awaiting migration + def count(self, classification: ValueClass | None) -> None: + if classification is None: + return + self.scanned += 1 + match classification: + case "migrated": + self.already_v2 += 1 + case "legacy": + self.legacy += 1 + case "undecryptable": + self.undecryptable += 1 + case _: + self.plaintext += 1 + def as_dict(self) -> dict[str, int]: return { "scanned": self.scanned, @@ -441,7 +457,7 @@ def _classify_callback_value(value: object) -> ValueClass: _COVERED_TABLE_SPECS: Final = [ ("model_table", "litellm_proxymodeltable", ("litellm_params",), ()), ("credentials", "litellm_credentialstable", ("credential_values",), ()), - ("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars"), ()), + ("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars", "static_headers", "env"), ()), ("mcp_user_credentials", "litellm_mcpusercredentials", (), ("credential_b64",)), ("mcp_user_env_vars", "litellm_mcpuserenvvars", (), ("values_b64",)), ] @@ -472,14 +488,18 @@ def _classify_into_report(report: LocationReport, value: str) -> None: names, base URLs, …) do not decrypt and fall through to ``plaintext``, so over-scanning a column is harmless to the residual count. """ - report.scanned += 1 - cls: Final = classify_value(value, key="scan") - if cls == "migrated": - report.already_v2 += 1 - elif cls == "legacy": - report.legacy += 1 - else: # plaintext / not-a-string - report.plaintext += 1 + report.count(classify_value(value, key="scan")) + + +def _classify_secret_map(value: object, key: str) -> ValueClass | None: + try: + decoded: Final = decode_secret_map(value, key=key) + except SecretMapDecodeError: + return "undecryptable" + if not decoded: + return None + ciphertext: Final = json.loads(value) if isinstance(value, str) and value.lstrip().startswith('"') else value + return "migrated" if is_migrated(ciphertext) else "legacy" async def _scan_one_table( @@ -503,6 +523,9 @@ async def _scan_one_table( raw = getattr(row, col, None) if raw is None: continue + if db_attr == "litellm_mcpservertable" and col in ("static_headers", "env"): + report.count(_classify_secret_map(raw, col)) + continue if isinstance(raw, str): try: raw = json.loads(raw) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ebcfab090b5..db4467dda5b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -196,6 +196,38 @@ class _ModelRowWhere(TypedDict): model_id: ReadOnly[str] +class _KeyUpdateResult(TypedDict): + token: ReadOnly[str] + data: ReadOnly[Mapping[str, object]] + + +class _KeyRowWhere(TypedDict): + token: ReadOnly[str] + + +class _BudgetRowWhere(TypedDict): + budget_id: ReadOnly[str] + + +class _BudgetRowSoftBudgetUpdate(TypedDict): + soft_budget: ReadOnly[float | None] + updated_by: ReadOnly[str] + + +class _BudgetRowSoftBudgetCreate(TypedDict): + soft_budget: ReadOnly[float] + created_by: ReadOnly[str] + updated_by: ReadOnly[str] + + +class _KeyUpdateTx(Protocol): + @property + def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ... + + @property + def litellm_budgettable(self) -> "TableActions[prisma_models.LiteLLM_BudgetTable]": ... + + class _ConfigTableActions(Protocol): """Config table surface this module needs; the shared repository seam exposes no ``update``.""" @@ -1812,11 +1844,7 @@ async def generate_key_fn( status_code=400, detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) - if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0): - raise HTTPException( - status_code=400, - detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, - ) + _validate_soft_budget_value(data.soft_budget) custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = ( _custom_key_generate_hook(proxy_server) @@ -2121,6 +2149,88 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ return non_default_values +def _validate_soft_budget_value(soft_budget: float | None) -> None: + if soft_budget is not None and (not math.isfinite(soft_budget) or soft_budget < 0): + raise HTTPException( + status_code=400, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {soft_budget}"}, + ) + + +async def _update_key_soft_budget( + db: _KeyUpdateTx, + existing_key_row: LiteLLM_VerificationToken, + soft_budget: float | None, + changed_by: str, +) -> str | None: + existing_budget_id: Final = existing_key_row.budget_id + if existing_budget_id is not None: + budget_update: Final[_BudgetRowSoftBudgetUpdate] = {"soft_budget": soft_budget, "updated_by": changed_by} + budget_where: Final[_BudgetRowWhere] = {"budget_id": existing_budget_id} + await db.litellm_budgettable.update(where=budget_where, data=budget_update) + return existing_budget_id + if soft_budget is None: + return None + budget_create: Final[_BudgetRowSoftBudgetCreate] = { + "soft_budget": soft_budget, + "created_by": changed_by, + "updated_by": changed_by, + } + created_budget: Final = await db.litellm_budgettable.create(data=budget_create) + return created_budget.budget_id + + +async def _apply_soft_budget_update( + data: UpdateKeyRequest, + non_default_values: Mapping[str, object], + db: _KeyUpdateTx, + existing_key_row: LiteLLM_VerificationToken, + changed_by: str, +) -> Mapping[str, object]: + remaining: Final = MappingProxyType({k: v for k, v in non_default_values.items() if k != "soft_budget"}) + updated_budget_id: Final = await _update_key_soft_budget( + db=db, + existing_key_row=existing_key_row, + soft_budget=data.soft_budget, + changed_by=changed_by, + ) + if updated_budget_id is not None and existing_key_row.budget_id is None: + return MappingProxyType({**remaining, "budget_id": updated_budget_id}) + return remaining + + +async def _update_key_row_with_soft_budget( + prisma_client: PrismaClient, + key: str, + data: UpdateKeyRequest, + non_default_values: Mapping[str, object], + existing_key_row: LiteLLM_VerificationToken, + changed_by: str, +) -> _KeyUpdateResult: + hashed_token: Final = _hash_token_if_needed(key) + key_where: Final[_KeyRowWhere] = {"token": hashed_token} + tx: _KeyUpdateTx + async with prisma_client.tx() as tx: + update_values: Final = await _apply_soft_budget_update( + data=data, + non_default_values=non_default_values, + db=tx, + existing_key_row=existing_key_row, + changed_by=changed_by, + ) + updated_row: Final = await tx.litellm_verificationtoken.update( + where=key_where, + data=with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ), + ) + updated_data: Final[Mapping[str, object]] = ( + updated_row.model_dump() if updated_row is not None else MappingProxyType({}) + ) + result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data} + return result + + async def prepare_key_update_data( data: UpdateKeyRequest | RegenerateKeyRequest, existing_key_row: LiteLLM_VerificationToken, @@ -2659,6 +2769,7 @@ async def _validate_update_key_data( (data.max_budget is not None and data.max_budget != existing_key_row.max_budget) or data.spend is not None or "budget_limits" in data.model_fields_set + or "soft_budget" in data.model_fields_set ) _existing_metadata: Final = getattr(existing_key_row, "metadata", None) @@ -2840,6 +2951,11 @@ async def update_key_fn( """ Update an existing API key's parameters. + The body is a merge patch: a field left out keeps its stored value, and on the key's own columns + an explicit null clears it. The metadata-backed fields below are the exception, merging into the + stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself + replaces the stored metadata wholesale. + Parameters: - key: Optional[str] - The key to update. Either key or key_alias must be provided. - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases) @@ -2857,7 +2973,7 @@ async def update_key_fn( - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) - - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. + - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget. - max_parallel_requests: Optional[int] - Rate limit for parallel requests - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} - tpm_limit: Optional[int] - Tokens per minute limit @@ -2913,6 +3029,7 @@ async def update_key_fn( """ from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, llm_router, premium_user, prisma_client, @@ -2928,6 +3045,8 @@ async def update_key_fn( detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) + _validate_soft_budget_value(data.soft_budget) + # get the row from db existing_key_row: Final = await _get_and_validate_existing_key( token=data.key, @@ -2984,10 +3103,22 @@ async def update_key_fn( existing_key_alias=existing_key_row.key_alias, ) - _data: Final = {**non_default_values, "token": key} if prisma_client is None: raise Exception("Not connected to DB!") - response: Final = await prisma_client.update_data(token=key, data=_data) + + changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name + response: Final = ( + await _update_key_row_with_soft_budget( + prisma_client=prisma_client, + key=key, + data=data, + non_default_values=non_default_values, + existing_key_row=existing_key_row, + changed_by=changed_by, + ) + if "soft_budget" in data.model_fields_set + else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key})) + ) # Delete - key from cache, since it's been updated! # key updated - a new model could have been added to this key. it should not block requests after this is done @@ -3598,7 +3729,7 @@ async def delete_key_fn( ) verbose_proxy_logger.debug( - "/keys/delete - cache after delete: %s", user_api_key_cache.in_memory_cache.cache_dict + "/keys/delete - cache after delete: %s", user_api_key_cache.key_object_cache.in_memory_cache.cache_dict ) asyncio.create_task( @@ -4428,6 +4559,23 @@ async def delete_verification_tokens( litellm_changed_by=litellm_changed_by, ) + # Snapshot before the delete: the FK cascade drops the mapping rows, but their + # cached jwt_key_mapping entries still resolve to the now-dead token (LIT-5380). + jwt_mapping_cache_keys: Final[tuple[str, ...]] = tuple( + cache_key + for keys_for_token in await asyncio.gather( + *( + get_jwt_key_mapping_cache_keys_for_token( + hashed_token=key.token, + prisma_client=prisma_client, + ) + for key in authorized_keys + if key.token is not None + ) + ) + for cache_key in keys_for_token + ) + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: deleted_tokens = await prisma_client.delete_data(tokens=tokens) if deleted_tokens is not None and len(deleted_tokens) != len(tokens): @@ -4440,6 +4588,8 @@ async def delete_verification_tokens( if len(deleted_tokens) != len(tokens): failed_tokens = [token for token in tokens if token not in deleted_tokens] + await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache) + else: raise Exception("DB not connected. prisma_client is None") except Exception as e: @@ -6427,7 +6577,7 @@ async def _list_key_helper( {"token": "desc"}, # fallback sort ] ), - include={"object_permission": True}, + include={"object_permission": True, "litellm_budget_table": True}, ) verbose_proxy_logger.debug("Fetched %s keys", len(keys)) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index ae266792391..4c97bbaf5de 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -193,6 +193,7 @@ if MCP_AVAILABLE: UpdateMCPServerRequest, UserAPIKeyAuth, UserMCPManagementMode, + is_per_server_oauth_discovery_eligible, ) from litellm.proxy.auth.user_api_key_auth import ( _user_api_key_auth_builder, @@ -2277,6 +2278,28 @@ if MCP_AVAILABLE: """Persist the OAuth2 access token obtained by the calling user.""" prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") await _authorize_and_fetch_mcp_server(prisma_client, user_api_key_dict, server_id) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # keep manager import lazy + global_mcp_server_manager as _manager, + ) + + # This endpoint accepts an opaque token with no upstream identity validation, so it must be + # closed for identity-bound servers or it becomes a bypass of the token-relay binding check. + registry_server: Final = _manager.get_mcp_server_by_id(server_id) + binding: Final = registry_server.oauth_identity_binding if registry_server else None + if binding is not None and binding.mode == "enforce": + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ # mutable-ok: FastAPI exception detail requires a JSON-serializable dictionary + "error": "oauth_identity_binding_enforced", + "error_description": ( + "Direct credential storage is disabled for this server: its OAuth identity " + "binding is enforced and this endpoint cannot validate the token's principal. " + "Complete the OAuth flow through the gateway instead." + ), + "server_id": server_id, + "credential_stored": False, + }, + ) user_id: Final = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( @@ -2672,6 +2695,8 @@ if MCP_AVAILABLE: """ Updates the MCP Server in the db. + Partial update: a field left out of the payload keeps its stored value, and a field sent as null is cleared. + Parameters: - payload: UpdateMCPServerRequest - Required. The updated mcp server data. ``` @@ -2714,6 +2739,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 @@ -3076,6 +3122,8 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: str | None = Header(None), ): + """Partial update: a field left out keeps its stored value, and a field sent as null is cleared, except + ``toolset_name`` and ``tools``, which a toolset always has; empty the tool selection with an explicit [].""" prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b77108911aa..94d2b773e14 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( _refresh_cached_team, @@ -109,6 +110,8 @@ 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.llms.bedrock import AwsSessionTag from litellm.types.proxy.management_endpoints.model_management_endpoints import ( AutoRouterClassifierDefaultPromptResponse, UpdateUsefulLinksRequest, @@ -349,6 +352,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 +399,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") @@ -850,6 +898,12 @@ async def patch_model( existing_litellm_params=db_model.litellm_params, ) + ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=patch_data.litellm_params, + user_api_key_dict=user_api_key_dict, + existing_litellm_params=db_model.litellm_params, + ) + _raise_on_strategy_router_write_violation( incoming_params=patch_data.litellm_params, existing_params=db_model.litellm_params, @@ -1603,6 +1657,10 @@ async def _update_existing_team_model_assignment( # No team_model_add/delete calls required; public name is already registered +def _canonical_session_tags(tags: Sequence[AwsSessionTag]) -> tuple[tuple[str, str], ...]: + return tuple(sorted((tag["Key"], tag["Value"]) for tag in tags)) + + class ModelManagementAuthChecks: """ Common auth checks for model management endpoints @@ -1657,6 +1715,28 @@ class ModelManagementAuthChecks: param="litellm_credential_name", ) + @staticmethod + def can_user_set_aws_session_tags( + litellm_params: GenericLiteLLMParams | None, + user_api_key_dict: UserAPIKeyAuth, + existing_litellm_params: GenericLiteLLMParams | None = None, + ) -> Literal[True]: + if litellm_params is None or litellm_params.aws_session_tags is None: + return True + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return True + existing_tags: Final = existing_litellm_params.aws_session_tags if existing_litellm_params is not None else None + if existing_tags is not None and _canonical_session_tags(existing_tags) == _canonical_session_tags( + litellm_params.aws_session_tags + ): + return True + raise ProxyException( + message=f"Only a proxy admin can set aws_session_tags on a model. Your role={user_api_key_dict.user_role}.", + type=ProxyErrorTypes.auth_error.value, + code=status.HTTP_403_FORBIDDEN, + param="aws_session_tags", + ) + @staticmethod async def allow_team_model_action( model_params: Deployment | updateDeployment, @@ -1990,6 +2070,11 @@ async def add_new_model( user_api_key_dict=user_api_key_dict, ) + ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=model_params.litellm_params, + user_api_key_dict=user_api_key_dict, + ) + _raise_on_strategy_router_write_violation( incoming_params=model_params.litellm_params, existing_params=None, @@ -2174,6 +2259,12 @@ async def update_model( existing_litellm_params=deployment.litellm_params, ) + ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=model_params.litellm_params, + user_api_key_dict=user_api_key_dict, + existing_litellm_params=deployment.litellm_params, + ) + _raise_on_strategy_router_write_violation( incoming_params=model_params.litellm_params, existing_params=deployment.litellm_params, diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 1e711b036d2..96e946424bd 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -80,7 +80,23 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_OrganizationTable as PrismaOrganizationTable from prisma.models import LiteLLM_UserTable as PrismaUserTable -router: Final = APIRouter() + +async def _enterprise_license_required( + _user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> None: + from litellm.proxy.proxy_server import premium_user + + if not premium_user: + raise HTTPException( + status_code=403, + detail={ + "error": "Organizations are only available for LiteLLM Enterprise users. " + f"{CommonProxyErrors.not_premium_user.value}" + }, + ) + + +router: Final = APIRouter(dependencies=[Depends(_enterprise_license_required)]) class _ObjectPermissionRow(Protocol): diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index fe658a13c24..1932e89717b 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -20,6 +20,7 @@ from litellm.proxy._types import ( LiteLLM_AuditLogs, LiteLLM_TeamTable, LitellmTableNames, + LitellmUserRoles, ProxyErrorTypes, ProxyException, TeamCallbackDeleteResponse, @@ -28,7 +29,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.callback_config_validation import callback_config_error +from litellm.proxy.common_utils.callback_config_validation import ( + callback_config_error, + cross_entry_family_error, +) from litellm.proxy.common_utils.callback_utils import ( _CALLBACK_VAR_ENCRYPTED_PREFIX, decrypt_callback_vars, @@ -230,6 +234,22 @@ def _callback_error(status_code: int, message: str) -> HTTPException: ) +def _unknown_team_error(team_id: str, user_api_key_dict: UserAPIKeyAuth, status_code: int) -> HTTPException: + """Report an unknown team without telling an unauthorized caller that it is unknown. + + These routes are reachable by any authenticated caller so that a team admin can + get as far as _verify_team_access. A distinct "does not exist" would therefore let + any valid key probe which team ids exist, so a caller who could not have managed + the team either way gets the same 403 body _verify_team_access raises. + """ + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return _callback_error(status_code, f"Team id = {team_id} does not exist.") + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You do not have access to this team", + ) + + @router.post( "/team/{team_id:path}/callback", tags=["team management"], @@ -304,10 +324,7 @@ async def add_team_callbacks( # Check if team_id exists already _existing_team = await prisma_client.get_data(team_id=team_id, table_name="team", query_type="find_unique") if _existing_team is None: - raise HTTPException( - status_code=400, - detail={"error": f"Team id = {team_id} does not exist. Please use a different team id."}, - ) + raise _unknown_team_error(team_id, user_api_key_dict, status.HTTP_400_BAD_REQUEST) # IDOR guard: only proxy admins / org admins / team admins of THIS # team may write callback credentials. Without this, any @@ -326,6 +343,28 @@ async def add_team_callbacks( if team_callback_settings is None or not isinstance(team_callback_settings, list): team_callback_settings = [] + # One entry has to own a credential family end to end. The entries are + # flattened into one dict before a request reads them, so an entry + # naming only a destination would pair with a key written on another + # entry and carry it to that destination -- a key a team admin can read + # back nowhere. Repeating a value the owning entry already stores is + # fine, which is how one integration covers both events. Proxy admins + # are exempt: they already hold every credential the proxy has. + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Decrypted, because the check compares the incoming values against + # the stored ones and the credentials are encrypted at rest. + decrypted_logging: Final = decrypt_callback_vars(team_metadata).get("logging") + stored_entries: Final = decrypted_logging if isinstance(decrypted_logging, list) else () + stored_entry_vars: Final = [ # mutable-ok: read-only input to the check, never stored + entry.get("callback_vars") or {} for entry in stored_entries + ] + family_error: Final = cross_entry_family_error(data.callback_vars, stored_entry_vars) + if family_error is not None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=family_error, + ) + ## check if it already exists, for the same callback event for callback in team_callback_settings: if ( @@ -452,7 +491,7 @@ async def delete_team_callback( team_id=team_id, table_name="team", query_type="find_unique" ) if _existing_team is None: - raise _callback_error(404, f"Team id = {team_id} does not exist.") + raise _unknown_team_error(team_id, user_api_key_dict, status.HTTP_404_NOT_FOUND) # IDOR guard: only proxy admins / org admins / team admins of THIS team may # deregister its callbacks, otherwise any authenticated key holder could @@ -726,10 +765,7 @@ async def get_team_callbacks( # Check if team_id exists _existing_team = await prisma_client.get_data(team_id=team_id, table_name="team", query_type="find_unique") if _existing_team is None: - raise HTTPException( - status_code=404, - detail={"error": f"Team id = {team_id} does not exist."}, - ) + raise _unknown_team_error(team_id, user_api_key_dict, status.HTTP_404_NOT_FOUND) # IDOR guard: callback metadata holds third-party API credentials # (Langfuse / Langsmith / GCS). Only proxy admins / org admins / diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 2ec68a10f65..c050368b3fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5601,7 +5601,7 @@ async def team_model_add( updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"updated_at": datetime.now(timezone.utc)}, - include={"object_permission": True}, + include={"litellm_model_table": True, "object_permission": True}, ) if updated_team is None: raise HTTPException( @@ -5688,7 +5688,7 @@ async def team_model_delete( updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"models": updated_models}, - include={"object_permission": True}, + include={"litellm_model_table": True, "object_permission": True}, ) if updated_team is None: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 3e6434a5afd..1ba90725eff 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -22,7 +22,6 @@ from html import escape from types import MappingProxyType from typing import ( TYPE_CHECKING, - Annotated, Any, Final, Literal, @@ -42,12 +41,13 @@ if TYPE_CHECKING: import jwt from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse -from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.constants import ( CLI_SSO_CLAIM_MAP, CLI_SSO_CLAIM_MAX_SCALAR_LENGTH, @@ -95,6 +95,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.auth.team_grants import TeamModelAliasTable from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.admin_ui_utils import ( admin_ui_disabled, @@ -209,31 +210,14 @@ def _team_detail_db(repo: TeamRepository) -> "TableActions[_TeamDetailRow]": return repo.table -_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str]) _SSO_TOKEN_CLAIMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -def _decode_model_aliases(value: object) -> object: - """``/team/new`` stores team model aliases as a JSON-encoded string in the Json column.""" - if not isinstance(value, str): - return value - try: - return _MODEL_ALIASES_ADAPTER.validate_json(value) - except ValidationError: - return None - - -class _TeamModelAliasTable(BaseModel): - model_config = ConfigDict(protected_namespaces=()) - - model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None - - class _TeamRowGrants(BaseModel): team_id: str team_alias: str | None = None models: tuple[str, ...] = () - litellm_model_table: _TeamModelAliasTable | None = None + litellm_model_table: TeamModelAliasTable | None = None class CliSsoTeamDetail(BaseModel): @@ -353,6 +337,16 @@ def _check_cli_sso_start_rate_limit( ) +def _read_cli_sso_flow(cache: DualCache, cache_key: str) -> object: + redis_cache: Final = cache.redis_cache + if redis_cache is None: + return cache.get_cache(key=cache_key) + try: + return redis_cache.get_cache(key=cache_key) + except RedisCircuitBreakerOpenError: + return None + + def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict: if isinstance(login_id, str) and login_id.startswith("sk-"): raise HTTPException( @@ -365,12 +359,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict: if not _is_valid_cli_sso_login_id(login_id): raise HTTPException(status_code=400, detail="Invalid CLI login session id") - cache_key: Final = _get_cli_sso_flow_cache_key(cast(str, login_id)) - redis_cache: Final = cache.redis_cache - if redis_cache is not None: - flow = redis_cache.get_cache(key=cache_key) - else: - flow = cache.get_cache(key=cache_key) + flow = _read_cli_sso_flow(cache, _get_cli_sso_flow_cache_key(cast(str, login_id))) if isinstance(flow, str): try: flow = _as_object(json.loads(flow)) diff --git a/litellm/proxy/middleware/per_request_root_path_middleware.py b/litellm/proxy/middleware/per_request_root_path_middleware.py new file mode 100644 index 00000000000..d5df91458d2 --- /dev/null +++ b/litellm/proxy/middleware/per_request_root_path_middleware.py @@ -0,0 +1,122 @@ +"""Per-request ``root_path`` resolution from ``SERVER_ROOT_PATHS``. + +``SERVER_ROOT_PATH`` is a startup scalar, so one deployment serves exactly one +client-visible URL path prefix; a request under any other prefix 404s before a +handler runs. When the ingress preserves several prefixes into one pod (e.g. +``/tenant-a/*`` and ``/tenant-b/*``), the matched prefix becomes that +request's ``scope["root_path"]`` instead: Starlette strips it during route +matching and rebuilds it into ``request.base_url``, so every emitted URL — +the MCP OAuth discovery ``resource`` (RFC 9728 §3) and the 401 challenges' +``resource_metadata`` among them — lands under the prefix the client called. +Opt-in: with ``SERVER_ROOT_PATHS`` unset the middleware is not added at all. +""" + +import os +from collections.abc import Sequence +from contextvars import ContextVar +from typing import Final + +from starlette.types import ASGIApp, Receive, Scope, Send + +from litellm._logging import verbose_proxy_logger + +SERVER_ROOT_PATHS_ENV: Final = "SERVER_ROOT_PATHS" + +# The effective ``root_path`` for the currently-handled request. Populated by +# ``PerRequestRootPathMiddleware`` from the (possibly-mutated) scope so code +# that emits URLs off the request path — the 401 challenges' resource_metadata +# and ``get_custom_url``'s SSO callbacks among them — can pick up the prefix +# the client actually called without threading scope through every call site. +# ``None`` means "middleware did not run" (the ``SERVER_ROOT_PATHS`` env is +# unset, so no per-request prefix exists); readers fall back to the scalar +# ``SERVER_ROOT_PATH`` in that case, which matches the pre-middleware behavior. +_request_root_path_var: Final[ContextVar[str | None]] = ContextVar("_request_root_path_var", default=None) + + +def get_request_root_path() -> str: + """Return the effective ``root_path`` for the current request. + + Reads the value ``PerRequestRootPathMiddleware`` stashed for this request; + falls back through :func:`~litellm.proxy.utils.get_server_root_path` (i.e. + the ``SERVER_ROOT_PATH`` env) when the middleware did not run — the + scalar-only deployment. Delegating to the existing helper keeps every + existing ``monkeypatch.setattr("litellm.proxy.utils.get_server_root_path"`` + test override working, and keeps a single source of truth for the scalar. + """ + value: Final = _request_root_path_var.get() + if value is not None: + return value + # Lazy import: utils.py imports this module (via the lazy import inside + # get_custom_url), so a top-level import would build a cycle at load time. + from litellm.proxy.utils import get_server_root_path # noqa: PLC0415 # lazy import breaks a two-way dep + + return get_server_root_path() + + +def normalize_root_paths(raw_paths: Sequence[str]) -> tuple[str, ...]: + """Strip whitespace and trailing slashes, dedupe, order longest-first; + warn and drop entries missing a leading ``/`` and the bare root.""" + kept: Final[list[str]] = [] # mutable-ok: local accumulator; escapes only as a tuple + for entry in raw_paths: + candidate = entry.strip() + if not candidate: + continue + if not candidate.startswith("/"): + verbose_proxy_logger.warning( + "%s entry %r does not start with '/' and will be ignored.", + SERVER_ROOT_PATHS_ENV, + entry, + ) + continue + candidate = candidate.rstrip("/") + if not candidate: + verbose_proxy_logger.warning( + "%s entry %r is the bare root and will be ignored; a root-mounted deployment needs no entry.", + SERVER_ROOT_PATHS_ENV, + entry, + ) + continue + if candidate not in kept: + kept.append(candidate) + return tuple(sorted(kept, key=len, reverse=True)) + + +def get_server_root_paths() -> tuple[str, ...]: + """The normalized ``SERVER_ROOT_PATHS`` prefixes, empty when unset.""" + configured: Final = os.getenv(SERVER_ROOT_PATHS_ENV, "") + if not configured.strip(): + return () + return normalize_root_paths(configured.split(",")) + + +class PerRequestRootPathMiddleware: + """Sets ``scope["root_path"]`` to the configured prefix matching the + request path on a whole-segment boundary. ``scope["path"]`` is left + untouched (Starlette strips ``root_path`` at route-match time). Must be + the outermost middleware so inner middlewares and the router see the + resolved value; a matched prefix overrides a scalar ``SERVER_ROOT_PATH`` + for that request. + """ + + def __init__(self, app: ASGIApp, root_paths: Sequence[str]) -> None: + self.app = app + self.root_paths: Final = normalize_root_paths(root_paths) + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] in ("http", "websocket"): + path: Final = scope.get("path", "") + for prefix in self.root_paths: + if path == prefix or path.startswith(prefix + "/"): + scope["root_path"] = prefix # rebind-ok: ASGI middleware contract; Router and base_url read it + break + # Stash the effective root_path (matched prefix, or the scope's + # existing value when nothing matched — i.e. FastAPI's scalar + # SERVER_ROOT_PATH) so code that emits URLs off the request path + # picks the same prefix the router will resolve the request under. + token: Final = _request_root_path_var.set(str(scope.get("root_path", ""))) + try: + await self.app(scope, receive, send) + finally: + _request_root_path_var.reset(token) + return + await self.app(scope, receive, send) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 15eeddbc489..b1f282a0978 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -15,6 +15,7 @@ from typing import ( runtime_checkable, ) +from litellm.batches.batch_utils import batch_cost_is_final from litellm.proxy._types import ProxyException from litellm.repositories.table_repositories import ( ManagedFileRepository, @@ -1357,12 +1358,7 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: enumerated the batch and none succeeded. A zero or unknown total means counts are unreported, so stay eligible and let the next poller pass revisit it. (#37713) """ - if response.output_file_id is not None: - return True - request_counts = response.request_counts - if request_counts is None: - return False - return request_counts.total > 0 and request_counts.completed == 0 + return batch_cost_is_final(response) async def update_batch_in_database( diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index bf07f4748ef..c315d30b8f3 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -45,6 +45,11 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.proxy.openai_files_endpoints.batch_file_validation import ( check_batch_file_upload, raise_batch_file_validation_failure, @@ -296,22 +301,22 @@ async def route_create_file( if managed_files_obj is None: raise ProxyException( message="Managed files hook not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if llm_router is None: raise ProxyException( message="LLM Router not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if not isinstance(managed_files_obj, BaseFileEndpoints): raise ProxyException( message="Managed files hook is not a BaseFileEndpoints", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) # Managed files internally calls llm_router.acreate_file() which includes loadbalancing @@ -713,17 +718,17 @@ async def create_file( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) finally: for spool in spools: @@ -812,22 +817,22 @@ async def get_file_content( if managed_files_obj is None: raise ProxyException( message="Managed files hook not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if llm_router is None: raise ProxyException( message="LLM Router not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if not isinstance(managed_files_obj, BaseFileEndpoints): raise ProxyException( message="Managed files hook is not a BaseFileEndpoints", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) @@ -1021,17 +1026,17 @@ async def get_file_content( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) @@ -1151,15 +1156,15 @@ async def get_file( if managed_files_obj is None: raise ProxyException( message="Managed files hook not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if not isinstance(managed_files_obj, BaseFileEndpoints): raise ProxyException( message="Managed files hook is not a BaseFileEndpoints", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) response = await managed_files_obj.afile_retrieve( @@ -1215,17 +1220,17 @@ async def get_file( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) @@ -1355,22 +1360,22 @@ async def delete_file( if managed_files_obj is None: raise ProxyException( message="Managed files hook not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if llm_router is None: raise ProxyException( message="LLM Router not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if not isinstance(managed_files_obj, BaseFileEndpoints): raise ProxyException( message="Managed files hook is not a BaseFileEndpoints", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) @@ -1427,17 +1432,17 @@ async def delete_file( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) @@ -1629,15 +1634,15 @@ async def list_files( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 32da2658b99..28a8bab1f24 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,6 +15,7 @@ import os import re from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast @@ -33,6 +34,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse @@ -52,6 +54,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_set_request_parsed_body, get_form_data, get_request_body, + is_json_content_type, ) from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, @@ -77,6 +80,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, ) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials +from litellm.types.router import LiteLLMParamsTypedDict from litellm.types.utils import LlmProviders from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager @@ -93,7 +97,6 @@ else: vertex_llm_base: Final = VertexBase() router: Final = APIRouter() -openai_passthrough_router: Final = APIRouter() default_vertex_config: Final = None passthrough_endpoint_router: Final = PassthroughEndpointRouter() @@ -119,6 +122,24 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li return False +class RelayRejection(TypedDict): + error: ReadOnly[str] + + +def _deployment_model_name(litellm_params: LiteLLMParamsTypedDict) -> str: + model: Final = litellm_params.get("model", "") + try: + return get_llm_provider(model=model, custom_llm_provider=litellm_params.get("custom_llm_provider"))[0] + except litellm.BadRequestError: + return model + + +def _models_served_by_group(llm_router: litellm.Router, model_group: str) -> frozenset[str]: + return frozenset( + _deployment_model_name(row["litellm_params"]) for row in llm_router.get_model_list(model_name=model_group) or () + ) + + def is_passthrough_request_streaming(request_body: object) -> bool: """ Returns True if the request is streaming. @@ -411,7 +432,7 @@ async def vllm_proxy_route( content=None, data=None, files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), + json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), params=None, headers=None, cookies=None, @@ -712,6 +733,10 @@ BEDROCK_ENDPOINT_ACTIONS: Final = { BEDROCK_STREAMING_ACTIONS: Final = {"invoke-with-response-stream", "converse-stream"} +def is_bedrock_count_tokens_endpoint(endpoint: str) -> bool: + return "count_tokens" in endpoint or "count-tokens" in endpoint + + def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: """ Extract model name from Bedrock endpoint path. @@ -942,7 +967,14 @@ async def handle_bedrock_count_tokens( except BedrockError as e: # Convert BedrockError to HTTPException for FastAPI verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e) - raise HTTPException(status_code=e.status_code, detail={"error": e.message}) + from litellm.litellm_core_utils.llm_response_utils.get_headers import get_response_headers + + provider_headers: Final = getattr(getattr(e, "response", None), "headers", None) + raise HTTPException( + status_code=e.status_code, + detail={"error": e.message}, + headers=get_response_headers(provider_headers) if provider_headers else None, + ) except HTTPException: # Re-raise HTTP exceptions as-is raise @@ -986,8 +1018,7 @@ async def bedrock_llm_proxy_route( request_body: Final = await _read_request_body(request=request) - # Special handling for count_tokens endpoints - if "count_tokens" in endpoint or "count-tokens" in endpoint: + if is_bedrock_count_tokens_endpoint(endpoint): return await handle_bedrock_count_tokens( endpoint=endpoint, request=request, @@ -1089,13 +1120,6 @@ async def bedrock_proxy_route( """ create_request_copy(request) - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - aws_region_name: Final = get_secret_str(secret_name="AWS_REGION_NAME") if not _is_bedrock_agent_runtime_route(endpoint=endpoint): return await bedrock_llm_proxy_route( @@ -1126,20 +1150,24 @@ async def bedrock_proxy_route( ) # Add or update query parameters + from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post from litellm.llms.bedrock.chat import BedrockConverseLLM bedrock_llm: Final = BedrockConverseLLM() - credentials: Final[Credentials] = bedrock_llm.get_credentials() - sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name) - headers: Final = {"Content-Type": "application/json"} # Assuming the body contains JSON data, parse it try: data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - _request: Final = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) - sigv4.add_auth(_request) - prepped: Final = _request.prepare() + prepped: Final = await run_aws_signing( + sign_aws_json_post, + get_credentials=bedrock_llm.get_credentials, + service_name="bedrock", + aws_region_name=aws_region_name, + url=str(updated_url), + body=json.dumps(data), + headers=MappingProxyType({"Content-Type": "application/json"}), + ) ## check for streaming is_streaming_request = False @@ -1197,13 +1225,6 @@ async def comprehend_medical_proxy_route( [Docs](https://docs.litellm.ai/docs/pass_through/comprehend_medical) """ - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call comprehendmedical. Run 'pip install boto3'.") - from .llm_provider_handlers.comprehend_medical_passthrough_logging_handler import ( COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS, ) @@ -1234,20 +1255,23 @@ async def comprehend_medical_proxy_route( if "stream" in data: raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member") - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post - credentials: Final[Credentials] = BaseAWSLLM().get_credentials(aws_region_name=aws_region_name) - sigv4: Final = SigV4Auth(credentials, "comprehendmedical", aws_region_name) - headers: Final = MappingProxyType( - { - "Content-Type": "application/x-amz-json-1.1", - "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", - } - ) target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" - _request: Final = AWSRequest(method="POST", url=target_url, data=json.dumps(data), headers=headers) - sigv4.add_auth(_request) - prepped: Final = _request.prepare() + prepped: Final = await run_aws_signing( + sign_aws_json_post, + get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), + service_name="comprehendmedical", + aws_region_name=aws_region_name, + url=target_url, + body=json.dumps(data), + headers=MappingProxyType( + { + "Content-Type": "application/x-amz-json-1.1", + "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", + } + ), + ) endpoint_func: Final = create_pass_through_route( endpoint=operation, @@ -1495,6 +1519,14 @@ async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> Async await upstream.aclose() +async def _relay_upstream_response(upstream: httpx.Response) -> Response: + return Response( + content=await upstream.aread(), + status_code=upstream.status_code, + headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), + ) + + async def _relay_azure_router_model( llm_router: litellm.Router, model: str, @@ -1504,30 +1536,37 @@ async def _relay_azure_router_model( is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, ) -> Response: - result: Final = await llm_router.allm_passthrough_route( - model=model, - method=request.method, - endpoint=endpoint, - request_query_params=request.query_params, - request_headers=_safe_get_request_headers(request), - stream=is_streaming_request, - content=None, - data=None, - files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), - params=None, - headers=None, - cookies=None, - litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + foreign_deployment: Final = foreign_azure_deployment( + endpoint, model, lambda: _models_served_by_group(llm_router, model) ) + if foreign_deployment is not None: + rejection: Final[RelayRejection] = { + "error": f"deployment '{foreign_deployment}' in the path is not served by model group '{model}'; " + "put the model group name in the deployments segment" + } + raise HTTPException(status_code=400, detail=rejection) + try: + result: Final = await llm_router.allm_passthrough_route( + model=model, + method=request.method, + endpoint=endpoint, + request_query_params=request.query_params, + request_headers=_safe_get_request_headers(request), + stream=is_streaming_request, + content=None, + data=None, + files=None, + json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), + params=None, + headers=None, + cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + ) + except httpx.HTTPStatusError as upstream_error: + return await _relay_upstream_response(upstream_error.response) if not is_streaming_request: - upstream: Final = cast(httpx.Response, result) - return Response( - content=await upstream.aread(), - status_code=upstream.status_code, - headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), - ) + return await _relay_upstream_response(cast(httpx.Response, result)) if inspect.isasyncgen(result): sse_headers: Final = {"content-type": "text/event-stream"} @@ -2257,11 +2296,6 @@ async def vertex_proxy_route( ) -@openai_passthrough_router.api_route( - "/openai_passthrough/{endpoint:path}", - methods=["GET", "POST", "PUT", "DELETE", "PATCH"], - tags=["OpenAI Pass-through", "pass-through"], -) @router.api_route( "/openai/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 5f6489a69ca..93fe3c5b31b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -4,6 +4,7 @@ OpenAI Passthrough Logging Handler Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions. """ +from collections.abc import Mapping, Sequence from datetime import datetime from typing import Final from urllib.parse import urlparse @@ -16,6 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, ) +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.openai.openai import OpenAIConfig as OpenAIConfigType from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig @@ -96,6 +98,47 @@ def _is_openai_compatible_url(url_route: str | None) -> bool: return False +def _is_remote_high_detail_image(part: object) -> bool: + if not isinstance(part, Mapping) or part.get("type") != "image_url": + return False + image_url: Final = part.get("image_url") + if not isinstance(image_url, Mapping): + return False + url: Final = image_url.get("url") + return ( + isinstance(url, str) and url.lower().startswith(("http://", "https://")) and image_url.get("detail") == "high" + ) + + +def _content_parts(message: Mapping[str, object]) -> Sequence[object]: + content: Final = message.get("content") + return content if isinstance(content, list) else () + + +def _without_remote_high_detail_images(message: Mapping[str, object]) -> Mapping[str, object]: + if not isinstance(message.get("content"), list): + return message + kept_parts: Final = [ # mutable-ok: token_counter reads message content only when it is a list + part for part in _content_parts(message) if not _is_remote_high_detail_image(part) + ] + return {**message, "content": kept_parts} # mutable-ok: token_counter rejects any message that is not a dict + + +def count_relayed_prompt_tokens(model: str, messages: Sequence[Mapping[str, object]] | None) -> int: + if messages is None: + return 0 + remote_high_detail_images: Final = sum( + 1 for message in messages for part in _content_parts(message) if _is_remote_high_detail_image(part) + ) + local_messages: Final = [ # mutable-ok: token_counter takes a list of messages + _without_remote_high_detail_images(message) for message in messages + ] + return ( + litellm.token_counter(model=model, messages=local_messages) + + high_detail_image_token_upper_bound() * remote_high_detail_images + ) + + class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): """ OpenAI-specific passthrough logging handler that provides cost tracking for /chat/completions endpoints. @@ -512,9 +555,10 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): def _build_complete_streaming_response( self, - all_chunks: list[str], + all_chunks: Sequence[str], litellm_logging_obj: LiteLLMLoggingObj, model: str, + messages: Sequence[Mapping[str, object]] | None = None, ) -> ModelResponse | TextCompletionResponse | None: """ Builds complete response from raw chunks for OpenAI streaming responses. @@ -558,7 +602,11 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return None # Build complete response from chunks - complete_streaming_response: Final = litellm.stream_chunk_builder(chunks=all_openai_chunks) + complete_streaming_response: Final = litellm.stream_chunk_builder( + chunks=all_openai_chunks, + messages=messages, + count_prompt_tokens=lambda: count_relayed_prompt_tokens(model, messages), + ) return complete_streaming_response diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 49ec18013b5..119a53c2411 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -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.""" diff --git a/litellm/proxy/pass_through_endpoints/openai_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/openai_passthrough_endpoints.py new file mode 100644 index 00000000000..f56a59dd560 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/openai_passthrough_endpoints.py @@ -0,0 +1,44 @@ +"""/openai_passthrough must be matched ahead of the native /{provider}/v1/files and +/{provider}/v1/batches routes, so unlike the other provider passthrough routes it is +registered at startup and defers to the lazily loaded handler per call.""" + +from typing import Final + +from fastapi import APIRouter, Depends, Request, Response + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router: Final = APIRouter() + + +@router.api_route( + "/openai_passthrough/{endpoint:path}", + methods=["GET", "POST", "PUT", "DELETE", "PATCH"], + tags=["OpenAI Pass-through", "pass-through"], +) +async def openai_passthrough_route( + endpoint: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> Response: + """ + Dedicated pass-through to the OpenAI API with no overlap with LiteLLM's native + implementations (e.g. the Responses API at /v1/responses). + + Examples: + - /openai_passthrough/v1/responses + - /openai_passthrough/v1/responses/{response_id} + - /openai_passthrough/v1/responses/{response_id}/input_items + + [Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough) + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import openai_proxy_route + + return await openai_proxy_route( + endpoint=endpoint, + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3cb9acc6110..ea4ede7e513 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -78,6 +78,11 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, ) +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) @@ -311,9 +316,9 @@ async def chat_completion_pass_through_endpoint( error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) @@ -322,7 +327,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): def get_response_headers( headers: httpx.Headers, litellm_call_id: str | None = None, - custom_headers: dict | None = None, + custom_headers: Mapping[str, str] | None = None, ) -> dict: # Exclude headers that uvicorn writes itself (server, date) and # encoding/length headers that don't survive re-serialization. @@ -1728,18 +1733,18 @@ async def pass_through_request( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(getattr(e, "detail", str(e)))), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), headers=custom_headers, ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), headers=custom_headers, ) diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 05df44242aa..a8ead86ac36 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -30,6 +30,23 @@ class PolicyAttachmentMatch(TypedDict): matched_via: str +def _attachment_specificity(attachment: PolicyAttachment) -> tuple[int, int]: + if attachment.is_global(): + return (0, 0) + + dims: Final = tuple( + specificity + for values, specificity in ( + (attachment.teams, 1), + (attachment.keys, 2), + (attachment.tags, 3), + (attachment.models, 4), + ) + if values + ) + return (max(dims, default=0), len(dims)) + + class AttachmentRegistry: """ In-memory registry for storing and managing policy attachments. @@ -116,31 +133,26 @@ class AttachmentRegistry: """ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - results: Final[list[PolicyAttachmentMatch]] = [] - seen_policies: Final[set[str]] = set() + matching_attachments: Final = sorted( + ( + attachment + for attachment in self._attachments + if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context) + ), + key=_attachment_specificity, + ) + unique_attachments: Final = tuple( + next(attachment for attachment in matching_attachments if attachment.policy == policy_name) + for policy_name in dict.fromkeys(attachment.policy for attachment in matching_attachments) + ) - for attachment in self._attachments: - scope = attachment.to_policy_scope() - if PolicyMatcher.scope_matches(scope=scope, context=context): - if attachment.policy not in seen_policies: - seen_policies.add(attachment.policy) - matched_via = self._describe_match_reason(attachment, context) - results.append( - { - "policy_name": attachment.policy, - "matched_via": matched_via, - } - ) - verbose_proxy_logger.debug( - "Attachment matched: policy=%s, matched_via=%s, context=(team=%s, key=%s, model=%s)", - attachment.policy, - matched_via, - context.team_alias, - context.key_alias, - context.model, - ) - - return results + return [ + { + "policy_name": attachment.policy, + "matched_via": self._describe_match_reason(attachment, context), + } + for attachment in unique_attachments + ] @staticmethod def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str: diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 4be0f556ed7..ed193c7f434 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -5,18 +5,27 @@ Runs guardrails sequentially per pipeline step definitions, handling pass/fail actions (allow, block, next, modify_response) and data forwarding. """ +import copy import time -from collections.abc import Sequence -from typing import Any, Final, Literal +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar + +from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import LOGS_GUARDRAIL_INFORMATION_MARKER from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import independent_snapshot +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, + independent_snapshot, +) +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -25,6 +34,14 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( PipelineStep, PipelineStepResult, ) +from litellm.types.utils import GenericGuardrailAPIInputs, StandardLoggingGuardrailInformation + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) + from litellm.proxy._types import UserAPIKeyAuth try: from fastapi.exceptions import HTTPException @@ -32,6 +49,221 @@ except ImportError: HTTPException = None +class UndeliverableStreamRewrite(Exception): + def __init__(self, guardrail_name: str) -> None: + super().__init__( + f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's " + "streaming pipeline cannot deliver" + ) + self.guardrail_name: Final = guardrail_name + + +def _tool_call_shape(tool_call: object) -> tuple[object, object]: + plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call + function: Final = plain.get("function") if isinstance(plain, Mapping) else None + if not isinstance(function, Mapping): + return (None, None) + return (function.get("name"), function.get("arguments")) + + +def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None: + return None if texts is None else tuple(texts) + + +def _scanned_texts(texts: Sequence[str] | None) -> tuple[str, ...]: + return tuple(texts or ()) + + +def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None: + return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls) + + +def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: + return sent is not None and returned is not None and returned != sent + + +def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: + return sent is not None and returned is not None and len(returned) != len(sent) + + +_GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object]) + + +def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT: + vars(method)[LOGS_GUARDRAIL_INFORMATION_MARKER] = True # rebind-ok: stamps the method the class body just defined + return method + + +class _StreamRewriteObserver(CustomGuardrail): + """Stand-in handed to the endpoint translation in place of a streaming pipeline step's + guardrail. It records whether the guardrail returned different output than it was given, + which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text and + tool-call rewrites are deliverable on translations that write them back across the + buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation, + and a rewrite that drops or adds a tool call on any translation, are discarded by the + executor, which releases the original chunks. + The inner guardrail's ``apply_guardrail`` already records the guardrail information + and span, so the observer's stays out of ``log_guardrail_information``.""" + + def __init__(self, inner: CustomGuardrail) -> None: + super().__init__(guardrail_name=inner.guardrail_name) + self.inner: Final = inner + self.rewrote_texts = False + self.rewrote_tool_calls = False + self.changed_tool_call_count = False + + def structured_messages_cover_full_request(self) -> bool: + return self.inner.structured_messages_cover_full_request() + + @_logged_by_inner_guardrail + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + sent_texts: Final = _text_snapshot(inputs.get("texts")) + sent_tool_shapes: Final = _tool_call_shapes(inputs.get("tool_calls")) + outputs: Final = await self.inner.apply_guardrail( + inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj + ) + returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls")) + self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) + self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes) + self.changed_tool_call_count = self.changed_tool_call_count or _changed_count( + sent_tool_shapes, returned_tool_shapes + ) + return outputs + + +class _ScannedTextRecorder(CustomGuardrail): + def __init__(self, guardrail_name: str) -> None: + super().__init__(guardrail_name=guardrail_name) + self.inputs: GenericGuardrailAPIInputs | None = None + + @_logged_by_inner_guardrail + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + self.inputs = inputs + return inputs + + +class _LegacyHookStreamAdapter(CustomGuardrail): + """Runs a guardrail that only implements the legacy post-call hook (no unified + ``apply_guardrail``, or ``use_native_lifecycle_hooks``) as a streaming pipeline step. The + endpoint translation hands it the texts it scanned plus the assembled response under + ``request_data["response"]``; the hook gets that response in the shape its route gives + non-streaming hooks, an exception it raises ends the stream through the executor's + fail/error classification, and the response it hands back, or the one it changed in place + and returned ``None`` for, is re-scanned by the same translation so its texts reach the + client through the translation's ended-stream write-back. A + replacement whose scanned texts do not line up with the originals, or whose tool calls + differ from them, is undeliverable, so the executor releases the original chunks. A stream + that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as + long as the hook left the tool calls alone.""" + + def __init__( + self, + inner: CustomGuardrail, + endpoint_translation: "BaseTranslation", + user_api_key_dict: "UserAPIKeyAuth", + ) -> None: + super().__init__(guardrail_name=inner.guardrail_name) + self.inner: Final = inner + self.endpoint_translation: Final = endpoint_translation + self.user_api_key_dict: Final = user_api_key_dict + + def structured_messages_cover_full_request(self) -> bool: + return self.inner.structured_messages_cover_full_request() + + @_logged_by_inner_guardrail + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response")) + replacement: Final = await self.inner.async_post_call_success_hook( + data=request_data, + user_api_key_dict=self.user_api_key_dict, + response=hooked, + ) + rewrite: Final = hooked if replacement is None else replacement + if rewrite is None: + return inputs + rescanned: Final = await self._rescan(rewrite, logging_obj) + if rescanned is None: + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + rewritten: Final = rescanned.get("texts") + if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))): + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + if not rewritten: + return inputs + rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} + return rewritten_inputs + + async def _rescan( + self, response: object, logging_obj: "LiteLLMLoggingObj | None" + ) -> GenericGuardrailAPIInputs | None: + recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown") + await self.endpoint_translation.process_output_response( + response=response, + guardrail_to_apply=recorder, + litellm_logging_obj=logging_obj, + user_api_key_dict=self.user_api_key_dict, + ) + return recorder.inputs + + +def _prepare_hook_input( + step: PipelineStep, + callback: CustomGuardrail, + data: dict, # mutable-ok: same request-payload shape the hooks mutate + raw_request_snapshot: dict | None, # mutable-ok: same request-payload shape as data +) -> tuple[dict, bool]: # mutable-ok: returns that same request-payload dict + """Inject the step's guardrail name into metadata so should_run_guardrail() allows it, + and pick the payload the step scans: a scan_raw_request step evaluates the pristine + pre-pipeline snapshot instead of `data` (which earlier pass_data steps in this same + pipeline may have already rewritten), same reason the normal sequential/parallel + guardrail loops do this.""" + if "metadata" not in data: + data["metadata"] = {} # mutable-ok: request metadata bucket, hooks mutate it + data["metadata"]["guardrails"] = [ + step.guardrail + ] # mutable-ok: guardrails list is part of the request-payload shape + + scans_raw_request: Final = callback.scan_raw_request + hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data + independent_snapshot(raw_request_snapshot) if scans_raw_request and raw_request_snapshot is not None else data + ) + if hook_input is not data: + hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] # mutable-ok: request metadata shape + return hook_input, scans_raw_request + + +def _release_original_chunks( + guardrail_name: str, + streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks, restored in place + originals: Sequence[object], +) -> None: + streaming_chunks[:] = originals # rebind-ok: the caller's buffer is the stream the client receives + verbose_proxy_logger.warning( + "Pipeline: guardrail '%s' rewrote the streamed response in a way this endpoint's streaming " + "pipeline cannot deliver yet; the rewrite was discarded and the original stream released", + guardrail_name, + ) + + class PipelineExecutor: """Executes guardrail pipelines with ordered, conditional step logic.""" @@ -44,6 +276,8 @@ class PipelineExecutor: call_type: str, policy_name: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + endpoint_translation: "BaseTranslation | None" = None, ) -> PipelineExecutionResult: """ Execute pipeline steps sequentially with conditional actions. @@ -60,6 +294,12 @@ class PipelineExecutor: step whose guardrail opted into ``scan_raw_request`` evaluates the original request instead of whatever an earlier ``pass_data`` step in this same pipeline already rewrote. + streaming_chunks: buffered chunks of a completed stream. When set + (with ``endpoint_translation``), post_call steps scan the + assembled streamed output through the endpoint translation + instead of calling ``async_post_call_success_hook``. + endpoint_translation: the guardrail translation for the streamed + endpoint, resolved by the caller. Returns: PipelineExecutionResult with terminal action and step results @@ -84,6 +324,8 @@ class PipelineExecutor: user_api_key_dict=user_api_key_dict, call_type=call_type, raw_request_snapshot=raw_request_snapshot, + streaming_chunks=streaming_chunks, + endpoint_translation=endpoint_translation, ) duration = time.perf_counter() - start_time @@ -109,8 +351,10 @@ class PipelineExecutor: action, ) - # Forward modified data to next step if pass_data is True - if step.pass_data and modified_data is not None: + # Forward modified data to the next step if pass_data is True; + # post_call response replacements always chain, matching the flat + # callback loop where each hook sees the previous hook's response + if modified_data is not None and (step.pass_data or mode == "post_call"): working_data = {**working_data, **modified_data} # Handle terminal actions @@ -118,18 +362,22 @@ class PipelineExecutor: return _allow_result(step_results=step_results, working_data=working_data, request_data=data) if action == "block": + _carry_working_guardrail_information(working_data=working_data, request_data=data) return PipelineExecutionResult( terminal_action="block", step_results=step_results, error_message=error_detail, original_exception=original_exception, + modified_data=working_data if working_data != data else None, ) if action == "modify_response": + _carry_working_guardrail_information(working_data=working_data, request_data=data) return PipelineExecutionResult( terminal_action="modify_response", step_results=step_results, modify_response_message=step.modify_response_message or error_detail, + modified_data=working_data if working_data != data else None, ) # action == "next" → continue to next step @@ -137,6 +385,65 @@ class PipelineExecutor: # Ran out of steps without a terminal action → default allow return _allow_result(step_results=step_results, working_data=working_data, request_data=data) + @staticmethod + async def _run_streaming_step( + step: PipelineStep, + callback: CustomGuardrail, + endpoint_translation: "BaseTranslation", + streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place + hook_input: dict[str, object], # mutable-ok: same request-payload shape as data + user_api_key_dict: "UserAPIKeyAuth", + litellm_logging_obj: "LiteLLMLoggingObj | None", + ) -> None: + """Run one streaming post_call step through the endpoint translation, delivering + text and tool-call rewrites on translations that support ended-stream write-back. A + guardrail without the unified interface runs its legacy post-call hook against the + assembled response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the + client yet (one on a translation without write-back, one that drops or adds a tool call, + or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is + discarded: the buffered chunks go back to the originals and the step passes, so the + client gets the stream the merge base sent, and the guardrail stays out of the + applied-guardrails header since its output never reached the client. The response an + earlier step's translation stored under ``request_data["response"]`` is dropped first, + so this step's hook sees the stream as the steps before it left it.""" + scanner: Final = ( + callback + if PipelineExecutor.supports_unified_execution(callback) + else _LegacyHookStreamAdapter(callback, endpoint_translation, user_api_key_dict) + ) + observer: Final = _StreamRewriteObserver(scanner) + deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites + originals: Final = copy.deepcopy(streaming_chunks) + hook_input.pop("response", None) # rebind-ok: an earlier step's stored response goes so this step's is stored + try: + if deliver_rewrites: + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=observer, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + deliver_ended_stream_rewrites=True, + ) + else: + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=observer, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + ) + except UndeliverableStreamRewrite: + _release_original_chunks(step.guardrail, streaming_chunks, originals) + return + if observer.changed_tool_call_count or ( + not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls) + ): + _release_original_chunks(step.guardrail, streaming_chunks, originals) + return + if not callback.records_own_guardrail_information: + add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) + @staticmethod async def _run_step( step: PipelineStep, @@ -145,6 +452,8 @@ class PipelineExecutor: user_api_key_dict: Any, call_type: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + endpoint_translation: "BaseTranslation | None" = None, ) -> tuple[ Literal["pass", "fail", "error"], dict | None, @@ -168,34 +477,17 @@ class PipelineExecutor: verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail) return ("error", None, f"Guardrail '{step.guardrail}' not found", None) + hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot) + snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input)) + + # Use unified_guardrail path if callback implements apply_guardrail + target: CustomLogger = callback + use_unified: Final = PipelineExecutor.supports_unified_execution(callback) + if use_unified and streaming_chunks is None: + hook_input["guardrail_to_apply"] = callback + target = UnifiedLLMGuardrails() + try: - # Inject guardrail name into metadata so should_run_guardrail() allows it - if "metadata" not in data: - data["metadata"] = {} - data["metadata"]["guardrails"] = [step.guardrail] - - # A scan_raw_request step evaluates the pristine pre-pipeline - # snapshot instead of `data` (which earlier pass_data steps in - # this same pipeline may have already rewritten), same reason - # the normal sequential/parallel guardrail loops do this. - scans_raw_request: Final = callback.scan_raw_request - hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data - independent_snapshot(raw_request_snapshot) - if scans_raw_request and raw_request_snapshot is not None - else data - ) - if hook_input is not data: - hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] - - # Use unified_guardrail path if callback implements apply_guardrail - target: CustomLogger = callback - use_unified: Final = ( - "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks - ) - if use_unified: - hook_input["guardrail_to_apply"] = callback - target = UnifiedLLMGuardrails() - if mode == "pre_call": response = await target.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -207,6 +499,24 @@ class PipelineExecutor: callback.mark_pre_call_hook_ran(data) if isinstance(response, dict): callback.mark_pre_call_hook_ran(response) + elif mode == "post_call" and streaming_chunks is not None: + if endpoint_translation is None: + return ( + "error", + None, + f"Guardrail '{step.guardrail}' cannot run on a stream without an endpoint translation", + None, + ) + await PipelineExecutor._run_streaming_step( + step=step, + callback=callback, + endpoint_translation=endpoint_translation, + streaming_chunks=streaming_chunks, + hook_input=hook_input, + user_api_key_dict=user_api_key_dict, + litellm_logging_obj=data.get("litellm_logging_obj"), + ) + response = None elif mode == "post_call": response = await target.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, @@ -220,11 +530,19 @@ class PipelineExecutor: # same contract as run_in_parallel/scan_raw_request elsewhere: any # data it returned is discarded, since applying it on top of the # raw snapshot would silently undo whatever an earlier step in - # this pipeline already did. - modified_data = None - if response is not None and isinstance(response, dict) and not scans_raw_request: - modified_data = response - return ("pass", modified_data, None, None) + # this pipeline already did. A post_call hook's non-None return is + # a replacement response (the flat callback-loop contract), carried + # under the same "response" key the step input uses. + if response is None or scans_raw_request: + return ("pass", None, None, None) + if mode == "post_call": + return ( + "pass", + {"response": response}, + None, + None, + ) # mutable-ok: modified-data contract is a plain dict + return ("pass", response if isinstance(response, dict) else None, None, None) except Exception as e: if CustomGuardrail._is_guardrail_intervention(e): @@ -233,6 +551,30 @@ class PipelineExecutor: else: verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e) return ("error", None, str(e), e) + finally: + if hook_input is not data: + _append_guardrail_information( + request_data=data, + entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:], + ) + + @staticmethod + def supports_unified_execution(callback: CustomGuardrail) -> bool: + """Whether this guardrail runs through the unified apply_guardrail path.""" + return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + + @staticmethod + def supports_streaming_execution(callback: CustomGuardrail) -> bool: + """Whether a streaming pipeline step can run this guardrail against the buffered + stream: through the unified path, or through its post-call hook on the assembled + response when that hook is its only streaming path. A guardrail with its own + streaming iterator hook, or with neither hook, keeps running on its own.""" + callback_type: Final = type(callback) + return PipelineExecutor.supports_unified_execution(callback) or ( + callback_type.async_post_call_success_hook is not CustomLogger.async_post_call_success_hook + and callback_type.async_post_call_streaming_iterator_hook + is CustomLogger.async_post_call_streaming_iterator_hook + ) @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: @@ -283,6 +625,40 @@ def _restore_request_guardrails( return {**working_data, "metadata": stripped} # mutable-ok: request dict +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +def _recorded_guardrail_information(source: Mapping[str, object]) -> list[StandardLoggingGuardrailInformation]: + bucket: Final = source.get(get_metadata_variable_name_from_kwargs(source)) + recorded: Final = bucket.get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None + return recorded if isinstance(recorded, list) else [] + + +def _append_guardrail_information( + request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data + entries: Sequence[StandardLoggingGuardrailInformation], +) -> None: + if not entries: + return + _, request_bucket = get_or_create_metadata_bucket(request_data) + existing: Final = request_bucket.get(_GUARDRAIL_INFORMATION_KEY) + if isinstance(existing, list): + existing.extend(entries) + return + request_bucket[_GUARDRAIL_INFORMATION_KEY] = list(entries) + + +def _carry_working_guardrail_information( + working_data: Mapping[str, object], + request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data +) -> None: + recorded: Final = _recorded_guardrail_information(working_data) + existing: Final = _recorded_guardrail_information(request_data) + if recorded is existing: + return + _append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing]) + + def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py new file mode 100644 index 00000000000..d284c44397e --- /dev/null +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -0,0 +1,152 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + add_policy_sources_to_metadata, + add_policy_to_applied_policies_header, +) +from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import resolve_model_group_alias +from litellm.types.proxy.policy_engine import PolicyMatchContext +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router + +PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...] + +_POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines) + + +@dataclass(frozen=True, slots=True) +class UngovernedRetrieval: + reason: Literal["no router", "response id names no deployment", "deployment no longer in the router"] + + +def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | UngovernedRetrieval: + if llm_router is None: + return UngovernedRetrieval("no router") + model_id: Final = ( + ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) if isinstance(response_id, str) else None + ) + if model_id is None: + return UngovernedRetrieval("response id names no deployment") + deployment: Final = llm_router.get_deployment(model_id) + if deployment is None: + return UngovernedRetrieval("deployment no longer in the router") + hidden_by: Final = _submit_model_hidden_by(deployment.model_name, llm_router.model_group_alias) + if hidden_by is not None: + verbose_proxy_logger.warning( + "Policy engine: background response %s re-matches policies on retrieval as model group %s (%s), " + "so a policy attached to the model name it was submitted as does not run on it", + response_id, + deployment.model_name, + hidden_by, + ) + return deployment.model_name + + +def _submit_model_hidden_by(model_group: str, model_group_alias: Mapping[str, object]) -> str | None: + if "*" in model_group: + return "a wildcard deployment" + aliases: Final = tuple( + alias for alias in model_group_alias if resolve_model_group_alias(model_group_alias, alias) == model_group + ) + if not aliases: + return None + return f"the target of model_group_alias {', '.join(aliases)}" + + +def _retrieval_context( + data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str +) -> PolicyMatchContext: + team_alias: Final = user_api_key_dict.team_alias + key_alias: Final = user_api_key_dict.key_alias + return PolicyMatchContext( + team_alias=team_alias if isinstance(team_alias, str) else None, + key_alias=key_alias if isinstance(key_alias, str) else None, + model=model_group, + tags=get_tags_from_request_body(data) or None, + ) + + +def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]: + matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context) + if not matches: + return (), MappingProxyType({}) + applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( + policy_names=[match["policy_name"] for match in matches], # mutable-ok: the matcher takes a list + context=context, + ) + post_call_pipelines: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context( + context=context, policy_names=applied_policy_names + ) + if pipeline.mode == "post_call" + ) + return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches}) + + +def attach_post_call_pipelines_to_retrieval( + data: dict[str, object], # mutable-ok: request-state dict the policy engine hooks all write in place + user_api_key_dict: "UserAPIKeyAuth", + llm_router: "Router | None", +) -> None: + if not get_policy_registry().is_initialized(): + return + model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router) + if isinstance(model_group, UngovernedRetrieval): + verbose_proxy_logger.warning( + "Policy engine: background response %s is retrieved without its post_call policy pipelines (%s)", + data.get("response_id"), + model_group.reason, + ) + return + context: Final = _retrieval_context(data, user_api_key_dict, model_group) + post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context) + _, bucket = get_or_create_metadata_bucket(data) + already_attached: Final = _POLICY_PIPELINES_ADAPTER.validate_python(bucket.get("_guardrail_pipelines") or ()) + attached_policy_names: Final = frozenset(policy_name for policy_name, _pipeline in already_attached) + added: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in post_call_pipelines + if policy_name not in attached_policy_names + ) + if not added: + return + pipelines: Final = (*already_attached, *added) + bucket["_guardrail_pipelines"] = pipelines + bucket["_pipeline_managed_guardrails"] = frozenset( + step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps + ) + for policy_name, _pipeline in added: + add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name) + for _policy_name, pipeline in added: + for step in pipeline.steps: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail) + add_policy_sources_to_metadata( + request_data=data, + policy_sources={ # mutable-ok: add_policy_sources_to_metadata takes a dict + policy_name: policy_sources[policy_name] for policy_name, _pipeline in added + }, + ) + verbose_proxy_logger.debug( + "Policy engine: attached post_call pipelines to the retrieval of background response %s (model group %s): %s", + data.get("response_id"), + model_group, + ", ".join(policy_name for policy_name, _pipeline in added), + ) diff --git a/litellm/proxy/prometheus_cleanup.py b/litellm/proxy/prometheus_cleanup.py index d9827723887..c65aaeedfaa 100644 --- a/litellm/proxy/prometheus_cleanup.py +++ b/litellm/proxy/prometheus_cleanup.py @@ -8,10 +8,13 @@ from __future__ import annotations import glob import os +import re from typing import Final from litellm._logging import verbose_proxy_logger +_LIVE_GAUGE_PID: Final = re.compile(r"gauge_live[a-z]*_(\d+)\.db$") + def wipe_directory(directory: str) -> None: """Delete all .db files in the directory. Called once before workers fork.""" @@ -38,3 +41,35 @@ def mark_worker_exit(worker_pid: int) -> None: verbose_proxy_logger.info("Prometheus cleanup: marked worker %s as dead", worker_pid) except Exception as e: verbose_proxy_logger.warning("Failed to mark prometheus worker %s as dead: %s", worker_pid, e) + + +def _is_running(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def mark_dead_workers(directory: str) -> tuple[int, ...]: + """Drop the live-gauge files of workers that no longer exist and return their pids. + + Uvicorn's multi-worker supervisor has no exit hook, so a replacement worker calls this at startup; without it + a crashed worker's in-flight gauges stay in the aggregate forever. + """ + owners: Final = frozenset( + int(match.group(1)) + for match in map(_LIVE_GAUGE_PID.search, glob.glob(os.path.join(directory, "gauge_live*_*.db"))) + if match is not None + ) + dead: Final = tuple(sorted(pid for pid in owners if pid != os.getpid() and not _is_running(pid))) + if not dead: + return dead + from prometheus_client import multiprocess + + for pid in dead: + multiprocess.mark_process_dead(pid, path=directory) + verbose_proxy_logger.info("Prometheus cleanup: marked dead workers %s in %s", dead, directory) + return dead diff --git a/litellm/proxy/prometheus_metrics_server.py b/litellm/proxy/prometheus_metrics_server.py index 4a9651d62e1..01479f9dd41 100644 --- a/litellm/proxy/prometheus_metrics_server.py +++ b/litellm/proxy/prometheus_metrics_server.py @@ -30,6 +30,7 @@ from litellm.integrations.prometheus_metrics_endpoint import make_metrics_asgi_a from litellm.llms.custom_httpx.http_handler import HTTPHandler METRICS_PATH: Final = "/metrics" +HEALTH_PATH: Final = "/health" PID_HEADER: Final = "x-litellm-metrics-pid" _PARENT_POLL_INTERVAL_SECONDS: Final = 1.0 _STARTUP_TIMEOUT_SECONDS: Final = 30.0 @@ -77,6 +78,10 @@ def build_metrics_app(multiproc_dir: str) -> FastAPI: app: Final = FastAPI(title="LiteLLM Prometheus metrics", docs_url=None, redoc_url=None, openapi_url=None) app.mount(METRICS_PATH, _add_pid_header(make_metrics_asgi_app(registry))) + @app.get(HEALTH_PATH) + def health() -> dict[str, str]: + return {"status": "healthy", "multiproc_dir": multiproc_dir} + return app diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e245367b1b4..c2d60cd5488 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -18,6 +18,12 @@ from pydantic import BaseModel, ConfigDict import litellm from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY +from litellm.proxy.db.pgbouncer import ( + PgBouncerError, + PgBouncerSettings, + export_pooled_database_url, + start_in_container_pgbouncer, +) from litellm.proxy.db.query_engine_reaper import start_query_engine_reaper if TYPE_CHECKING: @@ -1108,6 +1114,7 @@ def run_server( from litellm.proxy.db.token_auth import ( AZURE_POSTGRESQL_AUTH_ENV_VAR, IAM_TOKEN_DB_AUTH_ENV_VAR, + resolve_database_token_auth, token_auth_flag_enabled, ) @@ -1377,6 +1384,21 @@ def run_server( print( f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa: F541 ) + pgbouncer_settings: Final = PgBouncerSettings() + upstream_database_url: Final = os.getenv("DATABASE_URL") + if pgbouncer_settings.enabled and upstream_database_url is not None: + pooled_database_url: Final = start_in_container_pgbouncer( + pgbouncer_settings, upstream_database_url, token_auth=resolve_database_token_auth() + ) + if isinstance(pooled_database_url, PgBouncerError): + print( + f"\033[1;31mLiteLLM Proxy: LITELLM_PGBOUNCER_ENABLED is set but the in-container pgbouncer " + f"could not start: {pooled_database_url.reason}\033[0m", + file=sys.stderr, + flush=True, + ) + sys.exit(1) + export_pooled_database_url(pooled_database_url) if port == 4000 and ProxyInitializationHelpers._is_port_in_use(port): port = random.randint(1024, 49152) if prometheus_metrics_port == port: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1cd08fe27c0..c8fd7edfed6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16,7 +16,17 @@ import threading import time import traceback import warnings -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Collection, + Mapping, + MutableMapping, + Sequence, +) +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -50,6 +60,7 @@ from litellm.constants import ( AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, AUDIO_SPEECH_CHUNK_SIZE, + BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, BASE_MCP_ROUTE, DAILY_TAG_SPEND_BATCH_MULTIPLIER, DEFAULT_MAX_RECURSE_DEPTH, @@ -62,11 +73,13 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import ( UI_TEAM_ID, CallbackDelete, @@ -125,6 +138,13 @@ 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.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -132,11 +152,7 @@ from litellm.types.utils import ( TextCompletionResponse, TokenCountResponse, ) -from litellm.utils import ( - _invalidate_model_cost_lowercase_map, - load_credentials_from_list, - reapply_runtime_model_cost_registrations, -) +from litellm.utils import load_credentials_from_list if TYPE_CHECKING: from aiohttp import ClientSession @@ -145,6 +161,7 @@ if TYPE_CHECKING: from prisma import models as prisma_models from litellm.integrations.opentelemetry import OpenTelemetry + from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager Span = _Span | Any else: @@ -227,13 +244,14 @@ def generate_feedback_box(): import contextlib from collections import defaultdict from contextlib import asynccontextmanager -from functools import lru_cache +from functools import lru_cache, partial import litellm import litellm._redis from litellm import Router from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, @@ -262,16 +280,16 @@ from litellm.constants import ( WEEKLY_SPEND_REPORT_JOB_ID, ) from litellm.exceptions import RejectedRequestError -from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) -from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, + drop_params_flag, get_litellm_metadata_from_kwargs, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor @@ -286,7 +304,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.proxy._lazy_features import attach_lazy_features +from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot from litellm.proxy._types import * from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, @@ -300,6 +318,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.auth_utils import ( check_response_size_is_safe, is_request_body_safe, + log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check @@ -359,8 +378,11 @@ from litellm.proxy.common_utils.load_config_utils import ( ) from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations from litellm.proxy.common_utils.model_listing_utils import ( + ClaudeCodeRoutingNames, TeamModelNameTranslator, + claude_code_view_ids, configured_display_names, + is_claude_code_client, ) from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, @@ -405,6 +427,7 @@ from litellm.proxy.config_resolvers.alerting import ( ) from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( SPEND_LOG_CLEANUP_BOUND_SETTINGS, SpendLogCleanup, @@ -418,6 +441,7 @@ from litellm.proxy.db.exception_handler import ( ) from litellm.proxy.db.gateway_request_tracking import ( GatewayRequestAccumulator, + GatewayRequestRedisBuffer, flush_gateway_requests, ) from litellm.proxy.db.proxy_worker_heartbeat import ( @@ -446,7 +470,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( from litellm.proxy.hooks.prompt_injection_detection import ( _OPTIONAL_PromptInjectionDetection, ) -from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger +from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger, run_spend_event from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.list_api.common import ( PROBLEM_TYPE_BASE, @@ -569,6 +593,11 @@ from litellm.proxy.plugin_routes import ( from litellm.proxy.plugin_routes import ( router as plugin_router, ) +from litellm.proxy.spend_tracking.spend_event_producer import ( + CollectorSettings, + SpendEventProducer, + build_spend_event_producer, +) from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail try: @@ -592,6 +621,10 @@ from litellm.proxy.middleware.admission_control_middleware import ( from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) +from litellm.proxy.middleware.per_request_root_path_middleware import ( + PerRequestRootPathMiddleware, + get_server_root_paths, +) from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware from litellm.proxy.middleware.request_size_limit_middleware import ( RequestSizeLimitMiddleware, @@ -606,13 +639,8 @@ from litellm.proxy.openai_files_endpoints.files_endpoints import ( from litellm.proxy.openai_files_endpoints.files_endpoints import ( set_files_config, ) -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - openai_passthrough_router, - passthrough_endpoint_router, - vertex_ai_live_websocket_passthrough, -) -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - router as llm_passthrough_router, +from litellm.proxy.pass_through_endpoints.openai_passthrough_endpoints import ( + router as openai_passthrough_router, ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( initialize_pass_through_endpoints, @@ -620,12 +648,14 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( router as pass_through_router, ) +from litellm.proxy.prometheus_cleanup import mark_dead_workers, mark_worker_exit from litellm.proxy.public_endpoints import router as public_endpoints_router from litellm.proxy.public_endpoints.public_v1 import router as public_v1_router from litellm.proxy.rag_endpoints.endpoints import router as rag_router from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router from litellm.proxy.response_api_endpoints.endpoints import router as response_router from litellm.proxy.route_llm_request import route_request +from litellm.proxy.route_priority import hot_routes_first from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start @@ -892,7 +922,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 +942,18 @@ 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_counters_on_shutdown() -> None: + if prisma_client is None: + return + try: + await proxy_logging_obj.db_spend_update_writer.db_update_spend_transaction_handler( + prisma_client=prisma_client, n_retry_times=3, proxy_logging_obj=proxy_logging_obj + ) + except Exception as e: # noqa: BLE001 # shutdown must continue even if the commit fails + verbose_proxy_logger.exception("Error flushing spend counters on shutdown: %s", e) async def _flush_spend_logs_queue_on_shutdown() -> None: @@ -1046,6 +1089,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: init_verbose_loggers() + prometheus_multiproc_dir: Final = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + if prometheus_multiproc_dir: + mark_dead_workers(prometheus_multiproc_dir) + ## RUN WORKER STARTUP HOOKS (e.g., gflags initialization) ## _startup_hooks_env: Final = os.environ.get("LITELLM_WORKER_STARTUP_HOOKS", "") if _startup_hooks_env: @@ -1344,6 +1391,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: except Exception as e: verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e) + await _drain_spend_event_producer_on_shutdown() + + await flush_spend_counters_on_shutdown() + await _flush_spend_logs_queue_on_shutdown() await proxy_config.stop_config_sync_subscriber() @@ -1352,6 +1403,9 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await proxy_shutdown_event(worker_heartbeat=worker_heartbeat) + if prometheus_multiproc_dir: + mark_worker_exit(os.getpid()) + def _generate_stable_operation_id(route: "APIRoute") -> str: operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}") @@ -2267,6 +2321,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 @@ -2334,6 +2389,17 @@ open_telemetry_logger: OpenTelemetry | None = None gateway_request_accumulator: Final = GatewayRequestAccumulator() ### INITIALIZE GLOBAL LOGGING OBJECT ### proxy_logging_obj: ProxyLogging = ProxyLogging(user_api_key_cache=user_api_key_cache, premium_user=premium_user) + + +def _gateway_request_redis_buffer() -> GatewayRequestRedisBuffer | None: + """Shares the spend writer's transaction-buffer Redis and pod lock when use_redis_transaction_buffer is on.""" + writer: Final = proxy_logging_obj.db_spend_update_writer + redis_cache: Final = writer.redis_update_buffer.redis_cache + if redis_cache is None or not writer.redis_update_buffer._should_commit_spend_updates_to_redis(): + return None + return GatewayRequestRedisBuffer(redis_cache=redis_cache, pod_lock_manager=writer.pod_lock_manager) + + ### REDIS QUEUE ### async_result: Final = None celery_app_conn: Final = None @@ -2408,16 +2474,27 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): ) +spend_event_producer: SpendEventProducer | None = None + + def cost_tracking(): - global prisma_client + global prisma_client, spend_event_producer if prisma_client is not None: from litellm.integrations.shadow_eval_logger import ShadowEvalLogger - litellm.logging_callback_manager.add_litellm_callback(_ProxyDBLogger()) - litellm.logging_callback_manager.add_litellm_async_success_callback(_ProxyDBLogger()) + spend_event_producer = build_spend_event_producer(CollectorSettings(), fallback=run_spend_event) + litellm.logging_callback_manager.add_litellm_callback(_ProxyDBLogger(spend_event_producer)) + litellm.logging_callback_manager.add_litellm_async_success_callback(_ProxyDBLogger(spend_event_producer)) litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger()) +async def _drain_spend_event_producer_on_shutdown() -> None: + if spend_event_producer is None: + return + await spend_event_producer.close(drain_timeout=CollectorSettings().drain_timeout_seconds) + verbose_proxy_logger.info("collector: producer drained on shutdown. stats=%s", spend_event_producer.stats()) + + # Bounds authoritative DB re-reads when enforcing a budget against a # stale-low spend counter: at most one DB read per counter per window. SPEND_DB_FLOOR_CACHE_TTL_SECONDS: Final = 5 @@ -2560,10 +2637,11 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None ) -async def reseed_spend_counter_from_db(counter_key: str) -> None: +async def reseed_spend_counter_from_db(counter_key: str) -> bool: """Recover a counter that the reservation reconcile found in an inconsistent state (missing, or where applying the reconcile delta would drive it - negative) by reseeding it from the DB instead of deleting it. + negative) by reseeding it from the DB instead of deleting it. Returns + whether a DB row was found and the counter was reseeded. The DB row is a LAGGING authoritative floor, not post-request truth: the entity .spend column is flushed in batches (every PROXY_BATCH_WRITE_AT), so @@ -2578,8 +2656,9 @@ async def reseed_spend_counter_from_db(counter_key: str) -> None: """ db_spend: Final = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) if db_spend is None: - return + return False await _repair_stale_spend_counter(counter_key=counter_key, db_spend=db_spend) + return True async def _floor_spend_from_db( @@ -2641,21 +2720,14 @@ async def _authoritative_floor_spend( return db_spend -async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) -> tuple[float, bool]: - """Return (spend, authoritative). ``authoritative`` is True when the value - came from Redis or a fresh DB read (cross-pod truth), False when it came - from the per-pod in-memory copy or the caller's fallback. Only the - fail-closed path reads the flag; normal callers ignore it.""" - # 1. Redis first (cross-pod authoritative). On clean miss, skip - # in-memory: per-pod in-memory only has this pod's writes, so it - # would mask cross-pod increments. - redis_clean_miss = False +async def read_spend_counter_cache_value(counter_key: str) -> tuple[float | None, bool]: + """Return (value, authoritative) for the live counter, None when absent. A clean + Redis miss is final: the per-pod in-memory copy outlives the Redis TTL and only + holds this pod's writes, so it is consulted only when Redis is unreachable.""" if spend_counter_cache.redis_cache is not None: try: - val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) - if val is not None: - return float(val), True - redis_clean_miss = True + redis_val: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) + return (float(redis_val) if redis_val is not None else None), True except Exception as e: verbose_proxy_logger.debug( "get_current_spend: Redis read failed for %s, falling back to in-memory: %s", @@ -2663,13 +2735,20 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) e, ) - # 2. In-memory only when Redis is unreachable. - if not redis_clean_miss: - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val), False + in_memory_val: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + return (float(in_memory_val) if in_memory_val is not None else None), False - # 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass. + +async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) -> tuple[float, bool]: + """Return (spend, authoritative). ``authoritative`` is True when the value + came from Redis or a fresh DB read (cross-pod truth), False when it came + from the per-pod in-memory copy or the caller's fallback. Only the + fail-closed path reads the flag; normal callers ignore it.""" + cached_val, cached_authoritative = await read_spend_counter_cache_value(counter_key=counter_key) + if cached_val is not None: + return cached_val, cached_authoritative + + # Reseed from DB - fallback_spend lags cross-pod, would allow bypass. db_spend: Final = await SpendCounterReseed.coalesced( prisma_client=prisma_client, spend_counter_cache=spend_counter_cache, @@ -2682,6 +2761,12 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) return fallback_spend, False +@dataclass(frozen=True, slots=True) +class _PendingSpendIncrement: + counter_key: str + increment: float + + async def increment_spend_counters( token: str | None, team_id: str | None, @@ -2716,7 +2801,7 @@ async def increment_spend_counters( cost: Final[float] = response_cost - async def _key_scope(key_token: str) -> None: + async def _key_scope(key_token: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow # hashes raw "sk-..." keys before they reach the callback). The # startswith("sk-") check is a safety net matching update_cache — @@ -2727,30 +2812,29 @@ async def increment_spend_counters( hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token ) key_counter_key: Final = f"spend:key:{hashed_token}" - if key_counter_key not in reserved_counter_keys: - await _init_and_increment_spend_counter( - counter_key=key_counter_key, - source_cache_key=hashed_token, - increment=cost, + key_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + () + if key_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=key_counter_key, + source_cache_key=hashed_token, + increment=cost, + ), ) - - key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token) - if key_obj is None: - return - key_budget_limits = getattr(key_obj, "budget_limits", None) or ( - key_obj.get("budget_limits") if isinstance(key_obj, dict) else None ) - if isinstance(key_budget_limits, str): - key_budget_limits = json.loads(key_budget_limits) - if not isinstance(key_budget_limits, list): - return - for window in key_budget_limits: - duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration - key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at - key_window_counter = f"spend:key:{hashed_token}:window:{duration}" + + async def _key_window_increment(window: object) -> _PendingSpendIncrement | None: + duration = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + key_window_reset_at = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + key_window_counter: Final = f"spend:key:{hashed_token}:window:{duration}" key_window_start = get_budget_window_start(window) - if key_window_counter not in reserved_counter_keys: - await _init_and_increment_window_spend_counter( + pending_window: Final = ( + await _prepare_window_spend_counter_increment( counter_key=key_window_counter, entity_type="Key", entity_id=hashed_token, @@ -2758,6 +2842,9 @@ async def increment_spend_counters( window_start=key_window_start, increment=cost, ) + if key_window_counter not in reserved_counter_keys + else None + ) await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.KEY, entity_id=hashed_token, @@ -2767,33 +2854,48 @@ async def increment_spend_counters( increment=cost, request_started_at=request_started_at, ) + return pending_window - async def _team_scope(scope_team_id: str) -> None: - team_counter_key: Final = f"spend:team:{scope_team_id}" - if team_counter_key not in reserved_counter_keys: - await _init_and_increment_spend_counter( - counter_key=team_counter_key, - source_cache_key=f"team_id:{scope_team_id}", - increment=cost, - ) - - team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") - if team_obj is None: - return - team_budget_limits = getattr(team_obj, "budget_limits", None) or ( - team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token) + if key_obj is None: + return key_pending + key_budget_limits = getattr(key_obj, "budget_limits", None) or ( + key_obj.get("budget_limits") if isinstance(key_obj, dict) else None ) - if isinstance(team_budget_limits, str): - team_budget_limits = json.loads(team_budget_limits) - if not isinstance(team_budget_limits, list): - return - for window in team_budget_limits: - duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration - team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at - team_window_counter = f"spend:team:{scope_team_id}:window:{duration}" + if isinstance(key_budget_limits, str): + key_budget_limits = json.loads(key_budget_limits) + if not isinstance(key_budget_limits, list): + return key_pending + window_pending: Final = await asyncio.gather( + *(_key_window_increment(window) for window in key_budget_limits), return_exceptions=True + ) + return key_pending + tuple(item for item in window_pending if item is not None) + + async def _team_scope(scope_team_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + team_counter_key: Final = f"spend:team:{scope_team_id}" + team_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + () + if team_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=team_counter_key, + source_cache_key=f"team_id:{scope_team_id}", + increment=cost, + ), + ) + ) + + async def _team_window_increment(window: object) -> _PendingSpendIncrement | None: + duration = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + team_window_reset_at = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + team_window_counter: Final = f"spend:team:{scope_team_id}:window:{duration}" team_window_start = get_budget_window_start(window) - if team_window_counter not in reserved_counter_keys: - await _init_and_increment_window_spend_counter( + pending_window: Final = ( + await _prepare_window_spend_counter_increment( counter_key=team_window_counter, entity_type="Team", entity_id=scope_team_id, @@ -2801,6 +2903,9 @@ async def increment_spend_counters( window_start=team_window_start, increment=cost, ) + if team_window_counter not in reserved_counter_keys + else None + ) await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.TEAM, entity_id=scope_team_id, @@ -2810,25 +2915,47 @@ async def increment_spend_counters( increment=cost, request_started_at=request_started_at, ) + return pending_window - async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None: + team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") + if team_obj is None: + return team_pending + team_budget_limits = getattr(team_obj, "budget_limits", None) or ( + team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + ) + if isinstance(team_budget_limits, str): + team_budget_limits = json.loads(team_budget_limits) + if not isinstance(team_budget_limits, list): + return team_pending + window_pending: Final = await asyncio.gather( + *(_team_window_increment(window) for window in team_budget_limits), return_exceptions=True + ) + return team_pending + tuple(item for item in window_pending if item is not None) + + async def _team_member_scope( + scope_user_id: str, scope_team_id: str + ) -> tuple[_PendingSpendIncrement | BaseException, ...]: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" if team_member_counter_key in reserved_counter_keys: - return - await _init_and_increment_spend_counter( - counter_key=team_member_counter_key, - source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", - increment=cost, + return () + return ( + await _prepare_spend_counter_increment( + counter_key=team_member_counter_key, + source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", + increment=cost, + ), ) - async def _user_scope(scope_user_id: str) -> None: + async def _user_scope(scope_user_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: user_counter_key: Final = f"spend:user:{scope_user_id}" if user_counter_key in reserved_counter_keys: - return - await _init_and_increment_spend_counter( - counter_key=user_counter_key, - source_cache_key=scope_user_id, - increment=cost, + return () + return ( + await _prepare_spend_counter_increment( + counter_key=user_counter_key, + source_cache_key=scope_user_id, + increment=cost, + ), ) scope_coros: Final = tuple( @@ -2838,7 +2965,7 @@ async def increment_spend_counters( _team_scope(team_id) if team_id is not None else None, _team_member_scope(user_id, team_id) if user_id is not None and team_id is not None else None, _user_scope(user_id) if user_id is not None else None, - _increment_end_user_and_tag_spend_counters( + _prepare_end_user_and_tag_spend_increments( end_user_id=end_user_id, tags=tags, response_cost=cost, @@ -2846,14 +2973,14 @@ async def increment_spend_counters( ) if end_user_id is not None or tags is not None else None, - _increment_model_access_group_spend_counters( + _prepare_model_access_group_spend_increments( model_access_groups=model_access_groups, response_cost=cost, reserved_counter_keys=reserved_counter_keys, ) if model_access_groups else None, - _increment_org_spend_counter( + _prepare_org_spend_increment( org_id=org_id, response_cost=cost, reserved_counter_keys=reserved_counter_keys, @@ -2868,7 +2995,20 @@ async def increment_spend_counters( # as orphaned tasks that race the caller's reservation-counter invalidation; # all scopes settle, then the first error propagates as before. scope_results: Final = await asyncio.gather(*scope_coros, return_exceptions=True) - scope_errors: Final = [r for r in scope_results if isinstance(r, BaseException)] + scope_errors: Final = tuple( + item + for scope in scope_results + for item in (scope if isinstance(scope, tuple) else (scope,)) + if isinstance(item, BaseException) + ) + pending: Final = tuple( + item + for scope in scope_results + if not isinstance(scope, BaseException) + for item in scope + if not isinstance(item, BaseException) + ) + await _apply_spend_counter_increments(pending=pending) if scope_errors: raise scope_errors[0] @@ -2911,41 +3051,49 @@ async def _reconcile_budget_reservation_for_counter_update( return reserved_counter_keys -async def _increment_end_user_and_tag_spend_counters( +async def _prepare_end_user_and_tag_spend_increments( end_user_id: str | None, tags: list[str] | None, response_cost: float, reserved_counter_keys: set[str], -) -> None: - if end_user_id is not None: - await _init_and_increment_unreserved_spend_counter( - counter_key=f"spend:end_user:{end_user_id}", - source_cache_key=end_user_cache_key(end_user_id), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) - - if tags is None: - return - - seen_tags: Final[set[str]] = set() - for tag_name in tags: - if not tag_name or not isinstance(tag_name, str) or tag_name in seen_tags: - continue - seen_tags.add(tag_name) - await _init_and_increment_unreserved_spend_counter( - counter_key=f"spend:tag:{tag_name}", - source_cache_key=tag_cache_key(tag_name), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) +) -> tuple[_PendingSpendIncrement | BaseException, ...]: + unique_tags: Final = ( + tuple(dict.fromkeys(tag for tag in tags if tag and isinstance(tag, str))) if tags is not None else () + ) + results: Final = await asyncio.gather( + *( + coro + for coro in ( + _prepare_unreserved_spend_counter_increment( + counter_key=f"spend:end_user:{end_user_id}", + source_cache_key=end_user_cache_key(end_user_id), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + if end_user_id is not None + else None, + *( + _prepare_unreserved_spend_counter_increment( + counter_key=f"spend:tag:{tag_name}", + source_cache_key=tag_cache_key(tag_name), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + for tag_name in unique_tags + ), + ) + if coro is not None + ), + return_exceptions=True, + ) + return tuple(item for item in results if item is not None) -async def _increment_model_access_group_spend_counters( +async def _prepare_model_access_group_spend_increments( model_access_groups: Sequence[object], response_cost: float, reserved_counter_keys: set[str], -) -> None: +) -> tuple[_PendingSpendIncrement | BaseException, ...]: """Charge the model access groups that authorized this request. Without this the counter auth reads is written only by the reservation path, so @@ -2959,55 +3107,63 @@ async def _increment_model_access_group_spend_counters( unique_groups: Final = tuple( dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str)) ) - for group in unique_groups: - await _init_and_increment_unreserved_spend_counter( - counter_key=model_access_group_spend_counter_key(group), - source_cache_key=model_access_group_cache_key(group), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) + results: Final = await asyncio.gather( + *( + _prepare_unreserved_spend_counter_increment( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + for group in unique_groups + ), + return_exceptions=True, + ) + return tuple(item for item in results if item is not None) -async def _increment_org_spend_counter( +async def _prepare_org_spend_increment( org_id: str | None, response_cost: float, reserved_counter_keys: set[str], -) -> None: +) -> tuple[_PendingSpendIncrement, ...]: if org_id is None: - return + return () - await _init_and_increment_unreserved_spend_counter( + pending: Final = await _prepare_unreserved_spend_counter_increment( counter_key=f"spend:org:{org_id}", source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"], increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) + return (pending,) if pending is not None else () -async def _init_and_increment_unreserved_spend_counter( +async def _prepare_unreserved_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, reserved_counter_keys: set[str], -) -> None: +) -> _PendingSpendIncrement | None: if counter_key in reserved_counter_keys: - return + return None - await _init_and_increment_spend_counter( + return await _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key=source_cache_key, increment=increment, ) -async def _init_and_increment_spend_counter( +async def _prepare_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, -): +) -> _PendingSpendIncrement: """ Initialize counter from the authoritative DB spend value if not yet - set, then atomically increment in both in-memory and Redis. + set, then return the pending increment for the caller to apply in one + pipelined Redis call. On first access per pod: 1. Check spend_counter_cache (in-memory -> Redis via DualCache) @@ -3019,13 +3175,13 @@ async def _init_and_increment_spend_counter( the counter as absent and seed it. Using increment means the worst case is over-counting (conservative, blocks slightly early) rather than under-counting (would allow overspend). - 4. Increment atomically (both in-memory + Redis) + 4. Increment is returned for the caller to apply via pipeline """ await _ensure_spend_counter_initialized( counter_key=counter_key, source_cache_key=source_cache_key, ) - await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + return _PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _enqueue_window_spend_row_update( @@ -3077,20 +3233,20 @@ async def _enqueue_window_spend_row_update( ) -async def _init_and_increment_window_spend_counter( +async def _prepare_window_spend_counter_increment( counter_key: str, entity_type: str, entity_id: str, window_duration: str | None, window_start: datetime | None, increment: float, -): +) -> _PendingSpendIncrement | None: if window_start is None: verbose_proxy_logger.warning( "Skipping spend counter increment for invalid budget window %s", counter_key, ) - return + return None initialized: Final = await _ensure_window_spend_counter_initialized( counter_key=counter_key, @@ -3100,8 +3256,8 @@ async def _init_and_increment_window_spend_counter( window_start=window_start, ) if initialized is False: - return - await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + return None + return _PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _ensure_spend_counter_initialized( @@ -3234,6 +3390,34 @@ async def _invalidate_spend_counter(counter_key: str): ) +async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncrement]) -> None: + if not pending: + return + redis_cache: Final = spend_counter_cache.redis_cache + if redis_cache is None: + for item in pending: + await spend_counter_cache.async_increment_cache( + key=item.counter_key, + value=item.increment, + refresh_ttl=True, + ) + return + ttl: Final = redis_cache.get_ttl() + increment_list: Final = [ # mutable-ok: async_increment_pipeline signature requires list[RedisPipelineIncrementOperation] + RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl) + for item in pending + ] + try: + results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) + except Exception as e: + await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending)) + if isinstance(e, RedisCircuitBreakerOpenError): + return + raise + for item, current_value in zip(pending, results or ()): + spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value) + + async def update_cache( token: str | None, user_id: str | None, @@ -3605,13 +3789,50 @@ async def _run_direct_health_check_with_instrumentation( raise AssertionError("perform_health_check rejected every optional argument") +async def _window_gated_health_check_db_save( + save: Callable[[], Awaitable[bool]], + pod_lock_manager: PodLockManager | None, + lock_ttl: int | None, +) -> None: + """ + Persist at most once per window fleet-wide. A completed save keeps the lock as the + "this window's save is done" marker, so it is deliberately never released and expires + with the interval. A save that reports failure or is cancelled releases the lock so + another pod's cycle in the same window can retry, instead of the fleet going a whole + window without a write. + """ + if pod_lock_manager is None or pod_lock_manager.redis_cache is None: + await save() + return + acquired: Final = await pod_lock_manager.acquire_lock( + cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, + ttl=lock_ttl, + allow_reentrant=False, + ) + if not acquired: + verbose_proxy_logger.debug("background_health_check_db_save_skipped another pod persisted this window") + return + try: + persisted: Final = await save() + except BaseException: + await pod_lock_manager.release_lock(cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME) + raise + if not persisted: + verbose_proxy_logger.warning( + "background_health_check_db_save_incomplete released the window lock so another pod can retry" + ) + await pod_lock_manager.release_lock(cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME) + + def _schedule_background_health_check_db_save( - prisma_client, - shared_health_manager, + prisma_client: PrismaClient | None, + shared_health_manager: "SharedHealthCheckManager | None", model_list: list, healthy_endpoints: list, unhealthy_endpoints: list, -): + pod_lock_manager: PodLockManager | None = None, + lock_ttl: int | None = None, +) -> None: """Fire-and-forget: persist health check results to DB if prisma is available.""" if prisma_client is None: return @@ -3623,16 +3844,16 @@ def _schedule_background_health_check_db_save( checked_by: Final = shared_health_manager.pod_id if shared_health_manager is not None else "background_health_check" start_time: Final = time_module.time() - asyncio.create_task( - _save_background_health_checks_to_db( - prisma_client, - model_list, - healthy_endpoints, - unhealthy_endpoints, - start_time, - checked_by=checked_by, - ) + save: Final = partial( + _save_background_health_checks_to_db, + prisma_client, + model_list, + healthy_endpoints, + unhealthy_endpoints, + start_time, + checked_by=checked_by, ) + asyncio.create_task(_window_gated_health_check_db_save(save, pod_lock_manager, lock_ttl)) def _get_endpoint_exception_status(endpoint: dict, exceptions: dict) -> int: @@ -3939,6 +4160,8 @@ async def _run_background_health_check(): _llm_model_list, healthy_endpoints, unhealthy_endpoints, + pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, + lock_ttl=health_check_interval, ) # Write health state to router cache for health-check-driven routing @@ -4411,20 +4634,9 @@ def resolve_classifier_plugin( def _swap_in_model_cost_map(new_model_cost_map: dict) -> int: - """Adopt a freshly fetched cost map into this process's litellm state, return the model count""" - litellm.model_cost = new_model_cost_map - # Invalidate case-insensitive lookup map since model_cost was replaced - _invalidate_model_cost_lowercase_map() - # Repopulate provider model sets (e.g. litellm.anthropic_models) so that - # wildcard patterns like "anthropic/*" include any newly added models. - litellm.add_known_models(model_cost_map=new_model_cost_map) - # Counted before the re-apply below, which writes into this same dict, so the - # number reported describes the fetched price data alone. - fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0 - # The swap discards everything registered at runtime (deployment model_info, - # register_model overrides), so put it back on top of the fresh catalog. - reapply_runtime_model_cost_registrations() - return fetched_model_count + from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map + + return adopt_model_cost_map(new_model_cost_map) def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool: @@ -4775,7 +4987,9 @@ class ProxyConfig: if not isinstance(raw_params, dict): raise ValueError("general_settings.coordination_redis must be a mapping of Redis connection params") - coordination_params: Final = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(raw_params)) + coordination_params: Final = CoordinationRedisParams.model_validate( + _resolve_coordination_redis_env_refs(raw_params) + ) if not coordination_params.has_connection_target(): raise ValueError( "general_settings.coordination_redis needs a connection target: " @@ -5495,6 +5709,8 @@ class ProxyConfig: parse_budget_reset_time(value) setattr(litellm, key, value) + elif key == "drop_params": + litellm.drop_params = drop_params_flag(value, "litellm_settings.drop_params", verbose_proxy_logger) else: verbose_proxy_logger.debug( "%s setting litellm.%s=%s%s", @@ -5593,6 +5809,16 @@ class ProxyConfig: default_redis_ttl=ttl, ) + ### USER API KEY CACHE MAX SIZE (in-memory tier shared by keys, teams, users, end users, ...) ### + if "user_api_key_cache_max_size" in general_settings: + user_api_key_cache.update_in_memory_max_size( + ConfigGeneralSettings.model_validate( + MappingProxyType( + {"user_api_key_cache_max_size": general_settings["user_api_key_cache_max_size"]} + ) + ).user_api_key_cache_max_size + ) + ### PKCE MULTI-INSTANCE PREREQUISITE CHECK ### # PKCE verifiers are stored in redis_usage_cache when available so they can # be read back by any instance (not just the one that started the auth flow). @@ -5640,6 +5866,10 @@ class ProxyConfig: run_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)), ) + log_once_if_budget_reservation_disabled( + disabled=general_settings.get("disable_budget_reservation") is True, + ) + custom_key_generate: Final = general_settings.get("custom_key_generate", None) if custom_key_generate is not None: user_custom_key_generate = get_instance_fn(value=custom_key_generate, config_file_path=config_file_path) @@ -5823,6 +6053,10 @@ class ProxyConfig: set_files_config(config=files_config) ## default config for vertex ai routes + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + passthrough_endpoint_router, + ) + default_vertex_config: Final = config.get("default_vertex_config", None) passthrough_endpoint_router.set_default_vertex_config(config=default_vertex_config) @@ -6505,6 +6739,14 @@ class ProxyConfig: return parsed return None + async def get_hierarchical_router_settings( + self, + user_api_key_dict: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + proxy_logging_obj: ProxyLogging | None = None, + ) -> dict | None: + return await self._get_hierarchical_router_settings(user_api_key_dict, prisma_client, proxy_logging_obj) + async def _get_hierarchical_router_settings( self, user_api_key_dict: Optional["UserAPIKeyAuth"], @@ -6827,6 +7069,23 @@ class ProxyConfig: "enable_openai_websocket_passthrough" ) + if "user_api_key_cache_max_size" not in self._yaml_general_settings_keys: + db_cache_max_size: Final = _general_settings.get("user_api_key_cache_max_size") + try: + cache_max_size: Final = ConfigGeneralSettings.model_validate( + MappingProxyType({"user_api_key_cache_max_size": db_cache_max_size}) + ).user_api_key_cache_max_size + except ValidationError: + verbose_proxy_logger.warning( + "Ignoring invalid general_settings.user_api_key_cache_max_size=%r from the DB", db_cache_max_size + ) + else: + if cache_max_size is None: + general_settings.pop("user_api_key_cache_max_size", None) + else: + general_settings["user_api_key_cache_max_size"] = cache_max_size + user_api_key_cache.update_in_memory_max_size(cache_max_size) + ## STORE MODEL IN DB ## if "store_model_in_db" in _general_settings: value = _general_settings["store_model_in_db"] @@ -7091,11 +7350,10 @@ class ProxyConfig: ], ) - # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set) - if self._should_load_db_object(object_type="models"): - new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) - - # update llm router + load_models: Final = self._should_load_db_object(object_type="models") + new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) if load_models else None + await self.get_credentials(prisma_client=prisma_client) + if load_models: still_desired_ids = await self._update_llm_router( new_models=new_models, proxy_logging_obj=proxy_logging_obj ) @@ -7135,12 +7393,9 @@ class ProxyConfig: async def _resync_config_from_db() -> None: await self.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) - async def _resync_credentials_from_db() -> None: - await self.get_credentials(prisma_client=prisma_client) - subscriber: Final = ConfigSyncSubscriber( redis_cache=redis_cache, - resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db), + resync_callbacks=(_resync_config_from_db,), ) self.config_sync_subscriber = subscriber subscriber.start() @@ -7995,7 +8250,7 @@ class ProxyConfig: async def get_credentials(self, prisma_client: PrismaClient): try: - credentials = await CredentialsRepository(prisma_client).find_all() + credentials = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_all() credentials = [self.decrypt_credentials(cred) for cred in credentials] await self.delete_credentials(credentials) # delete credentials that are not in the all-up list CredentialAccessor.upsert_credentials(credentials) # upsert credentials that are in the all-up list @@ -8502,6 +8757,7 @@ _STREAM_KEEPALIVE: Final = object() _KEEPALIVE_MIN_SECONDS: Final = 1.0 _KEEPALIVE_MAX_SECONDS: Final = 300.0 _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) +_EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) async def _iter_with_keepalive( @@ -9092,7 +9348,9 @@ class ProxyStartupEvent: if persisted is None: return None - coordination_params: Final = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(persisted)) + coordination_params: Final = CoordinationRedisParams.model_validate( + _resolve_coordination_redis_env_refs(persisted) + ) if not coordination_params.has_connection_target(): verbose_proxy_logger.warning( "coordination_redis saved in the database names no connection target; ignoring it." @@ -9324,6 +9582,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 +9664,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 @@ -9445,7 +9774,7 @@ class ProxyStartupEvent: flush_gateway_requests, "interval", seconds=batch_writing_interval, - args=(prisma_client, gateway_request_accumulator), + args=(prisma_client, gateway_request_accumulator, _gateway_request_redis_buffer()), id="update_gateway_requests_job", replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, @@ -9508,19 +9837,6 @@ class ProxyStartupEvent: ) if store_model_in_db is True: - ### GET STORED CREDENTIALS ### - scheduler.add_job( - proxy_config.get_credentials, - "interval", - seconds=config_reload_interval_seconds, - # REMOVED jitter parameter - major cause of memory leak - args=[prisma_client], - id="get_credentials_job", - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - await proxy_config.get_credentials(prisma_client=prisma_client) - # MEMORY LEAK FIX: Increase interval from 10s to 30s minimum # Frequent polling was causing excessive memory allocations scheduler.add_job( @@ -9534,7 +9850,7 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - # this will load all existing models on proxy startup + # this will load all existing credentials and models on proxy startup await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) proxy_config.start_config_sync_subscriber( @@ -9573,6 +9889,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, @@ -10275,6 +10597,24 @@ async def model_list( wants_anthropic_format: Final = ( http_request is not None and http_request.headers.get("anthropic-version") is not None ) + client_headers: Final[Mapping[str, str]] = http_request.headers if http_request is not None else _EMPTY_HEADERS + view_router_settings: Final = ( + await proxy_config.get_hierarchical_router_settings(user_api_key_dict, prisma_client, proxy_logging_obj) + if wants_anthropic_format and is_claude_code_client(client_headers) + else None + ) + view_aliases: Final = ( + view_router_settings.get("model_group_alias") if isinstance(view_router_settings, Mapping) else None + ) + routing_names: Final = ClaudeCodeRoutingNames( + llm_router, + team_id or user_api_key_dict.team_id, + ( + user_api_key_dict.aliases, + user_api_key_dict.team_model_aliases, + view_aliases, + ), + ) # Validate scope parameter if provided if scope is not None and scope != "expand": @@ -10361,6 +10701,11 @@ async def model_list( return create_anthropic_model_list_response( admin_listing, display_names=configured_display_names(admin_entries, llm_router), + listed_ids=claude_code_view_ids( + admin_listing, + client_headers, + routing_names, + ), ) return dict( @@ -10409,6 +10754,11 @@ async def model_list( return create_anthropic_model_list_response( listing, display_names=configured_display_names(entries, llm_router), + listed_ids=claude_code_view_ids( + listing, + client_headers, + routing_names, + ), ) return dict( @@ -10507,14 +10857,14 @@ async def model_info( # Use the actual litellm model from the deployment to get provider info _, provider, _, _ = litellm.get_llm_provider(model=deployment.litellm_params.model) - response_id: Final = internal_to_public.get(resolved_model_id, model_id) - return create_model_info_response( - model_id=response_id, + response: Final = create_model_info_response( + model_id=resolved_model_id, provider=provider, include_metadata=False, fallback_type=None, llm_router=llm_router, ) + return {**response, "id": internal_to_public.get(resolved_model_id, model_id)} # mutable-ok: response id differs def _blocked_response_usage(original_response: object | None) -> "litellm.Usage": @@ -11440,6 +11790,10 @@ async def vertex_ai_live_passthrough_endpoint( This endpoint delegates to the WebSocket function defined in llm_passthrough_endpoints.py """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + vertex_ai_live_websocket_passthrough, + ) + return await vertex_ai_live_websocket_passthrough( websocket=websocket, model=model, @@ -12623,7 +12977,9 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer) + _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( + model=model_to_use, custom_tokenizer=custom_tokenizer + ) tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) @@ -12636,7 +12992,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - total_tokens: Final = await asyncify(litellm.token_counter)( + total_tokens: Final = await offload_token_count(litellm.token_counter)( model=model_to_use, text=prompt, messages=counted_messages, @@ -16642,6 +16998,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "cancel_on_disconnect": "Boolean", "disable_auto_add_proxy_admin_to_teams": "Boolean", "apply_user_budget_to_team_keys": "Boolean", + "user_api_key_cache_max_size": "Integer", } ) @@ -17459,6 +17816,66 @@ async def delete_callback( ) +def _normalize_callback_alias(callback_name: str) -> str: + callback_aliases: Final = ( + ("opentelemetry", "otel"), + ("s3_v2", "s3"), + ("aws_sqs", "sqs"), + ("custom_callback_api", "generic_api"), + ) + return next( + (canonical_name for alias, canonical_name in callback_aliases if alias == callback_name), + callback_name, + ) + + +def _callback_module_name(callback: CustomLogger | Callable[..., object]) -> str: + if inspect.ismethod(callback): + return callback.__func__.__module__ + if inspect.isfunction(callback): + return callback.__module__ + return type(callback).__module__ + + +def _is_litellm_internal_callback(callback_name: str, callback: CustomLogger | Callable[..., object]) -> bool: + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + module_owner: Final = _callback_module_name(callback).partition(".")[0] + is_registered_integration: Final = callback_name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE + return not is_registered_integration and module_owner in ("litellm", "litellm_enterprise") + + +def _is_instance_of_configured_callback( + callback_name: str, callback: CustomLogger | Callable[..., object], configured_classes: tuple[type, ...] +) -> bool: + """Self-naming OTel-family instances (`arize`, `weave_otel`) match by name, so a configured `logfire` (a bare + `OpenTelemetry`) does not hide YAML-configured siblings of the same class.""" + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + class_derived_name: Final = CustomLoggerRegistry.get_callback_str_from_class_type(type(callback)) + return isinstance(callback, configured_classes) and callback_name in (class_derived_name, type(callback).__name__) + + +def _hidden_runtime_callback_names(configured_callback_names: frozenset[str]) -> frozenset[str]: + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + configured_classes: Final = tuple( + CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE[name] + for name in configured_callback_names + if name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE + ) + configured_modules: Final = frozenset(name.rsplit(".", 1)[0] for name in configured_callback_names if "." in name) + internal_callback_names: Final = frozenset({"cache", "vector_store_pre_call_hook"}) + return internal_callback_names | frozenset( + callback_name + for callback_name, callback in litellm.logging_callback_manager.get_callback_objects() + if isinstance(callback, CustomGuardrail) + or _is_litellm_internal_callback(callback_name, callback) + or _is_instance_of_configured_callback(callback_name, callback, configured_classes) + or _callback_module_name(callback) in configured_modules + ) + + @router.get( "/get/config/callbacks", tags=["config.yaml"], @@ -17491,10 +17908,10 @@ async def get_config( # Normalize string callbacks to lists def normalize_callback(callback): if isinstance(callback, str): - return [callback] - elif callback is None: - return [] - return callback + return (callback,) + if callback is None: + return () + return tuple(callback) if isinstance(callback, (list, dict)) else () _success_callbacks = normalize_callback(_success_callbacks) _failure_callbacks = normalize_callback(_failure_callbacks) @@ -17525,6 +17942,30 @@ async def get_config( for _callback in _success_and_failure_callbacks: _data_to_return.append(process_callback(_callback, "success_and_failure", environment_variables)) + configured_callback_names: Final = frozenset( + _normalize_callback_alias(callback) + for callback in (_success_callbacks + _failure_callbacks + _success_and_failure_callbacks) + ) + runtime_callbacks_by_type: Final = litellm.logging_callback_manager.get_callbacks_by_type() + hidden_callback_names: Final = _hidden_runtime_callback_names(configured_callback_names) + runtime_callback_rows: Final = tuple( + (_normalize_callback_alias(callback_name), callback_type) + for callback_type, callback_names in ( + ("success", runtime_callbacks_by_type["success"]), + ("failure", runtime_callbacks_by_type["failure"]), + ("success_and_failure", runtime_callbacks_by_type["success_and_failure"]), + ) + for callback_name in callback_names + if callback_name not in hidden_callback_names + ) + runtime_only_rows: Final = sorted( + frozenset(row for row in runtime_callback_rows if row[0] not in configured_callback_names) + ) + _data_to_return.extend( + dict(process_callback(callback_name, callback_type, environment_variables), read_only=True) + for callback_name, callback_type in runtime_only_rows + ) + _data_to_return = _apply_callback_role_gate(_data_to_return, is_full_admin) # Check if slack alerting is on @@ -17659,6 +18100,7 @@ async def reload_model_cost_map( # Immediately reload the model cost map in the current pod from litellm.litellm_core_utils.get_model_cost_map import ( ModelCostMapReloadUnavailable, + get_model_cost_map_provenance, refetch_model_cost_map, ) @@ -17672,6 +18114,7 @@ async def reload_model_cost_map( models_count = _swap_in_model_cost_map(reload_result.model_cost_map) current_time = utc_now() proxy_config.model_cost_map_loaded_at = current_time + provenance: Final = get_model_cost_map_provenance() # Publish a new revision so every other pod reloads on its next poll; this pod has # already served it, so adopt it here rather than reloading again a tick later @@ -17686,6 +18129,7 @@ async def reload_model_cost_map( "status": "success", "models_count": models_count, "timestamp": current_time.isoformat(), + **provenance, } except HTTPException: raise @@ -17806,12 +18250,17 @@ async def get_model_cost_map_reload_status( try: global prisma_client + from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map_provenance, + ) + provenance: Final = get_model_cost_map_provenance() if prisma_client is None: verbose_proxy_logger.info("No database connection, returning not scheduled") - return reload_schedule_status(None) + return {**reload_schedule_status(None), **provenance} - return reload_schedule_status(await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME)) + schedule: Final = await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME) + return {**reload_schedule_status(schedule), **provenance} except Exception as e: verbose_proxy_logger.exception("Failed to get model cost map reload status: %s", e) raise HTTPException( @@ -17839,6 +18288,9 @@ async def get_model_cost_map_source( - url: the remote URL that was attempted (null when env-forced local) - is_env_forced: true if LITELLM_LOCAL_MODEL_COST_MAP=True forced local usage - fallback_reason: human-readable reason why remote failed (null on success) + - loaded_at: when this pod last loaded the map + - source_revision: git blob id of the loaded file, what git rev-parse : prints for it + - etag: the ETag of the remote fetch (null for the bundled backup) - model_count: number of models in the currently loaded cost map """ # Read-only source info — admin viewers can read. @@ -18248,7 +18700,7 @@ app.include_router(credential_router) app.include_router(openai_passthrough_router) app.include_router(batches_router) app.include_router(openai_files_router) -app.include_router(llm_passthrough_router) +reserve_lazy_slot(app, "llm_passthrough") app.include_router(pass_through_router) app.include_router(health_router) app.include_router(key_management_router) @@ -18288,6 +18740,7 @@ app.include_router(ui_discovery_endpoints_router) app.include_router(google_router) attach_lazy_features(app) +app.router.routes = hot_routes_first(app.router.routes) app.add_middleware( RequestSizeLimitMiddleware, get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"), @@ -18298,6 +18751,22 @@ app.add_middleware( get_settings=lambda: get_admission_control_settings(general_settings), state=admission_control_state, ) +# Added last on purpose - last-added is outermost, and the client-visible URL +# prefix must be resolved into scope["root_path"] before any inner middleware +# or the router inspects the path. Only added when SERVER_ROOT_PATHS is +# configured, so the default deployment's middleware stack is unchanged. +_server_root_paths: Final = get_server_root_paths() +if _server_root_paths: + if server_root_path and server_root_path != "/": + verbose_proxy_logger.warning( + "Both SERVER_ROOT_PATH=%r and SERVER_ROOT_PATHS=%r are set. A request " + "matching a SERVER_ROOT_PATHS prefix overrides the scalar root_path for " + "that request; unmatched requests keep SERVER_ROOT_PATH. Configure one " + "mechanism or the other.", + server_root_path, + _server_root_paths, + ) + app.add_middleware(PerRequestRootPathMiddleware, root_paths=_server_root_paths) async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse": @@ -18377,6 +18846,31 @@ async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "Streami ######################################################## +@app.api_route( + "/mcp/proxy", + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], # mutable-ok: FastAPI route methods +) +async def proxy_mcp_route(request: Request) -> Response: + """Serve the fixed three-tool MCP proxy surface.""" + from litellm.proxy._experimental.mcp_server.mcp_context import ( # pyright: ignore[reportPrivateUsage] # route-owned mode + _mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # route-owned mode + ) + from litellm.proxy._experimental.mcp_server.server import handle_streamable_http_mcp + from litellm.proxy._experimental.mcp_server.utils import is_mcp_available + + if not is_mcp_available(): + raise HTTPException(status_code=404, detail="Not Found") + + token: Final = _mcp_proxy_mode.set(True) + try: + scope: Final = dict(request.scope) # mutable-ok: ASGI scope rewrite + scope["_original_path"] = scope.get("path", "") + scope["path"] = BASE_MCP_ROUTE + return await _stream_mcp_asgi_response(handle_streamable_http_mcp, scope, request.receive) + finally: + _mcp_proxy_mode.reset(token) + + @app.api_route( BASE_MCP_ROUTE, methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], diff --git a/litellm/proxy/public_endpoints/autorouter_presets.json b/litellm/proxy/public_endpoints/autorouter_presets.json index 7d09db31127..7a251afc076 100644 --- a/litellm/proxy/public_endpoints/autorouter_presets.json +++ b/litellm/proxy/public_endpoints/autorouter_presets.json @@ -10,7 +10,12 @@ "REASONING": ["claude-opus-5"] }, "tier_model_configs": { - "REASONING": [{ "model_name": "claude-opus-5", "litellm_params": { "reasoning_effort": "high" } }] + "REASONING": [ + { + "model_name": "claude-opus-5", + "litellm_params": { "reasoning_effort": "high" } + } + ] }, "classifier_type": "heuristic_v2", "escalation_keywords": ["LITELLM ESCALATE"], @@ -23,16 +28,21 @@ }, "anthropic_family": { "label": "Anthropic Family", - "description": "Routes across the Claude model family: Haiku for simple queries, Sonnet for medium, Opus for complex, Opus at high thinking for reasoning.", + "description": "Routes across the Claude model family: Haiku for simple queries, Sonnet for medium, Opus for complex, Fable 5.1 at high thinking for reasoning.", "complexity_router_config": { "tiers": { "SIMPLE": ["claude-haiku-4-5"], "MEDIUM": ["claude-sonnet-5"], "COMPLEX": ["claude-opus-5"], - "REASONING": ["claude-opus-5"] + "REASONING": ["claude-fable-5-1"] }, "tier_model_configs": { - "REASONING": [{ "model_name": "claude-opus-5", "litellm_params": { "reasoning_effort": "high" } }] + "REASONING": [ + { + "model_name": "claude-fable-5-1", + "litellm_params": { "reasoning_effort": "high" } + } + ] }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], @@ -73,8 +83,18 @@ "REASONING": ["claude-opus-5"] }, "tier_model_configs": { - "MEDIUM": [{ "model_name": "muse-spark-1.2", "litellm_params": { "reasoning_effort": "xhigh" } }], - "COMPLEX": [{ "model_name": "kimi-k3", "litellm_params": { "reasoning_effort": "max" } }] + "MEDIUM": [ + { + "model_name": "muse-spark-1.2", + "litellm_params": { "reasoning_effort": "xhigh" } + } + ], + "COMPLEX": [ + { + "model_name": "kimi-k3", + "litellm_params": { "reasoning_effort": "max" } + } + ] }, "classifier_type": "llm", "classifier_llm_config": { @@ -93,16 +113,21 @@ }, "openai_family": { "label": "OpenAI Family", - "description": "Routes across the GPT model family: Luna for simple queries, Terra for medium, Sol for complex, Sol at xhigh thinking for reasoning.", + "description": "Routes across the GPT model family: Luna for simple queries, Terra for medium, Sol for complex, Astra at xhigh thinking for reasoning.", "complexity_router_config": { "tiers": { "SIMPLE": ["gpt-5.6-luna"], "MEDIUM": ["gpt-5.6-terra"], "COMPLEX": ["gpt-5.6-sol"], - "REASONING": ["gpt-5.6-sol"] + "REASONING": ["gpt-6-astra"] }, "tier_model_configs": { - "REASONING": [{ "model_name": "gpt-5.6-sol", "litellm_params": { "reasoning_effort": "xhigh" } }] + "REASONING": [ + { + "model_name": "gpt-6-astra", + "litellm_params": { "reasoning_effort": "xhigh" } + } + ] }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 66f8c2ea36f..cd781abee26 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -688,6 +688,13 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "CHATGPT", + "provider_display_name": "ChatGPT Subscription", + "litellm_provider": "chatgpt", + "credential_fields": [], + "default_model_placeholder": "chatgpt/gpt-5.4" + }, { "provider": "CLARIFAI", "provider_display_name": "Clarifai", diff --git a/litellm/proxy/realtime_endpoints/endpoints.py b/litellm/proxy/realtime_endpoints/endpoints.py index f41bf4dbd93..be2ac2ff33e 100644 --- a/litellm/proxy/realtime_endpoints/endpoints.py +++ b/litellm/proxy/realtime_endpoints/endpoints.py @@ -17,6 +17,11 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.types.realtime import ( RealtimeClientSecretRequest, RealtimeClientSecretResponse, @@ -304,15 +309,15 @@ async def create_realtime_client_secret( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", http_status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, http_status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, http_status.HTTP_400_BAD_REQUEST), ) raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) if upstream_resp.status_code != 200: @@ -495,15 +500,15 @@ async def proxy_realtime_calls( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", http_status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, http_status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, http_status.HTTP_400_BAD_REQUEST), ) raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) return Response( @@ -608,15 +613,15 @@ async def create_realtime_transcription_session( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "detail", getattr(e, "message", str(e))), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", http_status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, http_status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, http_status.HTTP_400_BAD_REQUEST), ) raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) if upstream_resp.status_code != 200: diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index dd5803796b7..16cd7368e4a 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -11,6 +11,11 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) router: Final = APIRouter() @@ -112,15 +117,15 @@ async def rerank( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 0fd242f2bc1..fac45d4391c 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Final, TypeAlias from fastapi import Request, Response from fastapi.responses import StreamingResponse +from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -74,6 +75,20 @@ class _StreamEventParser: parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) +async def _never_receive() -> Message: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + +def detach_request_from_client(request: Request) -> Request: + """Same scope (headers, parsed body, auth) but a receive() that never yields http.disconnect. + + The polling client closes its connection right after getting the polling id, so the + upstream call must not be cancelled by the client-disconnect guards. + """ + return Request(request.scope, _never_receive) + + async def background_streaming_task( polling_id: str, data: dict[str, object], @@ -123,7 +138,7 @@ async def background_streaming_task( # Pre-call checks (rate limits, guardrails, budget) were already run # before polling ID creation, so skip them here to avoid double-counting. response: Final[StreamingResponse] = await processor.base_process_llm_request( - request=request, + request=detach_request_from_client(request), fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, route_type="aresponses", diff --git a/litellm/proxy/route_priority.py b/litellm/proxy/route_priority.py new file mode 100644 index 00000000000..77815b678a6 --- /dev/null +++ b/litellm/proxy/route_priority.py @@ -0,0 +1,24 @@ +"""Starlette matches routes in registration order, so the routes that take the most traffic go first.""" + +from collections.abc import Sequence +from typing import Final + +from starlette.routing import BaseRoute, Route + +HOT_ROUTE_PATHS: Final[frozenset[str]] = frozenset( + ( + "/health/liveliness", + "/health/liveness", + "/v1/chat/completions", + "/chat/completions", + "/v1/messages", + ) +) + + +def _is_hot(route: BaseRoute) -> bool: + return isinstance(route, Route) and route.path in HOT_ROUTE_PATHS + + +def hot_routes_first(routes: Sequence[BaseRoute]) -> list[BaseRoute]: # mutable-ok: assigned to Router.routes, a list + return sorted(routes, key=lambda route: not _is_hot(route)) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 1c43668f227..7d521d54791 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + skills String[] @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -343,6 +344,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? @@ -491,7 +493,7 @@ model LiteLLM_JWTKeyMapping { updated_at DateTime @default(now()) @updatedAt updated_by String? - litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) @@unique([jwt_claim_name, jwt_claim_value]) @@index([jwt_claim_name, jwt_claim_value, is_active]) @@ -1035,6 +1037,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") updated_at DateTime @updatedAt @@ -1508,7 +1511,10 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") + baseline_models Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index ed2bc87597c..ef21551ca93 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -34,6 +34,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router +from litellm.rust_bridge.token_counter import count_anthropic_input_tokens, uses_anthropic_tokenizer from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.types.router import DeploymentTypedDict @@ -173,6 +174,29 @@ def _raise_counter_budget_exceeded( ) +_UNBILLED_ROUTES: Final[frozenset[str]] = frozenset( + { + "/models", + "/v1/models", + "/utils/token_counter", + "/responses/input_tokens", + "/v1/responses/input_tokens", + "/openai/v1/responses/input_tokens", + } +) +_TOKEN_COUNTING_SEGMENTS: Final[frozenset[str]] = frozenset({"count_tokens", "count-tokens"}) +_TOKEN_COUNTING_ACTION: Final = "countTokens" + + +def _is_token_counting_route(route: str) -> bool: + resource, _, action = route.rsplit("/", 1)[-1].partition(":") + return resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION + + +def _is_unbilled_route(route: str) -> bool: + return route in _UNBILLED_ROUTES or _is_token_counting_route(route) + + async def reserve_budget_for_request( request_body: dict, route: str, @@ -187,17 +211,11 @@ async def reserve_budget_for_request( end_user_object: object = None, apply_user_budget_to_team_keys: bool = False, fail_closed_budget_enforcement: bool = False, + raw_body: bytes | None = None, ) -> dict | None: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None - if route in { - "/models", - "/v1/models", - "/utils/token_counter", - "/responses/input_tokens", - "/v1/responses/input_tokens", - "/openai/v1/responses/input_tokens", - }: + if _is_unbilled_route(route): return None if get_model_from_request(request_body, route, llm_router=llm_router) is None: return None @@ -221,6 +239,7 @@ async def reserve_budget_for_request( request_body=request_body, route=route, llm_router=llm_router, + raw_body=raw_body, ) current_spend_by_counter_key: Final[dict[str, float]] = {} @@ -299,6 +318,7 @@ async def reserve_budget_for_request( "entries": applied_entries, "finalized": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), + "input_tokens": max(input_token_counts.values(), default=None), } @@ -905,13 +925,13 @@ async def _set_reserved_entry_actual_cost( increment=adjustment, ) elif reseed_on_inconsistent: - # Post-call reconcile / release: the counter was flushed or reseeded - # between reservation and reconcile (Redis restart / cross-pod reset), - # so the optimistic delta no longer applies. Recover by reseeding from - # the DB's lagging authoritative floor rather than deleting the counter - # and failing open — deleting it is what left budgets unenforced after a - # Redis reload. - await reseed_spend_counter_from_db(counter_key=counter_key) + # Post-call reconcile / release: the counter was flushed, expired or reseeded + # between reservation and reconcile, so the optimistic delta no longer applies. + # Reseed from the DB floor (which cannot include this request's cost yet) and + # add the settled cost, since increment_spend_counters skips reserved keys. + reseeded: Final = await reseed_spend_counter_from_db(counter_key=counter_key) + if reseeded and actual_cost > 0: + await _increment_spend_counter_cache(counter_key=counter_key, increment=actual_cost) else: # Pre-call admission resize: the in-flight reservation cost is not yet # persisted, so the DB floor would discard it. Keep the original @@ -925,18 +945,16 @@ async def _counter_can_apply_adjustment( counter_key: str, adjustment: float, ) -> bool: - from litellm.proxy.proxy_server import spend_counter_cache + from litellm.proxy.proxy_server import read_spend_counter_cache_value - current_value: Final = await spend_counter_cache.async_get_cache(key=counter_key) + try: + current_value, _ = await read_spend_counter_cache_value(counter_key=counter_key) + except (TypeError, ValueError): + return False if current_value is None: return False - try: - current_float: Final = float(current_value) - except (TypeError, ValueError): - return False - - return not (adjustment < 0 and current_float + adjustment < -1e-12) + return not (adjustment < 0 and current_value + adjustment < -1e-12) async def _release_applied_entries_best_effort( @@ -1258,7 +1276,7 @@ def _get_model_cost_info( llm_router: Router | None, ) -> Mapping[str, object] | None: if llm_router is not None: - model_group_info: Final = llm_router.get_model_group_info(model_group=model) + model_group_info: Final = llm_router.cached_model_group_info(model) if model_group_info is not None: return model_group_info.model_dump() return dict(litellm.get_model_info(model=model)) @@ -1300,7 +1318,7 @@ def _deployment_tiered_pricing_table( backend_model: Final = _get_value(_get_value(deployment, "litellm_params"), "model") if not isinstance(model_id, str) or not isinstance(backend_model, str): return None - deployment_model_info: Final = llm_router.get_deployment_model_info(model_id=model_id, model_name=backend_model) + deployment_model_info: Final = llm_router.cached_deployment_model_info(model_id, backend_model) if deployment_model_info is None: return None tiered_pricing: Final = deployment_model_info.get("tiered_pricing") @@ -1341,24 +1359,46 @@ async def count_request_input_tokens( request_body: dict, route: str, llm_router: Router | None, + raw_body: bytes | None = None, ) -> Mapping[str, int]: """Input-token count per candidate model, counted once per request. Tokenizing is the reservation path's dominant CPU cost and is O(prompt), so counting a large prompt inline stalls every other request on the worker. - Large prompts are counted in a worker thread, and the counts are reused by - both the max-cost and the input-cost estimate. + Models on the Anthropic tokenizer are counted from the raw body by the Rust + bridge when it is enabled, which parses and tokenizes with the GIL released. + Everything it declines is counted in Python, large prompts in a worker + thread. The counts are reused by both the max-cost and the input-cost + estimate. """ models: Final = _get_request_models(request_body=request_body, route=route, llm_router=llm_router) if not models: return MappingProxyType({}) - if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: - return _count_input_tokens_for_models(request_body=request_body, models=models) - return await asyncio.to_thread( - _count_input_tokens_for_models, - request_body=request_body, - models=models, + rust_count: Final = ( + await count_anthropic_input_tokens(raw_body) + if raw_body is not None and any(uses_anthropic_tokenizer(model) for model in models) + else None ) + rust_counts: Final = MappingProxyType( + { + model: rust_count.input_tokens + for model in models + if rust_count is not None and uses_anthropic_tokenizer(model) + } + ) + python_models: Final = tuple(model for model in models if model not in rust_counts) + if not python_models: + return rust_counts + python_counts: Final = ( + _count_input_tokens_for_models(request_body=request_body, models=python_models) + if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS + else await asyncio.to_thread( + _count_input_tokens_for_models, + request_body=request_body, + models=python_models, + ) + ) + return MappingProxyType({**rust_counts, **python_counts}) def _count_input_tokens_for_models( @@ -1389,12 +1429,15 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int: def _count_input_tokens(request_body: dict, model: str) -> int | None: try: if "messages" in request_body: - return litellm.token_counter( - model=model, - messages=request_body.get("messages") or [], - tools=request_body.get("tools"), - tool_choice=request_body.get("tool_choice"), - ) + try: + return litellm.token_counter( + model=model, + messages=request_body.get("messages") or (), + tools=request_body.get("tools"), + tool_choice=request_body.get("tool_choice"), + ) + except ValueError: + return _count_text_tokens(model=model, text=request_body.get("messages")) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: @@ -1442,11 +1485,7 @@ def _estimate_output_tokens( if _is_input_only_route(route=route): return 0 - requested: int | None = None - for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"): - requested = _to_int(request_body.get(key)) - if requested is not None: - break + requested: Final = _requested_output_tokens(request_body) # Clamp at min(requested-or-default, model_max-or-default). Two purposes: # (1) Without an explicit cap we still need a finite reservation so the @@ -1457,9 +1496,19 @@ def _estimate_output_tokens( # at the cap — the model can only physically emit max_output_tokens # anyway, so reserving more is both wasteful and a DoS surface. model_ceiling: Final = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK - if requested is None: - requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK - return min(requested, model_ceiling) + return min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None else requested, model_ceiling) + + +_OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens") + + +def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: + inference_config: Final = request_body.get("inferenceConfig") + candidates: Final = ( + *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), + inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None, + ) + return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) def _count_text_tokens(model: str, text: object) -> int: diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 7de18521edd..29688b61b3d 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,5 +1,7 @@ +import asyncio from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet +from datetime import datetime, timedelta from types import MappingProxyType from typing import Final, TypeVar @@ -7,6 +9,13 @@ from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import ( + SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, + SPEND_LOG_KEY_METADATA_CACHE_TTL, + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, + SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, +) from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.utils import PrismaClient from litellm.repositories.user_repository import UserRepository @@ -27,6 +36,33 @@ WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[]) ORDER BY token, deleted_at DESC """ +_SPEND_LOG_ALIAS_SQL: Final = """ +SELECT api_key AS digest, + MIN(key_alias) AS first_alias, + MAX(key_alias) AS last_alias, + MIN(team_id) AS first_team, + MAX(team_id) AS last_team, + MIN(user_id) AS first_owner, + MAX(user_id) AS last_owner +FROM ( + SELECT api_key, + NULLIF(metadata->>'user_api_key_alias', '') AS key_alias, + COALESCE(NULLIF(team_id, ''), NULLIF(metadata->>'user_api_key_team_id', '')) AS team_id, + COALESCE(NULLIF("user", ''), NULLIF(metadata->>'user_api_key_user_id', '')) AS user_id + FROM "LiteLLM_SpendLogs" + WHERE api_key = ANY($1::text[]) + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp +) named +WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL +GROUP BY api_key +""" + +_SPEND_LOG_STATEMENT_TIMEOUT_SQL: Final = f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}" +_SPEND_LOG_TRANSACTION_TIMEOUT: Final = timedelta(milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS) + +_HASHED_JWT_PREFIX: Final = "hashed-jwt-" + class KeyMetadataDict(TypedDict, total=False): key_alias: ReadOnly[str | None] @@ -42,7 +78,35 @@ class _TokenDigestRow(BaseModel): user_id: str | None = None +def _unanimous(first: str | None, last: str | None) -> str | None: + return first if first == last else None + + +class _SpendLogDigestRow(BaseModel): + digest: str + first_alias: str | None = None + last_alias: str | None = None + first_team: str | None = None + last_team: str | None = None + first_owner: str | None = None + last_owner: str | None = None + + def metadata(self) -> KeyMetadataDict: + return KeyMetadataDict( + key_alias=_unanimous(self.first_alias, self.last_alias), + team_id=_unanimous(self.first_team, self.last_team), + user_id=_unanimous(self.first_owner, self.last_owner), + ) + + _TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...]) +_SPEND_LOG_DIGEST_ROWS: Final = TypeAdapter(tuple[_SpendLogDigestRow, ...]) +_CACHED_KEY_METADATA: Final = TypeAdapter(KeyMetadataDict) +_SPEND_LOG_METADATA_CACHE: Final = InMemoryCache( + max_size_in_memory=SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, + default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL, +) +_SPEND_LOG_QUERY_LOCK: Final = asyncio.Lock() _EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({}) _EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({}) @@ -138,14 +202,6 @@ async def recover_double_hashed_key_metadata( prisma_client: PrismaClient, missing_keys: AbstractSet[str], ) -> Mapping[str, KeyMetadataDict]: - """ - Recover key_alias/team_id/user_id for DailyUserSpend.api_key values that - were double-hashed by the v1.99 spend-log provenance gate. - - Those rows store hash(VerificationToken.token) instead of the token, so the - exact join misses. Postgres hashes the token column itself, one pass over - active keys and one over deleted keys, so no key row crosses the wire. - """ sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) if not sha_missing: return _EMPTY_KEY_METADATA @@ -168,6 +224,117 @@ async def recover_double_hashed_key_metadata( return MappingProxyType({**from_active, **from_deleted}) +def _is_spend_log_digest(key: str) -> bool: + return is_valid_sha256_hash(key.removeprefix(_HASHED_JWT_PREFIX)) + + +def _spend_log_cache_key(digest: str, window: tuple[datetime, datetime]) -> str: + start, end = window + return f"spend_log_key_metadata:{digest}:{start.isoformat()}:{end.isoformat()}" + + +def _cached_spend_log_metadata( + cache: InMemoryCache, + digests: AbstractSet[str], + window: tuple[datetime, datetime], +) -> Mapping[str, KeyMetadataDict]: + return MappingProxyType( + { + digest: _CACHED_KEY_METADATA.validate_python(cached) + for digest in digests + for cached in (cache.get_cache(_spend_log_cache_key(digest, window)),) + if cached is not None + } + ) + + +async def _spend_log_rows_within_the_statement_timeout( + prisma_client: PrismaClient, + digests: AbstractSet[str], + window: tuple[datetime, datetime], +) -> Sequence[Mapping[str, object]]: + start, end = window + async with prisma_client.db.tx(timeout=_SPEND_LOG_TRANSACTION_TIMEOUT) as transaction: + await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) + return await transaction.query_raw(_SPEND_LOG_ALIAS_SQL, sorted(digests), start, end) + + +async def _query_spend_log_metadata( + prisma_client: PrismaClient, + digests: AbstractSet[str], + window: tuple[datetime, datetime], +) -> Mapping[str, KeyMetadataDict] | None: + rows: Final = await _db_or_empty( + lambda: _spend_log_rows_within_the_statement_timeout(prisma_client, digests, window), + "Failed spend-log alias recovery for %d missing keys: %s", + len(digests), + ) + if rows is None: + return None + return MappingProxyType( + { + row.digest: meta + for row in _SPEND_LOG_DIGEST_ROWS.validate_python(rows) + for meta in (row.metadata(),) + if row.digest in digests and any(meta.values()) + } + ) + + +def _remember_spend_log_metadata( + cache: InMemoryCache, digest: str, window: tuple[datetime, datetime], meta: KeyMetadataDict | None +) -> None: + key: Final = _spend_log_cache_key(digest, window) + if meta is not None: + cache.set_cache(key, meta) + return + missed_before: Final = f"{key}:missed-before" + if cache.get_cache(missed_before) is not None: + cache.set_cache(key, KeyMetadataDict()) + return + cache.set_cache(key, KeyMetadataDict(), ttl=SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL) + cache.set_cache(missed_before, True) + + +async def _spend_log_metadata_one_query_at_a_time( + prisma_client: PrismaClient, + cache: InMemoryCache, + lock: asyncio.Lock, + digests: AbstractSet[str], + window: tuple[datetime, datetime], +) -> Mapping[str, KeyMetadataDict]: + async with lock: + settled: Final = _cached_spend_log_metadata(cache, digests, window) + pending: Final = digests - frozenset(settled) + fresh: Final = ( + await _query_spend_log_metadata(prisma_client, pending, window) if pending else _EMPTY_KEY_METADATA + ) + found: Final = fresh if fresh is not None else _EMPTY_KEY_METADATA + for digest in pending: + _remember_spend_log_metadata(cache, digest, window, found.get(digest)) + return MappingProxyType({**settled, **found}) + + +async def recover_key_metadata_from_spend_logs( + prisma_client: PrismaClient, + missing_keys: AbstractSet[str], + window: tuple[datetime, datetime], + cache: InMemoryCache = _SPEND_LOG_METADATA_CACHE, + lock: asyncio.Lock = _SPEND_LOG_QUERY_LOCK, +) -> Mapping[str, KeyMetadataDict]: + digests: Final = frozenset(key for key in missing_keys if _is_spend_log_digest(key)) + if not digests: + return _EMPTY_KEY_METADATA + cached: Final = _cached_spend_log_metadata(cache, digests, window) + uncached: Final = digests - frozenset(cached) + settled: Final = ( + await _spend_log_metadata_one_query_at_a_time(prisma_client, cache, lock, uncached, window) + if uncached + else _EMPTY_KEY_METADATA + ) + return MappingProxyType({digest: meta for digest, meta in (*cached.items(), *settled.items()) if meta}) + + def _row_with_recovered_fields( row: Mapping[str, object], recovered: Mapping[str, KeyMetadataDict], diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 1d0eb12da75..950fcca2039 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -9,12 +9,17 @@ have been aggregated across models. """ from collections.abc import Callable, Mapping +from datetime import datetime from typing import TYPE_CHECKING, Final, NamedTuple import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY -from litellm.litellm_core_utils.llm_cost_calc.utils import _get_cost_per_unit, generic_cost_per_token +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + _get_cost_per_unit, + calculate_prompt_caching_savings, + generic_cost_per_token, +) from litellm.types.integrations.anthropic_cache_control_hook import ( GATEWAY_INJECTED_CACHE_METADATA_KEY, GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT, @@ -32,42 +37,13 @@ class SavingsSpend(NamedTuple): gateway_injected_caching: float = 0.0 -def _input_cache_read_and_write_cost(info: ModelInfo | None) -> tuple[float, float, float]: - """ - Return ``(input_cost, cache_read_cost, cache_write_cost)`` per token. - - ``info`` is whatever pricing the caller resolved -- deployment rates when the - request came through a router deployment, public rates otherwise -- so a - negotiated price is honoured here rather than silently replaced by the list rate. - ``None`` falls open to ``(0.0, 0.0, 0.0)`` so savings degrade to zero rather than - raising inside the spend writer. - - Prices are read through ``_get_cost_per_unit``, the same accessor the cost - calculator uses, which coerces the string prices a ``config.yaml`` can produce - (``"3e-7"``) and resolves service-tier suffixes. - - An absent cache price mirrors the input cost, which yields a zero discount on the - read leg and a zero premium on the write leg. Mirroring rather than taking - ``_get_cost_per_unit``'s 0.0 default is load-bearing on the write leg: a zero write - price would make the premium ``0 - input_cost``, turning a model that simply has no - write pricing into a spurious extra saving. - - The two legs then differ on an explicit ``0.0``, and the asymmetry is deliberate. A - free cache *write* does not exist -- entries carrying a literal zero (``deepseek-chat`` - does) mean "no separate price", so a falsy write price also mirrors input. A free - cache *read* is real: 15 models charge for input and serve reads for nothing, which - is the largest discount available, so the read leg keeps its literal zero. - """ - if info is None: - return 0.0, 0.0, 0.0 - input_cost: Final = _get_cost_per_unit(info, "input_cost_per_token") or 0.0 - cache_read_cost: Final = _get_cost_per_unit(info, "cache_read_input_token_cost", default_value=None) - cache_write_cost: Final = _get_cost_per_unit(info, "cache_creation_input_token_cost", default_value=None) - return ( - input_cost, - input_cost if cache_read_cost is None else cache_read_cost, - cache_write_cost if cache_write_cost else input_cost, - ) +def _coerce_billed_at(value: datetime | str | None) -> datetime | None: + if isinstance(value, datetime) or value is None: + return value + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None class _ModelIdentity(NamedTuple): @@ -323,6 +299,8 @@ def compute_autorouter_savings( selected_info: ModelInfo | None = None, baseline_info: ModelInfo | None = None, cost_breakdown: Mapping[str, object] | None = None, + baseline_deployment_id: str | None = None, + selected_deployment_id: str | None = None, ) -> float: """Net dollars the router saved, or cost, by serving this request on ``selected_model``. @@ -358,11 +336,12 @@ def compute_autorouter_savings( selected: Final = _resolve_model(selected_model, selected_provider) if baseline is None or selected is None: return 0.0 - # Same model is only the same cost when it is also the same deployment. Two - # deployments of one model can carry different negotiated rates, and routing from - # the dear one to the cheap one is a real saving that short-circuiting on the model - # name alone reports as zero. - if baseline == selected: + same_target: Final = ( + baseline_deployment_id == selected_deployment_id + if baseline_deployment_id and selected_deployment_id + else baseline == selected + ) + if same_target: return 0.0 basis: Final = _pricing_basis(cost_breakdown) effective_baseline_info: Final = baseline_info if baseline_info is not None else _model_info(baseline) @@ -541,6 +520,8 @@ def autorouter_savings_for_request( selected_info=_effective_model_info(router_instance, model_id, model or ""), baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""), cost_breakdown=cost_breakdown, + baseline_deployment_id=baseline_id, + selected_deployment_id=model_id, ) classifier_cost: Final = classifier_cost_from_decision(decision) return gross if classifier_cost is None else gross - classifier_cost @@ -586,6 +567,7 @@ def compute_savings_spend( llm_router: "Callable[[], Router | None] | None" = None, cost_breakdown: Mapping[str, object] | None = None, recorded_autorouter_savings: object = None, + billed_at: datetime | str | None = None, ) -> SavingsSpend: """ Dollar savings for one request, split by optimization driver. @@ -595,24 +577,10 @@ def compute_savings_spend( premium paid to write those entries, both derived here from ``usage_object`` so no caller can hand in a count that disagrees with the usage record. - The net form follows from what the request would have cost with caching off. The - provider reports ``prompt_tokens`` as the inclusive total of three disjoint - partitions (uncached text, cache reads, cache writes), so an uncached counterfactual - bills every one of those tokens at the flat input rate:: - - would_have_cost = (text + reads + writes) * input - actually_cost = text * input + reads * read_rate + writes * write_rate - savings = reads * (input - read_rate) - writes * (write_rate - input) - - So the write leg subtracts the write PREMIUM, not the whole write cost: those tokens - had to be sent either way, and the counterfactual already pays the input rate for - them. The premium stays signed, because a handful of models price writes below their - input rate and there the write is a genuine extra saving. - - A request that only writes cache and gets no hits therefore reports negative savings, - which is accurate: it really did cost more than the uncached call would have. The - daily rollup increments arithmetically, so those rows offset positive ones in the - same bucket. + The uncached counterfactual pays the ordinary input rate for the same prompt size + and tier. Cache writes subtract only the premium over that rate, split by TTL. + Savings stay signed: a write-only request can lose money, and daily rollups net + those losses against read savings. Caching is reported twice. ``prompt_caching`` is every net dollar caching saved, whoever caused it, which is what a customer means by "what did caching save me". @@ -638,12 +606,9 @@ def compute_savings_spend( calls this and only auto-routed ones need one, so looking it up eagerly at the call site would fetch and discard it on the rest. - ``cost_breakdown`` is what the cost calculator recorded for this request, and it - carries both what the request really cost and the tier and region it was priced on. - Only the auto-router driver reads it. Compression and prompt caching price a - hypothetical token delta off flat rate keys, so they are blind to tiered pricing in - the same way; that is pre-existing behaviour on two shipped drivers rather than - something introduced here, and moving those numbers is its own change. + ``cost_breakdown`` supplies the biller's tier and region to caching and auto-router + savings. Caching also uses the logged prompt size and TTL split. Compression retains + its flat input-rate estimate; changing that counterfactual is a separate concern. ``recorded_autorouter_savings`` is the figure the logging path stamped on the spend log's metadata, honoured over recomputation so the rollup, the turn table and the @@ -658,13 +623,24 @@ def compute_savings_spend( pricing: Final = _effective_model_info(router_instance, model_id, model or "") or ( _model_info(identity) if identity else None ) - input_cost, cache_read_cost, cache_write_cost = _input_cache_read_and_write_cost(pricing) + input_cost: Final = (_get_cost_per_unit(pricing, "input_cost_per_token") or 0.0) if pricing else 0.0 compression: Final = max(compression_saved_tokens, 0) * input_cost - cache_read_input_tokens: Final = extract_cache_read_tokens(usage_object) - cache_creation_input_tokens: Final = extract_cache_creation_tokens(usage_object) - read_discount: Final = max(cache_read_input_tokens, 0) * max(input_cost - cache_read_cost, 0.0) - write_premium: Final = max(cache_creation_input_tokens, 0) * (cache_write_cost - input_cost) - prompt_caching: Final = read_discount - write_premium + usage: Final = _usage_from_spend_log(usage_object) + basis: Final = _pricing_basis(cost_breakdown) + billed_at_datetime: Final = _coerce_billed_at(billed_at) + prompt_caching: Final = ( + calculate_prompt_caching_savings( + model_info=pricing, + usage=usage, + custom_llm_provider=identity.provider if identity else custom_llm_provider, + service_tier=basis.service_tier, + data_residency=basis.data_residency, + vertex_location=basis.vertex_location, + billed_at=billed_at_datetime, + ) + if pricing is not None and usage is not None + else 0.0 + ) gateway_injected_caching: Final = prompt_caching if gateway_injected_cache else 0.0 # The figure the logging path recorded wins, before the usage gate on purpose: a row diff --git a/litellm/proxy/spend_tracking/spend_event.py b/litellm/proxy/spend_tracking/spend_event.py new file mode 100644 index 00000000000..53f26346f85 --- /dev/null +++ b/litellm/proxy/spend_tracking/spend_event.py @@ -0,0 +1,418 @@ +"""Compact, typed success event handed from an inference worker to the collector. + +``build_spend_event`` runs on the inference worker right after ``Logging.async_success_handler`` +has built the ``standard_logging_object`` (so the cost is already known). It validates the success +callback's ``kwargs`` into the projection ``_PROXY_track_cost_callback`` and +``DBSpendUpdateWriter.update_database`` actually read: identities and metadata, timings, usage, the +standard logging payload without its prompt/response bodies, and the tool names. The request +messages, the raw ``proxy_server_request`` body and the full response travel only when spend logs +are configured to store prompts and responses. The cache key is the preset key the caching layer +already computed, never a fresh hash over the request body. + +``spend_event_callback_args`` rebuilds the ``(kwargs, response_obj, start_time, end_time)`` tuple +the existing cost pipeline consumes, so the sidecar runs the unchanged pipeline against the event. +""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.spend_log_tool_index import response_tool_call_names +from litellm.types.interactions import InteractionsAPIResponse +from litellm.types.utils import LiteLLMBatch, Usage + +SPEND_EVENT_VERSION: Final = 1 +CACHE_OFF_KEY: Final = "Cache OFF" + +ObjectMapping: TypeAlias = Mapping[str, object] + +_UNSERIALIZABLE_METADATA_KEYS: Final = frozenset({"user_api_key_auth", "litellm_parent_otel_span"}) +_STANDARD_LOGGING_BODY_KEYS: Final = frozenset({"messages", "response"}) +_STANDARD_LOGGING_DROPPED_KEYS: Final = frozenset({"model_parameters"}) +_NOT_OFFLOADED_RESPONSE_TYPES: Final = (LiteLLMBatch, InteractionsAPIResponse) + + +class _LitellmParams(TypedDict, total=False): + api_base: ReadOnly[str | None] + custom_llm_provider: ReadOnly[str | None] + litellm_call_id: ReadOnly[str | None] + user_api_key_end_user_id: ReadOnly[str | None] + metadata: ReadOnly[ObjectMapping | None] + litellm_metadata: ReadOnly[ObjectMapping | None] + proxy_server_request: ReadOnly[ObjectMapping | None] + preset_cache_key: ReadOnly[str | None] + + +class _DynamicParams(TypedDict, total=False): + turn_off_message_logging: ReadOnly[bool | None] + + +class _RequestBody(TypedDict, total=False): + tools: ReadOnly[Sequence[ObjectMapping] | None] + + +class _PassthroughPayload(TypedDict, total=False): + request_body: ReadOnly[_RequestBody | None] + + +class _ToolCallFunction(TypedDict): + name: ReadOnly[str] + arguments: ReadOnly[str] + + +class _ToolCall(TypedDict): + id: ReadOnly[str | None] + type: ReadOnly[Literal["function"]] + function: ReadOnly[_ToolCallFunction] + + +class _ToolCallMessage(TypedDict): + role: ReadOnly[Literal["assistant"]] + content: ReadOnly[None] + tool_calls: ReadOnly[Sequence[_ToolCall]] + + +class _ToolCallChoice(TypedDict): + index: ReadOnly[int] + finish_reason: ReadOnly[Literal["tool_calls"]] + message: ReadOnly[_ToolCallMessage] + + +class CompactResponse(TypedDict, total=False): + """What the spend pipeline reads off a response: its id, usage and which tools it called.""" + + id: ReadOnly[object] + model: ReadOnly[object] + usage: ReadOnly[object] + usage_info: ReadOnly[object] + status: ReadOnly[object] + background: ReadOnly[object] + choices: ReadOnly[Sequence[_ToolCallChoice]] + + +class _SuccessKwargs(TypedDict, total=False): + """The success callback's ``kwargs`` (``Logging.model_call_details``), validated and projected.""" + + litellm_call_id: ReadOnly[str | None] + call_type: ReadOnly[str | None] + model: ReadOnly[str | None] + custom_llm_provider: ReadOnly[str | None] + stream: ReadOnly[bool | None] + complete_streaming_response: ReadOnly[object] + cache_hit: ReadOnly[bool | None] + response_cost: ReadOnly[float | None] + completion_start_time: ReadOnly[datetime | None] + agent_id: ReadOnly[str | None] + litellm_trace_id: ReadOnly[str | None] + litellm_params: ReadOnly[_LitellmParams] + standard_logging_object: ReadOnly[ObjectMapping | None] + standard_callback_dynamic_params: ReadOnly[_DynamicParams | None] + combined_usage_object: ReadOnly[Usage | None] + realtime_tools: ReadOnly[Sequence[object] | None] + realtime_tool_calls: ReadOnly[Sequence[object] | None] + tools: ReadOnly[Sequence[ObjectMapping] | None] + passthrough_logging_payload: ReadOnly[_PassthroughPayload | None] + + +class _FunctionToolFunction(TypedDict): + name: ReadOnly[str] + + +class _FunctionTool(TypedDict): + type: ReadOnly[Literal["function"]] + function: ReadOnly[_FunctionToolFunction] + + +class SpendCallbackKwargs(TypedDict): + """The ``kwargs`` handed to ``_PROXY_track_cost_callback`` on the sidecar.""" + + litellm_call_id: ReadOnly[str | None] + call_type: ReadOnly[str | None] + model: ReadOnly[str | None] + custom_llm_provider: ReadOnly[str | None] + stream: ReadOnly[bool | None] + cache_hit: ReadOnly[bool | None] + response_cost: ReadOnly[float | None] + completion_start_time: ReadOnly[datetime | None] + agent_id: ReadOnly[str | None] + litellm_trace_id: ReadOnly[str | None] + litellm_params: ReadOnly[_LitellmParams] + standard_logging_object: ReadOnly[ObjectMapping | None] + standard_callback_dynamic_params: ReadOnly[_DynamicParams | None] + combined_usage_object: ReadOnly[Usage | None] + realtime_tools: ReadOnly[Sequence[object] | None] + realtime_tool_calls: ReadOnly[Sequence[object] | None] + tools: ReadOnly[Sequence[_FunctionTool] | None] + complete_streaming_response: NotRequired[ReadOnly[CompactResponse | None]] + + +_NO_LITELLM_PARAMS: Final[_LitellmParams] = {} +_SUCCESS_KWARGS: Final = TypeAdapter(_SuccessKwargs) +_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) +_COMPACT_RESPONSE: Final = TypeAdapter(CompactResponse) + + +class SpendEvent(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + version: Literal[1] + litellm_call_id: str | None + call_type: str | None + model: str | None + custom_llm_provider: str | None + stream: bool | None + complete_streaming_response: bool + cache_hit: bool | None + response_cost: float | None + start_time: datetime + end_time: datetime + completion_start_time: datetime | None + agent_id: str | None + litellm_trace_id: str | None + litellm_params: _LitellmParams + standard_logging_object: ObjectMapping | None + standard_callback_dynamic_params: _DynamicParams | None + response: CompactResponse | None + combined_usage: ObjectMapping | None + realtime_tools: Sequence[object] | None + realtime_tool_calls: Sequence[object] | None + request_tool_names: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class SpendEventCallbackArgs: + kwargs: SpendCallbackKwargs + response_obj: CompactResponse | None + start_time: datetime + end_time: datetime + + +@dataclass(frozen=True, slots=True) +class SpendEventBuildError: + reason: str + + +@dataclass(frozen=True, slots=True) +class SpendEventDecodeError: + reason: str + + +def is_offloadable_success(response_obj: object) -> bool: + """Batch retrieves and interaction polls branch on the concrete response class, so they stay in-process.""" + return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) + + +def _json_fallback(value: object) -> str: + return str(value) + + +def _mapping_or_none(value: object) -> ObjectMapping | None: + try: + return _OBJECT_MAPPING.validate_python(value) + except ValidationError: + return None + + +def _drop_keys(mapping: ObjectMapping, keys: frozenset[str]) -> ObjectMapping: + return MappingProxyType({key: value for key, value in mapping.items() if key not in keys}) + + +def _budget_reservation(metadata: ObjectMapping) -> ObjectMapping | None: + """The admission-time reservation, wherever the request setup left it, so the sidecar can reconcile it.""" + direct: Final = _mapping_or_none(metadata.get("user_api_key_budget_reservation")) + if direct is not None: + return direct + auth: Final = metadata.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth): + return auth.budget_reservation + auth_mapping: Final = _mapping_or_none(auth) + return _mapping_or_none(auth_mapping.get("budget_reservation")) if auth_mapping is not None else None + + +def _metadata_for_event( + metadata: ObjectMapping | None, budget_reservation: ObjectMapping | None +) -> ObjectMapping | None: + if metadata is None: + return None + kept: Final = _drop_keys(metadata, _UNSERIALIZABLE_METADATA_KEYS) + if budget_reservation is None: + return kept + return MappingProxyType({**kept, "user_api_key_budget_reservation": budget_reservation}) + + +def _litellm_params_for_event( + litellm_params: _LitellmParams, cache_key: str | None, store_bodies: bool +) -> _LitellmParams: + metadata: Final = litellm_params.get("metadata") + litellm_metadata: Final = litellm_params.get("litellm_metadata") + budget_reservation: Final = next( + ( + reservation + for source in (litellm_metadata, metadata) + if source is not None and (reservation := _budget_reservation(source)) is not None + ), + None, + ) + projected: Final[_LitellmParams] = { + "api_base": litellm_params.get("api_base"), + "custom_llm_provider": litellm_params.get("custom_llm_provider"), + "litellm_call_id": litellm_params.get("litellm_call_id"), + "user_api_key_end_user_id": litellm_params.get("user_api_key_end_user_id"), + "metadata": _metadata_for_event(metadata, budget_reservation), + "litellm_metadata": _metadata_for_event(litellm_metadata, budget_reservation), + "proxy_server_request": litellm_params.get("proxy_server_request") if store_bodies else None, + "preset_cache_key": cache_key, + } + return projected + + +def _standard_logging_for_event(sl_object: ObjectMapping | None, store_bodies: bool) -> ObjectMapping | None: + if sl_object is None: + return None + dropped: Final = ( + _STANDARD_LOGGING_DROPPED_KEYS if store_bodies else _STANDARD_LOGGING_DROPPED_KEYS | _STANDARD_LOGGING_BODY_KEYS + ) + return _drop_keys(sl_object, dropped) + + +def _tool_call(name: str) -> _ToolCall: + tool_call: Final[_ToolCall] = {"id": None, "type": "function", "function": {"name": name, "arguments": "{}"}} + return tool_call + + +def _tool_call_choice(names: Sequence[str]) -> _ToolCallChoice: + choice: Final[_ToolCallChoice] = { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": tuple(_tool_call(name) for name in names)}, + } + return choice + + +def _compact_response(response_obj: object) -> CompactResponse | None: + """Usage, identity and tool calls of the response, in chat-completions shape, without the content.""" + dumped: Final = response_obj.model_dump() if isinstance(response_obj, BaseModel) else _mapping_or_none(response_obj) + if dumped is None: + return None + scalars: Final = _COMPACT_RESPONSE.validate_python(_drop_keys(dumped, frozenset({"choices"}))) + tool_call_names: Final = response_tool_call_names(response_obj) + if not tool_call_names: + return scalars + with_tool_calls: Final[CompactResponse] = {**scalars, "choices": (_tool_call_choice(tool_call_names),)} + return with_tool_calls + + +def _tool_name(tool: ObjectMapping) -> str | None: + """Chat tools nest the name under ``function``; Anthropic and Responses API tools keep it at the top.""" + function: Final = _mapping_or_none(tool.get("function")) + name: Final = function.get("name") if function is not None else tool.get("name") + return name.strip() if isinstance(name, str) and name.strip() else None + + +def _request_tool_names(kwargs: _SuccessKwargs) -> tuple[str, ...]: + passthrough: Final = kwargs.get("passthrough_logging_payload") + request_body: Final = passthrough.get("request_body") if passthrough is not None else None + passthrough_tools: Final = request_body.get("tools") if request_body is not None else None + return tuple( + name + for source in (kwargs.get("tools"), passthrough_tools) + if source is not None + for tool in source + if (name := _tool_name(tool)) is not None + ) + + +def preset_spend_log_cache_key(litellm_params: _LitellmParams) -> str | None: + """The key the caching layer already stored in ``litellm_params``, or ``Cache OFF``; never hashes the body.""" + if litellm.cache is None: + return CACHE_OFF_KEY + return litellm_params.get("preset_cache_key") + + +def _function_tool(name: str) -> _FunctionTool: + tool: Final[_FunctionTool] = {"type": "function", "function": {"name": name}} + return tool + + +def build_spend_event( + raw_kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime, store_bodies: bool +) -> bytes | SpendEventBuildError: + """Validate the success callback's kwargs and serialize the event once, as a single JSON line.""" + try: + kwargs: Final = _SUCCESS_KWARGS.validate_python(raw_kwargs) + except ValidationError as error: + return SpendEventBuildError(reason=str(error)) + litellm_params: Final = kwargs.get("litellm_params", _NO_LITELLM_PARAMS) + sl_object: Final = kwargs.get("standard_logging_object") + cache_key: Final = preset_spend_log_cache_key(litellm_params) + response_cost: Final = sl_object.get("response_cost") if sl_object is not None else kwargs.get("response_cost") + combined_usage: Final = kwargs.get("combined_usage_object") + event: Final = SpendEvent( + version=SPEND_EVENT_VERSION, + litellm_call_id=kwargs.get("litellm_call_id"), + call_type=kwargs.get("call_type"), + model=kwargs.get("model"), + custom_llm_provider=kwargs.get("custom_llm_provider"), + stream=kwargs.get("stream"), + complete_streaming_response="complete_streaming_response" in kwargs, + cache_hit=kwargs.get("cache_hit"), + response_cost=response_cost if isinstance(response_cost, (int, float)) else None, + start_time=start_time, + end_time=end_time, + completion_start_time=kwargs.get("completion_start_time"), + agent_id=kwargs.get("agent_id"), + litellm_trace_id=kwargs.get("litellm_trace_id"), + litellm_params=_litellm_params_for_event(litellm_params, cache_key, store_bodies), + standard_logging_object=_standard_logging_for_event(sl_object, store_bodies), + standard_callback_dynamic_params=kwargs.get("standard_callback_dynamic_params"), + response=_compact_response(response_obj), + combined_usage=combined_usage.model_dump() if combined_usage is not None else None, + realtime_tools=kwargs.get("realtime_tools"), + realtime_tool_calls=kwargs.get("realtime_tool_calls"), + request_tool_names=_request_tool_names(kwargs), + ) + return event.model_dump_json(fallback=_json_fallback).encode() + b"\n" + + +def decode_spend_event(line: bytes) -> SpendEvent | SpendEventDecodeError: + try: + return SpendEvent.model_validate_json(line) + except ValidationError as error: + return SpendEventDecodeError(reason=str(error)) + + +def spend_event_callback_args(event: SpendEvent) -> SpendEventCallbackArgs: + """The ``(kwargs, response_obj, start_time, end_time)`` the in-process cost callback receives.""" + tools: Final = tuple(_function_tool(name) for name in event.request_tool_names) + kwargs: Final[SpendCallbackKwargs] = { + "litellm_call_id": event.litellm_call_id, + "call_type": event.call_type, + "model": event.model, + "custom_llm_provider": event.custom_llm_provider, + "stream": event.stream, + "cache_hit": event.cache_hit, + "response_cost": event.response_cost, + "completion_start_time": event.completion_start_time, + "agent_id": event.agent_id, + "litellm_trace_id": event.litellm_trace_id, + "litellm_params": event.litellm_params, + "standard_logging_object": event.standard_logging_object, + "standard_callback_dynamic_params": event.standard_callback_dynamic_params, + "combined_usage_object": Usage.model_validate(event.combined_usage) + if event.combined_usage is not None + else None, + "realtime_tools": event.realtime_tools, + "realtime_tool_calls": event.realtime_tool_calls, + "tools": tools or None, + } + if not event.complete_streaming_response: + return SpendEventCallbackArgs(kwargs, event.response, event.start_time, event.end_time) + streaming_kwargs: Final[SpendCallbackKwargs] = {**kwargs, "complete_streaming_response": event.response} + return SpendEventCallbackArgs(streaming_kwargs, event.response, event.start_time, event.end_time) diff --git a/litellm/proxy/spend_tracking/spend_event_producer.py b/litellm/proxy/spend_tracking/spend_event_producer.py new file mode 100644 index 00000000000..20c7f177af0 --- /dev/null +++ b/litellm/proxy/spend_tracking/spend_event_producer.py @@ -0,0 +1,338 @@ +"""Fire-and-forget push of serialized spend events from an inference worker to the pod-local sidecar. + +``LITELLM_COLLECTOR_ENABLED=true`` turns the push on in the gateway; the sidecar process sets +``LITELLM_JOB_ROLE=collector`` and always runs the pipeline in-process. Events queue in a bounded +in-memory buffer that a single writer task flushes over a unix socket or loopback TCP connection. +When the sidecar is unreachable, the buffer is full, or the connection breaks mid-write, each affected +event follows ``LITELLM_COLLECTOR_ON_UNAVAILABLE``: ``fallback`` runs the existing cost pipeline in +the worker, ``drop`` counts it and moves on. Transitions are logged with the counters, so a sidecar +outage is visible without scraping anything. + +Delivery is at-most-once: a sidecar crash loses the events the kernel already took from its socket. +A sidecar that stops gracefully half-closes each connection first (EOF towards the producer) and +keeps reading until the producer hangs up, so the producer switches to the unavailable policy without +losing the events in flight. A write that fails part-way follows the unavailable policy without double +counting: ``drain()`` only fails while part of the line is still buffered in this process, so the +sidecar can at most have read a truncated line, which it discards. When the gateway itself stops with +the writer stuck mid-send, only an event whose bytes are still in the producer's write buffer follows +the unavailable policy; the connection is aborted first so the sidecar discards the truncated line +instead of also counting it. Events from one uvicorn worker are handled in the order it produced them; +events from different workers interleave, exactly like the in-process callbacks do today. +""" + +import asyncio +import ipaddress +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias +from urllib.parse import urlsplit + +from pydantic import AliasChoices, Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from litellm._logging import verbose_proxy_logger + +COLLECTOR_ENV_PREFIX: Final = "LITELLM_COLLECTOR_" +COLLECTOR_JOB_ROLE: Final = "collector" +DEFAULT_COLLECTOR_ADDRESS: Final = "unix:///var/run/litellm/collector.sock" +RECONNECT_BACKOFF_SECONDS: Final = 1.0 +DROP_LOG_EVERY: Final = 1000 + +UnavailablePolicy: TypeAlias = Literal["fallback", "drop"] +PublishOutcome: TypeAlias = Literal["queued", "fallback", "dropped"] + + +class CollectorSettings(BaseSettings): + """``LITELLM_COLLECTOR_*`` env vars, shared by the gateway producer and the sidecar consumer.""" + + model_config = SettingsConfigDict( + env_prefix=COLLECTOR_ENV_PREFIX, case_sensitive=False, extra="ignore", frozen=True, populate_by_name=True + ) + + enabled: bool = False + address: str = DEFAULT_COLLECTOR_ADDRESS + buffer_size: int = Field(default=1000, ge=1) + on_unavailable: UnavailablePolicy = "fallback" + drain_timeout_seconds: float = Field(default=10.0, gt=0) + connect_timeout_seconds: float = Field(default=1.0, gt=0) + job_role: str | None = Field(default=None, validation_alias=AliasChoices("LITELLM_JOB_ROLE")) + + @property + def produces(self) -> bool: + return self.enabled and self.job_role != COLLECTOR_JOB_ROLE + + +@dataclass(frozen=True, slots=True) +class UnixAddress: + path: str + + +@dataclass(frozen=True, slots=True) +class TcpAddress: + host: str + port: int + + +@dataclass(frozen=True, slots=True) +class AddressError: + reason: str + + +CollectorAddress: TypeAlias = UnixAddress | TcpAddress + + +def _is_loopback(host: str) -> bool: + try: + return ipaddress.ip_address(host).is_loopback + except ValueError: + return host == "localhost" + + +def parse_collector_address(address: str) -> CollectorAddress | AddressError: + """``unix:///path/to.sock`` or ``tcp://127.0.0.1:port``; the socket carries unauthenticated spend events.""" + parsed: Final = urlsplit(address) + if parsed.scheme == "unix" and parsed.path: + return UnixAddress(path=parsed.path) + if parsed.scheme == "tcp" and parsed.hostname and parsed.port is not None: + if not _is_loopback(parsed.hostname): + return AddressError(reason=f"tcp collector address must be a loopback host, got {address!r}") + return TcpAddress(host=parsed.hostname, port=parsed.port) + return AddressError(reason=f"expected unix:///path or tcp://127.0.0.1:port, got {address!r}") + + +async def open_collector_connection( + address: CollectorAddress, timeout: float +) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + match address: + case UnixAddress(path=path): + return await asyncio.wait_for(asyncio.open_unix_connection(path), timeout) + case TcpAddress(host=host, port=port): + return await asyncio.wait_for(asyncio.open_connection(host, port), timeout) + + +def build_spend_event_producer( + settings: CollectorSettings, fallback: Callable[[bytes], Awaitable[None]] +) -> "SpendEventProducer | None": + """The gateway producer for these settings, or ``None`` when the pipeline stays in-process.""" + if not settings.produces: + return None + address: Final = parse_collector_address(settings.address) + if isinstance(address, AddressError): + verbose_proxy_logger.error("collector: %s; running the spend pipeline in-process", address.reason) + return None + verbose_proxy_logger.info( + "collector: offloading spend tracking to %s (buffer=%d, on_unavailable=%s)", + settings.address, + settings.buffer_size, + settings.on_unavailable, + ) + return SpendEventProducer( + address=address, + on_unavailable=settings.on_unavailable, + buffer_size=settings.buffer_size, + connect_timeout=settings.connect_timeout_seconds, + fallback=fallback, + ) + + +@dataclass(frozen=True, slots=True) +class _Connection: + reader: asyncio.StreamReader + writer: asyncio.StreamWriter + + @property + def alive(self) -> bool: + return not self.writer.is_closing() and not self.reader.at_eof() + + +@dataclass(frozen=True, slots=True) +class SpendEventProducerStats: + queued: int + sent: int + fallback: int + dropped: int + connected: bool + + +class SpendEventProducer: + """Bounded buffer plus one writer task per process; see the module docstring for the contract.""" + + def __init__( + self, + address: CollectorAddress, + on_unavailable: UnavailablePolicy, + buffer_size: int, + connect_timeout: float, + fallback: Callable[[bytes], Awaitable[None]], + clock: Callable[[], float] = time.monotonic, + open_connection: Callable[ + [CollectorAddress, float], Awaitable[tuple[asyncio.StreamReader, asyncio.StreamWriter]] + ] = open_collector_connection, + ) -> None: + self._address = address + self._on_unavailable = on_unavailable + self._buffer_size = buffer_size + self._connect_timeout = connect_timeout + self._fallback = fallback + self._clock = clock + self._open_connection = open_connection + self._queue: asyncio.Queue[bytes] | None = None + self._writer_task: asyncio.Task[None] | None = None + self._connection: _Connection | None = None + self._in_flight: bytes | None = None + self._closing = False + self._next_connect_at = 0.0 + self._queued = 0 + self._sent = 0 + self._fallback_count = 0 + self._dropped = 0 + + def stats(self) -> SpendEventProducerStats: + return SpendEventProducerStats( + queued=self._queued, + sent=self._sent, + fallback=self._fallback_count, + dropped=self._dropped, + connected=self._connection is not None, + ) + + async def publish(self, line: bytes) -> PublishOutcome: + """Hand one serialized event to the writer task, or apply the unavailable policy right away.""" + if self._closing or self._clock() < self._next_connect_at: + return await self._unavailable(line, "sidecar unreachable") + queue: Final = self._ensure_writer() + try: + queue.put_nowait(line) + except asyncio.QueueFull: + return await self._unavailable(line, "buffer full") + self._queued += 1 + return "queued" + + async def close(self, drain_timeout: float) -> None: + """Flush the buffer for up to ``drain_timeout`` seconds, then apply the unavailable policy to the rest.""" + self._closing = True + queue: Final = self._queue + task: Final = self._writer_task + if queue is None or task is None: + return + try: + await asyncio.wait_for(queue.join(), drain_timeout) + except asyncio.TimeoutError: + verbose_proxy_logger.warning( + "collector: %s events still buffered after %.1fs drain timeout", queue.qsize(), drain_timeout + ) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + unsent: Final = self._take_unsent() + await self._disconnect() + if unsent is not None: + await self._unavailable(unsent, "shutdown") + while not queue.empty(): + await self._unavailable(queue.get_nowait(), "shutdown") + + def _take_unsent(self) -> bytes | None: + """The in-flight event if any of its bytes never left this process, aborting the half-written connection.""" + in_flight: Final = self._in_flight + self._in_flight = None + connection: Final = self._connection + if in_flight is None: + return None + if connection is None: + return in_flight + if connection.writer.transport.get_write_buffer_size() == 0: + return None + connection.writer.transport.abort() + return in_flight + + def _ensure_writer(self) -> asyncio.Queue[bytes]: + if self._queue is None: + self._queue = asyncio.Queue(maxsize=self._buffer_size) + if self._writer_task is None or self._writer_task.done(): + self._writer_task = asyncio.get_running_loop().create_task(self._run_writer(self._queue)) + return self._queue + + async def _run_writer(self, queue: asyncio.Queue[bytes]) -> None: + while True: + line = await queue.get() + try: + await self._send(line) + finally: + queue.task_done() + + async def _send(self, line: bytes) -> None: + self._in_flight = line + connection: Final = await self._connect() + if connection is None: + self._in_flight = None + await self._unavailable(line, "sidecar unreachable") + return + try: + connection.writer.write(line) + await connection.writer.drain() + except (ConnectionError, OSError, RuntimeError) as error: # uvloop: RuntimeError on a closed transport + self._in_flight = None + await self._disconnect() + self._next_connect_at = self._clock() + RECONNECT_BACKOFF_SECONDS + await self._unavailable(line, f"write failed: {error}") + return + self._in_flight = None + self._sent += 1 + + async def _connect(self) -> _Connection | None: + if self._connection is not None and self._connection.alive: + return self._connection + await self._disconnect() + if self._clock() < self._next_connect_at: + return None + try: + reader, writer = await self._open_connection(self._address, self._connect_timeout) + except (ConnectionError, OSError, asyncio.TimeoutError) as error: + self._next_connect_at = self._clock() + RECONNECT_BACKOFF_SECONDS + verbose_proxy_logger.warning( + "collector: cannot reach %s (%s); applying %s policy for %.0fs. stats=%s", + self._address, + error, + self._on_unavailable, + RECONNECT_BACKOFF_SECONDS, + self.stats(), + ) + return None + self._connection = _Connection(reader=reader, writer=writer) + verbose_proxy_logger.info("collector: connected to %s. stats=%s", self._address, self.stats()) + return self._connection + + async def _disconnect(self) -> None: + connection: Final = self._connection + self._connection = None + if connection is None: + return + connection.writer.close() + try: + await connection.writer.wait_closed() + except (ConnectionError, OSError): + pass + + async def _unavailable(self, line: bytes, reason: str) -> PublishOutcome: + if self._on_unavailable == "fallback": + self._fallback_count += 1 + fallback: Final = asyncio.ensure_future(self._run_fallback(line, reason)) + try: + await asyncio.shield(fallback) + except asyncio.CancelledError: + await fallback + raise + return "fallback" + self._dropped += 1 + if self._dropped % DROP_LOG_EVERY == 1: + verbose_proxy_logger.warning("collector: dropping spend event (%s). stats=%s", reason, self.stats()) + return "dropped" + + async def _run_fallback(self, line: bytes, reason: str) -> None: + try: + await self._fallback(line) + except Exception: # noqa: BLE001 # one failing event must not kill the writer task + verbose_proxy_logger.exception("collector: in-process fallback failed (%s)", reason) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index f8831ca4152..8cfb6354dd0 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -26,7 +26,12 @@ from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME +from litellm.constants import ( + EMPTY_MAPPING, + LITELLM_TRUNCATED_PAYLOAD_FIELD, + LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, +) +from litellm.litellm_core_utils.classifier_logging import classifier_audit_fields, classifier_input_snapshot from litellm.proxy._types import * from litellm.proxy._types import ProviderBudgetResponse, ProviderBudgetResponseObject from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -2920,6 +2925,32 @@ async def _fetch_session_representatives( return [rep_by_key[key] for key in session_keys if key in rep_by_key] # mutable-ok: rows are enriched in place +async def _count_grouped_sessions( + prisma_client: "PrismaClient", + where_clause: str, + sql_params: Sequence[object], + next_param_index: int, +) -> tuple[int, bool]: + """Count the sessions matching the filter, returning ``(total, total_is_capped)`` bounded by the count cap.""" + count_query: Final = f""" + SELECT COUNT(*) AS total_count + FROM ( + SELECT 1 + FROM "LiteLLM_SpendLogs" + WHERE {where_clause} + GROUP BY {_SESSION_GROUP_KEY_SQL} + LIMIT ${next_param_index} + ) AS bounded_sessions + """ + count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw( + prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 + ) + raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0 + return ( + (SPEND_LOGS_PAGINATION_COUNT_CAP, True) if raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP else (raw_total, False) + ) + + async def _ui_session_grouped_spend_logs( prisma_client: "PrismaClient", sql_conditions: Sequence[str], @@ -2939,11 +2970,19 @@ async def _ui_session_grouped_spend_logs( next ``page_size`` sessions ordered by ``(MAX(startTime), session_key, api_key)``, resumed from the ``session_cursor`` keyset ``'||'`` instead of an OFFSET, so - page depth does not degrade the query plan. Each session is represented + page depth does not degrade the query plan. A request for ``page > 1`` + without a cursor (the UI jumping straight to the last page, or back to a + page it never walked through) falls back to ``OFFSET (page - 1) * + page_size``, trimmed to the end of the ``SPEND_LOGS_PAGINATION_COUNT_CAP`` + window the capped ``total`` promises, so a page never runs past that total + and one starting at or past it returns no rows without a query. Each session is represented by its newest non-MCP row, enriched by ``_build_ui_spend_logs_response`` exactly like the flat listing, and the response carries ``next_session_cursor`` / ``has_more`` while ``total`` counts sessions - (capped like the flat total). + (capped like the flat total). A page that runs out of sessions while still + holding some is itself the end of the list, so its ``total`` is + ``offset + len(page)`` and the grouped count query is skipped; a page that + starts past the end says nothing about the total, so that one is counted. """ where_clause: Final = " AND ".join(sql_conditions) if sql_conditions else "TRUE" cmp_op: Final = "<" if sort_desc else ">" @@ -2958,6 +2997,10 @@ async def _ui_session_grouped_spend_logs( ) cursor_params: Final[tuple[object, ...]] = cursor if cursor else () limit_index: Final = next_param_index + len(cursor_params) + offset: Final = (page - 1) * page_size if cursor is None else 0 + page_limit: Final = min(page_size, SPEND_LOGS_PAGINATION_COUNT_CAP - offset) + offset_params: Final[tuple[int, ...]] = (offset,) if offset and page_limit > 0 else () + offset_clause: Final = f"OFFSET ${limit_index + 1}" if offset_params else "" page_query: Final = f""" SELECT {_SESSION_KEY_EXPR} AS session_key, @@ -2968,36 +3011,29 @@ async def _ui_session_grouped_spend_logs( GROUP BY {_SESSION_GROUP_KEY_SQL} {having_clause} ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction} - LIMIT ${limit_index} + LIMIT ${limit_index} {offset_clause} """ - page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw( - prisma_client, page_query, *sql_params, *cursor_params, page_size + 1 + page_rows: Final[Sequence[_SessionPageRow]] = ( + () + if page_limit <= 0 + else await _query_raw(prisma_client, page_query, *sql_params, *cursor_params, page_limit + 1, *offset_params) ) - has_more: Final = len(page_rows) > page_size - visible_rows: Final = page_rows[:page_size] + has_more: Final = len(page_rows) > page_limit + visible_rows: Final = page_rows[:page_limit] next_cursor: Final = ( f"{visible_rows[-1]['last_activity']}|{visible_rows[-1]['api_key']}|{visible_rows[-1]['session_key']}" if has_more and visible_rows else None ) - count_query: Final = f""" - SELECT COUNT(*) AS total_count - FROM ( - SELECT 1 - FROM "LiteLLM_SpendLogs" - WHERE {where_clause} - GROUP BY {_SESSION_GROUP_KEY_SQL} - LIMIT ${next_param_index} - ) AS bounded_sessions - """ - count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw( - prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 + page_starts_inside_the_list: Final = offset == 0 or len(page_rows) > 0 + page_ends_the_list: Final = cursor is None and page_limit > 0 and not has_more and page_starts_inside_the_list + total_records, total_is_capped = ( + (offset + len(page_rows), False) + if page_ends_the_list + else await _count_grouped_sessions(prisma_client, where_clause, sql_params, next_param_index) ) - raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0 - total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP - total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total session_keys: Final = tuple((row["session_key"], row["api_key"]) for row in visible_rows) data: Final[list[dict[str, object]]] = ( # mutable-ok: _build_ui_spend_logs_response writes onto each row @@ -3099,7 +3135,11 @@ async def _resolve_request_response_payload( proxy_server_request: Final = row.get("proxy_server_request") pg_payload: Final = RequestResponsePayload(messages, response, proxy_server_request) - if ( + stored_request: Final = classifier_input_snapshot(proxy_server_request) + truncated_audit: Final = bool(stored_request and classifier_audit_fields(stored_request)) and ( + LITELLM_TRUNCATED_PAYLOAD_FIELD in str(proxy_server_request) + ) + if not truncated_audit and ( _spend_log_field_has_content(messages) or _spend_log_field_has_content(response) or _spend_log_field_has_content(proxy_server_request) @@ -3124,10 +3164,22 @@ async def _resolve_request_response_payload( if payload is None: return pg_payload + cold_audit: Final = classifier_audit_fields(payload) + resolved_request: Final = ( + { + **(classifier_input_snapshot(payload.get("proxy_server_request")) or stored_request or EMPTY_MAPPING), + **cold_audit, + } + if cold_audit + else payload.get("proxy_server_request") + ) + if truncated_audit: + return RequestResponsePayload(messages, response, resolved_request if cold_audit else proxy_server_request) + return RequestResponsePayload( messages=payload.get("messages"), response=payload.get("response"), - proxy_server_request=payload.get("proxy_server_request"), + proxy_server_request=resolved_request, ) @@ -4633,16 +4685,16 @@ async def _assert_user_can_view_request_id( Verify the requesting non-admin user is allowed to view this spend-log row. Allowed when the log belongs to the user directly, or to one of their permitted teams (admin or ``/spend/logs`` permission). - Raises HTTP 403 if not. + Raises HTTP 403 if not, including when no spend-log row exists for the + request_id (e.g. it was pruned by retention), so a missing row can't be + used to read a payload out of cold storage via the detail endpoint. """ row: Final = await _find_spend_log_row(prisma_client, request_id) - if row is None: + + if row is not None and row.user is not None and row.user == user_api_key_dict.user_id: return - if row.user is not None and row.user == user_api_key_dict.user_id: - return - - if row.team_id: + if row is not None and row.team_id: can_view: Final = await _can_team_member_view_log( prisma_client=prisma_client, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8a06bf68b81..01398d38687 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -4,6 +4,7 @@ import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt +from types import MappingProxyType from typing import Final, Literal, Protocol, cast, runtime_checkable from pydantic import BaseModel @@ -11,17 +12,22 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + EMPTY_MAPPING, LITELLM_PROXY_MASTER_KEY_ALIAS, LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, + MAX_SPEND_LOG_MODEL_NAME_LENGTH, + MCP_SPEND_LOG_MODEL_PREFIX, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, + UNKNOWN_MODEL_SPEND_LOG_MODEL, ) from litellm.constants import ( MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB, ) +from litellm.litellm_core_utils.classifier_logging import classifier_audit_fields, without_classifier_audit from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, reconstruct_model_name, @@ -34,6 +40,7 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload, SpendLogsRouterMetadata +from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token from litellm.types.utils import ( @@ -331,10 +338,16 @@ def _sl_attribution_fallback( return standard_logging_payload.get(field) or "" +def _looks_like_model_name(model: str) -> bool: + candidate: Final = model.removeprefix(MCP_SPEND_LOG_MODEL_PREFIX) + return len(candidate) <= MAX_SPEND_LOG_MODEL_NAME_LENGTH and not any(char.isspace() for char in candidate) + + def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogsPayload: if kwargs is None: kwargs = {} + rejected_as_unknown_model: Final = isinstance(response_obj, ProxyModelNotFoundError) if response_obj is None: response_obj = {} elif not isinstance(response_obj, BaseModel) and not isinstance(response_obj, dict): @@ -432,9 +445,19 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs or None ) raw_model: Final = cast(str, kwargs.get("model") or "") - model_name: Final = ( + resolved_model: Final = ( standard_logging_payload.get("model") if standard_logging_payload is not None else None ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) + failed_with_prompt_shaped_model: Final = ( + _get_status_for_spend_log(metadata=metadata) == "failure" + and not _model_group + and not _looks_like_model_name(resolved_model) + ) + model_name: Final = ( + UNKNOWN_MODEL_SPEND_LOG_MODEL + if rejected_as_unknown_model or failed_with_prompt_shaped_model + else resolved_model + ) litellm_call_id: Final = cast( str | None, kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), @@ -533,10 +556,12 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs additional_usage_values["cache_creation_input_tokens"] = cache_write_tokens clean_metadata["additional_usage_values"] = additional_usage_values - if litellm.cache is not None: - cache_key = litellm.cache.get_cache_key(**kwargs) - else: + if litellm.cache is None: cache_key = "Cache OFF" + elif litellm_params.get("preset_cache_key") is not None: + cache_key = litellm_params["preset_cache_key"] + else: + cache_key = litellm.cache.get_cache_key(**kwargs) if cache_hit is True: import time @@ -590,6 +615,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs metadata=metadata, standard_logging_payload=standard_logging_payload, omit_when_missing=_omits_session_id_when_missing(metadata), + batch_trace_session_id=_get_batch_trace_session_id(call_type=call_type, request_id=id), ), request_duration_ms=_get_request_duration_ms(start_time, end_time), status=_get_status_for_spend_log( @@ -628,20 +654,44 @@ def _omits_session_id_when_missing(metadata: Mapping[str, object] | None) -> boo return general_settings.get("missing_session_id") == "omit" +_BATCH_TRACE_CALL_TYPES: Final = frozenset( + { + CallTypes.create_batch.value, + CallTypes.acreate_batch.value, + CallTypes.retrieve_batch.value, + CallTypes.aretrieve_batch.value, + } +) + + +def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -> str | None: + """A batch's create row and its poller-written cost row both derive their request id + from the same batch id (the cost row appends BATCH_COST_REQUEST_ID_SUFFIX), so using + that id as the session groups the batch lifecycle into one trace on the logs UI. The + poller builds its own logging context, so per-request trace ids can never link them.""" + if call_type not in _BATCH_TRACE_CALL_TYPES or not request_id: + return None + return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX) + + def _get_session_id_for_spend_log( kwargs: Mapping[str, object], metadata: Mapping[str, object] | None, standard_logging_payload: StandardLoggingPayload | None, omit_when_missing: bool, + batch_trace_session_id: str | None = None, ) -> str | None: """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may - be a copied trace id.""" + be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks + the per-request trace ids because those differ between the create call and the cost poller's row.""" if omit_when_missing: session_id: Final = metadata.get("session_id") if metadata else None return str(session_id) if session_id else None from litellm._uuid import uuid + if batch_trace_session_id is not None: + return batch_trace_session_id if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: return str(standard_logging_payload.get("trace_id")) if kwargs.get("litellm_trace_id") is not None: @@ -818,7 +868,7 @@ def _get_messages_for_spend_logs_payload( standard_logging_payload: StandardLoggingPayload | None, metadata: dict | None = None, ) -> str: - if _should_store_prompts_and_responses_in_spend_logs(): + if should_store_prompts_and_responses_in_spend_logs(): if standard_logging_payload is not None: call_type: Final = standard_logging_payload.get("call_type", "") if call_type == "_arealtime": @@ -835,7 +885,7 @@ _SENSITIVE_REQUEST_BODY_KEYS: Final = frozenset({"secret_fields"}) def _sanitize_request_body_for_spend_logs_payload( - request_body: dict, + request_body: Mapping[str, object], visited: set | None = None, max_string_length_prompt_in_db: int | None = None, ) -> dict: @@ -1068,7 +1118,7 @@ def _sanitize_guardrail_information_for_spend_logs( here to match OTEL's defensive read pattern; otherwise iteration would yield the dict's keys and crash the whole spend-log write. """ - if guardrail_information is None or _should_store_prompts_and_responses_in_spend_logs(): + if guardrail_information is None or should_store_prompts_and_responses_in_spend_logs(): return guardrail_information entries: Final = [guardrail_information] if isinstance(guardrail_information, dict) else guardrail_information return [_redact_prompt_fields_in_guardrail_entry(entry) for entry in entries if isinstance(entry, dict)] @@ -1140,7 +1190,7 @@ def _sanitize_error_information_for_spend_logs( sanitized = cast(dict, {**error_information}) - if not _should_store_prompts_and_responses_in_spend_logs(): + if not should_store_prompts_and_responses_in_spend_logs(): for field in ("error_message", "traceback"): value = sanitized.get(field) if isinstance(value, str): @@ -1217,14 +1267,18 @@ def _get_proxy_server_request_for_spend_logs_payload( kwargs: dict | None = None, ) -> str: """ - Only store if _should_store_prompts_and_responses_in_spend_logs() is True + Only store if should_store_prompts_and_responses_in_spend_logs() is True If turn_off_message_logging is enabled, redact messages in the request body. """ - if _should_store_prompts_and_responses_in_spend_logs(): - _proxy_server_request: Final = cast(dict | None, litellm_params.get("proxy_server_request", {})) + if should_store_prompts_and_responses_in_spend_logs(): + _proxy_server_request: Final = cast(dict | None, litellm_params.get("proxy_server_request", EMPTY_MAPPING)) if _proxy_server_request is not None: - _request_body = _proxy_server_request.get("body", {}) or {} + _request_body = _proxy_server_request.get("body", EMPTY_MAPPING) or EMPTY_MAPPING + + standard_payload: Final = (kwargs or EMPTY_MAPPING).get("standard_logging_object") + if isinstance(standard_payload, Mapping): + _request_body = MappingProxyType({**_request_body, **classifier_audit_fields(standard_payload)}) if kwargs is not None: realtime_tools: Final = kwargs.get("realtime_tools") @@ -1247,7 +1301,7 @@ def _get_proxy_server_request_for_spend_logs_payload( # If redaction is enabled, convert to serializable dict before redacting if should_redact_message_logging(model_call_details=model_call_details): - _request_body = _convert_mapping_to_json_serializable(_request_body) + _request_body = _convert_mapping_to_json_serializable(without_classifier_audit(_request_body)) perform_redaction(model_call_details=_request_body, result=None) _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body) @@ -1267,7 +1321,7 @@ def _get_vector_store_request_for_spend_logs_payload( """ If user does not want to store prompts and responses, then remove the content from the vector store request metadata """ - if _should_store_prompts_and_responses_in_spend_logs(): + if should_store_prompts_and_responses_in_spend_logs(): return vector_store_request_metadata # if user does not want to store prompts and responses, then remove the content from the vector store request metadata @@ -1291,7 +1345,7 @@ def _get_response_for_spend_logs_payload( ) -> str: if payload is None: return "{}" - if _should_store_prompts_and_responses_in_spend_logs(): + if should_store_prompts_and_responses_in_spend_logs(): response_obj: object = payload.get("response") if response_obj is None: return "{}" @@ -1339,7 +1393,7 @@ def _get_response_for_spend_logs_payload( return "{}" -def _should_store_prompts_and_responses_in_spend_logs() -> bool: +def should_store_prompts_and_responses_in_spend_logs() -> bool: from litellm.proxy.proxy_server import general_settings from litellm.secret_managers.main import get_secret_bool diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b4e4dfeae67..8e6142090a4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -11,7 +11,7 @@ import sys import threading import time import traceback -from collections.abc import AsyncGenerator, Awaitable, Callable, Collection, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart @@ -37,6 +37,7 @@ from litellm.proxy._types import ( SpendLogsMetadata, SpendLogsPayload, ) +from litellm.proxy.common_utils.openai_error_payload import openai_error_param from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse @@ -67,6 +68,7 @@ except ImportError: raise ImportError("backoff is not installed. Please install it via 'pip install backoff'") from fastapi import HTTPException, status +from pydantic import TypeAdapter import litellm import litellm.litellm_core_utils @@ -93,12 +95,14 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, + get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms import load_guardrail_translation_mappings from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( @@ -121,7 +125,13 @@ from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, ) +from litellm.proxy.db.health_check_latest import ( + LatestHealthCheckRow, + fetch_latest_health_checks, + fetch_latest_health_checks_for_models, +) from litellm.proxy.db.log_db_metrics import log_db_metrics +from litellm.proxy.db.pgbouncer import database_url_is_pooled from litellm.proxy.db.prisma_client import ( PrismaWrapper, parse_iam_endpoint_from_url, @@ -139,6 +149,7 @@ from litellm.proxy.db.token_auth import ( ) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, + resolve_endpoint_translation, ) from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck @@ -152,9 +163,11 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.hooks.sensitive_data_routing import ( _PROXY_SensitiveDataRoutingHandler, ) -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( @@ -170,6 +183,7 @@ from litellm.repositories.verification_token_repository import ( ) from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.mcp import ( MCPDuringCallResponseObject, MCPPreCallRequestObject, @@ -177,6 +191,9 @@ from litellm.types.mcp import ( ) from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams +from litellm.utils import ( + _add_custom_logger_callback_to_specific_event, # pyright: ignore[reportPrivateUsage] # only string-to-logger helper +) if TYPE_CHECKING: from mcp.types import CallToolResult @@ -188,6 +205,7 @@ if TYPE_CHECKING: from prisma.types import HttpConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -446,12 +464,296 @@ def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardrail ) -def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[str]: - managed: Final = _policy_state_metadata(data).get("_pipeline_managed_guardrails") +def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> frozenset[str]: + return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps) + + +def pipeline_managed_guardrail_names( + data: Mapping[str, object], mode: Literal["pre_call", "post_call"] +) -> frozenset[str]: + return _pipeline_step_guardrail_names( + tuple((policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == mode) + ) + + +def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]: + resolved: Final = tuple( + litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( + cast( # cast-ok: the resolver returns None for unknown names, filtered below + _custom_logger_compatible_callbacks_literal, callback + ) + ) + if isinstance(callback, str) + else callback + for callback in litellm.callbacks + ) + present: Final = tuple(callback for callback in resolved if callback is not None) + guardrails: Final = tuple(callback for callback in present if isinstance(callback, CustomGuardrail)) + others: Final = cast( # cast-ok: mirrors the legacy loop, which treated every non-guardrail entry as a CustomLogger + "tuple[CustomLogger, ...]", + tuple(callback for callback in present if not isinstance(callback, CustomGuardrail)), + ) + return (guardrails, others) + + +def _merge_pipeline_metadata_bucket( + data: dict, bucket_key: str, modified_bucket_value: object +) -> None: # mutable-ok: request payload dict, written in place + if not isinstance(modified_bucket_value, dict): + return + modified_bucket: Final = cast("dict[str, object]", modified_bucket_value) # cast-ok: metadata buckets are str-keyed + surviving_writes: Final = { + key: value for key, value in modified_bucket.items() if key != "guardrails" + } # mutable-ok: merged into the live request metadata bucket in place + existing_bucket: Final = data.get(bucket_key) + if isinstance(existing_bucket, dict): + cast("dict[str, object]", existing_bucket).update(surviving_writes) # cast-ok: metadata buckets are str-keyed + else: + data[bucket_key] = surviving_writes + + +def _merge_pipeline_metadata_writes( + data: dict, modified_data: Mapping[str, object] +) -> None: # mutable-ok: request payload dict, written in place + """ + Copy metadata-bucket writes from a pipeline's working copy back onto the request. + + Post_call pipelines run step hooks against a copied request dict so the payload + already sent upstream stays untouched, but hooks record proxy-internal logging + state in the metadata buckets (``applied_guardrails`` for response headers, + ``standard_logging_guardrail_information`` for spend logs), and those writes + must reach the request dict the proxy keeps reading after the pipeline returns. + + The ``guardrails`` key is the executor's per-step activation flag for + ``should_run_guardrail``, not a hook write, so it stays in the working copy. + """ + for bucket_key in ("metadata", "litellm_metadata"): + _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) + + +def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool: + callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) + if callback is None: + return False + if PipelineExecutor.supports_unified_execution(callback): + return True return ( - frozenset(cast("Collection[str]", managed)) # cast-ok: the policy engine wrote these guardrail names - if managed - else frozenset() + translation is not None + and type(translation).assembles_streamed_response + and PipelineExecutor.supports_streaming_execution(callback) + ) + + +def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]: + return tuple( + (policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" + ) + + +_PENDING_BACKGROUND_RESPONSE_STATUSES: Final = frozenset(("queued", "in_progress")) + + +def _is_pending_background_response(response: LLMResponseTypes) -> bool: + return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES + + +def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline") -> frozenset[str]: + resolved: Final = PolicyResolver.resolve_policy_guardrails( + policy_name=policy_name, policies=get_policy_registry().get_all_policies() + ) + return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps) + + +def _guardrails_run_standalone_pre_call(data: Mapping[str, object]) -> frozenset[str]: + return frozenset( + callback.guardrail_name + for callback in litellm.callbacks + if isinstance(callback, CustomGuardrail) + and callback.guardrail_name is not None + and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + ) + + +def _without_names( + bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write + slot: str, + names: frozenset[str], +) -> None: + claimed: Final = bucket.get(slot) + if not isinstance(claimed, list): + return + remaining: Final = [ # mutable-ok: the slot stays a list, the shape every applied_* header writer appends to + name for name in claimed if name not in names + ] + if remaining: + bucket[slot] = remaining # rebind-ok: the slot lives in the shared request-state dict, rewritten in place + else: + bucket.pop(slot) + + +def _withdraw_deferred_claims( + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data + deferred: Sequence[tuple[str, "GuardrailPipeline"]], +) -> None: + outside_by_policy: Final = MappingProxyType( + {policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred} + ) + running_elsewhere: Final = pipeline_managed_guardrail_names(data, "pre_call").union( + _guardrails_run_standalone_pre_call(data), *outside_by_policy.values() + ) + withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside) + withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere + _, bucket = get_or_create_metadata_bucket(data) + _without_names(bucket, "applied_policies", withdrawn_policies) + _without_names(bucket, "applied_guardrails", withdrawn_guardrails) + sources: Final = bucket.get("policy_sources") + if not isinstance(sources, dict): + return + remaining_sources: Final = { # mutable-ok: policy_sources stays a dict, the shape its writer updates in place + name: reason for name, reason in sources.items() if name not in withdrawn_policies + } + if remaining_sources: + bucket["policy_sources"] = remaining_sources + else: + bucket.pop("policy_sources") + + +def _defer_post_call_pipelines( + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data + response: ResponsesAPIResponse, +) -> None: + deferred: Final = _post_call_pipelines(data) + if not deferred: + return + verbose_proxy_logger.debug( + "Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s", + response.id, + response.status, + ", ".join(policy_name for policy_name, _pipeline in deferred), + ) + tag_matched: Final = _tag_matched_deferrals(data, deferred) + if tag_matched: + verbose_proxy_logger.warning( + "Policy engine: background response %s matched post_call policies through a request tag at submit; " + "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " + "does not govern the completed response: %s", + response.id, + ", ".join(tag_matched), + ) + body_selected: Final = _body_selected_deferrals(data, deferred) + if body_selected: + verbose_proxy_logger.warning( + "Policy engine: background response %s matched post_call policies through the request body's policies " + "list at submit; retrieval carries no request body, so those policies do not govern the completed " + "response: %s", + response.id, + ", ".join(body_selected), + ) + _withdraw_deferred_claims(data, deferred) + + +def _tag_matched_deferrals( + data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] +) -> tuple[str, ...]: + sources: Final = _policy_state_metadata(data).get("policy_sources") + if not isinstance(sources, dict): + return () + return tuple( + policy_name + for policy_name, _pipeline in deferred + if policy_name in sources and "tag:" in str(sources[policy_name]) + ) + + +def _body_selected_deferrals( + data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] +) -> tuple[str, ...]: + sources: Final = _policy_state_metadata(data).get("policy_sources") + attributed: Final = frozenset(sources) if isinstance(sources, dict) else frozenset() + return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed) + + +def _pipeline_unsupported_streaming_guardrails( + pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" +) -> tuple[str, ...]: + return tuple( + dict.fromkeys( + step.guardrail + for step in pipeline.steps + if not _pipeline_step_supports_streaming(step.guardrail, translation) + ) + ) + + +def _pipeline_is_streamable( + policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" +) -> bool: + unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation) + if not unsupported: + return True + verbose_proxy_logger.warning( + "Policy '%s' has post_call pipeline guardrails a streaming pipeline cannot run on this route yet; they " + "need the unified apply_guardrail interface, or a post-call hook without a streaming iterator hook on a " + "route whose translation assembles the streamed response. The stream skips the pipeline and its " + "guardrails run on their own: %s", + policy_name, + ", ".join(unsupported), + ) + return False + + +def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseTranslation | None": + resolved: Final = resolve_endpoint_translation(user_api_key_dict, None) + return None if resolved is None else resolved[1] + + +def stream_gated_guardrail_names( + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> frozenset[str]: + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if translation is None: + return frozenset() + return _pipeline_step_guardrail_names( + tuple( + (policy_name, pipeline) + for policy_name, pipeline in _post_call_pipelines(request_data) + if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) + ) + ) + + +def _streamable_post_call_pipelines( + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> tuple[tuple[str, "GuardrailPipeline"], ...]: + """ + The post_call pipelines a streaming response can be gated through. + + Streaming pipelines scan the buffered stream through the endpoint guardrail + translation of the request route, so every step's guardrail needs either the + unified apply_guardrail interface or, on a route whose translation assembles + the streamed response, a post-call hook that is its only streaming path, and + the route needs a translation. A pipeline that + cannot be run that way yet is left out and its guardrails run on the stream + on their own, the way they did before pipelines ran on streams at all, with + a warning naming the pipeline. + """ + post_call_pipelines: Final = _post_call_pipelines(request_data) + if not post_call_pipelines: + return () + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if translation is None: + verbose_proxy_logger.warning( + "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet " + "(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run " + "on their own: %s", + user_api_key_dict.request_route, + ", ".join(policy_name for policy_name, _pipeline in post_call_pipelines), + ) + return () + return tuple( + (policy_name, pipeline) + for policy_name, pipeline in post_call_pipelines + if _pipeline_is_streamable(policy_name, pipeline, translation) ) @@ -575,9 +877,10 @@ _EMPTY_LIFT: Final = MappingProxyType({}) def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]: """Failure-path callbacks run after ``litellm_logging_obj`` is popped from request_data (it is not serialisable), so the caller merges these fields - onto request_data first: the first-handoff instant for preprocessing - latency, recovered or estimated usage for token counts, and the standard - logging object for deployment attribution on failed-request spend logs.""" + onto request_data first: the request start and first-handoff instants for + duration and preprocessing latency, the call type, recovered or estimated + usage for token counts, and the standard logging object for deployment + attribution on failed-request spend logs.""" _logging_obj: Final = request_data.get("litellm_logging_obj") if _logging_obj is None: return _EMPTY_LIFT @@ -589,7 +892,9 @@ def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, dispatched=_first_handoff is not None, ) _entries: Final = ( + ("start_time", _model_call_details.get("start_time")), ("first_api_call_start_time", _first_handoff), + ("call_type", _model_call_details.get("call_type")), ("combined_usage_object", None if _usage_to_lift is None else _usage_to_lift[0]), ("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)), ("standard_logging_object", _model_call_details.get("standard_logging_object")), @@ -857,6 +1162,14 @@ class ProxyLogging: litellm.logging_callback_manager.add_litellm_async_success_callback(callback) litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + # Runs after load_config applied every litellm_settings key: logger __init__s read e.g. s3_callback_params + success_callbacks: Final = tuple(cb for cb in litellm.success_callback if isinstance(cb, str)) + failure_callbacks: Final = tuple(cb for cb in litellm.failure_callback if isinstance(cb, str)) + for callback in success_callbacks: + _add_custom_logger_callback_to_specific_event(callback, "success") + for callback in failure_callbacks: + _add_custom_logger_callback_to_specific_event(callback, "failure") + async def update_request_status(self, litellm_call_id: str, status: Literal["success", "fail"]): # only use this if slack alerting is being used if self.alerting is None: @@ -924,7 +1237,13 @@ class ProxyLogging: "incoming_bearer_token": kwargs.get("incoming_bearer_token"), "metadata": {"headers": kwargs.get("headers") or {}}, } - + user_api_key_auth: Final = kwargs.get("user_api_key_auth") + if isinstance(user_api_key_auth, UserAPIKeyAuth): + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_auth, + data=synthetic_data, + metadata_variable_name="metadata", + ) return synthetic_data def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None: @@ -1579,7 +1898,8 @@ class ProxyLogging: call_type: str, event_hook: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data - ) -> dict: + response: LLMResponseTypes | None = None, + ) -> tuple[dict, LLMResponseTypes | None]: # mutable-ok: returns the request-payload dict onward """ Execute guardrail pipelines if any are configured for this request. @@ -1591,20 +1911,27 @@ class ProxyLogging: ``scan_raw_request`` evaluates the pristine request, not whatever an earlier ``pass_data`` step in the same pipeline already rewrote. - Returns the (possibly modified) data dict. + Returns the (possibly modified) data dict, plus the replacement + response when a post_call pipeline step returned one (None when the + response is unchanged), matching the flat callback-loop contract. """ pipelines: Final = _policy_pipelines(data) if not pipelines: - return data + return data, None + current_response = response # rebind-ok: chains each pipeline's replacement response into the next for policy_name, pipeline in pipelines: if pipeline.mode != event_hook: continue + step_input: dict = ( + {**data, "response": current_response} if current_response is not None else data + ) # mutable-ok: same request-payload shape as data + result: PipelineExecutionResult = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode=pipeline.mode, - data=data, + data=step_input, user_api_key_dict=user_api_key_dict, call_type=call_type, policy_name=policy_name, @@ -1615,26 +1942,46 @@ class ProxyLogging: result=result, data=data, policy_name=policy_name, + original_response=current_response, ) - return data + if current_response is not None and result.modified_data is not None: + current_response = result.modified_data.get("response", current_response) + + return data, current_response if current_response is not response else None @staticmethod def _handle_pipeline_result( result: PipelineExecutionResult, data: dict, policy_name: str, + original_response: "LLMResponseTypes | Sequence[object] | None" = None, ) -> dict: """ Handle a PipelineExecutionResult — allow, block, or modify_response. Returns data dict if allowed, raises on block/modify_response. + ``original_response`` is set on the post_call path, where the request + payload (already sent upstream) must stay untouched; a replacement + response carried in ``modified_data`` is adopted by the caller, and + metadata-bucket writes (applied guardrails, guardrail logging info) + are merged back so headers and spend logs still see them, on block + and modify_response too, so failure spend records keep guardrail + cost and status. On the + streaming path it is the buffered chunk list, carried into + ``ModifyResponseException.original_response`` for usage reporting. """ if result.terminal_action == "allow": if result.modified_data is not None: - data.update(result.modified_data) + if original_response is None: + data.update(result.modified_data) + else: + _merge_pipeline_metadata_writes(data, result.modified_data) return data + if result.modified_data is not None: + _merge_pipeline_metadata_writes(data, result.modified_data) + if result.terminal_action == "block": original_exception: Final = result.original_exception if original_exception is not None and not _exception_changes_request_flow(original_exception): @@ -1672,6 +2019,7 @@ class ProxyLogging: request_data=data, guardrail_name=f"pipeline:{policy_name}", detection_info=None, + original_response=original_response, ) return data @@ -1789,7 +2137,7 @@ class ProxyLogging: try: # Execute guardrail pipelines before the normal callback loop - data = await self._maybe_execute_pipelines( + data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below data=data, user_api_key_dict=user_api_key_dict, call_type=call_type, @@ -1798,7 +2146,7 @@ class ProxyLogging: ) # Get pipeline-managed guardrails to skip in normal loop - pipeline_managed: Final = _pipeline_managed_guardrail_names(data) + pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call") caps: Final = ProxyLogging._callback_capabilities() # Skip the per-request callback walk entirely when nothing in @@ -2563,7 +2911,7 @@ class ProxyLogging: original_exception=original_exception, ) - request_data.update(_failure_fields_to_lift(request_data)) + request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data)) # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) @@ -2757,6 +3105,24 @@ class ProxyLogging: daemon=True, ).start() + async def _run_post_call_pipelines( + self, + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data + user_api_key_dict: UserAPIKeyAuth, + response: LLMResponseTypes, + ) -> LLMResponseTypes | None: + if _is_pending_background_response(response): + _defer_post_call_pipelines(data, response) + return None + _, pipeline_response = await self._maybe_execute_pipelines( + data=data, + user_api_key_dict=user_api_key_dict, + call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", + event_hook="post_call", + response=response, + ) + return pipeline_response + async def post_call_success_hook( self, data: dict, @@ -2776,36 +3142,33 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks - guardrail_callbacks: Final[list[CustomGuardrail]] = [] - other_callbacks: Final[list[CustomLogger]] = [] + pipeline_response: Final = await self._run_post_call_pipelines( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + ) + if pipeline_response is not None: + response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below + + pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call") + guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() try: - for callback in litellm.callbacks: - _callback: CustomLogger | None = None - if isinstance(callback, str): - _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( - cast(_custom_logger_compatible_callbacks_literal, callback) - ) - else: - _callback = callback - - if _callback is not None: - if isinstance(_callback, CustomGuardrail): - guardrail_callbacks.append(_callback) - else: - other_callbacks.append(_callback) - ############## Handle Guardrails ######################################## - ############################################################################# - # Merge model-level guardrails before checking which guardrails to run guardrail_data: Final = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple( - callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False) + callback + for callback in guardrail_callbacks + if getattr(callback, "run_in_parallel", False) + and not (callback.guardrail_name and callback.guardrail_name in pipeline_managed) ) for callback in guardrail_callbacks: # Main - V2 Guardrails implementation + if callback.guardrail_name and callback.guardrail_name in pipeline_managed: + continue + if getattr(callback, "run_in_parallel", False): continue @@ -3102,11 +3465,16 @@ class ProxyLogging: # dict lookups + llm_router.get_deployment() per callback per chunk. _cached_guardrail_data: dict | None = None _guardrail_data_computed = False + pipeline_gated: Final = ( + stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset() + ) for callback in litellm.callbacks: try: _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): + if callback.guardrail_name in pipeline_gated: + continue # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks @@ -3163,12 +3531,13 @@ class ProxyLogging: 1. /chat/completions """ caps: Final = ProxyLogging._callback_capabilities() + post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict) # Fast path: no real overrides. Internal proxy CustomLogger callbacks # (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default # ``async for chunk: yield chunk`` body, so wrapping the iterator # through each of them adds N pass-through trampolines per chunk for # zero behavior change. Skip the chain entirely and stream through. - if not caps.iterator_overrides: + if not caps.iterator_overrides and not post_call_pipelines: try: async for chunk in response: yield chunk @@ -3188,8 +3557,11 @@ class ProxyLogging: current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) + pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines) for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): + if resolved_callback.guardrail_name in pipeline_gated_names: + continue if ( resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True @@ -3229,6 +3601,18 @@ class ProxyLogging: ), ) + pipeline_translation: Final = ( + resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None + ) + if pipeline_translation is not None: + current_response = self._pipeline_gated_stream( + response=current_response, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + pipelines=post_call_pipelines, + translation=pipeline_translation, + ) + try: async for chunk in current_response: yield chunk @@ -3244,6 +3628,71 @@ class ProxyLogging: # we reach this point the metadata is fully populated. ProxyLogging._fire_deferred_stream_logging(request_data) + async def _pipeline_gated_stream( + self, + response: "AsyncGenerator[object, None]", + user_api_key_dict: UserAPIKeyAuth, + request_data: dict, # mutable-ok: same request-payload shape the hooks mutate + pipelines: "tuple[tuple[str, GuardrailPipeline], ...]", + translation: "tuple[str, BaseTranslation]", + ) -> "AsyncGenerator[Any, None]": + """ + Execute post_call policy pipelines against a streamed response. + + Buffers the whole stream (nothing reaches the client until every + pipeline allows it), then runs each pipeline's steps against the + assembled output through the endpoint guardrail translation, the same + machinery flat post_call guardrails use at end of stream. An allow + releases the buffered chunks: verbatim when no guardrail rewrote the + output, rewritten in place when one rewrote text or a tool call and the + translation delivers ended-stream rewrites (later steps then re-scan the + rewritten chunks, so rewrites chain). A rewrite the translation cannot + deliver yet (one on a route without write-back, or a shape the route + refuses) is discarded by the executor and the original chunks are + released; a block or modify_response terminates with the translation's + block chunks or the raised error. + """ + buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict + async for item in response: + buffered.append(item) + if not buffered: + return + + call_type, endpoint_translation = translation + + for policy_name, pipeline in pipelines: + result: PipelineExecutionResult = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode="post_call", + data=request_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + policy_name=policy_name, + streaming_chunks=buffered, + endpoint_translation=endpoint_translation, + ) + try: + ProxyLogging._handle_pipeline_result( + result, data=request_data, policy_name=policy_name, original_response=buffered + ) + except ModifyResponseException as e: + if e.original_response is None: + e.original_response = buffered + async for block_chunk in unified_guardrail.handle_streaming_block( + e, endpoint_translation, stream_started=False, responses_so_far=() + ): + yield block_chunk + return + except HTTPException as e: + async for error_chunk in unified_guardrail.emit_streaming_http_error( + e, call_type, buffered, request_data + ): + yield error_chunk + return + + for buffered_item in buffered: + yield buffered_item + @staticmethod def _fire_deferred_stream_logging(request_data: dict) -> None: """ @@ -3463,6 +3912,11 @@ async def prefetch_config_params(prisma_client: "PrismaClient | None", param_nam ) +_WRITER_WRITABILITY_PROBE_SQL: Final = "SELECT current_setting('transaction_read_only') AS transaction_read_only" +_WRITER_WRITABILITY_PROBE_ROWS: Final = TypeAdapter(list[dict[str, object]]) +_READ_ONLY_RECREATE_BACKOFF_CAP_SECONDS: Final = 600 + + class _ForcedRecreateDeclined(Exception): """A forced recreate was declined by the engine-generation guard. @@ -3522,7 +3976,7 @@ class _StaleReadEngine: class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() - spend_log_flush_requested: ClassVar[asyncio.Event] = asyncio.Event() + spend_log_flush_requested: "asyncio.Event | None" = None spend_log_queue_bytes: ClassVar[int] = 0 spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None tool_usage_transactions: list["ToolUsageTransaction"] = [] @@ -3557,6 +4011,7 @@ class PrismaClient: verbose_proxy_logger.error("Please run 'prisma generate' to generate the Prisma client.") raise Exception("Unable to find Prisma binaries. Please run 'prisma generate' first.") token_auth: Final = self.token_auth + writer_token_auth: Final = None if database_url_is_pooled() else token_auth # When read-replica routing is on, tag log lines with [writer]/[reader] # so the two wrappers' interleaved token refresh logs can be told apart. # Single-DB deployments get an empty prefix (logs unchanged). @@ -3565,13 +4020,13 @@ class PrismaClient: if http_client is not None: writer_wrapper = PrismaWrapper( original_prisma=Prisma(http=http_client), - token_auth=token_auth, + token_auth=writer_token_auth, log_prefix=writer_log_prefix, ) else: writer_wrapper = PrismaWrapper( original_prisma=Prisma(), - token_auth=token_auth, + token_auth=writer_token_auth, log_prefix=writer_log_prefix, ) @@ -3640,6 +4095,8 @@ class PrismaClient: self._db_health_watchdog_task: asyncio.Task | None = None self._db_last_reconnect_attempt_ts: float = 0.0 self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15"))) + self._db_read_only_recreate_ts: float = 0.0 + self._db_read_only_recreate_streak: int = 0 self._db_health_watchdog_interval_seconds: int = max( 5, int(os.getenv("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", "30")) ) @@ -5357,15 +5814,20 @@ class PrismaClient: writer: Final = self.writer_db if force_recreate is False: try: - await writer.query_raw("SELECT 1") - verbose_proxy_logger.info( - "Writer healthy on probe; skipping recreate (engine " - "likely already replaced by a token refresh)." - ) - if isinstance(self.db, RoutingPrismaWrapper): - self.db.mark_writer_recovered() - await self._start_engine_watcher() - return + if await self._writer_is_read_only(writer): + verbose_proxy_logger.warning( + "Writer answers the probe but its session is read-only " + "(writes fail with SQLSTATE 25006); recreating Prisma client." + ) + else: + verbose_proxy_logger.info( + "Writer healthy on probe; skipping recreate (engine " + "likely already replaced by a token refresh)." + ) + if isinstance(self.db, RoutingPrismaWrapper): + self.db.mark_writer_recovered() + await self._start_engine_watcher() + return except Exception as probe_err: verbose_proxy_logger.warning( "Writer probe failed (%s); recreating Prisma client.", @@ -5685,6 +6147,18 @@ class PrismaClient: reason="db_health_watchdog_writer_unavailable", timeout_seconds=self._db_watchdog_reconnect_timeout_seconds, ) + continue + if await asyncio.wait_for( + self._writer_is_read_only(self.writer_db), + timeout=self._db_health_watchdog_probe_timeout_seconds, + ): + await self.recreate_read_only_writer( + reason="db_health_watchdog_writer_read_only", + timeout_seconds=self._db_watchdog_reconnect_timeout_seconds, + ) + continue + self._db_read_only_recreate_streak = 0 + self._db_read_only_recreate_ts = 0.0 except asyncio.CancelledError: break except Exception as e: @@ -5696,6 +6170,39 @@ class PrismaClient: else: verbose_proxy_logger.debug("Prisma DB health watchdog observed non-DB error: %s", e) + async def recreate_read_only_writer(self, reason: str, timeout_seconds: float | None = None) -> bool: + """Force-recreate the client behind a writer session that rejects writes + (SQLSTATE 25006). Each recreate doubles the wait before the next one + until the watchdog sees a writable session again, so a database that is + read-only as a whole (replica, failover in progress) does not get its + engine killed on every watchdog cycle or failed write.""" + backoff_seconds: Final = min( + self._db_reconnect_cooldown_seconds * 2 ** min(self._db_read_only_recreate_streak, 10), + _READ_ONLY_RECREATE_BACKOFF_CAP_SECONDS, + ) + if time.time() - self._db_read_only_recreate_ts < backoff_seconds: + verbose_proxy_logger.debug( + "Writer session still read-only after %s recreate(s); backing off %ss. reason=%s", + self._db_read_only_recreate_streak, + backoff_seconds, + reason, + ) + return False + verbose_proxy_logger.warning( + "Writer session is read-only (writes fail with SQLSTATE 25006); recreating Prisma client. reason=%s", + reason, + ) + self._db_read_only_recreate_ts = time.time() + self._db_read_only_recreate_streak += 1 + return await self.attempt_db_reconnect(reason=reason, timeout_seconds=timeout_seconds, force_recreate=True) + + async def _writer_is_read_only(self, writer: PrismaWrapper) -> bool: + """True iff the pooled writer session answers reads but rejects writes (SQLSTATE 25006).""" + rows: Final = _WRITER_WRITABILITY_PROBE_ROWS.validate_python( + await writer.query_raw(_WRITER_WRITABILITY_PROBE_SQL) + ) + return any(row.get("transaction_read_only") == "on" for row in rows) + def _probe_target_wrapper(self) -> PrismaWrapper: """The Prisma wrapper a `SELECT 1` health probe actually reaches. @@ -5965,48 +6472,13 @@ class PrismaClient: verbose_proxy_logger.error("Error getting health check history: %s", e) return [] - async def get_all_latest_health_checks(self) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]": - """ - Get the latest health check for each model. + async def get_all_latest_health_checks(self) -> tuple[LatestHealthCheckRow, ...]: + """Latest health check per (model_id, model_name), deduplicated in Postgres.""" + return await fetch_latest_health_checks(self) - Uses DB-level DISTINCT ON (model_id, model_name) with ORDER BY checked_at DESC - (via Prisma ``distinct`` + ``order``) so we never load the full history into memory. - """ - try: - return await HealthCheckRepository(self).table.find_many( - distinct=["model_id", "model_name"], - order=[ - {"model_id": "asc"}, - {"model_name": "asc"}, - {"checked_at": "desc"}, - ], - ) - except Exception as e: - verbose_proxy_logger.error("Error getting all latest health checks: %s", e) - return [] - - async def get_latest_health_checks_for_models( - self, model_names: "Sequence[str]" - ) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]": - """ - Get the latest health check for each of the named models. - - Same DISTINCT ON as ``get_all_latest_health_checks``, bounded to the models asked - about, so a paged caller reads health for its page instead of for the whole table. - """ - if not model_names: - return () - latest_first: Final = (("model_id", "asc"), ("model_name", "asc"), ("checked_at", "desc")) - order: Final = [{field: direction} for field, direction in latest_first] # mutable-ok: prisma order is a list - try: - return await HealthCheckRepository(self).table.find_many( - where={"model_name": {"in": list(model_names)}}, # mutable-ok: prisma filters are dicts and lists - distinct=["model_id", "model_name"], # mutable-ok: prisma distinct takes a list - order=order, - ) - except Exception as e: # noqa: BLE001 # health decorates a list; a driver error must not fail the page - verbose_proxy_logger.error("Error getting latest health checks for models: %s", e) - return () + async def get_latest_health_checks_for_models(self, model_names: Sequence[str]) -> tuple[LatestHealthCheckRow, ...]: + """Same as ``get_all_latest_health_checks``, bounded to the named models.""" + return await fetch_latest_health_checks_for_models(self, model_names) ### HELPER FUNCTIONS ### @@ -6248,23 +6720,27 @@ async def enqueue_spend_logs( ) -def request_spend_log_flush() -> None: - """Wake the queue monitor now rather than leaving the rows for its next poll. +def request_spend_log_flush(prisma_client: PrismaClient) -> None: + """Wake this client's queue monitor now rather than leaving the rows for its next poll. The Responses API hands the client an id it can chain from straight away, and that lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. Repeated requests coalesce into the monitor's next pass, so the batching holds. + A request made before the monitor is running is dropped, and loses nothing: the + monitor reads the queue on its first pass, before it ever waits on a request. """ - PrismaClient.spend_log_flush_requested.set() + flush_requested: Final = prisma_client.spend_log_flush_requested + if flush_requested is not None: + flush_requested.set() -async def _wait_for_spend_log_flush_request(interval: float) -> bool: +async def _wait_for_spend_log_flush_request(flush_requested: asyncio.Event, interval: float) -> bool: """Wait out ``interval``, returning early and True when a flush was requested.""" try: - await asyncio.wait_for(PrismaClient.spend_log_flush_requested.wait(), timeout=interval) + await asyncio.wait_for(flush_requested.wait(), timeout=interval) except asyncio.TimeoutError: return False - PrismaClient.spend_log_flush_requested.clear() + flush_requested.clear() return True @@ -6691,6 +7167,8 @@ async def _monitor_spend_logs_queue( max_backoff: Final = 30.0 # Maximum backoff interval in seconds backoff_multiplier: Final = 1.5 # Exponential backoff multiplier current_interval = base_interval + flush_requested: Final = asyncio.Event() + prisma_client.spend_log_flush_requested = flush_requested # rebind-ok: the client owns its monitor's flush signal verbose_proxy_logger.info( "Starting spend logs queue monitor (threshold: %s, poll_interval: %ss)", threshold, base_interval @@ -6729,7 +7207,7 @@ async def _monitor_spend_logs_queue( # Exponential backoff when no logs to process current_interval = min(current_interval * backoff_multiplier, max_backoff) - if await _wait_for_spend_log_flush_request(current_interval): + if await _wait_for_spend_log_flush_request(flush_requested, current_interval): current_interval = base_interval except Exception as e: spend_log_error("Error in spend logs queue monitor: %s", str(e), exc=e) @@ -7102,6 +7580,12 @@ def _get_openapi_url() -> str | None: return "/openapi.json" +def _recreate_writer_on_read_only_transaction(prisma_client: "PrismaClient | None") -> None: + if prisma_client is None: + return + asyncio.create_task(prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction")) + + def handle_exception_on_proxy(e: Exception) -> ProxyException: """ Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible @@ -7109,12 +7593,16 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: from fastapi import status verbose_proxy_logger.exception("Exception: %s", e) + if PrismaDBExceptionHandler.is_read_only_transaction_error(e): + from litellm.proxy.proxy_server import prisma_client + + _recreate_writer_on_read_only_transaction(prisma_client) if isinstance(e, HTTPException): return ProxyException( message=getattr(e, "detail", f"error({e})"), type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), + param=openai_error_param(e), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), ) elif isinstance(e, ProxyException): @@ -7123,7 +7611,7 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: return ProxyException( message=str(e), type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), + param=openai_error_param(e), code=_status_code, ) @@ -7206,7 +7694,19 @@ def get_custom_url(request_base_url: str, route: str | None = None) -> str: else: base_url = request_base_url - server_root_path: Final = get_server_root_path() + # get_request_root_path() returns the prefix the router is actually + # resolving this request under: the matched SERVER_ROOT_PATHS entry when + # PerRequestRootPathMiddleware ran, otherwise the SERVER_ROOT_PATH scalar. + # This keeps the emitted URL under one prefix — the one the client called — + # instead of stacking the scalar onto a request already living under a + # dynamic prefix (which would produce /tenant-a/legacy/... — a path that + # doesn't exist). join_paths()'s tail-dedup then collapses the append when + # base_url (i.e. request.base_url) already ends in the same prefix. + from litellm.proxy.middleware.per_request_root_path_middleware import ( # noqa: PLC0415 # lazy: middleware imports utils + get_request_root_path, + ) + + server_root_path: Final = get_request_root_path() if route is not None: if server_root_path != "": # First join base_url with server_root_path, then with route @@ -7498,6 +7998,88 @@ async def get_available_models_for_user( return all_models +def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None: + try: + return get_model_info(model) + except Exception as e: + verbose_proxy_logger.debug( + "create_model_info_response: cost map lookup failed for %s: %s", + model, + e, + ) + return None + + +def _resolve_listing_model_info( + deployment_model: str | None, + listed_model: str, + listed_info: ModelInfo | None, + get_model_info: Callable[[str], ModelInfo], +) -> tuple[ModelInfo, ...]: + """ + Cost-map entries describing one deployment behind a listed model, best source first. + + The name a model is listed under is an arbitrary public alias, so it often misses the + cost map and lands on a fallback-generalization rule that answers with a conservative + family baseline instead of the real model's limits; the deployment's underlying model + is what the request actually reaches. Both names are kept because either can + generalize, and because a deployment's own model is registered into the cost map as a + stub that carries no limits of its own. Exact entries are consulted before generalized + ones, and each field is then taken from the first entry that has it. + + ``listed_info`` is resolved once by the caller, since a group with several distinct + underlying models resolves the same alias for each of them. + """ + # Fast path, and the only one a wildcard-expanded name takes: with a single name + # there is nothing to order, so skip the generalization test entirely. This keeps + # the per-model cost of the listing on the hot path #33721 exists to protect. + if deployment_model is None or deployment_model == listed_model: + return () if listed_info is None else (listed_info,) + + deployment_info: Final = _safe_get_model_info(deployment_model, get_model_info) + if deployment_info is None: + return () if listed_info is None else (listed_info,) + if listed_info is None: + return (deployment_info,) + + from litellm.utils import is_generalized_model_info + + # Both names resolved: the deployment's model leads unless it only generalized + # while the listed name is an exact cost-map entry. + if is_generalized_model_info(deployment_info) and not is_generalized_model_info(listed_info): + return (listed_info, deployment_info) + return (deployment_info, listed_info) + + +def _first_token_limit(candidates: tuple[ModelInfo, ...], field: str) -> int | None: + return next( + (limit for limit in (coerce_token_limit(info.get(field)) for info in candidates) if limit is not None), + None, + ) + + +def _group_token_limit(candidate_sets: tuple[tuple[ModelInfo, ...], ...], field: str) -> int | None: + """The widest limit any deployment behind the listed name declares for ``field``. + + A model group is normally one model behind several interchangeable deployments, so + there is a single value to report and the choice of aggregate does not arise. + + When a group genuinely mixes models no single number is right, and the widest is the + deliberate pick over the narrowest for two reasons. It is what ``/model_group/info`` + has long reported to the Admin UI, so the two surfaces agree; disagreeing is the very + complaint this resolution path exists to fix. And of the two ways to be wrong, + under-advertising is worse: a client that trusts a narrowed window silently refuses + prompts the group would have served, while an over-long prompt that reaches a smaller + deployment comes back as a legible context-length error -- and does not reach one at + all when ``enable_pre_call_checks`` is set, which filters deployments the prompt does + not fit. + """ + limits: Final = tuple( + limit for limit in (_first_token_limit(candidates, field) for candidates in candidate_sets) if limit is not None + ) + return max(limits) if limits else None + + def create_model_info_response( model_id: str, provider: str, @@ -7522,31 +8104,48 @@ def create_model_info_response( "owned_by": provider, } - try: - model_cost_info: ModelInfo | None = get_model_info(model_id) - except Exception as e: - verbose_proxy_logger.debug( - "create_model_info_response: cost map lookup failed for %s: %s", - model_id, - e, - ) - model_cost_info = None + listing_info: Final = llm_router.get_model_listing_info(model_id) if llm_router is not None else None - max_input_tokens: int | None = None - max_output_tokens: int | None = None - if model_cost_info is not None: - max_input_tokens = coerce_token_limit(model_cost_info.get("max_input_tokens")) - max_output_tokens = coerce_token_limit(model_cost_info.get("max_output_tokens")) - mode: Final = model_cost_info.get("mode") - if isinstance(mode, str): - base["mode"] = mode + # One entry per distinct model behind the listed name; (None,) when the router knows + # nothing about it, so the listed name is resolved on its own as before. + deployment_models: Final[tuple[str | None, ...]] = ( + listing_info.cost_map_keys if listing_info is not None and listing_info.cost_map_keys else (None,) + ) + listed_info: Final = _safe_get_model_info(model_id, get_model_info) + candidate_sets: Final = tuple( + _resolve_listing_model_info( + deployment_model=deployment_model, + listed_model=model_id, + listed_info=listed_info, + get_model_info=get_model_info, + ) + for deployment_model in deployment_models + ) + + max_input_tokens: int | None = _group_token_limit(candidate_sets, "max_input_tokens") + max_output_tokens: int | None = _group_token_limit(candidate_sets, "max_output_tokens") + mode: Final = next( + ( + m + for m in ( + cast("Mapping[str, object]", info).get("mode") # cast-ok: an entry need not carry "mode" + for candidates in candidate_sets + for info in candidates + ) + if isinstance(m, str) + ), + None, + ) + if mode is not None: + base["mode"] = mode + + if listing_info is not None: + if listing_info.max_input_tokens is not None: + max_input_tokens = listing_info.max_input_tokens + if listing_info.max_output_tokens is not None: + max_output_tokens = listing_info.max_output_tokens if llm_router is not None: - configured_input, configured_output = llm_router.get_configured_token_limits(model_id) - if configured_input is not None: - max_input_tokens = configured_input - if configured_output is not None: - max_output_tokens = configured_output configured_mode: Final = llm_router.get_configured_mode(model_id) if isinstance(configured_mode, str): base["mode"] = configured_mode diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index f2070e6604c..8363aaee99a 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -161,6 +161,8 @@ async def can_user_access_vector_store( this vector store id. 5. The caller's team_id matches the vector store's team_id. + A dashboard session credential is evaluated against the same effective + contexts as listing (its own grants plus each real team of the user). Otherwise access is denied. """ if _is_proxy_admin(user_api_key_dict): @@ -169,7 +171,8 @@ async def can_user_access_vector_store( if vector_store.get("team_id") is None: return True - return await _is_vector_store_granted(vector_store, user_api_key_dict) + auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict) + return await _is_vector_store_granted_to_any(vector_store, auth_contexts) async def _is_vector_store_granted( @@ -219,7 +222,7 @@ async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> ) -async def _vector_store_listing_auth_contexts( +async def _vector_store_auth_contexts( user_api_key_dict: UserAPIKeyAuth, ) -> tuple[UserAPIKeyAuth, ...]: if not is_ui_session_credential(user_api_key_dict): @@ -250,7 +253,7 @@ async def filter_listable_vector_stores( if _is_proxy_admin(user_api_key_dict): return tuple(vector_stores) - auth_contexts: Final = await _vector_store_listing_auth_contexts(user_api_key_dict) + auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict) return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)]) diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py index 881f7a66cea..dcf9ddfc32a 100644 --- a/litellm/repositories/__init__.py +++ b/litellm/repositories/__init__.py @@ -2,6 +2,7 @@ Repository classes for database operations. """ +from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.credentials_repository import CredentialsRepository @@ -92,6 +93,7 @@ __all__ = [ "AdaptiveRouterStateRepository", "AgentsRepository", "AuditLogRepository", + "AutoRouterSessionRepository", "BatchTable", "BudgetCascadeUnitOfWork", "BudgetRepository", diff --git a/litellm/repositories/autorouter_session_repository.py b/litellm/repositories/autorouter_session_repository.py new file mode 100644 index 00000000000..d05ef9421ca --- /dev/null +++ b/litellm/repositories/autorouter_session_repository.py @@ -0,0 +1,34 @@ +""" +Repository for the auto-router per-session rollup (LiteLLM_AutoRouterSession). +""" + +from typing import TYPE_CHECKING, Final + +from litellm.models.autorouter_session import LiteLLM_AutoRouterSession +from litellm.repositories.base_repository import BaseRepository +from litellm.repositories.prisma_protocols import TableActions + +if TYPE_CHECKING: + from prisma import models as prisma_models + + +class AutoRouterSessionRepository(BaseRepository[LiteLLM_AutoRouterSession]): + @property + def table(self) -> TableActions["prisma_models.LiteLLM_AutoRouterSession"]: + return self.prisma_client.db.litellm_autoroutersession + + @property + def model_class(self) -> type[LiteLLM_AutoRouterSession]: + return LiteLLM_AutoRouterSession + + async def find_latest_for_key(self, api_key: str, session_id: str) -> LiteLLM_AutoRouterSession | None: + """The session's most recently active router row under exactly this key hash, or None. + + The key is the row's own partition, not a filter over a wider read: the spend writer keyed the + row under the caller's api_key, so a key can only ever see what it wrote itself. + """ + record: Final = await self.table.find_first( + where={"api_key": api_key, "session_id": session_id}, # mutable-ok: Prisma where filter must be a dict + order={"last_turn_at": "desc"}, # mutable-ok: Prisma order clause must be a dict + ) + return self._to_model(record) diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index acc7c8dcda8..d24eb8ffc62 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -3,7 +3,7 @@ Model repository for database operations on LiteLLM_ProxyModelTable. """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.models.model import LiteLLM_ProxyModelTable @@ -105,6 +105,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={"model_id": {"not": model_id}} # mutable-ok: Prisma requires plain dicts for query serialization + ) + return tuple(self._to_model_list(records)) + async def find_by_team_id(self, team_id: str) -> list[LiteLLM_ProxyModelTable]: """Find models associated with a specific team. diff --git a/litellm/repositories/object_permission_repository.py b/litellm/repositories/object_permission_repository.py index 6b1f9c68e47..7736939c696 100644 --- a/litellm/repositories/object_permission_repository.py +++ b/litellm/repositories/object_permission_repository.py @@ -40,6 +40,7 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): blocked_tools: list[str] | None = None, mcp_toolsets: list[str] | None = None, search_tools: list[str] | None = None, + skills: list[str] | None = None, ) -> LiteLLM_ObjectPermissionTable: """Create a new object permission record.""" data: Final[dict[str, Any]] = {} @@ -63,6 +64,8 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): data["mcp_toolsets"] = mcp_toolsets if search_tools is not None: data["search_tools"] = search_tools + if skills is not None: + data["skills"] = skills return await self.create(data) @@ -79,6 +82,7 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): blocked_tools: list[str] | None = None, mcp_toolsets: list[str] | None = None, search_tools: list[str] | None = None, + skills: list[str] | None = None, ) -> LiteLLM_ObjectPermissionTable | None: """Update an object permission record.""" data: Final[dict[str, Any]] = {} @@ -102,6 +106,8 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): data["mcp_toolsets"] = mcp_toolsets if search_tools is not None: data["search_tools"] = search_tools + if skills is not None: + data["skills"] = skills return await self.update(object_permission_id, data, id_field="object_permission_id") diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index d7f8cd8f8bd..660dd8f0c92 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -437,14 +437,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): response_created_event_data["temperature"] = self.responses_api_request["temperature"] if "text" in self.responses_api_request: response_created_event_data["text"] = self.responses_api_request["text"] - if "tool_choice" in self.responses_api_request: - # Transform tool_choice from dict format (e.g., {"type": "auto"}) to string format - response_created_event_data["tool_choice"] = ( - LiteLLMCompletionResponsesConfig._transform_tool_choice(self.responses_api_request["tool_choice"]) - or "auto" + response_created_event_data["tool_choice"] = ( + LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + self.responses_api_request.get("tool_choice") ) - else: - response_created_event_data["tool_choice"] = "auto" + ) if "tools" in self.responses_api_request: response_created_event_data["tools"] = self.responses_api_request["tools"] else: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b2d1a69e0d8..fca5b0d11cf 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -27,8 +27,10 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( ) from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam +from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam +from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam from openai.types.responses.tool_param import FunctionToolParam -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger @@ -68,6 +70,7 @@ from litellm.types.llms.openai import ( ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ResponsesAPIStatus, + ToolChoice, ValidChatCompletionMessageContentTypes, ValidChatCompletionMessageContentTypesLiteral, ) @@ -126,6 +129,7 @@ _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) +_RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice) @runtime_checkable @@ -267,6 +271,27 @@ class LiteLLMCompletionResponsesConfig: # Return as-is for unknown formats return tool_choice + @staticmethod + def _transform_tool_choice_for_responses_api_response(tool_choice: object) -> ToolChoice: + if tool_choice is None: + return "auto" + try: + return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(tool_choice) + except ValidationError: + return LiteLLMCompletionResponsesConfig._chat_tool_choice_as_responses_api_tool_choice(tool_choice) + + @staticmethod + def _chat_tool_choice_as_responses_api_tool_choice(tool_choice: object) -> ToolChoice: + match tool_choice, LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice): + case {"type": "custom"}, {"function": {"name": str(custom_name)}}: + return ToolChoiceCustomParam(type="custom", name=custom_name) + case _, {"type": "function", "function": {"name": str(function_name)}}: + return ToolChoiceFunctionParam(type="function", name=function_name) + case _, "none" | "auto" | "required" as normalized: + return normalized + case _, _: + return "auto" + @staticmethod def _should_drop_derived_web_search_options(model: str, custom_llm_provider: str | None) -> bool: """ @@ -2263,7 +2288,9 @@ class LiteLLMCompletionResponsesConfig: ), parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False), temperature=getattr(chat_completion_response, "temperature", 0), - tool_choice=getattr(chat_completion_response, "tool_choice", "auto"), + tool_choice=LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + responses_api_request.get("tool_choice") + ), tools=getattr(chat_completion_response, "tools", []), top_p=getattr(chat_completion_response, "top_p", None), max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 5e74b7324b4..a68cd02e61b 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -4,10 +4,11 @@ from collections.abc import Coroutine, Generator, Iterable, Mapping from contextlib import contextmanager from dataclasses import dataclass from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast import httpx from pydantic import BaseModel +from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger @@ -17,6 +18,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i from litellm.constants import request_timeout from litellm.integrations.anthropic_cache_control_hook import CARRY_UNMATCHED_MESSAGE_POINTS from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.common_utils import ( update_responses_input_with_model_file_ids, @@ -177,7 +179,7 @@ async def aresponses_api_with_mcp( ( mcp_tools_with_litellm_proxy, other_tools, - ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform) # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata) @@ -236,6 +238,7 @@ async def aresponses_api_with_mcp( "timeout": timeout, "custom_llm_provider": custom_llm_provider, **kwargs, + "_skip_mcp_handler": True, } # Handle MCP streaming if requested @@ -405,6 +408,37 @@ def _bridges_to_chat_completions( return responses_api_provider_config is None or use_chat_completions_api is True +_ResponsesCompatibilityFailure: TypeAlias = Literal["encrypted_task_unsupported"] + + +def _encrypted_task_support_failure( + responses_api_provider_config: BaseResponsesAPIConfig | None, use_chat_completions_api: bool +) -> _ResponsesCompatibilityFailure | None: + if ( + responses_api_provider_config is None + or _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api) + or not responses_api_provider_config.supports_encrypted_agent_messages() + ): + return "encrypted_task_unsupported" + return None + + +def _raise_responses_compatibility_failure( + failure: _ResponsesCompatibilityFailure, model: str, custom_llm_provider: str | None +) -> NoReturn: + match failure: + case "encrypted_task_unsupported": + raise litellm.exception_type( + model=model, + custom_llm_provider=custom_llm_provider, + original_exception=ValueError( + "Encrypted task classification requires a compatible native Responses deployment" + ), + ) + case _: + assert_never(failure) + + def _deployment_passes_through_responses(model_info: object) -> bool: """Whether ``model_info.supported_endpoints`` opts the deployment into native ``{api_base}/responses``.""" if not isinstance(model_info, dict): @@ -898,13 +932,14 @@ def _responses_try_dispatch_mcp_gateway( custom_llm_provider: str | None, kwargs: dict[str, object], _is_async: bool, + skip_mcp_handler: bool, ) -> Any | None: """Return a response when MCP gateway handles the call; otherwise None.""" from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) - if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): + if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): return None mcp_call_kwargs: Final = { "input": input, @@ -1074,6 +1109,8 @@ def responses( litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) _is_async: Final = kwargs.pop("aresponses", False) is True + skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) + require_encrypted_task_support: Final = kwargs.pop("_require_encrypted_task_support", False) is True use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs) client_headers: Final = kwargs.get("headers") @@ -1168,6 +1205,7 @@ def responses( custom_llm_provider=custom_llm_provider, kwargs=kwargs, _is_async=_is_async, + skip_mcp_handler=skip_mcp_handler, ) if _mcp_dispatch is not None: return _mcp_dispatch @@ -1181,6 +1219,17 @@ def responses( model, custom_llm_provider, deployment_model_info ) + if ( + require_encrypted_task_support + and ( + compatibility_failure := _encrypted_task_support_failure( + responses_api_provider_config, use_chat_completions_api + ) + ) + is not None + ): + _raise_responses_compatibility_failure(compatibility_failure, model, custom_llm_provider) + local_vars.update(kwargs) # Map reasoning_effort (from litellm_params/proxy config) to reasoning when not set if reasoning is None and "reasoning_effort" in local_vars: @@ -1253,7 +1302,7 @@ def responses( responses_api_provider_config=responses_api_provider_config, response_api_optional_params=response_api_optional_params, allowed_openai_params=allowed_openai_params, - drop_params=request_drop_params if isinstance(request_drop_params, bool) else None, + drop_params=normalize_drop_params(request_drop_params), ) litellm_logging_obj.update_from_kwargs( @@ -2081,7 +2130,7 @@ def compact_responses( responses_api_provider_config=responses_api_provider_config, response_api_optional_params=response_api_optional_params, allowed_openai_params=None, - drop_params=request_drop_params if isinstance(request_drop_params, bool) else None, + drop_params=normalize_drop_params(request_drop_params), ) # Pre Call logging diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index a75b3768636..ae18d5f6f1b 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -106,7 +106,7 @@ async def acompletion_with_mcp( ( mcp_tools_with_litellm_proxy, other_tools, - ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_tools_with_litellm_proxy: # No MCP tools, proceed with regular completion @@ -114,6 +114,7 @@ async def acompletion_with_mcp( model=model, messages=messages, tools=tools, + _skip_mcp_handler=True, **kwargs, ) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 15434bedbb7..a5021e2f777 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,6 +1,6 @@ import re import traceback -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload @@ -11,6 +11,7 @@ from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.utils import ( + iter_known_server_prefixes, logging_safe_mcp_headers, split_server_prefix_from_name, strip_known_server_prefix, @@ -23,6 +24,7 @@ from litellm.types.llms.openai import ( ResponsesAPIStreamingResponse, ) from litellm.types.llms.openai import ToolParam as ResponsesToolParam +from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.utils import ( CallTypes, ChatCompletionMessageCustomToolCall, @@ -45,6 +47,7 @@ else: # NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling ToolParam: TypeAlias = Mapping[str, object] +SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]] class MCPToolResult(TypedDict): @@ -56,14 +59,74 @@ class MCPToolResult(TypedDict): LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" -# Matches any URL whose path ends with /mcp/ — covers both root-path -# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments. -# A false-positive match (e.g. an external URL that happens to end with /mcp/) results -# in a "server not found" error from the internal gateway, not a silent failure or data leak, -# so this broad pattern is intentional and preferred over anchoring to localhost only. _PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") +def _mcp_server_url(tool: ToolParam) -> str | None: + if not isinstance(tool, dict) or tool.get("type") != "mcp": + return None + server_url: Final = tool.get("server_url") + return server_url if isinstance(server_url, str) else None + + +def _names_gateway_explicitly(tool: ToolParam) -> bool: + return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) + + +def _proxy_path_mcp_name(tool: ToolParam) -> str | None: + server_url: Final = _mcp_server_url(tool) + match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) + return None if match is None else match.group(1) + + +def _registered_mcp_servers() -> Collection[MCPServer]: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + return global_mcp_server_manager.get_registry().values() + + +def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool: + requested: Final = name.lower() + return any( + requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name)) + or name in (server.access_groups or ()) + for server in servers + ) + + +async def _toolset_exists(name: str) -> bool: + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return False + return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None + except Exception as e: + verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e) + return False + + +async def _gateway_served_names( + names: Collection[str], + servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers, + toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists, +) -> frozenset[str]: + registered: Final = tuple(servers()) if names else () + return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)]) + + +async def _served_mcp_path_names( + tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] +) -> frozenset[str]: + names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None) + return await served_names(names) if names else frozenset[str]() + + class LiteLLM_Proxy_MCP_Handler: """ Helper class with static methods for MCP integration with Responses API. @@ -87,57 +150,41 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: - """ - Returns True if any MCP tool should be handled via the litellm proxy MCP gateway. - This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/. - """ - if tools: - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "mcp": - server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): - return True - if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(server_url): - return True - return False + """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in + /mcp/. `_split_mcp_tools` then settles which of the latter the gateway actually serves.""" + return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ()) @staticmethod - def _parse_mcp_tools( + def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: + items: Final = tuple(tools or ()) + gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] + other_tools: Final[list[Any]] = [tool for tool in items if not _names_gateway_explicitly(tool)] + return gateway_tools, other_tools + + @staticmethod + async def _split_mcp_tools( tools: Iterable[Mapping[str, object]] | None, - ) -> tuple[list[ToolParam], list[Any]]: - """ - Parse tools and separate MCP tools with litellm_proxy from other tools. + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> SplitTools: + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) + return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( + [ + {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} + if (name := _proxy_path_mcp_name(tool)) in served + else tool + for tool in items + ] + ) - Returns: - Tuple of (mcp_tools_with_litellm_proxy, other_tools) - """ - mcp_tools_with_litellm_proxy: Final[list[ToolParam]] = [] - other_tools: Final[list[Any]] = [] - - if tools: - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "mcp": - server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): - mcp_tools_with_litellm_proxy.append(tool) - elif isinstance(server_url, str): - # Also intercept URLs like http://localhost:4000/mcp/atlassian_test - # by rewriting them to the internal litellm_proxy format. - m = _PROXY_MCP_PATH_RE.match(server_url) - if m: - rewritten = { - **tool, - "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}", - } - mcp_tools_with_litellm_proxy.append(rewritten) - else: - other_tools.append(tool) - else: - other_tools.append(tool) - else: - other_tools.append(tool) - - return mcp_tools_with_litellm_proxy, other_tools + @staticmethod + async def routes_through_gateway( + tools: Iterable[Mapping[str, object]] | None, + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> tuple[bool, ...]: + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) + return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items) @staticmethod async def _apply_toolset_permissions( diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9f9016c5a7f..40ff88fc557 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti import httpx from openai._streaming import SSEDecoder +from pydantic import BaseModel, ValidationError from typing_extensions import TypeIs import litellm @@ -438,18 +439,7 @@ class BaseResponsesAPIStreamingIterator: if self._persist_completed_response_before_logging: self._persist_completed_response_to_cache(is_async=is_async) - # Create a copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) - # to chat completion format (prompt_tokens/completion_tokens) for internal logging - # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with - # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) - logging_response = self.completed_response - if self.completed_response is not None and hasattr(self.completed_response, "model_dump"): - try: - logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump()) - except Exception: - # Fallback to original if serialization fails - pass + logging_response: Final[object] = _logging_copy(self.completed_response) self._restore_provider_response_headers(logging_response) end_time: Final = datetime.now() @@ -488,10 +478,10 @@ class BaseResponsesAPIStreamingIterator: def _restore_provider_response_headers(self, logging_response: object) -> None: """Re-apply the provider's response headers to the copy handed to logging callbacks. - ``model_validate(model_dump())`` above drops pydantic private attributes, so the + ``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the ``_hidden_params`` the provider transform set on the nested response are lost. Returns early - when that copy fell back to the original event, so logging-only state never lands on the - object the caller is iterating. + when the event was not a pydantic model and logging got the original, so logging-only state + never lands on the object the caller is iterating. """ if logging_response is self.completed_response: return @@ -544,7 +534,7 @@ class BaseResponsesAPIStreamingIterator: def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None)) if usage_obj is None: return try: @@ -1293,14 +1283,46 @@ def _add_text_like_part_events( ) +def _logging_copy(event: object) -> object: + """Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never + reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the + deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow + copies of the event and its nested response still keep the caller's ``usage`` attribute separate.""" + if not isinstance(event, BaseModel): + return event + try: + return type(event).model_validate(event.model_dump()) + except Exception: + return _detached_shallow_copy(event) + + +def _detached_shallow_copy(event: BaseModel) -> BaseModel: + nested: Final[object] = getattr(event, "response", None) + if isinstance(nested, BaseModel): + return event.model_copy(update={"response": nested.model_copy()}) + return event.model_copy() + + +def _usage_as_model(usage: object) -> ResponseAPIUsage | None: + if isinstance(usage, ResponseAPIUsage): + return usage + if not isinstance(usage, dict): + return None + try: + return ResponseAPIUsage.model_validate(usage) + except ValidationError: + return None + + def _stamp_responses_usage_cost( response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None ) -> None: if response_obj is None or logging_obj is None: return - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None)) if usage_obj is None: return + response_obj.usage = usage_obj # rebind-ok: the stamped cost has to ride on the response the client receives if isinstance(getattr(usage_obj, "cost", None), (int, float)): return try: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 3ca7b0503bf..599e978df6a 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -25,9 +25,15 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, SpecialEnums, Usage, + text_tokens_without_nested_reasoning, ) +def _output_token_detail(details: object, field: str) -> int | None: + value: Final = getattr(details, field, None) + return value if isinstance(value, int) else None + + def _is_object_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a list is a Sequence of anything return isinstance(value, list) @@ -537,6 +543,49 @@ class ResponsesAPIRequestUtils: return request_input + @staticmethod + def strip_encrypted_reasoning_from_input(request_input: object) -> None: + """Drop reasoning items the routed deployment cannot decrypt, keeping their readable summary. + + Mutates ``request_input`` in place: the router's fallback snapshot shares this + list object, so a rebound list would replay the stripped items on the fallback hop. + """ + if not isinstance(request_input, list): + return + items: Final = cast(list[object], request_input) # cast-ok: untyped client json + stripped: Final = tuple(ResponsesAPIRequestUtils._without_encrypted_reasoning(item) for item in items) + items[:] = (item for item in stripped if item is not None) # rebind-ok: list shared with fallback snapshot + + @staticmethod + def _without_encrypted_reasoning(item: object) -> object | None: + if not isinstance(item, dict): + return item + reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json + if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"): + return reasoning + readable: Final = any( + ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content") + ) + if not readable: + return None + kept: Final[dict[str, object]] = { # mutable-ok: request item rebuilt without the undecryptable keys + key: value for key, value in reasoning.items() if key not in ("encrypted_content", "id") + } + return kept + + @staticmethod + def _has_readable_text(value: object) -> bool: + """A reasoning item's ``summary``/``content`` carries readable text: a non-empty string, or a + list holding at least one block with a non-empty ``text`` field (summary_text / output_text).""" + if isinstance(value, str): + return bool(value.strip()) + if isinstance(value, list): + return any( + isinstance(block, dict) and bool(cast(Mapping[str, object], block).get("text")) # cast-ok: untyped json + for block in value + ) + return False + @staticmethod def _build_responses_api_response_id( custom_llm_provider: str | None, @@ -1137,11 +1186,22 @@ class ResponseAPILoggingUtils: response_api_usage, "output_tokens_details", None ) if output_tokens_details: + reasoning_tokens: Final = _output_token_detail(output_tokens_details, "reasoning_tokens") + image_tokens: Final = _output_token_detail(output_tokens_details, "image_tokens") + audio_tokens: Final = _output_token_detail(output_tokens_details, "audio_tokens") + reported_text_tokens: Final = _output_token_detail(output_tokens_details, "text_tokens") completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None), - image_tokens=getattr(output_tokens_details, "image_tokens", None), - text_tokens=getattr(output_tokens_details, "text_tokens", None), - audio_tokens=getattr(output_tokens_details, "audio_tokens", None), + reasoning_tokens=reasoning_tokens, + image_tokens=image_tokens, + text_tokens=None + if reported_text_tokens is None + else text_tokens_without_nested_reasoning( + completion_tokens=completion_tokens, + text_tokens=reported_text_tokens, + reasoning_tokens=reasoning_tokens or 0, + other_modality_tokens=(audio_tokens or 0) + (image_tokens or 0), + ), + audio_tokens=audio_tokens, ) extra_usage_fields: Final = { diff --git a/litellm/router.py b/litellm/router.py index 6d6efad9f42..597bcfaa20f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -21,7 +21,16 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Callable, + Generator, + Iterator, + Mapping, + MutableMapping, + Sequence, +) from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -44,18 +53,22 @@ from litellm.caching.caching import ( RedisCache, RedisClusterCache, ) +from litellm.caching.redis_cache import log_redis_failure from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, INTERNAL_CALL_ORIGIN_METADATA_KEY, + OUTPUT_TOKEN_CEILING_PARAMS, + ROUTING_REQUEST_TAGS_METADATA_KEY, RUNTIME_UPDATABLE_ROUTER_SETTINGS, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.asyncify import asyncify, run_async_function +from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, coerce_token_limit, @@ -86,6 +99,8 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_credentials_in_payload, mask_sensitive_structure, ) +from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.llms.base_llm.passthrough.transformation import replace_path_segment from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, @@ -136,6 +151,8 @@ from litellm.router_utils.common_utils import ( _is_proxy_admin_request, filter_team_based_models, filter_web_search_deployments, + get_request_team_id, + provider_for_generic_call, resolve_model_group_alias, truncate_fallback_error_detail, warn_on_provider_credential_mismatch, @@ -214,6 +231,7 @@ from litellm.types.router import ( CredentialLiteLLMParams, CustomRoutingStrategyBase, Deployment, + DeploymentModelListingInfo, DeploymentTypedDict, FallbackAccessCheck, GuardrailTypedDict, @@ -403,6 +421,10 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: + return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else () + + def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: """ Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still @@ -612,6 +634,50 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( RETRY_BREADCRUMB_LIMIT: Final = 4 +class FallbackAwareStreamWrapper(CustomStreamWrapper): + """Base for the Router's chat-completion stream wrappers, which are built around the + attempt the Router picked first and have to repoint themselves when a fallback takes over.""" + + fallback_headers_adopted: bool = False + + def adopt_fallback_response_headers( + self, + fallback_response: object, + prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]], + ) -> None: + """Repoint this wrapper at the deployment that served the stream. + + Replaces rather than merges, so the failed attempt's `x-request-id`, rate limit + counters, `model_id` and `api_base` cannot reach the proxy's response headers or + its callbacks. + """ + self._response_headers = getattr(fallback_response, "_response_headers", None) + fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params + if fallback_hidden_params: + self._hidden_params = { # mutable-ok: the rest of litellm writes into _hidden_params + **fallback_hidden_params, + # dict() because add_retry_fallback_headers mutates additional_headers in place + "additional_headers": dict(fallback_headers), # mutable-ok: see above + } + self._base_hidden_params = { # mutable-ok: CustomStreamWrapper keeps this snapshot as a dict + **self._hidden_params, + "response_cost": None, + } + self.fallback_headers_adopted = True + + +def as_output_cap(value: object) -> int | None: + """A client-sent output cap coerced to an int: ints, floats and numeric strings, never bools + or negatives.""" + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + return None + try: + cap: Final = int(float(value)) + except (ValueError, OverflowError): + return None + return cap if cap >= 0 else None + + class Router: model_names: set = set() cache_responses: bool | None = False @@ -879,6 +945,9 @@ class Router: # ``id()``-reuse risk after GC). See # ``litellm.proxy.auth.auth_checks._is_model_cost_zero``. self._zero_cost_cache: dict[str, bool] = {} + self.cached_deployment_model_info = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)( + self.get_deployment_model_info + ) self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None self._init_routing_groups(None) self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = () @@ -919,7 +988,6 @@ class Router: DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER ) self.health_state_cache = DeploymentHealthCache(cache=self.cache, staleness_threshold=float(_staleness)) - self.failed_calls = InMemoryCache() # cache to track failed call per deployment, if num failed calls within 1 minute > allowed fails, then add it to cooldown if num_retries is not None: self.num_retries = num_retries @@ -1212,7 +1280,7 @@ class Router: selector = LeastBusyLoggingHandler(router_cache=self.cache) if register_callbacks: if isinstance(litellm.input_callback, list): - litellm.input_callback.append(selector) + litellm.logging_callback_manager.add_litellm_input_callback(selector) else: litellm.input_callback = [selector] case RoutingStrategy.USAGE_BASED_ROUTING.value: @@ -1258,6 +1326,43 @@ class Router: if isinstance(litellm.input_callback, list): litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids] + def _apply_updated_routing_strategy_args(self) -> None: + """ + Re-link the default group's selector to the current `routing_strategy_args`. + + Selectors freeze their `RoutingArgs` at construction, so a runtime args + update would otherwise keep serving the boot-time values until restart. + Latency/usage state survives the rebuild: it lives in the shared router + cache, not on the selector. + """ + strategy: Final = self._normalize_strategy(self.routing_strategy) + if strategy == "lar1": + from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy + + apply_lar1_routing_strategy(self, self.routing_strategy_args) + return + + attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "") + current: Final = getattr(self, attr, None) if attr is not None else None + if attr is None or current is None: + return + + try: + rebuilt: Final = self._build_strategy_selector( + strategy=strategy or "", + routing_strategy_args=self.routing_strategy_args, + ) + except (TypeError, ValidationError): + verbose_router_logger.exception( + "Invalid routing_strategy_args %s for '%s'; keeping the previous ones", + self.routing_strategy_args, + strategy, + ) + return + + self._unregister_router_selectors((current,)) + setattr(self, attr, rebuilt) + def routing_strategy_init(self, routing_strategy: RoutingStrategy | str, routing_strategy_args: dict): verbose_router_logger.info("Routing strategy: %s", routing_strategy) self._validate_routing_strategy(routing_strategy) @@ -1487,9 +1592,33 @@ class Router: self._override_selectors[strategy] = self._build_strategy_selector( strategy=strategy, routing_strategy_args={}, + register_callbacks=False, ) return self._override_selectors[strategy] + def _override_selector_pre_call_check( + self, strategy: str | None, selector: RouterStrategySelector | None, deployment: dict + ) -> None: + """ + Override selectors are not in `litellm.callbacks`, so the pre-call check that + `routing_strategy_pre_call_checks` runs for the router's own selectors (rpm + accounting for `usage-based-routing-v2`) runs here, for the overriding request only. + """ + if selector is None or strategy is None or selector is not self._override_selectors.get(strategy): + return + selector.pre_call_check(deployment) + + async def _async_override_selector_pre_call_check( + self, + strategy: str | None, + selector: RouterStrategySelector | None, + deployment: dict, + parent_otel_span: Span | None, + ) -> None: + if selector is None or strategy is None or selector is not self._override_selectors.get(strategy): + return + await selector.async_pre_call_check(deployment, parent_otel_span) + def _get_routing_context( self, model: str, request_kwargs: dict | None = None ) -> tuple[str | None, RouterStrategySelector | None]: @@ -2572,6 +2701,18 @@ class Router: return fallback_hidden_params, {} return fallback_hidden_params, cast("dict[str, object]", fallback_headers) + @staticmethod + def _adopt_fallback_response_headers( + wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]", + fallback_response: object, + ) -> tuple[dict[str, object], dict[str, object]]: + """Repoint the wrapper at `fallback_response`, returning its prepared hidden params.""" + prepared: Final = Router._prepare_fallback_hidden_params(fallback_response) + adopting_wrapper: Final = wrapper_ref() + if adopting_wrapper is not None: + adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared) + return prepared + @staticmethod def _apply_fallback_hidden_params_to_item( fallback_item: object, @@ -2611,7 +2752,7 @@ class Router: held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack() - class FallbackStreamWrapper(CustomStreamWrapper): + class FallbackStreamWrapper(FallbackAwareStreamWrapper): def __init__(self, async_generator: AsyncGenerator): # Copy attributes from the original model_response super().__init__( @@ -2619,6 +2760,7 @@ class Router: model=model_response.model, custom_llm_provider=model_response.custom_llm_provider, logging_obj=model_response.logging_obj, + _response_headers=getattr(model_response, "_response_headers", None), ) self._async_generator = async_generator inner_chunks: Final[object] = getattr(model_response, "chunks", None) @@ -2695,8 +2837,17 @@ class Router: # If fallback returns a streaming response, iterate over it if hasattr(fallback_response, "__aiter__"): - prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) + prepared_fallback_hidden_params = Router._adopt_fallback_response_headers( + wrapper_ref, fallback_response + ) + fallback_headers_are_settled = False async for fallback_item in fallback_response: + if not fallback_headers_are_settled: + fallback_headers_are_settled = True # rebind-ok: one-shot latch + # a fallback that failed over again only repoints itself once it yields + prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields + Router._adopt_fallback_response_headers(wrapper_ref, fallback_response) + ) Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item @@ -2738,7 +2889,11 @@ class Router: e, ) - return FallbackStreamWrapper(stream_with_fallbacks()) + wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) + # weak, so the generator closing over it does not keep the wrapper out of + # refcount teardown and delay the `finally` that releases the deployment slot + wrapper_ref: Final = weakref.ref(wrapped_response) + return wrapped_response @staticmethod def _extract_partial_responses_usage( @@ -3167,13 +3322,14 @@ class Router: """ from litellm.exceptions import MidStreamFallbackError - class SyncFallbackStreamWrapper(CustomStreamWrapper): + class SyncFallbackStreamWrapper(FallbackAwareStreamWrapper): def __init__(self, sync_generator: Generator): super().__init__( completion_stream=sync_generator, model=model_response.model, custom_llm_provider=model_response.custom_llm_provider, logging_obj=model_response.logging_obj, + _response_headers=getattr(model_response, "_response_headers", None), ) self._sync_generator = sync_generator if hasattr(model_response, "_hidden_params"): @@ -3229,8 +3385,17 @@ class Router: ) if hasattr(fallback_response, "__iter__"): - prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) + prepared_fallback_hidden_params = Router._adopt_fallback_response_headers( + wrapper_ref, fallback_response + ) + fallback_headers_are_settled = False for fallback_item in fallback_response: + if not fallback_headers_are_settled: + fallback_headers_are_settled = True # rebind-ok: one-shot latch + # a fallback that failed over again only repoints itself once it yields + prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields + Router._adopt_fallback_response_headers(wrapper_ref, fallback_response) + ) Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item @@ -3268,7 +3433,10 @@ class Router: close_err, ) - return SyncFallbackStreamWrapper(stream_with_fallbacks()) + wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks()) + # weak, for the same reason as the async twin + wrapper_ref: Final = weakref.ref(wrapped_response) + return wrapped_response async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): """ @@ -3534,6 +3702,13 @@ class Router: effective_model_info: Final = kwargs.get("model_info") or deployment.get("model_info") or MappingProxyType({}) self._set_failed_deployment_id_on_exception(exception, MappingProxyType({"model_info": effective_model_info})) + @staticmethod + def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None: + effective_model_info: Final = kwargs.get("model_info") + deployment_id: Final = effective_model_info.get("id") if isinstance(effective_model_info, Mapping) else None + if isinstance(deployment_id, str) and deployment_id: + exception.retry_skip_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # dynamic stamp, read by _deployment_ids_to_skip_on_retry + def _update_kwargs_with_default_litellm_params( self, kwargs: dict, metadata_variable_name: str | None = "metadata" ) -> None: @@ -3651,6 +3826,11 @@ class Router: refund_stale_reservation_before_retry(self.cache, kwargs) set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=deployment_has_io_token_limits(deployment)) + kwargs[metadata_variable_name].setdefault( + ROUTING_REQUEST_TAGS_METADATA_KEY, + tuple(_get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name)), + ) + ## DEPLOYMENT-LEVEL TAGS deployment_tags: Final = deployment.get("litellm_params", {}).get("tags") if deployment_tags: @@ -4139,10 +4319,12 @@ class Router: } ) litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) - prompt_management_deployment: Final = self.get_available_deployment( + specific_deployment: Final = kwargs.pop("specific_deployment", None) + prompt_management_deployment: Final = await self.async_get_available_deployment( model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), + messages=cast(list[dict[str, str]], messages), # cast-ok: selection reads messages structurally + specific_deployment=specific_deployment, + request_kwargs=kwargs, ) self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs) @@ -4229,6 +4411,7 @@ class Router: model=model, messages=[{"role": "user", "content": "prompt"}], specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) data: Final = deployment["litellm_params"].copy() @@ -4259,6 +4442,7 @@ class Router: verbose_router_logger.info("litellm.image_generation(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) if model_name is not None: self.fail_calls[model_name] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e async def aimage_generation(self, prompt: str, model: str, **kwargs): @@ -4343,6 +4527,7 @@ class Router: verbose_router_logger.info("litellm.aimage_generation(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) if model_name is not None: self.fail_calls[model_name] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e async def atranscription(self, file: FileTypes, model: str, **kwargs): @@ -4447,9 +4632,10 @@ class Router: verbose_router_logger.info("litellm.atranscription(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) if model_name is not None: self.fail_calls[model_name] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e - async def aspeech(self, model: str, input: str, voice: str, **kwargs): + async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): """ Example Usage: @@ -4501,7 +4687,7 @@ class Router: ) raise e - async def _aspeech(self, model: str, input: str, voice: str, **kwargs): + async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): model_name: Final = model try: verbose_router_logger.debug("Inside _aspeech()- model: %s; kwargs: %s", model, kwargs) @@ -4525,7 +4711,7 @@ class Router: **{ **data, "input": input, - "voice": voice, + "voice": data.get("voice") if voice is None else voice, "client": model_client, **kwargs, } @@ -4561,6 +4747,7 @@ class Router: verbose_router_logger.info("litellm.aspeech(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) if model_name is not None: self.fail_calls[model_name] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e async def arerank(self, model: str, **kwargs): @@ -4619,6 +4806,7 @@ class Router: verbose_router_logger.info("litellm.arerank(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) if model_name is not None: self.fail_calls[model_name] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e def text_completion( @@ -4753,6 +4941,7 @@ class Router: verbose_router_logger.info("litellm.atext_completion(model=%s)\x1b[31m Exception %s\x1b[0m", model, e) if model is not None: self.fail_calls[model] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e async def aadapter_completion( @@ -4843,6 +5032,7 @@ class Router: verbose_router_logger.info("litellm.aadapter_completion(model=%s)\x1b[31m Exception %s\x1b[0m", model, e) if model is not None: self.fail_calls[model] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e async def _asearch_with_fallbacks(self, original_function: Callable, **kwargs): @@ -5016,7 +5206,7 @@ class Router: # If get_llm_provider fails, fall back to using model_name as-is replacement_model_name = model_name - kwargs["endpoint"] = kwargs["endpoint"].replace(model, replacement_model_name) + kwargs["endpoint"] = replace_path_segment(kwargs["endpoint"], model, replacement_model_name) return kwargs async def _ageneric_api_call_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): @@ -5052,16 +5242,7 @@ class Router: kwargs=kwargs, model=model, model_name=model_name ) - # Get custom_llm_provider from deployment params - try: - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - except Exception: - custom_llm_provider = None + custom_llm_provider: Final = provider_for_generic_call(data) response_kwargs: Final = { **data, @@ -5572,15 +5753,7 @@ class Router: # Perform pre-call checks for routing strategy self.routing_strategy_pre_call_checks(deployment=deployment) - try: - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - except Exception: - custom_llm_provider = None + custom_llm_provider: Final = provider_for_generic_call(data) response: Final = original_function( **{ @@ -5625,6 +5798,7 @@ class Router: model=model, input=input, specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) data: Final = deployment["litellm_params"].copy() @@ -5663,6 +5837,7 @@ class Router: verbose_router_logger.info("litellm.embedding(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) if model_name is not None: self.fail_calls[model_name] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e async def aembedding( @@ -5750,6 +5925,7 @@ class Router: verbose_router_logger.info("litellm.aembedding(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) if model_name is not None: self.fail_calls[model_name] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e #### FILES API #### @@ -6123,6 +6299,7 @@ class Router: ) if model is not None: self.fail_calls[model] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e async def aretrieve_batch( @@ -6343,6 +6520,7 @@ class Router: ) if model is not None: self.fail_calls[model] += 1 + self._stamp_retry_skip_deployment_id(e, kwargs) raise e async def alist_batches( @@ -7458,6 +7636,23 @@ class Router: Context_Policy_Fallbacks={content_policy_fallbacks}", ) + @staticmethod + def _deployment_ids_to_skip_on_retry(exception: Exception, already_skipped: object) -> tuple[str, ...]: + failed_deployment_id: Final[str | None] = getattr(exception, "retry_skip_deployment_id", None) or getattr( + exception, "failed_deployment_id", None + ) + status_code: Final = getattr(exception, "status_code", None) + if not failed_deployment_id or not isinstance(status_code, int): + return () + if litellm._should_retry(status_code): # pyright: ignore[reportPrivateUsage] # as in should_retry_this_error + return () + already_skipped_ids: Final = _as_retry_skipped_deployment_ids(already_skipped) + skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id)))) + verbose_router_logger.debug( + "Retry skips deployments that already answered %s to this request: %s", status_code, skipped + ) + return skipped + @tracer.wrap() async def async_function_with_retries(self, *args, **kwargs): verbose_router_logger.debug("Inside async function with retries.") @@ -7553,6 +7748,12 @@ class Router: ## LOGGING if num_retries > 0: kwargs = self.log_retry(kwargs=kwargs, e=original_exception) + first_skipped_ids: Final = self._deployment_ids_to_skip_on_retry( + exception=original_exception, + already_skipped=kwargs.get("_retry_skipped_deployment_ids"), + ) + if first_skipped_ids: + kwargs["_retry_skipped_deployment_ids"] = first_skipped_ids # rebind-ok: the next attempt reads it else: raise @@ -7622,6 +7823,12 @@ class Router: except Exception: raise e + skipped_ids = self._deployment_ids_to_skip_on_retry( + exception=e, + already_skipped=kwargs.get("_retry_skipped_deployment_ids"), + ) + if skipped_ids: + kwargs["_retry_skipped_deployment_ids"] = skipped_ids # rebind-ok: the next attempt reads it _timeout = self._time_to_sleep_before_retry( e=e, remaining_retries=remaining_retries, @@ -9075,10 +9282,12 @@ class Router: if prefs_raw is not None: model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) - # `input_cost_per_token` is a LiteLLM_Params field per types/router.py. + # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set. lp = d.get("litellm_params") if isinstance(d, dict) else d.litellm_params lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) cost = lp_dict.get("input_cost_per_token") + if cost is None: + cost = mi_dict.get("input_cost_per_token") if cost is not None: model_to_cost[name] = float(cost) @@ -9262,6 +9471,12 @@ class Router: #### VALIDATE MODEL ######## # Check if this is a prompt management model before validating as LLM provider litellm_model: Final = deployment.litellm_params.model + if isinstance(deployment.litellm_params.drop_params, str): + verbose_router_logger.warning( + "model=%s drop_params=%r is not a flag value, treating it as unset", + deployment.model_name, + deployment.litellm_params.drop_params, + ) is_prompt_management_model = False if "/" in litellm_model: @@ -9832,6 +10047,7 @@ class Router: model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.custom_llm_provider, ) + self._invalidate_model_group_info_cache() def delete_deployment(self, id: str) -> Deployment | None: """ @@ -10000,15 +10216,71 @@ class Router: return None return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable + def get_model_listing_info(self, model_name: str) -> DeploymentModelListingInfo | None: + """ + Return what the concrete deployments behind model_name contribute to its + /v1/models entry: the cost-map keys for their underlying models, plus the widest + token limits explicitly configured in their model_info. Resolved via O(1) index + lookup. + + Returns None for wildcard-expanded or unknown names, where the listed name is the + real model name and no deployment-specific information exists, and treats a + malformed configured limit as absent rather than failing the listing. + + The whole group is read rather than just its first deployment, so a group that + mixes models does not advertise a window that depends on config order; the widest + one is reported, which is what get_model_group_info already shows the Admin UI. + Keys are deduplicated, so the ordinary group of interchangeable deployments of one + model still costs the caller a single cost-map lookup. Unlike get_model_group_info, + this never triggers pattern matching or deep copies, so it is safe to call per + listed model on the /v1/models hot path. + """ + indices: Final = self.model_name_to_deployment_indices.get(model_name) + if not indices: + return None + + deployments: Final = tuple(self.model_list[index] for index in indices) + model_infos: Final = tuple(deployment.get("model_info") or MappingProxyType({}) for deployment in deployments) + params: Final = tuple(deployment.get("litellm_params") or MappingProxyType({}) for deployment in deployments) + # base_model resolution mirrors get_router_model_info: unset or blank means the + # deployment's own model name is the cost-map key. + cost_map_keys: Final = tuple( + dict.fromkeys( # deduplicates while preserving config order + key + for key in ( + model_info.get("base_model") or litellm_params.get("base_model") or litellm_params.get("model") + for model_info, litellm_params in zip(model_infos, params) + ) + if isinstance(key, str) and key + ) + ) + return DeploymentModelListingInfo( + cost_map_keys=cost_map_keys, + max_input_tokens=self._widest_configured_limit(model_infos, "max_input_tokens"), + max_output_tokens=self._widest_configured_limit(model_infos, "max_output_tokens"), + ) + + @staticmethod + def _widest_configured_limit(model_infos: Sequence[Mapping[str, Any]], field: str) -> int | None: + """The largest usable value of ``field`` across a group's configured model_info blocks.""" + limits: Final = tuple( + limit + for limit in (coerce_token_limit(model_info.get(field)) for model_info in model_infos) + if limit is not None + ) + return max(limits) if limits else None + def get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]": """ Return (max_input_tokens, max_output_tokens) explicitly configured in a concrete deployment's model_info for model_name, via O(1) index lookup. Returns (None, None) for wildcard-expanded or unknown names, and treats a - malformed configured value as absent rather than failing the listing. Unlike - get_model_group_info, this never triggers pattern matching or deep copies, so it - is safe to call per listed model on the /v1/models hot path. + malformed configured value as absent rather than failing the caller. + + Deliberately reads one deployment rather than aggregating the group the way + get_model_listing_info does: its caller truncates an embedding input to this + value, so the widest window in a mixed group would be the wrong answer there. """ deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) if deployment is None: @@ -10793,6 +11065,9 @@ class Router: """ return self.get_model_group_info(model_group) + def cached_model_group_info(self, model_group: str) -> ModelGroupInfo | None: + return self._cached_get_model_group_info(model_group) + async def get_remaining_model_group_usage(self, model_group: str) -> dict[str, int]: model_group_info: Final = self._cached_get_model_group_info(model_group) @@ -10944,6 +11219,43 @@ class Router: return ids + def get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: + """ + Deployment ids that could serve ``model`` for ``team_id``, following the same + precedence ``_common_checks_available_deployment`` uses to build a candidate pool: + ``model_group_alias``, then a routing group, then the first matching early-resolve + path for a name that is not a ``model_name`` (team route, wildcard pattern via + ``get_deployments_by_pattern``, team pattern router, default deployment), then the + ``model_name`` and team indexes. Delegating to the router's own resolvers keeps this + aligned with how a route actually resolves rather than re-deriving it, and unlike + ``_common_checks_available_deployment`` it is read-only: it does not apply request + fallbacks and (with ``include_team_models`` left off) does not raise. Lets a pre-call + check tell a genuine cross-group route from same-group unavailability without leaking + deployment ids into request kwargs bound for the provider. + """ + resolved: Final = self._get_model_from_alias(model=model) or model + routing_group_members: Final = self._get_routing_group_deployments(model=resolved, team_id=team_id) + if routing_group_members is not None: + return self._deployment_ids(routing_group_members) + early: Final = self._try_early_resolve_deployments_for_model_not_in_names( + model=resolved, request_team_id=team_id + ) + if early is not None: + early_deployments: Final = early[1] + return self._deployment_ids( + (early_deployments,) if isinstance(early_deployments, Mapping) else early_deployments + ) + return self._deployment_ids(self._get_all_deployments(model_name=resolved, team_id=team_id)) + + @staticmethod + def _deployment_ids(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]: + return frozenset( + str(model_info["id"]) + for deployment in deployments + for model_info in (deployment.get("model_info"),) + if isinstance(model_info, Mapping) and model_info.get("id") is not None + ) + def has_model_id(self, candidate_id: str) -> bool: """ O(1) membership check for a deployment ID without allocating large lists. @@ -11533,6 +11845,7 @@ class Router: result and bypass budget enforcement. """ self._cached_get_model_group_info.cache_clear() + self.cached_deployment_model_info.cache_clear() self._zero_cost_cache.clear() self._routing_group_rows = None @@ -11661,7 +11974,7 @@ class Router: _existing_router_settings: Final = self.get_settings() rebuild_routing_groups = False - relink_lar1_from_args = False + routing_args_updated = False for var in kwargs: if var in RUNTIME_UPDATABLE_ROUTER_SETTINGS: if var in _int_settings: @@ -11700,15 +12013,13 @@ class Router: ) rebuild_routing_groups = True elif var == "routing_strategy_args": - relink_lar1_from_args = True + routing_args_updated = True setattr(self, var, value) else: verbose_router_logger.debug("Setting %s is not allowed", var) - if relink_lar1_from_args and self._normalize_strategy(self.routing_strategy) == "lar1": - from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - - apply_lar1_routing_strategy(self, self.routing_strategy_args) + if routing_args_updated: + self._apply_updated_routing_strategy_args() if rebuild_routing_groups: self._init_routing_groups(self._routing_groups_input) @@ -11854,7 +12165,7 @@ class Router: try: if not self._pre_call_checks_need_token_count(model, healthy_deployments): return None - return await asyncify(self._count_pre_call_check_tokens)( + return await offload_token_count(self._count_pre_call_check_tokens)( messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter request_kwargs=request_kwargs, @@ -12082,27 +12393,7 @@ class Router: if team_deployments: return model, team_deployments elif include_team_models: - team_deployments = [ - self.model_list[index] - for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() - if public_model_name == model - for index in indices - ] - team_ids: Final = { - team_id - for deployment in team_deployments - for team_id in [(deployment.get("model_info") or {}).get("team_id")] - if team_id is not None - } - if len(team_ids) > 1: - raise litellm.BadRequestError( - message=( - f"Model name '{model}' matches deployments from multiple teams. " - "Specify the deployment ID directly to disambiguate." - ), - model=model, - llm_provider="", - ) + team_deployments = self._team_deployments_across_teams(model) if team_deployments: return model, team_deployments @@ -12129,6 +12420,45 @@ class Router: return None + def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]: + """Every team's deployments under public name `model`, for a proxy admin calling without a team.""" + team_deployments: Final = [ + self.model_list[index] + for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() + if public_model_name == model + for index in indices + ] + team_ids: Final = { + team_id + for deployment in team_deployments + for team_id in [(deployment.get("model_info") or {}).get("team_id")] + if team_id is not None + } + if len(team_ids) > 1: + raise litellm.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + return team_deployments + + def deployments_for_request( + self, model: str, request_kwargs: Mapping[str, object] + ) -> Sequence[DeploymentTypedDict]: + """The deployments `model` names for this caller, through the same alias, then team-first, then + global, then admin-across-teams resolution `_common_checks_available_deployment` applies, so + strategy selection and compression policy can never disagree with deployment selection about + which marker a name means.""" + registered_name: Final = self._get_model_from_alias(model=model) or model + team_id: Final = get_request_team_id(request_kwargs) + deployments: Final = self._get_all_deployments(model_name=registered_name, team_id=team_id) + if deployments or team_id is not None or not _is_proxy_admin_request(request_kwargs): + return deployments + return self._team_deployments_across_teams(registered_name) + @staticmethod def _is_strategy_marker_deployment(deployment: Mapping[str, object]) -> bool: litellm_params: Final = deployment.get("litellm_params") @@ -12156,11 +12486,7 @@ class Router: - Dict, if specific model chosen """ - request_team_id: str | None = None - if request_kwargs is not None: - metadata: Final = request_kwargs.get("metadata") or {} - litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {} - request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + request_team_id: Final = get_request_team_id(request_kwargs) # check if aliases set on litellm model alias map if specific_deployment is True: return model, self._get_deployment_by_litellm_model(model=model) @@ -12185,7 +12511,9 @@ class Router: include_team_models=_is_proxy_admin_request(request_kwargs), ) if early is not None: - return early + if not isinstance(early[1], list): + return early + return early[0], self._drop_strategy_markers(early[0], early[1]) ## get healthy deployments ### get all deployments @@ -12262,19 +12590,22 @@ class Router: model ] # update the model to the actual value if an alias has been passed in - marker_flags: Final = tuple(self._is_strategy_marker_deployment(d) for d in healthy_deployments) - if not any(marker_flags): - return model, healthy_deployments - selectable: Final = [ # mutable-ok: matches this function's list contract expected by downstream filters - d for d, is_marker in zip(healthy_deployments, marker_flags, strict=True) if not is_marker + return model, self._drop_strategy_markers(model, healthy_deployments) + + def _drop_strategy_markers( + self, model: str, deployments: Sequence[DeploymentTypedDict] + ) -> list[DeploymentTypedDict]: + """A strategy marker is never a callable deployment, whichever resolution arm produced it.""" + selectable: Final = [ # mutable-ok: matches _common_checks_available_deployment's list contract + d for d in deployments if not self._is_strategy_marker_deployment(d) ] - if not selectable: + if deployments and not selectable: raise litellm.BadRequestError( message=f"You passed in model={model}. {RouterErrors.only_strategy_marker_deployments.value}", model=model, llm_provider="", ) - return model, selectable + return selectable def _filter_deployments_by_model_access_groups( self, @@ -12338,6 +12669,7 @@ class Router: input: str | list | None = None, specific_deployment: bool | None = False, parent_otel_span: Span | None = None, + health_check_probe: bool = False, ) -> list[dict] | dict: """ Get the healthy deployments for a model. @@ -12387,6 +12719,7 @@ class Router: healthy_deployments = await self._async_filter_health_check_unhealthy_deployments( healthy_deployments=healthy_deployments, parent_otel_span=parent_otel_span, + health_check_probe=health_check_probe, ) cooldown_deployments: Final = await _async_get_cooldown_deployments( @@ -12452,7 +12785,7 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - healthy_deployments = litellm.utils._get_order_filtered_deployments( + healthy_deployments = litellm.utils.get_order_filtered_deployments( cast(list[dict], healthy_deployments), target_order=_target_order ) @@ -12460,11 +12793,24 @@ class Router: ## this request via weighted-failover. Always honored, regardless of the ## router-level flag, so a stale exclusion key on kwargs cannot escape. _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm.utils._get_excluded_filtered_deployments( + healthy_deployments = litellm.utils.get_excluded_filtered_deployments( cast(list[dict], healthy_deployments), excluded_deployment_ids=_excluded_deployment_ids, ) + ## RETRY SKIP ## -> drop deployments that already refused this request with a + ## non-retryable status, unless that leaves nothing, so the caller still gets + ## the provider's own error instead of a no-deployments error. + _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( + request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None + ) + healthy_deployments = ( + litellm.utils.get_excluded_filtered_deployments( + healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids + ) + or healthy_deployments + ) + if len(healthy_deployments) == 0: exception: Final = await async_raise_no_deployment_exception( litellm_router_instance=self, @@ -12487,7 +12833,83 @@ class Router: request_kwargs.pop(carrier, None) @staticmethod - def _drop_client_effort_carriers_a_tier_pin_supersedes( + def _tier_ceiling_under_the_surface_name( + tier_litellm_params: Mapping[str, object], responses_call: bool + ) -> Mapping[str, object]: + """``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are one + ceiling under three names, and each surface reads exactly one of them: the + Responses bridge builds its internal ``max_tokens`` from ``max_output_tokens`` + and would overwrite the tier's, chat and /v1/messages never read + ``max_output_tokens``, and litellm already renames ``max_tokens`` to + ``max_completion_tokens`` for the OpenAI models that require it. Collapse + whatever the tier carries onto the surface's own name, preferring a value the + operator already wrote under that name.""" + surface_key: Final = "max_output_tokens" if responses_call else "max_tokens" + carried: Final = tuple( + key + for key in (surface_key, "max_tokens", "max_completion_tokens", "max_output_tokens") + if key in tier_litellm_params + ) + if not carried: + return tier_litellm_params + return MappingProxyType( + { + **{k: v for k, v in tier_litellm_params.items() if k not in OUTPUT_TOKEN_CEILING_PARAMS}, + surface_key: tier_litellm_params[carried[0]], + } + ) + + def _pin_tier_params_onto_request( + self, + model: str, + tier_litellm_params: Mapping[str, object] | None, + request_kwargs: dict, + responses_call: bool, + ) -> bool: + """Apply a routing strategy's per-tier litellm_params on top of the request and report + whether they pinned an output ceiling, so the caller can hand the request its own ceiling + back on a routing pass that pins none.""" + if not tier_litellm_params: + return False + accepted_tier_params: Final = self._tier_params_the_target_accepts(model, tier_litellm_params, request_kwargs) + surface_tier_params: Final = self._tier_ceiling_under_the_surface_name( + accepted_tier_params, responses_call=responses_call + ) + self._drop_client_carriers_a_tier_pin_supersedes(request_kwargs, surface_tier_params) + request_kwargs.update(surface_tier_params) + return not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(surface_tier_params) + + @staticmethod + def _restore_client_ceiling_no_tier_pins(request_kwargs: MutableMapping[str, object]) -> None: + """A model-group fallback re-enters routing with the kwargs an earlier auto-router pass + already rewrote, so a ceiling sized for that pass's tier would ride onto a group no tier + chose. When this pass pins none, hand the request back exactly the carriers the caller + sent, which the first pinning pass stamped. The stamp lives in a metadata bucket a + caller can also write, so the proxy strips the key at ingestion and this read takes + nothing but the three ceiling carriers as integers: no other key ever reaches kwargs.""" + stamped: Final = next( + ( + bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) + for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket + ), + None, + ) + if not isinstance(stamped, dict): + return + callers_ceiling: Final = MappingProxyType( + { + carrier: cap + for carrier, value in stamped.items() + if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None + } + ) + for carrier in OUTPUT_TOKEN_CEILING_PARAMS: + request_kwargs.pop(carrier, None) + request_kwargs.update(callers_ceiling) + + @staticmethod + def _drop_client_carriers_a_tier_pin_supersedes( request_kwargs: dict[str, object], tier_litellm_params: Mapping[str, object], ) -> None: @@ -12497,7 +12919,22 @@ class Router: the ``reasoning_effort`` alias, so a pinned effort only reaches the wire if the client's other encodings are removed before the merge. Non-effort fields a carrier also holds (``output_config.format``, - ``reasoning.summary``) are kept.""" + ``reasoning.summary``) are kept. An output ceiling has the same shape: + ``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are + one setting under three names, and a provider handed two of them either + rejects the request or picks one by iteration order.""" + if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params): + _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) + metadata_bucket.setdefault( + CLIENT_OUTPUT_CEILING_METADATA_KEY, + { + carrier: request_kwargs[carrier] + for carrier in OUTPUT_TOKEN_CEILING_PARAMS + if carrier in request_kwargs + }, + ) + for carrier in OUTPUT_TOKEN_CEILING_PARAMS: + request_kwargs.pop(carrier, None) if "reasoning_effort" not in tier_litellm_params: return request_kwargs.pop("thinking", None) @@ -12538,6 +12975,7 @@ class Router: # Execute Pre-Routing Hooks # this hook can modify the model, messages before the routing decision is made ######################################################### + responses_call: Final = input is not None and messages is None pre_routing_hook_response: Final = await self.async_pre_routing_hook( model=model, request_kwargs=request_kwargs, @@ -12549,12 +12987,14 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - if pre_routing_hook_response.litellm_params: - accepted_tier_params: Final = self._tier_params_the_target_accepts( - model, pre_routing_hook_response.litellm_params, request_kwargs - ) - self._drop_client_effort_carriers_a_tier_pin_supersedes(request_kwargs, accepted_tier_params) - request_kwargs.update(accepted_tier_params) + tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + model=model, + tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, + request_kwargs=request_kwargs, + responses_call=responses_call, + ) + if not tier_pins_ceiling: + self._restore_client_ceiling_no_tier_pins(request_kwargs) ######################################################### # Resolve the strategy and logger AFTER the pre-routing hook, since @@ -12571,10 +13011,16 @@ class Router: parent_otel_span=parent_otel_span, ) if isinstance(healthy_deployments, dict): + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, healthy_deployments, parent_otel_span + ) return healthy_deployments # When encrypted content affinity pins to a specific deployment, if request_kwargs.get("_encrypted_content_affinity_pinned") and len(healthy_deployments) == 1: + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, healthy_deployments[0], parent_otel_span + ) return healthy_deployments[0] start_time: Final = time.time() @@ -12600,6 +13046,9 @@ class Router: parent_otel_span=parent_otel_span, ) raise exception + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, deployment, parent_otel_span + ) verbose_router_logger.info( "get_available_deployment for model: %s, Selected deployment: %s for model: %s", model, @@ -12654,6 +13103,7 @@ class Router: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(request_kwargs) # 1. Execute pre-routing hook + responses_call: Final = input is not None and messages is None pre_routing_hook_response: Final = await self.async_pre_routing_hook( model=model, request_kwargs=request_kwargs, @@ -12665,12 +13115,14 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - if pre_routing_hook_response.litellm_params: - accepted_tier_params: Final = self._tier_params_the_target_accepts( - model, pre_routing_hook_response.litellm_params, request_kwargs - ) - self._drop_client_effort_carriers_a_tier_pin_supersedes(request_kwargs, accepted_tier_params) - request_kwargs.update(accepted_tier_params) + tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + model=model, + tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, + request_kwargs=request_kwargs, + responses_call=responses_call, + ) + if not tier_pins_ceiling: + self._restore_client_ceiling_no_tier_pins(request_kwargs) # 2. Get healthy deployments healthy_deployments: Final = await self.async_get_healthy_deployments( @@ -12682,6 +13134,8 @@ class Router: parent_otel_span=parent_otel_span, ) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) + # 3. If specific deployment returned, verify if it supports pass-through if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: @@ -12692,6 +13146,9 @@ class Router: ) litellm_params: Final = healthy_deployments.get("litellm_params", {}) if litellm_params.get("use_in_pass_through"): + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, healthy_deployments, parent_otel_span + ) return healthy_deployments else: raise litellm.BadRequestError( @@ -12712,7 +13169,6 @@ class Router: # 5. Apply load balancing strategy start_time: Final = time.perf_counter() - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if strategy == "simple-shuffle": return simple_shuffle( llm_router_instance=self, @@ -12736,6 +13192,9 @@ class Router: parent_otel_span=parent_otel_span, ) raise exception + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, deployment, parent_otel_span + ) verbose_router_logger.info( "async_get_available_deployment_for_pass_through model: %s, selected deployment: %s", @@ -12838,12 +13297,8 @@ class Router: return filtered - def _model_name_has_plain_deployments(self, model: str) -> bool: - indices: Final = self.model_name_to_deployment_indices.get(model) or () - return any(not self._is_strategy_marker_deployment(self.model_list[idx]) for idx in indices) - def _select_pre_routing_strategy( - self, model: str, request_kwargs: dict + self, model: str, request_kwargs: Mapping[str, object] ) -> "TaggedPreRoutingStrategy[PreRoutingStrategy] | None": """ Resolve the pre-routing strategy for `model`, disambiguating deployments @@ -12854,6 +13309,12 @@ class Router: deployment the strategy was registered from via its (model_name, tags) pair. + The registries are keyed by the marker deployment's own `model_name`, which + for a team-scoped router is the internal `model_name_{team}_{uuid}` while + the caller sends the team's public name. So the names looked up are the + `model_name`s of whatever deployments this caller's request resolves `model` + to, and `model` itself when it resolves to none. + With tag filtering enabled, router-wide or by the request's enable_tag_filtering (which the proxy sets from key/team router_settings), strategies that all carry real tags matching none of @@ -12861,12 +13322,14 @@ class Router: deployments: returning None hands the request to ordinary tag-aware deployment selection. """ - candidates: Final[list[TaggedPreRoutingStrategy[PreRoutingStrategy]]] = [ - *self.auto_routers.get(model, []), - *self.complexity_routers.get(model, []), - *self.adaptive_routers.get(model, []), - *self.quality_routers.get(model, []), - ] + registries: Final = (self.auto_routers, self.complexity_routers, self.adaptive_routers, self.quality_routers) + if not any(registries): + return None + deployments: Final = self.deployments_for_request(model, request_kwargs) + registered_names: Final = tuple(dict.fromkeys(str(d["model_name"]) for d in deployments)) or (model,) + candidates: Final = tuple( + tagged for registry in registries for name in registered_names for tagged in registry.get(name, []) + ) if not candidates: return None @@ -12884,7 +13347,7 @@ class Router: if ( (self.enable_tag_filtering or request_scoped_filtering) and all(tagged.tags for tagged in candidates) - and self._model_name_has_plain_deployments(model) + and any(not self._is_strategy_marker_deployment(d) for d in deployments) ): return None return candidates[0] @@ -12935,8 +13398,10 @@ class Router: return await session_cache.async_get_cache(key=cache_key) return await session_cache.redis_cache.async_get_cache(key=cache_key) except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis - verbose_router_logger.warning( - "Failed to read Claude Code session router binding; using the requested model: %s", + log_redis_failure( + verbose_router_logger, + logging.WARNING, + "Failed to read Claude Code session router binding; using the requested model", e, ) return None @@ -12996,11 +13461,12 @@ class Router: Used for the litellm auto-router to modify the request before the routing decision is made. - `model` is whatever the caller asked for, which may be a `model_group_alias` key, while the - strategy registries and the marker deployment are keyed by the marker's own `model_name`, so - every lookup below resolves the alias first. Only the lookups: the caller-facing name stays - the alias, since spend metadata is stamped before routing and the response carries the tier - group the strategy picked. + `model` is whatever the caller asked for, which may be a `model_group_alias` key or a team's + public model name, while the strategy registries and the marker deployment are keyed by the + marker's own `model_name`, so every lookup below resolves the alias first and the team name + through the deployment path. Only the lookups: the caller-facing name stays the alias, since + spend metadata is stamped before routing and the response carries the tier group the + strategy picked. """ requested_registered_model_name: Final = self._get_model_from_alias(model=model) or model registered_model_name: Final = await self._resolve_claude_code_session_router( @@ -13037,7 +13503,6 @@ class Router: messages_for_routing, model_hop_compression_armed, policy_for_model, - team_id_from_request, ) # Same tag-aware lookup the proxy's pre-call arming used, so an alias with @@ -13045,7 +13510,7 @@ class Router: compression_policy: Final = policy_for_model( llm_router=self, model_alias=registered_model_name, - team_id=team_id_from_request(request_kwargs), + request_kwargs=request_kwargs, request_tags=_get_tags_from_request_kwargs(request_kwargs), ) # Shared compression already ran in the pre-call hook, so reuse it rather than @@ -13114,7 +13579,9 @@ class Router: # Per-tier `litellm_params` on the hook response are deliberate overrides # the caller applies on top, so those keys are never forwarded here. marker_params: Final = ( - self._forwardable_alias_marker_params(model=registered_model_name, strategy_tags=selected_strategy.tags) + self._forwardable_alias_marker_params( + model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs + ) if pre_routing_hook_response is not None else () ) @@ -13132,13 +13599,14 @@ class Router: return pre_routing_hook_response def _forwardable_alias_marker_params( - self, model: str, strategy_tags: tuple[str, ...] + self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] ) -> tuple[tuple[str, object], ...]: marker_params: Final = tuple( litellm_params - for idx in self.model_name_to_deployment_indices.get(model, ()) - if isinstance(litellm_params := self.model_list[idx].get("litellm_params", {}), dict) - and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + for deployment in self.deployments_for_request(model, request_kwargs) + if str((litellm_params := deployment["litellm_params"]).get("model", "")).startswith( + AUTO_ROUTER_MODEL_PREFIX + ) ) tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags @@ -13311,6 +13779,7 @@ class Router: specific_deployment=specific_deployment, request_kwargs=request_kwargs, ) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: @@ -13319,6 +13788,7 @@ class Router: model=model, llm_provider="", ) + self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments) return healthy_deployments parent_otel_span: Final[Span | None] = _get_parent_otel_span_from_kwargs(request_kwargs) @@ -13357,7 +13827,7 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - healthy_deployments = litellm.utils._get_order_filtered_deployments( + healthy_deployments = litellm.utils.get_order_filtered_deployments( healthy_deployments, target_order=_target_order ) @@ -13365,11 +13835,22 @@ class Router: ## this request via weighted-failover. See async counterpart in ## async_get_healthy_deployments for details. _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm.utils._get_excluded_filtered_deployments( + healthy_deployments = litellm.utils.get_excluded_filtered_deployments( healthy_deployments, excluded_deployment_ids=_excluded_deployment_ids, ) + ## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments. + _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( + request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None + ) + healthy_deployments = ( + litellm.utils.get_excluded_filtered_deployments( + healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids + ) + or healthy_deployments + ) + if len(healthy_deployments) == 0: model_ids = self.get_model_ids(model_name=model) _cooldown_time = self.cooldown_cache.get_min_cooldown( @@ -13383,7 +13864,6 @@ class Router: cooldown_list=_cooldown_list, ) - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if strategy == "simple-shuffle": # if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm ############## Check 'weight' param set for weighted pick ################# @@ -13415,6 +13895,7 @@ class Router: enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, ) + self._override_selector_pre_call_check(strategy, strategy_selector, deployment) verbose_router_logger.info( "get_available_deployment for model: %s, Selected deployment: %s for model: %s", model, @@ -13458,6 +13939,8 @@ class Router: specific_deployment=specific_deployment, ) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) + # 2. If the returned is a specific deployment (Dict), verify and return directly if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: @@ -13468,6 +13951,7 @@ class Router: ) litellm_params: Final = healthy_deployments.get("litellm_params", {}) if litellm_params.get("use_in_pass_through"): + self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments) return healthy_deployments else: # Specific deployment does not support pass-through @@ -13527,7 +14011,6 @@ class Router: ) # 6. Apply load balancing strategy - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if strategy == "simple-shuffle": return simple_shuffle( llm_router_instance=self, @@ -13559,6 +14042,7 @@ class Router: enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, ) + self._override_selector_pre_call_check(strategy, strategy_selector, deployment) verbose_router_logger.info( "get_available_deployment_for_pass_through model: %s, selected deployment: %s", @@ -13618,6 +14102,7 @@ class Router: self, healthy_deployments: list[dict], parent_otel_span: Span | None = None, + health_check_probe: bool = False, ) -> list[dict]: """ Filter out deployments marked unhealthy by background health checks. @@ -13654,8 +14139,7 @@ class Router: ] if not filtered: - verbose_router_logger.warning("All deployments marked unhealthy by health checks, bypassing health filter") - return healthy_deployments + return [] if health_check_probe else healthy_deployments # mutable-ok: empty list signals unavailable probe return filtered diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 1a33ea23bd4..7f376b46a8d 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -123,7 +123,12 @@ class AdaptiveRouter: self._cells[(rt, model)] = initial_cell(prefs, rt) async def load_state_from_db(self, prisma_client: Any) -> None: - """Override cold-start cells with persisted state. Called once at startup.""" + """Add each row's persisted delta to a freshly computed cold-start prior. + + A row holds an accumulated delta, not a full posterior, and can be one-sided + (e.g. beta=0) - assigning it straight into the cell would zero out a Beta shape + parameter and crash thompson_sample() on every later draw for that cell. + """ if prisma_client is None: return try: @@ -139,7 +144,12 @@ class AdaptiveRouter: continue if row.model_name not in self.config.available_models: continue - self._cells[(rt, row.model_name)] = BanditCell(alpha=row.alpha, beta=row.beta) + prefs = self.model_to_prefs.get(row.model_name) or _default_prefs() + prior = initial_cell(prefs, rt) + self._cells[(rt, row.model_name)] = BanditCell( + alpha=prior.alpha + row.alpha, + beta=prior.beta + row.beta, + ) loaded += 1 verbose_router_logger.info( "AdaptiveRouter[%s]: loaded %d cells from DB", @@ -451,6 +461,8 @@ class AdaptiveRouter: if d_alpha == 0 and d_beta == 0: continue cell_key = (attribution_type, target_model) + if cell_key not in self._cells: + continue self._cells[cell_key] = apply_delta( self._cells[cell_key], d_alpha, diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index 6b443026f61..d08afa8c1f6 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -2,6 +2,7 @@ Auto-Routing Strategy that works with a Semantic Router Config """ +import asyncio from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional @@ -75,6 +76,8 @@ class AutoRouter(CustomLogger): self.auto_sync_value = self.DEFAULT_AUTO_SYNC_VALUE self.loaded_routes: list[Route] = self._load_semantic_routing_routes() self.routelayer: SemanticRouter | None = None + self._routelayer_lock = asyncio.Lock() + self._routelayer_build_task: asyncio.Task[SemanticRouter] | None = None self.default_model = default_model self.embedding_model: str = embedding_model self.max_input_chars: int = max_input_chars @@ -115,6 +118,45 @@ class AutoRouter(CustomLogger): ) return auto_router_routes + def _build_routelayer(self) -> "SemanticRouter": + """Synchronous (embeds every route's utterances); run only via `_ensure_routelayer`.""" + if self.routelayer is not None: + return self.routelayer + + from semantic_router.routers import SemanticRouter + + routelayer: Final = SemanticRouter( + routes=self.loaded_routes, + encoder=self.encoder, + auto_sync=self.auto_sync_value, + ) + self.routelayer = routelayer + return routelayer + + def _clear_build_task_on_failure(self, build_task: "asyncio.Task[SemanticRouter]") -> None: + """Runs even with no caller left awaiting, so a failure never stays cached forever.""" + if build_task is self._routelayer_build_task and not build_task.cancelled() and build_task.exception(): + self._routelayer_build_task = None + + async def _ensure_routelayer(self) -> "SemanticRouter": + """Build the route layer once, off the event loop, shared across concurrent callers. + + A shared task (not a bare `asyncio.to_thread` awaited under the lock) survives one + caller's cancellation, so `cancel_on_disconnect` can't free a second caller into + starting a duplicate build. + """ + if self.routelayer is not None: + return self.routelayer + async with self._routelayer_lock: + if self.routelayer is not None: + return self.routelayer + build_task = self._routelayer_build_task + if build_task is None: + build_task = asyncio.ensure_future(asyncio.to_thread(self._build_routelayer)) + build_task.add_done_callback(self._clear_build_task_on_failure) + self._routelayer_build_task = build_task + return await asyncio.shield(build_task) + @staticmethod def _extract_text_from_messages(messages: list[dict[str, Any]]) -> str: """ @@ -151,8 +193,6 @@ class AutoRouter(CustomLogger): Used for the litellm auto-router to modify the request before the routing decision is made. """ - from semantic_router.routers import SemanticRouter - from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages from litellm.types.router import PreRoutingHookResponse @@ -164,17 +204,7 @@ class AutoRouter(CustomLogger): if resolved_messages is None: return None - routelayer = self.routelayer - if routelayer is None: - ####################### - # Create the route layer - ####################### - routelayer = SemanticRouter( - routes=self.loaded_routes, - encoder=self.encoder, - auto_sync=self.auto_sync_value, - ) - self.routelayer = routelayer + routelayer = await self._ensure_routelayer() message_content: Final = self._extract_text_from_messages(resolved_messages) route_name: Final = await self._matched_route_name(routelayer, message_content, request_kwargs) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 8235761ca98..686d57e2b77 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -3,12 +3,13 @@ Base class across routing strategies to abstract commmon functions like batch in """ import asyncio +import logging from abc import ABC from typing import Final from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.caching.redis_cache import RedisPipelineIncrementOperation, log_redis_failure from litellm.constants import DEFAULT_REDIS_SYNC_INTERVAL @@ -147,7 +148,7 @@ class BaseRoutingStrategy(ABC): return return_result except Exception as e: - verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e) + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) self.redis_increment_operation_queue = [] def add_to_in_memory_keys_to_update(self, key: str): diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index a8d51f95e45..3e094df7ac8 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -20,6 +20,7 @@ anthropic: import asyncio import builtins +import logging from collections.abc import Mapping from datetime import datetime, timedelta, timezone from typing import Any, Final @@ -27,7 +28,7 @@ from typing import Any, Final import litellm from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, @@ -92,6 +93,13 @@ class _LiteLLMParamsDictView: return dict(self._params) +async def _push_increments_to_redis(redis_cache: RedisCache, queued: list[RedisPipelineIncrementOperation]) -> None: + try: + await redis_cache.async_increment_pipeline(increment_list=queued) + except Exception as e: + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) + + class RouterBudgetLimiting(CustomLogger): def __init__( self, @@ -536,17 +544,13 @@ class RouterBudgetLimiting(CustomLogger): "Pushing Redis Increment Pipeline for queue: %s", self.redis_increment_operation_queue, ) - if len(self.redis_increment_operation_queue) > 0: - asyncio.create_task( - self.dual_cache.redis_cache.async_increment_pipeline( - increment_list=self.redis_increment_operation_queue, - ) - ) - + queued: Final = self.redis_increment_operation_queue self.redis_increment_operation_queue = [] + if queued: + asyncio.create_task(_push_increments_to_redis(self.dual_cache.redis_cache, queued)) except Exception as e: - verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e) + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) async def _sync_in_memory_spend_with_redis(self): """ @@ -601,7 +605,7 @@ class RouterBudgetLimiting(CustomLogger): verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, value) except Exception as e: - verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e) + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) def _get_budget_config_for_deployment( self, diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 93dddfb3d20..4aa342ea59f 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -195,6 +195,40 @@ model_list: session_affinity_ttl_seconds: 300 ``` +## Custom dimensions + +Add `custom_dimensions` under `complexity_router_config` to give domain keywords or regex patterns their own weighted signal + +```yaml +custom_dimensions: + - name: internalFrameworks + weight: 0.9 + keywords: [orbitmesh, fluxgate] + - name: sqlMigration + weight: 0.7 + patterns: ['\b(create|alter|drop)\s{1,4}table\b'] + - name: dataPipeline + weight: 0.4 + scoring_mode: match_count + keywords: [airflow, dbt, snowflake] +``` + +Each dimension contributes its weight once when any matcher hits the current ask. Repeated matches do not increase it. The built-in score and tier boundaries are unchanged, and the total score is not renormalized. Keywords use the existing case-insensitive word-boundary and CJK rules. Regexes search the first 2048 characters case-insensitively and compile during configuration validation and router initialization, never per request + +`scoring_mode` is optional and defaults to `binary`, the behavior above. `match_count` grades the dimension by how many distinct matchers hit: none scores 0 and emits no signal, one scores half the weight, two or more score the full weight. Repeated occurrences of one matcher never raise the count, keywords are distinct case-insensitively, patterns are distinct by source, and a keyword and a pattern are always distinct from each other. Matching stops as soon as the selected mode's maximum is reached, so a binary dimension still stops at its first hit. Existing configurations without the field keep binary scoring and the same tuning fingerprint, so the field only counts as a tuning change when set to `match_count` + +### Weights through the API versus the dashboard + +The API and YAML store exactly the weights written. A `dimension_weights` map and inline custom weights are read literally, missing recognized built-in names score zero, and nothing renormalizes the vector, so a total other than 1 is legal and scores accordingly. The dashboard's heuristic scoring editor is the one place that rebalances: editing one weight there holds it and redistributes the remainder across the other active dimensions in the draft, then Save sends the resulting explicit values, which the backend stores and scores as written. Opening a router, applying a preset, editing matchers, changing `scoring_mode`, or saving unrelated fields never normalizes existing weights + +Only `heuristic`, `heuristic_first` and `hybrid` accept custom dimensions. Each name must be a unique ASCII identifier starting with a letter, at most 64 characters, and cannot reuse a built-in dimension name or a key in `dimension_weights`. Set its weight inline, greater than zero and at most one + +Patterns are checked at configuration time against a grammar whose worst case stays a few milliseconds on 2048 characters. Every quantifier needs an explicit upper bound of at most 64 and must repeat a single character or character class, so `\s{1,4}` is accepted while `\s+`, `(a|aa){0,12}` and `(?:ab){0,64}` are refused. Backreferences, lookarounds, atomic groups and possessive quantifiers are refused as well. Each pattern is then costed: alternation branches and repeat lengths multiply the ways the engine can retry, and every later piece of the pattern is charged once per path that can reach it, so `a?a?a?a?a?a?a?a?` followed by a long fixed tail is refused even though each quantifier is small. The budget is 2048 work units per pattern and 8192 across the router. An invalid or over-budget pattern fails the write with a message naming the pattern and the rule it broke + +Limits are 16 dimensions, 32 combined keywords/patterns per dimension, 256 characters per matcher and 4096 matcher characters per dimension. Matching runs inline on the request path with no timeout and no worker thread, because the grammar is what bounds the cost. These are routing hints, not security enforcement rules + +The existing heuristic-v1 tuning quota covers custom dimensions, their weights and their scoring mode: one changed router without an auto-router license, unlimited with the entitlement. Omitting `custom_dimensions` preserves existing scoring. Routing decisions and spend logs include signals such as `custom (sqlMigration)` without recording the configured pattern or matched text. The field is configured through YAML, the model API, or the dashboard's heuristic scoring editor + ## Usage Once configured, use the model name like any other: @@ -236,6 +270,18 @@ change or default takeover records `cause: modality_escalation` with the displac pinned by session affinity, and by default a KEPT session pin bypasses the gate: a session pinned to a text-only model keeps it even when an image arrives. +Context-window and modality recovery take priority over the default model. If a compatible tier +cannot serve, the router checks the remaining compatible recovery tiers before using `default_model`. +A capacity failure without those constraints tries the selected tier's peers, then the default + +The default must fit the context and accept the request's modality. It cannot bypass routing plugins +or a plan-mode floor. Context fit uses the auto-router's existing buffer even when Router-wide pre-call +checks are off. Missing context metadata retains the existing unknown-window behavior + +Health fallback records `cause: health_default_fallback` and `health_displaced:` in `signals`. +It does not replace the session's tier pin. Adaptive feedback retains the model that actually served, +but a default outside the adaptive candidate pool does not become a normal candidate + Add `modality_pin_override: true` to lift that last exemption. The image turn is then re-placed the same way every other decision is, and records `cause: modality_pin_override` whether or not the tier moved, since the model left the pin either way. The pin itself is untouched: the session @@ -327,6 +373,19 @@ model_list: keep the classifier deployment or provider default, or set a supported value such as `none` or `low` to override that call. +When the current ask is a Responses API `agent_message` containing `encrypted_content`, LLM +classification preserves the encrypted task and uses native Responses. This also bypasses the +local scoring shortcut in `heuristic_first` and `hybrid` modes. The configured classifier must use +a native OpenAI or Azure OpenAI Responses deployment with access to the encrypted content. The +provider handles the encrypted task, and the classifier still chooses the tier dynamically + +Compatibility is checked after normal deployment selection. A paused incompatible member of the +classifier group does not prevent an eligible compatible deployment from classifying the task + +Unsupported classifier deployments and provider decryption errors use the existing +`classifier_fallback` policy. No fixed tier is introduced for encrypted tasks. Plaintext asks and +requests carrying only historical encrypted reasoning retain the existing classifier path + Classifier calls have a one-attempt hard deadline. After a timeout, the router opens a process-local circuit for that classifier and sends every session through `classifier_fallback` for `classifier_llm_config.circuit_breaker_cooldown_seconds` (30 seconds by default). When the cooldown @@ -408,11 +467,26 @@ If 2+ reasoning markers are detected in the user message, the request is promote Reasoning markers in the system prompt do **not** trigger the reasoning override. This prevents system prompts like "Think step by step before answering" from forcing all requests to the reasoning tier. +For requests identified by a `claude-cli/` or `claude-code/` user agent, the LLM classifier omits caller system +text to avoid classifying environment, agent, and skill catalogs. The current ask, configured prior-turn context, +and trajectory signal remain unchanged. The routed completion still receives the original system text. This +also excludes genuine task constraints supplied only in Claude Code system messages. Other clients keep the +existing system-context behavior. The browser routing preview has no client-identity field and retains that +generic behavior; use the real client when checking Claude Code routing. + ### Harness Reminder Blocks Agent harnesses inject their own context into the conversation as ordinary message text. That text is plumbing, not something a human asked for, so the router strips complete reminder blocks before classifying and picking a tier. A turn that is nothing but a reminder block strips to empty and is skipped, and the router falls back to the last real ask instead -By default a block is anything between `` and ``. `reminder_markers` replaces that with your harness's own delimiters. Many harnesses use a different envelope per agent type, so list every pair you emit: +By default the router strips complete `` blocks. For requests with a Codex user agent, it also strips complete ``, ``, ``, and `` blocks, plus repository instructions from the fixed heading prefix `# AGENTS.md instructions for ` through ``, regardless of the repository path. Other clients keep those tags and their contents + +The proxy records the incoming user agent in request metadata. SDK callers can supply `metadata.user_agent` (or `litellm_metadata.user_agent` on Responses requests), or configure `reminder_markers` explicitly when their client identity is unavailable + +The Codex `Message Type: NEW_TASK` wrapper and its delegated-task payload remain available for classification. Cleanup applies to the current ask and quoted prior turns; the routed request retains its original content + +In `classification_mode: user_turn`, complete text-only reminder tails leave the preceding fresh ask eligible for classification. Assistant turns and tool results still mark continuations, including tool results carried alongside reminder text + +`reminder_markers` replaces these defaults with your harness's own delimiters. Many harnesses use a different envelope per agent type, so list every pair you emit: ```yaml model_list: @@ -427,7 +501,7 @@ model_list: close: "[[SUBAGENT_CONTEXT_END]]" ``` -Setting `reminder_markers` replaces the built-in `` pair rather than adding to it, so list that pair too if your harness also emits it. Matching is case-insensitive. Blocks that nest or overlap across pairs are stripped whole. An unclosed delimiter is not a block and is left in place, which keeps prose that merely mentions a delimiter from being eaten +Setting `reminder_markers` replaces all built-in pairs, including the Codex heading pair, so include every default your harness still needs. Matching is case-insensitive. Blocks that nest or overlap across pairs are stripped whole. An unclosed delimiter is not a block and is left in place, which keeps prose that merely mentions a delimiter from being eaten ### Code Detection diff --git a/litellm/router_strategy/complexity_router/classification_rubrics.py b/litellm/router_strategy/complexity_router/classification_rubrics.py index 9f168eabbc4..1dae2902fad 100644 --- a/litellm/router_strategy/complexity_router/classification_rubrics.py +++ b/litellm/router_strategy/complexity_router/classification_rubrics.py @@ -108,6 +108,11 @@ _CALIBRATION_EXAMPLES: Final[Mapping[ClassificationRubric, str]] = MappingProxyT BUSINESS_TIER_CRITERIA: Final[Mapping[ComplexityTier, str]] = MappingProxyType( { + ComplexityTier.NON_REASONING: ( + "operational requests whose whole job is to pass information along or put it in a requested " + "shape: relaying or reformatting tool or system output, acknowledging a completed action, or " + "extracting a stated field. Use it only when no judgment about the content is asked for." + ), ComplexityTier.SIMPLE: ( "greetings, chitchat, or lookups of a fact, policy, price, or date with a short known answer. " "Never for analysis, strategy, or non-trivial work, even if the request is only one sentence." diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 98a1eb7ac9e..9deccc9a468 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -20,24 +20,27 @@ import random import re import time from collections.abc import Callable, Iterator, Mapping, Sequence -from itertools import accumulate, islice, takewhile +from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast -from pydantic import BaseModel, create_model +from pydantic import BaseModel, TypeAdapter, ValidationError, create_model from litellm._logging import verbose_router_logger from litellm.constants import ( EMPTY_MAPPING, INTERNAL_CALL_ORIGIN_METADATA_KEY, + OUTPUT_TOKEN_CEILING_PARAMS, RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.classifier_logging import masked_originating_request from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_metadata_variable_name_from_kwargs, + is_codex_user_agent, ) from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -45,6 +48,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( request_contains_image_content, ) from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload +from litellm.llms.anthropic.common_utils import is_claude_code_user_agent from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.complexity_router.tier_predictor import ( @@ -55,6 +59,7 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionImageObject, ChatCompletionTextObject, + ResponsesAPIResponse, ) from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, @@ -67,6 +72,7 @@ from litellm.types.utils import ( from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section from .config import ( CALIBRATION_EXAMPLES_HEADING, + CUSTOM_PATTERN_SCAN_CHARS, DEFAULT_CLASSIFICATION_RUBRIC, DEFAULT_CODE_KEYWORDS, DEFAULT_ESCALATION_KEYWORDS, @@ -81,6 +87,7 @@ from .config import ( ClassificationRubric, ComplexityRouterConfig, ComplexityTier, + CustomDimension, TierDefinition, ) from .stall_detector import detect_stalled_task @@ -115,8 +122,20 @@ def _tier_name(tier: ComplexityTier | str) -> str: return tier.value if isinstance(tier, ComplexityTier) else tier +def _built_in_tier_or_none(tier_name: str) -> ComplexityTier | None: + """The built-in tier a `tiers` key names, or None when the key is an operator-defined name.""" + return ComplexityTier.__members__.get(tier_name) + + _CLASSIFICATION_TIER_CRITERIA: Final[Mapping[ComplexityTier, str]] = MappingProxyType( { + ComplexityTier.NON_REASONING: ( + "operational requests whose whole job is to pass information along or put it in a " + "requested shape: relaying or reformatting tool output, acknowledging a completed action, " + "or extracting a stated value. Use it only when no judgment about the content is asked for; " + "the moment the request is to summarize, compare, explain, debug, or decide, it belongs " + "in a higher tier however short it is." + ), ComplexityTier.SIMPLE: ( "greetings, chitchat, or factual lookups with a short known answer. Do not use this tier for " "unsolved problems, proofs, deep theory, multi-step analysis, or non-trivial code, even if the " @@ -349,7 +368,7 @@ def _parent_session_kwargs(request_kwargs: Mapping[str, Any] | None) -> Mapping[ return {k: kwargs[k] for k in ("litellm_session_id", "litellm_trace_id") if kwargs.get(k) is not None} -def _response_cost_or_none(response: ModelResponse) -> float | None: +def _response_cost_or_none(response: ModelResponse | ResponsesAPIResponse) -> float | None: hidden_params: Final = response._hidden_params if not isinstance(hidden_params, dict): return None @@ -372,6 +391,13 @@ def _effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | N _REMINDER_OPEN: Final = "" _REMINDER_CLOSE: Final = "" _DEFAULT_REMINDER_MARKERS: Final = ((_REMINDER_OPEN, _REMINDER_CLOSE),) +_CODEX_REMINDER_MARKERS: Final = _DEFAULT_REMINDER_MARKERS + ( + ("", ""), + ("", ""), + ("", ""), + ("", ""), + ("# agents.md instructions for ", ""), +) _TRUNCATION_MARKER: Final = "..." _TRUNCATION_HEAD_FRACTION: Final = 0.3 @@ -471,6 +497,42 @@ def _human_text(content: object, marker_pairs: tuple[tuple[str, str], ...] = _DE return _strip_reminder_blocks(_message_text(content), marker_pairs) +def _encrypted_classifier_task( + request_kwargs: Mapping[str, object] | None, + marker_pairs: tuple[tuple[str, str], ...], +) -> dict[str, object] | None: + from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages + + raw_input: Final = (request_kwargs or EMPTY_MAPPING).get("input") + if not isinstance(raw_input, list) or (request_kwargs or EMPTY_MAPPING).get("messages"): + return None + try: + items: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(raw_input) + except ValidationError: + return None + current: Final = next( + ( + item + for item in reversed(items) + if (messages := resolve_structured_messages(messages=None, request_kwargs={"input": [item]})) + and any(_iter_human_asks_newest_first(messages, marker_pairs)) + ), + None, + ) + if current is None or current.get("type") != "agent_message" or not isinstance(current.get("content"), list): + return None + try: + parts: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(current["content"]) + except ValidationError: + return None + if not any(part.get("type") == "encrypted_content" and part.get("encrypted_content") for part in parts): + return None + return { + **current, + "content": [part for part in parts if part.get("type") in ("input_text", "encrypted_content")], + } + + def _iter_human_asks_newest_first( messages: Sequence[Mapping[str, object]], marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, @@ -576,6 +638,18 @@ def _last_human_ask_index( ) +def _is_reminder_only_turn(message: Mapping[str, object], marker_pairs: tuple[tuple[str, str], ...]) -> bool: + if message.get("role") != "user": + return False + content: Final = message.get("content") + if not isinstance(content, str) and not ( + isinstance(content, list) and all(isinstance(part, Mapping) and part.get("type") == "text" for part in content) + ): + return False + text: Final = _message_text(content) + return bool(text.strip()) and not _strip_reminder_blocks(text, marker_pairs) + + def _newest_turn_is_human_ask( messages: Sequence[Mapping[str, object]] | None, marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, @@ -586,21 +660,23 @@ def _newest_turn_is_human_ask( Anchored on `_last_human_ask_index` so every surface's plumbing reads as a continuation: chat-completions tool turns are role=tool, Messages-surface tool_result turns flatten to empty human text, and a hybrid turn carrying an ask alongside a tool_result still counts as an ask. - Compared against the newest non-system message rather than the raw tail, because Claude Code - appends a system-role reminder after the human turn; that trailing plumbing is neither an ask - nor loop traffic and must not turn a fresh ask into a continuation. An unreadable request (no - messages) is treated as a continuation: there is no ask to classify, which is the same reading - `_extract_current_ask_and_system_prompt` gives it downstream. + Trailing system messages and complete text-only reminders do not turn a fresh ask into a + continuation. Assistant turns and non-text content, including tool results alongside reminders, + still form continuation boundaries. An unreadable request has no ask to classify. """ if not messages: return False - newest_non_system: Final = next( - (index for index in range(len(messages) - 1, -1, -1) if messages[index].get("role") != "system"), + newest_activity: Final = next( + ( + index + for index in range(len(messages) - 1, -1, -1) + if messages[index].get("role") != "system" and not _is_reminder_only_turn(messages[index], marker_pairs) + ), None, ) - if newest_non_system is None: + if newest_activity is None: return False - return _last_human_ask_index(messages, marker_pairs) == newest_non_system + return _last_human_ask_index(messages, marker_pairs) == newest_activity def _iter_system_scope_texts( @@ -861,6 +937,7 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo "modality_escalation", "modality_pin_override", "health_failover", + "health_default_fallback", ) and not decision.get("context_escalated") and _CLASSIFIER_CIRCUIT_OPEN_SIGNAL not in (decision.get("signals") or ()) @@ -878,6 +955,15 @@ class DimensionScore: self.signal = signal +class _CustomDimensionMatchers(NamedTuple): + """One custom dimension's distinct matchers and the number of hits that saturates its score.""" + + dimension: CustomDimension + keywords: tuple[str, ...] + patterns: tuple[re.Pattern[str], ...] + saturation: int + + class KeywordOverride(NamedTuple): """A keyword_tier_rules match: the winning tier and, on the lexical path, the keyword that fired.""" @@ -1013,6 +1099,15 @@ def _group_provably_fits(facts: tuple[int | None, bool], needed: int, buffer: fl return window is not None and not has_unknown and needed <= int(window * buffer) +class _RequestContextFit(NamedTuple): + facts: Mapping[str, tuple[int | None, bool]] + needed: int | None + buffer: float + + def accepts(self, model: str) -> bool: + return self.needed is None or _window_can_hold(self.facts.get(model, (None, True))[0], self.needed, self.buffer) + + class _ContextWindowPlacement(NamedTuple): """Where the context-window gate placed the request: the placement tier, the subset of its pool the pick may use, and every configured group not provably misfit (the adaptive filter).""" @@ -1119,6 +1214,15 @@ class ComplexityRouter(CustomLogger): self.config.custom_technical_keywords, ) self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS + self._custom_dimensions = tuple( + _CustomDimensionMatchers( + dimension, + tuple(dict.fromkeys(keyword.lower() for keyword in dimension.keywords)), + tuple(re.compile(pattern, re.IGNORECASE) for pattern in dict.fromkeys(dimension.patterns)), + 2 if dimension.scoring_mode == "match_count" else 1, + ) + for dimension in self.config.custom_dimensions + ) if self.config.has_custom_tiers: self.escalation_keywords: tuple[str, ...] = () elif self.config.escalation_keywords is not None: @@ -1223,7 +1327,7 @@ class ComplexityRouter(CustomLogger): """ if self.config.has_custom_tiers: return tuple(dict.fromkeys(model for models in self._tier_pools().values() for model in models)) - for tier in reversed(TIER_SEVERITY_ORDER): + for tier in reversed(self.config.active_tier_severity_order()): models = self.config.tiers.get(tier.value) if models: return tuple(models) if isinstance(models, list) else (models,) @@ -1320,6 +1424,28 @@ class ComplexityRouter(CustomLogger): score: Final = score_high if match_count >= high_threshold else score_low return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count + def _count_custom_hits(self, matchers: _CustomDimensionMatchers, user_text: str, scanned: str) -> int: + hits: Final = chain( + (self._keyword_matches(user_text, keyword) for keyword in matchers.keywords), + (pattern.search(scanned) is not None for pattern in matchers.patterns), + ) + return sum(islice((1 for hit in hits if hit), matchers.saturation)) + + def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]: + if not self._custom_dimensions: + return () + scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS] + return tuple( + ( + DimensionScore( + matchers.dimension.name, hits / matchers.saturation, f"custom ({matchers.dimension.name})" + ), + matchers.dimension.weight, + ) + for matchers in self._custom_dimensions + if (hits := self._count_custom_hits(matchers, user_text, scanned)) + ) + def _score_multi_step(self, text: str) -> DimensionScore: """Score based on multi-step patterns.""" hits: Final = sum(1 for p in self._multi_step_patterns if p.search(text)) @@ -1415,12 +1541,13 @@ class ComplexityRouter(CustomLogger): self._score_question_complexity(prompt), ] - # Collect signals - signals: Final = [d.signal for d in dimensions if d.signal is not None] + custom_dimensions: Final = self._score_custom_dimensions(prompt, user_text) + signals: Final = [d.signal for d in (*dimensions, *(d for d, _ in custom_dimensions)) if d.signal is not None] - # Compute weighted score weights: Final = self.config.dimension_weights - weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + sum( + dimension.score * weight for dimension, weight in custom_dimensions + ) boundaries: Final = self._effective_tier_boundaries() clears_override_floor: Final = weighted_score >= self._effective_reasoning_override_min_score() @@ -1574,6 +1701,10 @@ class ComplexityRouter(CustomLogger): return self._classify_with_heuristic_v2(prompt) if self.config.classifier_type == "custom": return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages) + if self.config.classifier_type in ("heuristic_first", "hybrid") and _encrypted_classifier_task( + request_kwargs, self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) + ): + return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages) if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None: return await self._classify_heuristic_first(prompt, system_prompt, request_kwargs, messages) if self.config.classifier_type == "hybrid" and self.config.classifier_llm_config is not None: @@ -1846,7 +1977,11 @@ class ComplexityRouter(CustomLogger): default_model: Final = self.config.default_model pools: Final = self._tier_pools() tier: Final = next( - (candidate for candidate in TIER_SEVERITY_ORDER if default_model in pools.get(candidate.value, ())), + ( + candidate + for candidate in self.config.active_tier_severity_order() + if default_model in pools.get(candidate.value, ()) + ), ComplexityTier.MEDIUM, ) return ClassificationOutcome( @@ -1864,16 +1999,12 @@ class ComplexityRouter(CustomLogger): Call the configured classifier model with a system/user role split and prior-turn context. Builds a structured classification prompt with: - - System message: the stable classifier rubric AND the caller's own system prompt (task - constraints). This is the largest, most repeated part of the call, so keeping it in the - system role lets the provider prompt-cache it across a session's classifier calls. - - User message: the variable payload -- a few prior user turns for context and the current - ask to classify. + - System message: the stable classifier rubric. + - User message: the caller's system prompt quoted as task context, prior turns, and the current ask. Args: prompt: The current user ask text (already extracted as the real human ask, not tool results) - system_prompt: The caller's system prompt (task constraints), always included so later - turns never lose it + system_prompt: Caller task constraints, omitted from classification for Claude Code requests request_kwargs: Request metadata for spend attribution messages: Full message history for extracting prior turns and the trajectory signal """ @@ -1884,6 +2015,7 @@ class ComplexityRouter(CustomLogger): raise ValueError("classifier_llm_config is not set") include_assistant: Final = self.config.classifier_context_include_assistant_turns + marker_pairs: Final = self._reminder_markers_for_request(request_kwargs or {}) context_enabled: Final = bool(messages) and self.config.classifier_context_window_size > 0 prior_turns: Final = ( _extract_prior_turns( @@ -1893,26 +2025,30 @@ class ComplexityRouter(CustomLogger): budget_chars=self.config.classifier_context_budget_chars, per_turn_chars=self.config.classifier_context_per_turn_chars, include_assistant=include_assistant, - marker_pairs=self._reminder_markers, + marker_pairs=marker_pairs, ) if context_enabled else () ) has_prior_conversation: Final = ( context_enabled - and len( - tuple( - islice( - _iter_context_turns_newest_first(messages or (), include_assistant, self._reminder_markers), 2 - ) - ) - ) + and len(tuple(islice(_iter_context_turns_newest_first(messages or (), include_assistant, marker_pairs), 2))) > 1 ) + encrypted_task: Final = _encrypted_classifier_task(request_kwargs, marker_pairs) + caller_system_prompt: Final = ( + None + if any( + is_claude_code_user_agent(user_agent) + for metadata in (self._iter_metadata_dicts(request_kwargs) if request_kwargs is not None else ()) + if isinstance(user_agent := metadata.get("user_agent"), str) + ) + else system_prompt + ) user_payload: Final = self._build_classifier_user_payload( - prompt=prompt, - system_prompt=system_prompt, + prompt="The delegated task in the following agent_message." if encrypted_task is not None else prompt, + system_prompt=caller_system_prompt, prior_turns=prior_turns, messages=messages, has_prior_conversation=has_prior_conversation, @@ -1944,34 +2080,42 @@ class ComplexityRouter(CustomLogger): if llm_config.reasoning_effort is not None: classifier_call_params = MappingProxyType({"reasoning_effort": llm_config.reasoning_effort}) + payload: Final = ( + self._native_classifier_payload(messages_for_call, response_format, encrypted_task) + if encrypted_task is not None + else MappingProxyType( + {"messages": messages_for_call, "response_format": response_format, **classifier_call_params} + ) + ) proxy_server_request: Final = { - "body": { - "model": llm_config.model, - "messages": messages_for_call, - "response_format": response_format, - **classifier_call_params, - } + "originating_request_masked": masked_originating_request(request_kwargs), + "body": {"model": llm_config.model, **payload}, } + classify: Final = ( + self.litellm_router_instance.aresponses + if encrypted_task is not None + else self.litellm_router_instance.acompletion + ) classifier_timeout_s: Final[float] = llm_config.timeout_ms / 1000 - response: Final[ModelResponse] = await asyncio.wait_for( - self.litellm_router_instance.acompletion( + response: Final[ModelResponse | ResponsesAPIResponse] = await asyncio.wait_for( + classify( model=llm_config.model, - messages=messages_for_call, stream=False, - response_format=response_format, timeout=classifier_timeout_s, num_retries=0, disable_fallbacks=True, metadata=metadata, proxy_server_request=proxy_server_request, turn_off_message_logging=turn_off_message_logging, - **classifier_call_params, + **payload, **_parent_session_kwargs(request_kwargs), ), timeout=classifier_timeout_s, ) - content: Final = response.choices[0].message.content + content: Final = ( + response.output_text if isinstance(response, ResponsesAPIResponse) else response.choices[0].message.content + ) if not content: raise ValueError("LLM classifier returned empty content") raw_tier: Final = _LabeledTierClassification.model_validate_json(content).tier @@ -1980,6 +2124,33 @@ class ComplexityRouter(CustomLogger): raise ValueError(f"LLM classifier returned an unrecognized tier: {raw_tier!r}") return tier, _response_cost_or_none(response) + def _native_classifier_payload( + self, + messages: list[AllMessageValues], # mutable-ok: existing transformation accepts the SDK message list + response_format: Mapping[str, object], + encrypted_task: Mapping[str, object], + ) -> Mapping[str, object]: + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + transformation: Final = LiteLLMResponsesTransformationHandler() + input_items, instructions = transformation.convert_chat_completion_messages_to_responses_api(messages) + llm_config: Final = self.config.classifier_llm_config + reasoning: Final = ( + {"reasoning": {"effort": llm_config.reasoning_effort}} + if llm_config is not None and llm_config.reasoning_effort is not None + else {} + ) + return { + "input": [*input_items, encrypted_task], + "instructions": instructions, + "text": transformation.transform_response_format_to_text_format(dict(response_format)), + "store": False, + "_require_encrypted_task_support": True, + **reasoning, + } + @staticmethod def _build_classifier_user_payload( prompt: str, @@ -2067,11 +2238,15 @@ class ComplexityRouter(CustomLogger): raise ValueError(f"No model configured for tier {tier_key} and no default_model set") def _litellm_params_for_model(self, tier: ComplexityTier | str | None, model: str) -> Mapping[str, object]: - if tier is None: - return MappingProxyType({}) - entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) + entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) if tier is not None else () entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None) - return entry.litellm_params if entry is not None else MappingProxyType({}) + explicit: Final = entry.litellm_params if entry is not None else MappingProxyType({}) + if not self.config.max_tokens_from_tier_model or not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(explicit): + return explicit + ceiling: Final = self._group_output_ceiling(model) + if ceiling is None: + return explicit + return MappingProxyType({**explicit, "max_tokens": ceiling}) @staticmethod def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str: @@ -2174,9 +2349,12 @@ class ComplexityRouter(CustomLogger): else: model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[]) + # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set. lp = deployment.get("litellm_params") if isinstance(deployment, dict) else deployment.litellm_params lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) cost = lp_dict.get("input_cost_per_token") + if cost is None: + cost = mi_dict.get("input_cost_per_token") model_to_cost[name] = float(cost) if cost is not None else 0.0 self.adaptive_router = AdaptiveRouter( @@ -2228,7 +2406,8 @@ class ComplexityRouter(CustomLogger): return self._fitting_tier_fallback(classified_tier, fit_filter) request_type: Final = classify_prompt(user_message) - classified_idx: Final = TIER_SEVERITY_ORDER.index(classified_tier) + severity_order: Final = self.config.active_tier_severity_order() + classified_idx: Final = severity_order.index(classified_tier) pools: Final = self._tier_pools() classified_candidates: Final = _allowed(tuple(pools.get(_tier_name(classified_tier), ())), fit_filter) cold_start_candidates: Final = tuple( @@ -2292,9 +2471,7 @@ class ComplexityRouter(CustomLogger): distance = 0 else: model_tiers = self._model_tiers.get(model, (classified_tier,)) - distance = min( - abs(TIER_SEVERITY_ORDER.index(model_tier) - classified_idx) for model_tier in model_tiers - ) + distance = min(abs(severity_order.index(model_tier) - classified_idx) for model_tier in model_tiers) score = quality_weight * quality_sample + cost_weight * cost_score - penalty_weight * distance candidate_scores.append( { @@ -2361,7 +2538,7 @@ class ComplexityRouter(CustomLogger): body if isinstance(body, Mapping) else None, resolved_messages, tuple(self.config.plan_mode_patterns or ()), - self._reminder_markers, + self._reminder_markers_for_request(request_kwargs), ) def _matched_housekeeping_sentinel(self, newest_ask: str | None) -> str | None: @@ -2403,12 +2580,15 @@ class ComplexityRouter(CustomLogger): return name if self.config.has_custom_tiers else ComplexityTier(name) def _deployment_window(self, group: str, deployment: Mapping[str, object]) -> int | None: + return self._deployment_limit(group, deployment, "max_input_tokens") + + def _deployment_limit( + self, group: str, deployment: Mapping[str, object], key: Literal["max_input_tokens", "max_output_tokens"] + ) -> int | None: from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider deployment_model_info: Final = deployment.get("model_info") - declared: Final = ( - deployment_model_info.get("max_input_tokens") if isinstance(deployment_model_info, Mapping) else None - ) + declared: Final = deployment_model_info.get(key) if isinstance(deployment_model_info, Mapping) else None if isinstance(declared, int): return declared litellm_params: Final = deployment.get("litellm_params") @@ -2425,18 +2605,34 @@ class ComplexityRouter(CustomLogger): deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts received_model_name=group, ) - window: Final = model_info.get("max_input_tokens") + limit: Final = model_info.get(key) except Exception: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others return None - return window if isinstance(window, int) else None + return limit if isinstance(limit, int) else None + + def _group_deployments(self, group: str) -> Sequence[Mapping[str, object]]: + list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None) + deployments: Final = list_models(model_name=group) if callable(list_models) else None + return tuple(deployments) if isinstance(deployments, list) else () + + def _group_output_ceiling(self, group: str) -> int | None: + """Smallest max_output_tokens across the group's deployments, or None when any deployment + declares none: the core router picks within the group without a fit check, and a ceiling + above an unmapped member's real limit is a provider 400 on that member.""" + deployments: Final = self._group_deployments(group) + ceilings: Final = tuple( + ceiling + for deployment in deployments + if (ceiling := self._deployment_limit(group, deployment, "max_output_tokens")) is not None + ) + return min(ceilings) if ceilings and len(ceilings) == len(deployments) else None def _group_window_facts(self, group: str) -> tuple[int | None, bool]: """(smallest declared context window across the group's deployments, whether any deployment declares none). The core router picks a deployment within the group without a fit check, so the group is only as safe as its smallest member.""" - list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None) - deployments: Final = list_models(model_name=group) if callable(list_models) else None - if not isinstance(deployments, list) or not deployments: + deployments: Final = self._group_deployments(group) + if not deployments: return (None, True) windows: Final = tuple( window for deployment in deployments if (window := self._deployment_window(group, deployment)) is not None @@ -2483,24 +2679,44 @@ class ComplexityRouter(CustomLogger): """Real-tokenizer count of the resolved messages plus the out-of-band carriers, off the event loop; None when counting fails, and the gate then leaves the placement alone.""" import litellm - from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.token_counter import offload_token_count out_of_band: Final = self._out_of_band_request_text(request_kwargs) try: - counted: Final = await asyncify(litellm.token_counter)( + counted: Final = await offload_token_count(litellm.token_counter)( messages=cast(list, resolved_messages) # cast-ok: token_counter only iterates the sequence ) - return counted + (await asyncify(litellm.token_counter)(text=out_of_band) if out_of_band else 0) + return counted + (await offload_token_count(litellm.token_counter)(text=out_of_band) if out_of_band else 0) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.debug("ComplexityRouter: context-window token count failed. Got - %s", e) return None + async def _request_context_fit( + self, + resolved_messages: Sequence[Mapping[str, object]] | None, + request_kwargs: Mapping[str, object], + ) -> _RequestContextFit: + if not self.config.enable_context_window_escalation or not resolved_messages: + return _RequestContextFit(EMPTY_MAPPING, None, self.config.context_window_escalation_buffer) + names: Final = frozenset(model for pool in self._tier_pools().values() for model in pool) | frozenset( + (self.config.default_model,) if self.config.default_model else () + ) + facts: Final = MappingProxyType({name: self._group_window_facts(name) for name in names}) + known: Final = tuple(window for window, _ in facts.values() if window is not None) + buffer: Final = self.config.context_window_escalation_buffer + needs_count: Final = known and self._request_byte_upper_bound(resolved_messages, request_kwargs) > int( + min(known) * buffer + ) + needed: Final = await self._counted_request_tokens(resolved_messages, request_kwargs) if needs_count else None + return _RequestContextFit(facts=facts, needed=needed, buffer=buffer) + async def _context_window_placement( self, tier: ComplexityTier | str, resolved_messages: Sequence[Mapping[str, object]] | None, request_kwargs: Mapping[str, object], pool_override: tuple[str, ...] | None = None, + context_fit: _RequestContextFit | None = None, ) -> _ContextWindowPlacement | None: """Correct a decided placement whose models provably cannot hold the prompt, or None (the placement stands). Only a real tokenizer count ever moves a request, escalation @@ -2512,17 +2728,10 @@ class ComplexityRouter(CustomLogger): pool: Final = pool_override if pool_override is not None else tuple(pools.get(_tier_name(tier), ())) if not pool: return None - facts: Final = MappingProxyType({group: self._group_window_facts(group) for group in pool}) - known_windows: Final = tuple(window for window, _ in facts.values() if window is not None) - if not known_windows: + fit: Final = context_fit or await self._request_context_fit(resolved_messages, request_kwargs) + if fit.needed is None: return None - buffer: Final = self.config.context_window_escalation_buffer - if self._request_byte_upper_bound(resolved_messages, request_kwargs) <= int(min(known_windows) * buffer): - return None - needed: Final = await self._counted_request_tokens(resolved_messages, request_kwargs) - if needed is None: - return None - return self._placement_for_tokens(tier=tier, pool=pool, pools=pools, facts=facts, needed=needed) + return self._placement_for_tokens(tier=tier, pool=pool, pools=pools, facts=fit.facts, needed=fit.needed) def _placement_for_tokens( self, @@ -2534,14 +2743,16 @@ class ComplexityRouter(CustomLogger): needed: int, ) -> _ContextWindowPlacement | None: buffer: Final = self.config.context_window_escalation_buffer - in_tier: Final = tuple(group for group in pool if _window_can_hold(facts[group][0], needed, buffer)) + in_tier: Final = tuple( + group for group in pool if _window_can_hold(facts.get(group, (None, True))[0], needed, buffer) + ) if in_tier and len(in_tier) == len(pool): return None holdable: Final = frozenset( group for tier_pool in pools.values() for group in tier_pool - if _window_can_hold(self._group_window_facts(group)[0], needed, buffer) + if _window_can_hold(facts.get(group, (None, True))[0], needed, buffer) ) if in_tier: return _ContextWindowPlacement(tier=tier, allowed_models=in_tier, holdable_models=holdable) @@ -2549,7 +2760,7 @@ class ComplexityRouter(CustomLogger): proven = tuple( group for group in pools.get(name, ()) - if _group_provably_fits(self._group_window_facts(group), needed, buffer) + if _group_provably_fits(facts.get(group, (None, True)), needed, buffer) ) if proven: return _ContextWindowPlacement( @@ -2590,10 +2801,15 @@ class ComplexityRouter(CustomLogger): def _tier_for_model(self, model: str) -> ComplexityTier | None: """Return the most-severe configured tier whose pool contains this model.""" pools: Final = self._tier_pools() - matched: Final = tuple(ComplexityTier(tier_name) for tier_name, models in pools.items() if model in models) + order: Final = self.config.active_tier_severity_order() + matched: Final = tuple( + tier + for tier_name, models in pools.items() + if model in models and (tier := _built_in_tier_or_none(tier_name)) is not None and tier in order + ) if not matched: return None - return max(matched, key=TIER_SEVERITY_ORDER.index) + return max(matched, key=order.index) def _escalate_tier(self, tier: ComplexityTier | str) -> ComplexityTier | str: """Bump a tier one step up to the next-higher configured tier. @@ -2608,9 +2824,10 @@ class ComplexityRouter(CustomLogger): if self.config.has_custom_tiers: return tier configured: Final = frozenset(self.config.tiers) - current_index: Final = TIER_SEVERITY_ORDER.index(tier) + order: Final = self.config.active_tier_severity_order() + current_index: Final = order.index(tier) higher_tiers: Final = tuple( - candidate for candidate in TIER_SEVERITY_ORDER[current_index + 1 :] if candidate.value in configured + candidate for candidate in order[current_index + 1 :] if candidate.value in configured ) return higher_tiers[0] if higher_tiers else tier @@ -2689,6 +2906,7 @@ class ComplexityRouter(CustomLogger): messages: list[dict[str, Any]] | None, # mutable-ok: forwarded verbatim to the list-typed re-pick resolved_messages: Sequence[Mapping[str, object]] | None, request_kwargs: dict, # mutable-ok: same shape the hook receives + context_fit: _RequestContextFit | None = None, ) -> PreRoutingHookResponse: """Replace a routed model that cannot accept this request's image input. @@ -2719,7 +2937,8 @@ class ComplexityRouter(CustomLogger): or self._model_accepts_image_input(response.model) ): return response - eligible: Final = self._modality_eligible_models() + fit: Final = context_fit or await self._request_context_fit(resolved_messages, request_kwargs) + eligible: Final = frozenset(name for name in self._modality_eligible_models() if fit.accepts(name)) names: Final = self.config.tier_names() pools: Final = self._tier_pools() decided: Final = decision.get("tier") if decision is not None else None @@ -2832,17 +3051,21 @@ class ComplexityRouter(CustomLogger): where the prompt never arrives as messages. Probed on a COPY of request_kwargs because the owner pops routing bookkeeping off the - dict it is handed (`_target_order`, `_excluded_deployment_ids`), and this is a - speculative question about a model that may never be picked. + dict it is handed (`_target_order`, `_excluded_deployment_ids`, + `_retry_skipped_deployment_ids`), and this is a speculative question about a model + that may never be picked. Every way the owner says "nothing here can serve this" is a negative verdict: no healthy deployment for the group at all (BadRequestError, which ContextWindowExceededError - subclasses), every deployment filtered out (RouterRateLimitError), and every deployment - over its RPM (RouterRateLimitErrorBasic). Anything else is unknown rather than negative, - so it reads as capacity: absent information must never decide the verdict. + subclasses), every deployment filtered out (RouterRateLimitError), every deployment over + its RPM (RouterRateLimitErrorBasic), and every deployment refused by a filter that reports + exhaustion as a bare ValueError naming a RouterErrors marker -- provider and deployment + budgets, and tag routing, which have no typed error of their own. Anything else is unknown + rather than negative, so it reads as capacity: absent information must never decide the + verdict. """ from litellm.exceptions import BadRequestError - from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic + from litellm.types.router import RouterErrors, RouterRateLimitError, RouterRateLimitErrorBasic probe_kwargs: Final = dict(request_kwargs) # mutable-ok: the owner pops routing keys off the dict it is handed try: @@ -2852,10 +3075,15 @@ class ComplexityRouter(CustomLogger): messages=messages, input=input, parent_otel_span=_get_parent_otel_span_from_kwargs(request_kwargs), + health_check_probe=True, ) - except (RouterRateLimitError, RouterRateLimitErrorBasic, BadRequestError): + except (RouterRateLimitError, RouterRateLimitErrorBasic, BadRequestError) as exc: + verbose_router_logger.debug("health probe unavailable model=%s error=%s", model_name, type(exc).__name__) return False except Exception as exc: # noqa: BLE001 # a speculative eligibility read must fail open on unknown faults + if isinstance(exc, ValueError) and any(marker.value in str(exc) for marker in RouterErrors): + verbose_router_logger.debug("health probe exhausted model=%s error=%s", model_name, exc) + return False verbose_router_logger.debug( "ComplexityRouter: eligibility probe for %s failed, treating the group as live: %s", model_name, exc ) @@ -2869,76 +3097,124 @@ class ComplexityRouter(CustomLogger): input: str | list | None, # mutable-ok: mirrors the owner's own input parameter, which this forwards verbatim resolved_messages: Sequence[Mapping[str, object]] | None, request_kwargs: dict, # mutable-ok: same shape the hook receives + context_fit: _RequestContextFit | None = None, ) -> PreRoutingHookResponse: - """Replace a decided model group that has no serving capacity with a live peer in the same tier. - - Applied to the decided response at the hook's exits, so every arm that can place a request - is covered by one owner: a fresh classification, a replayed or escalated session pin, a - plan-mode floor, a context-window escalation, an adaptive pick, and whatever arm is added - next. Peers come from the DECIDED tier only; climbing to another tier is deliberately not - done here, since a higher tier costs more than the classifier asked for. - - Serving capacity is one question asked of one owner (`_model_group_can_serve`), so the - substitute is only ever a group the pipeline would actually accept for this request. The - pick then runs through `_pick_model_for_tier`, so routing plugins decide the substitute - exactly as they decided the original. - - Fails open everywhere it cannot be sure: an unreadable eligibility view, a decision - carrying no tier (default_model), or a tier whose every peer is unusable too. It fails - CLOSED on a plugin that empties the pool, leaving the original decision to fail rather - than serving a model the plugin excluded. - """ + """Try compatible tier recovery before the default, preserving request policy and fit.""" decision: Final = response.routing_decision decided_tier: Final = decision.get("tier") if decision is not None else None if decision is None or not isinstance(decided_tier, str): return response - peers: Final = tuple(self._tier_pools().get(decided_tier, ())) - if len(peers) < 2: - return response - if await self._model_group_can_serve(response.model, messages, input, request_kwargs): + fit: Final = context_fit or await self._request_context_fit(resolved_messages, request_kwargs) + if fit.accepts(response.model) and await self._model_group_can_serve( + response.model, messages, input, request_kwargs + ): return response eligible: Final = ( self._modality_eligible_models() if self.config.modality_routing and resolved_messages and request_contains_image_content(resolved_messages) else None ) - candidates: Final = tuple( - peer for peer in peers if peer != response.model and (eligible is None or peer in eligible) + pools: Final = self._tier_pools() + context_recovery: Final = bool(decision.get("context_escalated")) or any( + not fit.accepts(model) for model in pools.get(decided_tier, ()) ) - if not candidates: - return response - servable: Final = await asyncio.gather( - *(self._model_group_can_serve(peer, messages, input, request_kwargs) for peer in candidates) + modality_recovery: Final = eligible is not None + names: Final = self.config.tier_names() + tiers: Final = ( + tuple(names[names.index(decided_tier) :]) + if (context_recovery or modality_recovery) and decided_tier in names + else (decided_tier,) ) - live: Final = tuple(peer for peer, can_serve in zip(candidates, servable) if can_serve) - if not live: - return response - repick_messages: Final = ( - list(resolved_messages) if resolved_messages else None # mutable-ok: the pick's param is list-typed - ) - try: - new_model: Final = await self._pick_model_for_tier( - decided_tier if self.config.has_custom_tiers else ComplexityTier(decided_tier), - messages, - repick_messages, # pyright: ignore[reportArgumentType] # hook-resolved message dicts; the pick only reads them - request_kwargs, - allowed_models=live, + + async def recover_tier(candidate_tier: str) -> PreRoutingHookResponse | None: + peers: Final = tuple( + model + for model in pools.get(candidate_tier, ()) + if not context_recovery + or candidate_tier == decided_tier + or fit.needed is None + or _group_provably_fits(fit.facts.get(model, (None, True)), fit.needed, fit.buffer) ) - except ValueError as exc: - verbose_router_logger.debug( - "ComplexityRouter: health failover found no candidate the routing plugins allow: %s", exc + candidates: Final = tuple( + peer + for peer in peers + if peer != response.model and fit.accepts(peer) and (eligible is None or peer in eligible) ) + servable: Final = await asyncio.gather( + *(self._model_group_can_serve(peer, messages, input, request_kwargs) for peer in candidates) + ) + live: Final = tuple(peer for peer, can_serve in zip(candidates, servable) if can_serve) + if live: + repick_messages: Final = ( + list(resolved_messages) if resolved_messages else None # mutable-ok: the pick's param is list-typed + ) + try: + new_model: Final = await self._pick_model_for_tier( + candidate_tier if self.config.has_custom_tiers else ComplexityTier(candidate_tier), + messages, + repick_messages, # pyright: ignore[reportArgumentType] # hook-resolved message dicts; the pick only reads them + request_kwargs, + allowed_models=live, + ) + except ValueError as exc: + verbose_router_logger.debug( + "ComplexityRouter: health failover found no candidate the routing plugins allow: %s", exc + ) + else: + self._restamp_adaptive_choice(request_kwargs, response.model, new_model) + verbose_router_logger.info( + "ComplexityRouter: routing decision cause=health_failover, routed_model=%s, displaced=%s", + new_model, + response.model, + ) + new_decision: Final = self._build_routing_decision( + routed_model=new_model, + cause="health_failover", + tier=candidate_tier, + score=decision.get("score"), + signals=(*(decision.get("signals") or ()), f"health_displaced:{response.model}"), + matched_keyword=decision.get("matched_keyword"), + escalation_keyword=decision.get("escalation_keyword"), + escalated=bool(decision.get("escalated", False)), + classifier_model=decision.get("classifier_model"), + classifier_cost=decision.get("classifier_cost"), + conversation_continuing=bool(decision.get("conversation_continuing", True)), + tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model), + context_escalation_original_tier=decision.get("context_escalation_original_tier"), + ) + return response.model_copy( + update={ # mutable-ok: model_copy types update as a plain dict + "model": new_model, + "litellm_params": self._litellm_params_for_model(candidate_tier, new_model), + "routing_decision": new_decision, + } + ) + return None + + for candidate_tier in tiers: + if (recovered := await recover_tier(candidate_tier)) is not None: + return recovered + default_model: Final = self.config.default_model + plan_mode_active: Final = self._matched_plan_mode_signal(request_kwargs, resolved_messages) is not None + if ( + plan_mode_active + or self.config.plugins + or not default_model + or default_model == response.model + or not fit.accepts(default_model) + or (eligible is not None and default_model not in eligible) + or not await self._model_group_can_serve(default_model, messages, input, request_kwargs) + ): return response - self._restamp_adaptive_choice(request_kwargs, response.model, new_model) + self._restamp_adaptive_choice(request_kwargs, response.model, default_model) verbose_router_logger.info( - "ComplexityRouter: routing decision cause=health_failover, routed_model=%s, displaced=%s", - new_model, + "ComplexityRouter: routing decision cause=health_default_fallback, routed_model=%s, displaced=%s", + default_model, response.model, ) - new_decision: Final = self._build_routing_decision( - routed_model=new_model, - cause="health_failover", - tier=decision.get("tier"), + default_decision: Final = self._build_routing_decision( + routed_model=default_model, + cause="health_default_fallback", score=decision.get("score"), signals=(*(decision.get("signals") or ()), f"health_displaced:{response.model}"), matched_keyword=decision.get("matched_keyword"), @@ -2947,14 +3223,14 @@ class ComplexityRouter(CustomLogger): classifier_model=decision.get("classifier_model"), classifier_cost=decision.get("classifier_cost"), conversation_continuing=bool(decision.get("conversation_continuing", True)), - tier_litellm_params=self._litellm_params_for_model(decided_tier, new_model), + tier_litellm_params=self._litellm_params_for_model(None, default_model), context_escalation_original_tier=decision.get("context_escalation_original_tier"), ) return response.model_copy( update={ # mutable-ok: model_copy types update as a plain dict - "model": new_model, - "litellm_params": self._litellm_params_for_model(decided_tier, new_model), - "routing_decision": new_decision, + "model": default_model, + "litellm_params": self._litellm_params_for_model(None, default_model), + "routing_decision": default_decision, } ) @@ -3156,6 +3432,18 @@ class ComplexityRouter(CustomLogger): """ return _extract_current_ask_and_system_prompt(messages) + def _reminder_markers_for_request(self, request_kwargs: Mapping[str, object]) -> tuple[tuple[str, str], ...]: + if self.config.reminder_markers is not None: + return self._reminder_markers + if any( + is_codex_user_agent(user_agent) + for metadata_key in ("litellm_metadata", "metadata") + if isinstance(metadata := request_kwargs.get(metadata_key), Mapping) + if isinstance(user_agent := metadata.get("user_agent"), str) + ): + return _CODEX_REMINDER_MARKERS + return _DEFAULT_REMINDER_MARKERS + @staticmethod def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: """Metadata may land on `metadata` or `litellm_metadata` depending on the @@ -3257,6 +3545,8 @@ class ComplexityRouter(CustomLogger): # chat-completions messages, so it is real work on every non-chat surface, and # both the conversation shape and the classifier read the same list. resolved_messages: Final = self._resolve_messages(messages, request_kwargs) + context_fit: Final = await self._request_context_fit(resolved_messages, request_kwargs) + marker_pairs: Final = self._reminder_markers_for_request(request_kwargs) conversation_continuing: Final = _conversation_is_continuing(resolved_messages) use_session_affinity: Final = self._uses_tier_pin @@ -3266,7 +3556,7 @@ class ComplexityRouter(CustomLogger): # In 'user_turn' mode a held pin is replayed only on continuation turns; a new human # ask falls through and re-classifies. session_affinity restores pin-first for asks too. pin_replay_allowed: Final = bool(self.config.session_affinity) or not _newest_turn_is_human_ask( - resolved_messages, self._reminder_markers + resolved_messages, marker_pairs ) if cache_key is not None and pin_replay_allowed: @@ -3277,7 +3567,7 @@ class ComplexityRouter(CustomLogger): pin_escalation_keyword: str | None = None if self.escalation_keywords: user_message: Final = ( - _newest_turn_ask(resolved_messages, self._reminder_markers) if resolved_messages else None + _newest_turn_ask(resolved_messages, marker_pairs) if resolved_messages else None ) if user_message is not None: pin_escalation_keyword = self._matched_escalation_keyword(user_message) @@ -3306,7 +3596,11 @@ class ComplexityRouter(CustomLogger): pin_source_tier: Final = self._tier_for_model(routed_model) pin_placement: Final = ( await self._context_window_placement( - pin_source_tier, resolved_messages, request_kwargs, pool_override=(routed_model,) + pin_source_tier, + resolved_messages, + request_kwargs, + pool_override=(routed_model,), + context_fit=context_fit, ) if pin_source_tier is not None else None @@ -3376,11 +3670,13 @@ class ComplexityRouter(CustomLogger): messages, resolved_messages, request_kwargs, + context_fit, ), messages, input, resolved_messages, request_kwargs, + context_fit, ) ) @@ -3392,14 +3688,18 @@ class ComplexityRouter(CustomLogger): specific_deployment=specific_deployment, conversation_continuing=conversation_continuing, resolved_messages=resolved_messages, + context_fit=context_fit, ) response: Final = ( await self._gate_response_health( - await self._gate_response_modality(routed_response, messages, resolved_messages, request_kwargs), + await self._gate_response_modality( + routed_response, messages, resolved_messages, request_kwargs, context_fit + ), messages, input, resolved_messages, request_kwargs, + context_fit, ) if routed_response is not None else None @@ -3434,6 +3734,7 @@ class ComplexityRouter(CustomLogger): specific_deployment: bool | None = False, conversation_continuing: bool = True, resolved_messages: Sequence[Mapping[str, object]] | None = None, + context_fit: _RequestContextFit | None = None, ) -> PreRoutingHookResponse | None: """ Classifies the request by complexity and returns the appropriate model. @@ -3465,7 +3766,8 @@ class ComplexityRouter(CustomLogger): # Determine whether the original request used messages directly has_original_messages: Final = messages is not None and len(messages) > 0 - user_message, system_prompt = _extract_current_ask_and_system_prompt(resolved_messages, self._reminder_markers) + marker_pairs: Final = self._reminder_markers_for_request(request_kwargs) + user_message, system_prompt = _extract_current_ask_and_system_prompt(resolved_messages, marker_pairs) classifier_images: Final = self._classifier_image_parts(resolved_messages) if user_message is None and not classifier_images: @@ -3484,6 +3786,7 @@ class ComplexityRouter(CustomLogger): ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs ) fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM + default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -3492,11 +3795,13 @@ class ComplexityRouter(CustomLogger): cause="default_fallback", tier=fallback_tier, conversation_continuing=conversation_continuing, + tier_litellm_params=default_tier_params, ), + litellm_params=default_tier_params, ) ask: Final = user_message or "" - newest_ask: Final = _newest_turn_ask(resolved_messages, self._reminder_markers) + newest_ask: Final = _newest_turn_ask(resolved_messages, marker_pairs) escalation_keyword: Final = self._matched_escalation_keyword(newest_ask) if newest_ask is not None else None # Resolved here rather than beside the classifier because the keyword-override path below # returns before any classification runs, and a forced tier gets stuck for the same reason @@ -3519,6 +3824,7 @@ class ComplexityRouter(CustomLogger): _tier_name(plan_floor), routed_model, ) + plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -3530,7 +3836,9 @@ class ComplexityRouter(CustomLogger): matched_keyword=plan_mode_sentinel, escalation_keyword=escalation_keyword, escalated=False, + tier_litellm_params=plan_tier_params, ), + litellm_params=plan_tier_params, ) override: Final = await self._resolve_keyword_tier_override(ask, request_kwargs) @@ -3598,7 +3906,9 @@ class ComplexityRouter(CustomLogger): plan_floored: Final = tier != pre_floor_tier if plan_floored: signals = (*signals, "plan_mode_floor") - context_placement: Final = await self._context_window_placement(tier, resolved_messages, request_kwargs) + context_placement: Final = await self._context_window_placement( + tier, resolved_messages, request_kwargs, context_fit=context_fit + ) tier, signals, context_original_tier = _apply_context_placement(tier, signals, context_placement) score_repr: Final = f"{score:.3f}" if score is not None else "n/a" fallback_model: Final = self.config.default_model if not self.config.plugins else None @@ -3623,6 +3933,7 @@ class ComplexityRouter(CustomLogger): outcome.signals, fallback_model, ) + fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model) return PreRoutingHookResponse( model=fallback_model, messages=messages if has_original_messages else None, @@ -3633,7 +3944,9 @@ class ComplexityRouter(CustomLogger): signals=outcome.signals, escalation_keyword=escalation_keyword, escalated=False, + tier_litellm_params=fallback_tier_params, ), + litellm_params=fallback_tier_params, ) if self.config.adaptive: # hard_floor rather than a hard pick, and passed whenever the sentinel is present diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index c483a0b7073..1f1b5a5cc4b 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -5,13 +5,21 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas All values are configurable via proxy config.yaml. """ -from collections.abc import Mapping +import math +import re +import warnings +from collections.abc import Iterable, Mapping from enum import Enum from types import MappingProxyType -from typing import Annotated, Final, Literal +from typing import Annotated, Final, Literal, NamedTuple from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator +with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + import sre_constants + import sre_parse + from litellm.types.llms.openai import REASONING_EFFORT from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin @@ -21,6 +29,7 @@ from .tier_predictor import TrainedTierArtifact class ComplexityTier(str, Enum): """Complexity tiers for routing decisions.""" + NON_REASONING = "NON_REASONING" SIMPLE = "SIMPLE" MEDIUM = "MEDIUM" COMPLEX = "COMPLEX" @@ -54,6 +63,16 @@ TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = ( ComplexityTier.REASONING, ) +NON_REASONING_TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = ( + ComplexityTier.NON_REASONING, + *TIER_SEVERITY_ORDER, +) + + +def tier_severity_order(non_reasoning_enabled: bool) -> tuple[ComplexityTier, ...]: + return NON_REASONING_TIER_SEVERITY_ORDER if non_reasoning_enabled else TIER_SEVERITY_ORDER + + DEFAULT_TIER_DISTANCE_PENALTY: Final[float] = 0.5 DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE: Final[int] = 3 @@ -134,6 +153,9 @@ def normalize_classification_examples(value: str | None) -> str | None: return _normalize_operator_section(value, "classification_examples", MAX_CLASSIFICATION_EXAMPLES_CHARS) +_BUILT_IN_TIER_NAMES: Final[str] = ", ".join(ComplexityTier.__members__) + + class TierDefinition(BaseModel): """An operator-defined tier: the name the LLM classifier must return and its rubric description.""" @@ -144,7 +166,7 @@ class TierDefinition(BaseModel): default=None, description=( "What belongs in this tier; rendered as this tier's bullet in the classifier rubric. " - "Required unless the name is a built-in tier (SIMPLE/MEDIUM/COMPLEX/REASONING), which " + f"Required unless the name is a built-in tier ({_BUILT_IN_TIER_NAMES}), which " "inherits the built-in criteria when omitted" ), ) @@ -166,7 +188,7 @@ class TierDefinition(BaseModel): if description is None and name.upper() not in ComplexityTier.__members__: raise ValueError( f"tier_definitions entry {name!r} must have a description: only the built-in tiers " - "(SIMPLE, MEDIUM, COMPLEX, REASONING) carry one the rubric can inherit" + f"({_BUILT_IN_TIER_NAMES}) carry one the rubric can inherit" ) rendered_on_one_line: Final = (name, description or "") if any("\n" in part or "\r" in part for part in rendered_on_one_line): @@ -569,6 +591,125 @@ class ClassifierLLMConfig(BaseModel): return self +MAX_CUSTOM_PATTERN_REPEAT: Final[int] = 64 +MAX_CUSTOM_PATTERN_WORK: Final[int] = 2048 +MAX_CUSTOM_DIMENSIONS_WORK: Final[int] = 8192 +MAX_CUSTOM_PATTERN_DEPTH: Final[int] = 16 +CUSTOM_PATTERN_SCAN_CHARS: Final[int] = 2048 + +_ATOM_OPCODES: Final = frozenset( + {sre_constants.LITERAL, sre_constants.NOT_LITERAL, sre_constants.ANY, sre_constants.IN, sre_constants.CATEGORY} +) +_REPEAT_OPCODES: Final = frozenset({sre_constants.MAX_REPEAT, sre_constants.MIN_REPEAT}) + + +class _PatternCost(NamedTuple): + paths: int + steps: int + + +def _atom_steps(node: object) -> int: + if isinstance(node, tuple) and len(node) == 2 and node[0] is sre_constants.IN: + return 1 + len(node[1]) + return 1 + + +def _repeat_cost(argument: object) -> _PatternCost | str: + if not isinstance(argument, tuple) or len(argument) != 3: + return "unsupported repeat structure" + low, high, body = argument + if high > MAX_CUSTOM_PATTERN_REPEAT or len(body) != 1 or body[0][0] not in _ATOM_OPCODES: + return "requires a single character or class repeated at most 64 times; use {n,m} instead of *, + or {n,}" + choices: Final = high - low + 1 + return _PatternCost(choices, 1 + high * _atom_steps(body[0]) + choices) + + +def _node_cost(node: object, depth: int) -> _PatternCost | str: + if not isinstance(node, tuple) or len(node) != 2: + return "unsupported regex structure" + opcode, argument = node + if opcode in _ATOM_OPCODES or opcode is sre_constants.AT: + return _PatternCost(1, _atom_steps(node)) + if opcode is sre_constants.SUBPATTERN: + return _sequence_cost(argument[-1], depth + 1) + if opcode is sre_constants.BRANCH: + costs: Final = tuple(_sequence_cost(branch, depth + 1) for branch in argument[1]) + refused: Final = next((cost for cost in costs if isinstance(cost, str)), None) + if refused is not None: + return refused + return _PatternCost( + sum(cost.paths for cost in costs if isinstance(cost, _PatternCost)), + len(costs) + sum(cost.steps for cost in costs if isinstance(cost, _PatternCost)), + ) + if opcode in _REPEAT_OPCODES: + return _repeat_cost(argument) + return "contains an unsupported regex construct" + + +def _sequence_cost(nodes: Iterable[object], depth: int) -> _PatternCost | str: + if depth > MAX_CUSTOM_PATTERN_DEPTH: + return "nests deeper than 16 levels" + costs: Final = tuple(_node_cost(node, depth) for node in nodes) + refused: Final = next((cost for cost in costs if isinstance(cost, str)), None) + if refused is not None: + return refused + valid: Final = tuple(cost for cost in costs if isinstance(cost, _PatternCost)) + # Choices multiply across a sequence; every continuation can execute once per preceding path. + total: Final = _PatternCost( + math.prod(cost.paths for cost in valid), + 1 + sum(cost.steps * math.prod(prior.paths for prior in valid[:index]) for index, cost in enumerate(valid)), + ) + if total.steps > MAX_CUSTOM_PATTERN_WORK: + return "exceeds the per-pattern regex work budget" + return total + + +def custom_pattern_work(pattern: str) -> int | str: + try: + re.compile(pattern, re.IGNORECASE) + parsed: Final = sre_parse.parse(pattern, re.IGNORECASE) + except (re.error, RecursionError, OverflowError): + return "is not a valid regex" + cost: Final = _sequence_cost(tuple(parsed), 0) + return cost if isinstance(cost, str) else cost.steps + + +class CustomDimension(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + name: str = Field(min_length=1, max_length=64, pattern=r"^[A-Za-z][A-Za-z0-9_]*$") + weight: float = Field(gt=0, le=1, allow_inf_nan=False) + keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32) + patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32) + scoring_mode: Literal["binary", "match_count"] = Field( + default="binary", + description=( + "'binary' scores 1 when any matcher hits. 'match_count' scores 0.5 when one distinct matcher hits and 1 " + "when two or more do; repeated occurrences of one matcher never raise it. Keywords are distinct " + "case-insensitively, patterns by source, and a keyword and a pattern are always distinct from each other." + ), + ) + + @model_validator(mode="after") + def _validate_matchers(self) -> "CustomDimension": + matchers: Final = (*self.keywords, *self.patterns) + if not matchers or any(not matcher.strip() for matcher in matchers): + raise ValueError("custom dimensions require nonblank keywords and/or patterns") + if len(matchers) > 32 or sum(map(len, matchers)) > 4096: + raise ValueError("custom dimensions allow at most 32 matchers and 4096 matcher characters each") + costs: Final = tuple((pattern, custom_pattern_work(pattern)) for pattern in self.patterns) + rejected: Final = tuple(f"pattern {pattern!r} {work}" for pattern, work in costs if isinstance(work, str)) + if rejected: + raise ValueError("custom dimension " + "; ".join(rejected)) + return self + + def pattern_work(self) -> int: + """Combined work estimate of the validated patterns.""" + return sum( + work for work in (custom_pattern_work(pattern) for pattern in self.patterns) if isinstance(work, int) + ) + + class ComplexityRouterConfig(BaseModel): """Configuration for the ComplexityRouter.""" @@ -584,6 +725,20 @@ class ComplexityRouterConfig(BaseModel): default_factory=dict, ) + enable_non_reasoning_tier: bool = Field( + default=False, + description=( + "Add NON_REASONING as a fifth built-in tier below SIMPLE, for operational agent traffic " + "that relays or reformats information rather than reasoning about it. Off by default: " + "turning it on adds a rung to this router's ladder, a bullet to the LLM classifier's " + "rubric, and a value the classifier may return, all of which move tier decisions and " + "spend on an already-deployed router. Requires an LLM classifier or a custom classifier " + "plugin, since the heuristic scorers cannot produce the tier, and a model in `tiers` " + "under the NON_REASONING key. Escalation still walks up from it, and it is never the " + "savings baseline or a `heuristic_v2` prediction." + ), + ) + tier_definitions: tuple[TierDefinition, ...] | None = Field( default=None, description=( @@ -671,6 +826,20 @@ class ComplexityRouterConfig(BaseModel): description="Weights for each scoring dimension", ) + custom_dimensions: tuple[CustomDimension, ...] = Field( + default=(), + max_length=16, + description=( + "Named dimensions added to the heuristic-v1 score. Each contributes its inline weight once " + "when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters; " + "scoring_mode 'match_count' instead grades half weight for one distinct matcher and full for two or more. " + "Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, " + "backreferences and lookarounds are rejected. Conservative work limits include alternation paths, " + "repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. " + "Only heuristic, heuristic_first and hybrid accept this field. Uses the existing heuristic tuning quota." + ), + ) + # Keyword lists (overridable) code_keywords: list[str] | None = Field( default=None, @@ -802,9 +971,11 @@ class ComplexityRouterConfig(BaseModel): "classified against what it refers to. Counts turns of both roles when " "classifier_context_include_assistant_turns is enabled. These turns are sent to the classifier " "model, which may " - "be a different deployment or provider than the routed completion model; that call already " - "carries the current user ask and the caller's system prompt in full. Set to 0 to send neither " - "prior turns nor any conversation context beyond the current ask. Only applies when " + "be a different deployment or provider than the routed completion model; that call carries " + "the current user ask and, except for Claude Code requests, the extracted system-role text in full. " + "Claude Code system text is omitted to avoid classifying harness instructions; the routed " + "completion still receives it. Set to 0 to send neither prior turns nor " + "any conversation context beyond the current ask. Only applies when " "classifier_type is 'llm'." ), ) @@ -816,9 +987,9 @@ class ComplexityRouterConfig(BaseModel): "context window, per classification call. Turns are taken newest first and quoted whole " "while they fit, so a conversation small enough to quote entirely is never cut; once the " "budget runs out the older turns are dropped whole and only the turn straddling the " - "boundary is truncated, into whatever space is left. The current ask and the caller's " - "system prompt sit outside this budget and are always sent in full, as does the numbering " - "each quoted turn carries. A budget under 120 leaves no room to quote a turn and " + "boundary is truncated, into whatever space is left. The current ask and, except for Claude " + "Code requests, the extracted system-role text sit outside this budget and are sent in full, as does " + "the numbering each quoted turn carries. A budget under 120 leaves no room to quote a turn and " "suppresses the block; set classifier_context_window_size to 0 to turn context off " "deliberately. Only applies when classifier_type is 'llm'." ), @@ -948,6 +1119,20 @@ class ComplexityRouterConfig(BaseModel): "wording the built-ins don't cover, or after a client release changes its strings." ), ) + max_tokens_from_tier_model: bool = Field( + default=True, + description=( + "Set max_tokens on every routed request to the output ceiling of the tier model it " + "lands on, replacing whatever the caller sent. A caller behind an auto-router cannot " + "pick one value that fits every tier: the smallest tier's ceiling starves a bigger " + "tier's thinking budget, and a bigger tier's ceiling is rejected by the smallest. The " + "ceiling is the smallest max_output_tokens across the tier model's deployments, read " + "from each deployment's model_info and then the model cost map; a tier model with a " + "deployment whose ceiling is unknown keeps the caller's value. A max_tokens, " + "max_completion_tokens or max_output_tokens in the tier's own litellm_params still " + "wins. Set false to forward the caller's value unchanged." + ), + ) route_housekeeping_to_cheapest_tier: bool = Field( default=True, description=( @@ -1109,8 +1294,9 @@ class ComplexityRouterConfig(BaseModel): "Override the delimiter pairs used to recognize and strip harness-injected reminder " "blocks before classification. A harness that wraps injected context differently per " "agent type (main, subagent, cron) lists every pair it emits. Replaces, rather than " - "adds to, the built-in default of ('', ''), so a " - "harness that also emits that pair lists it too. Matching is case-insensitive." + "adds to, the built-in system-reminder pair and the Codex envelope pairs enabled " + "for Codex user agents, so list every built-in pair your harness also emits. " + "Matching is case-insensitive." ), ) @@ -1245,6 +1431,27 @@ class ComplexityRouterConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_custom_dimensions(self) -> "ComplexityRouterConfig": + if not self.custom_dimensions: + return self + if self.classifier_type not in ("heuristic", "heuristic_first", "hybrid"): + raise ValueError("custom_dimensions requires classifier_type heuristic, heuristic_first or hybrid") + names: Final = tuple(dimension.name.casefold() for dimension in self.custom_dimensions) + reserved: Final = frozenset(name.casefold() for name in DEFAULT_DIMENSION_WEIGHTS) + weighted: Final = frozenset(name.casefold() for name in self.dimension_weights) + if len(frozenset(names)) != len(names) or frozenset(names) & reserved: + raise ValueError("custom dimension names must be unique and must not shadow built-in dimensions") + if frozenset(names) & weighted: + raise ValueError("custom dimension weights must be inline, not in dimension_weights") + work: Final = sum(dimension.pattern_work() for dimension in self.custom_dimensions) + if work > MAX_CUSTOM_DIMENSIONS_WORK: + raise ValueError( + f"custom_dimensions regex work estimate is {work}; the limit across the router is " + f"{MAX_CUSTOM_DIMENSIONS_WORK}" + ) + return self + @field_validator("heuristic_first_max_tier", mode="before") @classmethod def _coerce_heuristic_first_max_tier(cls, value: object) -> object: @@ -1338,11 +1545,15 @@ class ComplexityRouterConfig(BaseModel): which still makes it a dependency on every one of those requests.""" return self.classifier_type in LLM_CLASSIFIER_TYPES + def active_tier_severity_order(self) -> tuple[ComplexityTier, ...]: + """This router's built-in ladder, ascending; not meaningful for a custom tier set.""" + return tier_severity_order(self.enable_non_reasoning_tier) + def tier_names(self) -> tuple[str, ...]: """The active tier names: the defined names, or the built-in set in severity order.""" if self.tier_definitions is not None: return tuple(definition.name for definition in self.tier_definitions) - return tuple(tier.value for tier in TIER_SEVERITY_ORDER) + return tuple(tier.value for tier in self.active_tier_severity_order()) def classifier_wire_labels(self) -> tuple[str, ...]: """The tier names the classifier is told to emit: defined names, or the display labels.""" @@ -1434,6 +1645,36 @@ class ComplexityRouterConfig(BaseModel): if present ) + @model_validator(mode="after") + def _validate_non_reasoning_tier(self) -> "ComplexityRouterConfig": + """Require a classifier that can emit the opt-in tier and a model to route it to.""" + non_reasoning_key: Final = ComplexityTier.NON_REASONING.value + if not self.enable_non_reasoning_tier: + if not self.has_custom_tiers and non_reasoning_key in self.tiers: + raise ValueError( + f"tiers names {non_reasoning_key} but enable_non_reasoning_tier is False, so no request " + "can route there; set enable_non_reasoning_tier: true or drop the tier" + ) + return self + if self.has_custom_tiers: + raise ValueError( + "enable_non_reasoning_tier cannot be combined with tier_definitions: a custom tier set " + f"replaces the built-in ladder, so name a tier {non_reasoning_key} in tier_definitions instead" + ) + if self.classifier_type not in ("llm", "custom"): + raise ValueError( + f"enable_non_reasoning_tier requires classifier_type 'llm' or 'custom', got " + f"{self.classifier_type!r}: the heuristic scorers only produce the four tiers from SIMPLE up, " + f"so nothing would ever classify as {non_reasoning_key}" + ) + if not self.tiers.get(non_reasoning_key): + raise ValueError( + f"enable_non_reasoning_tier requires tiers to map {non_reasoning_key} to at least one model: " + "the tier exists to send operational traffic somewhere cheaper, and an unconfigured tier " + "would fall through to the default model" + ) + return self + @model_validator(mode="after") def _validate_tier_definitions(self) -> "ComplexityRouterConfig": if self.tier_definitions is None: @@ -1456,7 +1697,7 @@ class ComplexityRouterConfig(BaseModel): if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first", "hybrid"): raise ValueError( "tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only " - "produces the four built-in tiers, as does heuristic_v2" + "produces the built-in tiers from SIMPLE up, as does heuristic_v2" ) conflicts: Final = self._tier_definition_conflicts() if conflicts: @@ -1610,7 +1851,7 @@ class ComplexityRouterConfig(BaseModel): def labeled_tiers(self) -> tuple[tuple[ComplexityTier, str], ...]: """Every tier paired with its display name, in ascending severity order.""" - return tuple((tier, self.tier_label(tier)) for tier in TIER_SEVERITY_ORDER) + return tuple((tier, self.tier_label(tier)) for tier in self.active_tier_severity_order()) def tier_for_label(self, label: str) -> ComplexityTier | None: """Resolve a display name back to its tier, case-insensitively, then canonical names.""" @@ -1618,7 +1859,7 @@ class ComplexityRouterConfig(BaseModel): labeled: Final = self.labeled_tiers() return next( (tier for tier, tier_label in labeled if tier_label.casefold() == folded), - next((tier for tier in TIER_SEVERITY_ORDER if tier.value.casefold() == folded), None), + next((tier for tier, _ in labeled if tier.value.casefold() == folded), None), ) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1433e8ba4d4..0b73f4e31a7 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,17 +1,111 @@ -#### What this does #### -# identifies least busy deployment -# How is this achieved? -# - Before each call, have the router print the state of requests {"deployment": "requests_in_flight"} -# - use litellm.input_callbacks to log when a request is just about to be made to a model - {"deployment-id": traffic} -# - use litellm.success + failure callbacks to log when a request completed -# - in get_available_deployment, for a given model group name -> pick based on traffic - -import random +import logging +from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger +IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 + + +class _ModelInfo(TypedDict, total=False): + id: ReadOnly[str | int | None] + + +class _Metadata(TypedDict, total=False): + model_group: ReadOnly[str | None] + + +class _LitellmParams(TypedDict, total=False): + metadata: ReadOnly[_Metadata | None] + model_info: ReadOnly[_ModelInfo | None] + + +class _CallKwargs(TypedDict, total=False): + litellm_params: ReadOnly[_LitellmParams | None] + + +class _DeploymentModelInfo(TypedDict): + id: ReadOnly[str | int] + + +class _Deployment(TypedDict): + model_info: ReadOnly[_DeploymentModelInfo] + + +_CALL_KWARGS: Final = TypeAdapter(_CallKwargs) +_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) +_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) + + +def _request_count_key(model_group: str, deployment_id: str) -> str: + return f"{model_group}_request_count:{deployment_id}" + + +def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: + try: + call: Final = _CALL_KWARGS.validate_python(kwargs) + except ValidationError: + return None + litellm_params: Final = call.get("litellm_params") + metadata: Final = litellm_params.get("metadata") if litellm_params else None + model_info: Final = litellm_params.get("model_info") if litellm_params else None + model_group: Final = metadata.get("model_group") if metadata else None + deployment_id: Final = model_info.get("id") if model_info else None + if model_group is None or deployment_id is None: + return None + return model_group, str(deployment_id) + + +def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + return tuple( + _request_count_key(model_group, str(deployment["model_info"]["id"])) + for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) + ) + + +def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: + return tuple(0 if value is None else int(value) for value in values) + + +def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: + values: Final = _MEMORY_COUNTS.validate_python(raw) + if values is None or len(values) != len(keys): + return (0,) * len(keys) + return _as_counts(values) + + +def _least_busy( + healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...] +) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])] + + +def _warn_unreadable(model_group: str, error: Exception) -> None: + log_redis_failure( + verbose_router_logger, + logging.WARNING, + f"least-busy routing could not read the shared in-flight counts for {model_group}, " + "falling back to this worker's own counts", + error, + ) + + +def _warn_unwritable(key: str, error: Exception) -> None: + log_redis_failure( + verbose_router_logger, + logging.WARNING, + f"least-busy routing could not update the in-flight count under {key}", + error, + ) + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False @@ -20,195 +114,101 @@ class LeastBusyLoggingHandler(CustomLogger): def __init__(self, router_cache: DualCache): self.router_cache = router_cache + self.router_cache_id = str(id(router_cache)) - def log_pre_api_call(self, model, messages, kwargs): - """ - Log when a model is being used. + def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: + self._increment(kwargs, 1) - Caching based on model group. - """ - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - request_count_api_key: Final = f"{model_group}_request_count" - # update cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_dict[id] = request_count_dict.get(id, 0) + 1 + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - except Exception: - pass + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - def log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - def _get_available_deployments( - self, - healthy_deployments: list, - all_deployments: dict, - ): - """ - Helper to get deployments using least busy strategy - """ - for d in healthy_deployments: - ## if healthy deployment not yet used - if d["model_info"]["id"] not in all_deployments: - all_deployments[d["model_info"]["id"]] = 0 - # map deployment to id - # pick least busy deployment - min_traffic = float("inf") - min_deployment = None - for k, v in all_deployments.items(): - if v < min_traffic: - min_traffic = v - min_deployment = k - if min_deployment is not None: - ## check if min deployment is a string, if so, cast it to int - for m in healthy_deployments: - if m["model_info"]["id"] == min_deployment: - return m - min_deployment = random.choice(healthy_deployments) - else: - min_deployment = random.choice(healthy_deployments) - return min_deployment + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 def get_available_deployments( - self, - model_group: str, - healthy_deployments: list, - ): - """ - Sync helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _as_counts(redis_cache.batch_get_counts(list(keys))) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) - async def async_get_available_deployments(self, model_group: str, healthy_deployments: list): - """ - Async helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + async def async_get_available_deployments( + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(keys))) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) + + def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + local: Final = self.router_cache.increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local < 0: + self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + if redis_cache is None: + return + redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) + except Exception as e: + _warn_unwritable(key, e) + + async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + local: Final = await self.router_cache.async_increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local is not None and local < 0: + await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + if redis_cache is None: + return + await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) + except Exception as e: + _warn_unwritable(key, e) diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index b927df0c438..d271349914e 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -1,6 +1,6 @@ #### What this does #### # picks based on response time (for streaming, this is time to first token) -from datetime import datetime, timedelta +from datetime import datetime from typing import Final import litellm @@ -39,7 +39,7 @@ class LowestCostLoggingHandler(CustomLogger): # ------------ """ { - {model_group}_map: { + cost_map:{model_group}: { id: { f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } @@ -50,18 +50,14 @@ class LowestCostLoggingHandler(CustomLogger): current_hour: Final = datetime.now().strftime("%H") current_minute: Final = datetime.now().strftime("%M") 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 + cost_key: Final = f"cost_map:{model_group}" 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 @@ -116,33 +112,27 @@ class LowestCostLoggingHandler(CustomLogger): # ------------ """ { - {model_group}_map: { + cost_map:{model_group}: { id: { - "cost": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } } """ - cost_key: Final = f"{model_group}_map" + cost_key: Final = f"cost_map:{model_group}" current_date: Final = datetime.now().strftime("%Y-%m-%d") current_hour: Final = datetime.now().strftime("%H") 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 # ------------ @@ -185,7 +175,7 @@ class LowestCostLoggingHandler(CustomLogger): """ Returns a deployment with the lowest cost """ - cost_key: Final = f"{model_group}_map" + cost_key: Final = f"cost_map:{model_group}" request_count_dict: Final = await self.router_cache.async_get_cache(key=cost_key) or {} diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index a1b67eaeaf9..e902192811c 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -1,9 +1,13 @@ #### 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 math import ceil from typing import TYPE_CHECKING, Any, Final +from pydantic import Field + import litellm from litellm import ModelResponse, token_counter, verbose_logger from litellm.caching.caching import DualCache @@ -23,6 +27,25 @@ class RoutingArgs(LiteLLMPydanticObjectBase): ttl: float = 1 * 60 * 60 # 1 hour lowest_latency_buffer: float = 0 max_latency_list_size: int = 10 + ttft_percentile: float | None = Field(default=None, gt=0, le=1) + + +def _average_latency(samples: Sequence[float]) -> float: + if not samples: + return 0.0 + return sum(samples) / len(samples) + + +def _percentile_latency(samples: Sequence[float], percentile: float) -> float: + values: Final = sorted(samples) + index: Final = ceil(len(values) * percentile) - 1 + return values[index] + + +def _ttft_seconds(elapsed: timedelta | float) -> float: + if isinstance(elapsed, timedelta): + return elapsed.total_seconds() + return float(elapsed) class LowestLatencyLoggingHandler(CustomLogger): @@ -79,14 +102,13 @@ class LowestLatencyLoggingHandler(CustomLogger): # breaks JSON serialization when the router cache syncs to # Redis (issue #33169) response_ms = response_ms.total_seconds() - time_to_first_token_response_time = None + time_to_first_token: float | None = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time + time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time) final_value: float = response_ms - time_to_first_token: float | None = None total_tokens = 0 if isinstance(response_obj, ModelResponse): @@ -104,13 +126,6 @@ class LowestLatencyLoggingHandler(CustomLogger): else: final_value = response_seconds - if time_to_first_token_response_time is not None: - if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() - else: - ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) - # ------------ # Update usage # ------------ @@ -131,14 +146,14 @@ class LowestLatencyLoggingHandler(CustomLogger): ## Time to first token if time_to_first_token is not None: if ( - len(request_count_dict[id].get("time_to_first_token", [])) + len(request_count_dict[id].get("time_to_first_token_seconds", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ - 1: - ] + [time_to_first_token] + request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][ + "time_to_first_token_seconds" + ][1:] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -245,7 +260,7 @@ class LowestLatencyLoggingHandler(CustomLogger): {model_group}_map: { id: { "latency": [..] - "time_to_first_token": [..] + "time_to_first_token_seconds": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } @@ -266,14 +281,13 @@ class LowestLatencyLoggingHandler(CustomLogger): # breaks JSON serialization when the router cache syncs to # Redis (issue #33169) response_ms = response_ms.total_seconds() - time_to_first_token_response_time = None + time_to_first_token: float | None = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time + time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time) final_value: float = response_ms total_tokens = 0 - time_to_first_token: float | None = None if isinstance(response_obj, ModelResponse): _usage: Final = getattr(response_obj, "usage", None) @@ -289,13 +303,6 @@ class LowestLatencyLoggingHandler(CustomLogger): final_value = float(normalized_value) else: final_value = response_seconds - - if time_to_first_token_response_time is not None: - if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() - else: - ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) # ------------ # Update usage # ------------ @@ -321,14 +328,14 @@ class LowestLatencyLoggingHandler(CustomLogger): ## Time to first token if time_to_first_token is not None: if ( - len(request_count_dict[id].get("time_to_first_token", [])) + len(request_count_dict[id].get("time_to_first_token_seconds", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ - 1: - ] + [time_to_first_token] + request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][ + "time_to_first_token_seconds" + ][1:] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -426,28 +433,21 @@ class LowestLatencyLoggingHandler(CustomLogger): or float("inf") ) item_latency = item_map.get("latency", []) - item_ttft_latency = item_map.get("time_to_first_token", []) + item_ttft_latency = item_map.get("time_to_first_token_seconds", []) item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) 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) + selected_latency = ( + _percentile_latency(item_ttft_latency, self.routing_args.ttft_percentile) + if use_ttft and self.routing_args.ttft_percentile is not None + else _average_latency(item_ttft_latency if use_ttft else item_latency) + ) # -------------- # # Debugging Logic @@ -456,7 +456,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] = selected_latency # -------------- # # End of Debugging Logic # -------------- # @@ -466,7 +466,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, selected_latency)) if len(potential_deployments) == 0: return None diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index 65bcce0e532..860e89cea22 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -41,8 +41,7 @@ def simple_shuffle( ############## Check if 'weight' or 'rpm' or 'tpm' param set for a weighted pick ################# for weight_by in ["weight", "rpm", "tpm"]: - weight = healthy_deployments[0].get("litellm_params").get(weight_by, None) - if weight is not None: + if any(m["litellm_params"].get(weight_by) is not None for m in healthy_deployments): weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments] verbose_router_logger.debug("\nweight %s", weights) total_weight = sum(weights) diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index eabd9278cf6..d4f46e94579 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -13,7 +13,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload from litellm._logging import verbose_logger -from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY +from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.types.router import ConsumedRequestTagsStamp, DeploymentTypedDict, RouterErrors @@ -461,7 +461,10 @@ def _request_tags_after_router_consumption(metadata: object, model: str) -> Sequ if not isinstance(metadata, Mapping): return None typed_metadata: Final[Mapping[str, object]] = metadata - request_tags: Final = _tags_in_metadata(typed_metadata) + request_tags: Final = _tags_in_metadata( + typed_metadata, + key=ROUTING_REQUEST_TAGS_METADATA_KEY if ROUTING_REQUEST_TAGS_METADATA_KEY in typed_metadata else "tags", + ) stamp: Final = typed_metadata.get(CONSUMED_REQUEST_TAGS_METADATA_KEY) if not isinstance(stamp, ConsumedRequestTagsStamp) or stamp.model_group != model: return request_tags @@ -646,7 +649,7 @@ async def get_deployments_for_tag( return healthy_deployments -def _tags_in_metadata(metadata: object) -> list[str]: +def _tags_in_metadata(metadata: object, key: str = "tags") -> list[str]: """ Tags out of a metadata bucket the caller controls the shape of. @@ -657,7 +660,7 @@ def _tags_in_metadata(metadata: object) -> list[str]: if not isinstance(metadata, Mapping): return [] typed_metadata: Final[Mapping[str, object]] = metadata - tags: Final = typed_metadata.get("tags") + tags: Final = typed_metadata.get(key) if isinstance(tags, str) or not isinstance(tags, Sequence): return [] typed_tags: Final[Sequence[object]] = tags diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 3ec92ad226a..3251ea457cf 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -50,6 +50,12 @@ def prepare_response_for_header_attachment(response: object) -> object | None: return response +def response_has_hidden_params(response: object) -> bool: + if isinstance(response, dict): + return "_hidden_params" in response + return hasattr(response, "_hidden_params") + + def ensure_response_additional_headers(response: object) -> dict[str, object]: hidden_params: Final = get_hidden_params_dict(response, create=isinstance(response, dict)) _write_hidden_params(response, hidden_params) diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py new file mode 100644 index 00000000000..9699ab886b9 --- /dev/null +++ b/litellm/router_utils/auto_router_tuning_baseline.py @@ -0,0 +1,174 @@ +"""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_v2" + +HEURISTIC_V1_TUNING_FIELDS: Final = ( + "tiers", + "tier_model_configs", + "classifier_type", + "tier_boundaries", + "reasoning_override_min_score", + "token_thresholds", + "dimension_weights", + "custom_dimensions", + "code_keywords", + "reasoning_keywords", + "technical_keywords", + "custom_technical_keywords", + "simple_keywords", + "escalation_keywords", + "keyword_tier_rules", +) + +_TUNING_FIELD_SET: Final = frozenset(HEURISTIC_V1_TUNING_FIELDS) + +_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 explicitly supplied heuristic-v1 tuning fields, normalized, or None when invalid.""" + raw: Final = _mapping(complexity_router_config) + try: + validated: Final = ComplexityRouterConfig.model_validate(raw) + except ValidationError: + return None + supplied: Final = ((_TUNING_FIELD_SET - frozenset(("tier_model_configs",))) & frozenset(raw)) | ( + frozenset(("tier_model_configs",)) if validated.tier_model_configs else frozenset() + ) + payload: Final = validated.model_dump( + mode="json", + include=supplied, + exclude={ + "custom_dimensions": { + index: {"scoring_mode"} + for index, dimension in enumerate(validated.custom_dimensions) + if dimension.scoring_mode == "binary" + } + }, + ) + return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +DEFAULT_TUNING_FINGERPRINT: Final = tuning_fingerprint(_EMPTY) + + +def uses_heuristic_v1(complexity_router_config: object) -> bool: + """Whether a config's primary classifier path is the heuristic-v1 scorer.""" + return _mapping(complexity_router_config).get("classifier_type", "heuristic") in _V1_SCORING_CLASSIFIER_TYPES + + +def router_identity(deployment: Mapping[str, object]) -> str | None: + """Stable identity for a complexity-router deployment, across tuning edits.""" + model_info: Final = _mapping(deployment.get("model_info")) + model_id: Final = model_info.get("id") + if model_info.get("db_model") is True and isinstance(model_id, str) and model_id: + return f"db:{model_id}" + model_name: Final = deployment.get("model_name") + if not isinstance(model_name, str) or not model_name: + return None + litellm_params: Final = _mapping(deployment.get("litellm_params")) + tags: Final = litellm_params.get("tags") + normalized_tags: Final = ( + tuple(sorted(str(tag) for tag in tags)) + if isinstance(tags, Iterable) and not isinstance(tags, str) + else _EMPTY_TAGS + ) + return f"yaml:{json.dumps((model_name, normalized_tags), separators=(',', ':'))}" + + +def heuristic_v1_router_fingerprint(deployment: Mapping[str, object]) -> tuple[str, str] | None: + """The identity/fingerprint pair for a heuristic-v1 complexity router, else None.""" + litellm_params: Final = _mapping(deployment.get("litellm_params")) + model: Final = litellm_params.get("model") + config: Final = litellm_params.get("complexity_router_config") + if ( + not isinstance(model, str) + or not model.startswith(_AUTO_ROUTER_COMPLEXITY_PREFIX) + or not uses_heuristic_v1(config) + ): + return None + identity: Final = router_identity(deployment) + fingerprint: Final = tuning_fingerprint(config) + if identity is None or fingerprint is None: + return None + return identity, fingerprint + + +def snapshot_tuning_baselines(deployments: Iterable[Mapping[str, object]]) -> Mapping[str, str]: + """One immutable first-observation baseline for every heuristic-v1 complexity router.""" + return MappingProxyType( + { + identity: fingerprint + for deployment in deployments + if (pair := heuristic_v1_router_fingerprint(deployment)) is not None + for identity, fingerprint in (pair,) + } + ) # mutable-ok: MappingProxyType owns the completed immutable snapshot + + +def is_mutable_tuned_candidate(candidate: Mapping[str, object], baselines: Mapping[str, str]) -> bool: + pair: Final = heuristic_v1_router_fingerprint(candidate) + if pair is None: + return False + identity, fingerprint = pair + return fingerprint != baselines.get(identity, DEFAULT_TUNING_FINGERPRINT) + + +def mutable_tuned_identities( + deployments: Iterable[Mapping[str, object]], baselines: Mapping[str, str] +) -> frozenset[str]: + """Heuristic-v1 routers whose current tuning differs from their baseline or shipped default.""" + return frozenset( + identity + for deployment in deployments + if (pair := heuristic_v1_router_fingerprint(deployment)) is not None + for identity, _ in (pair,) + if is_mutable_tuned_candidate(deployment, baselines) + ) + + +def tuning_limit_violation(*, held: int, limit: int | None) -> str | None: + if limit is None or held <= limit: + return None + return ( + f"At most {limit} auto-router(s) with changed heuristic scorer settings or tier models can be modified " + "without an auto-router license. Keep this router on its recorded settings, or revert the other changed " + "router to its baseline, or remove one of them." + ) + + +def tuning_quota_violation( + *, + candidate: Mapping[str, object], + others: Iterable[Mapping[str, object]], + baselines: Mapping[str, str], + limit: int | None, +) -> str | None: + """Why a change to candidate tuning exceeds the baseline-relative free quota.""" + if limit is None: + return None + pair: Final = heuristic_v1_router_fingerprint(candidate) + if pair is None: + return None + identity, _ = pair + if not is_mutable_tuned_candidate(candidate, baselines): + return None + held: Final = mutable_tuned_identities(others, baselines) - frozenset((identity,)) + return tuning_limit_violation(held=len(held) + 1, limit=limit) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 280a7defcf8..49cca8ee99e 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject +import litellm from litellm._logging import verbose_logger, verbose_router_logger from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError @@ -26,6 +27,18 @@ def _is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool return getattr(user_api_key_auth, "user_role", None) == "proxy_admin" +def get_request_team_id(request_kwargs: Mapping[str, object] | None) -> str | None: + """The caller's team id, from whichever metadata bucket this surface writes to.""" + if request_kwargs is None: + return None + for bucket_name in ("metadata", "litellm_metadata"): + bucket = request_kwargs.get(bucket_name) + team_id = bucket.get("user_api_key_team_id") if isinstance(bucket, Mapping) else None + if isinstance(team_id, str) and team_id: + return team_id + return None + + def resolve_model_group_alias(model_group_alias: object, model: str) -> str | None: """ Resolve ``model`` through a ``model_group_alias`` map. @@ -110,7 +123,7 @@ def filter_team_based_models( metadata: Final = request_kwargs.get("metadata") or {} litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {} - request_team_id: Final = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + request_team_id: Final = get_request_team_id(request_kwargs) if request_team_id is None and _is_proxy_admin_request(request_kwargs) and isinstance(healthy_deployments, list): requested_model: Final = ( request_kwargs.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group") @@ -244,6 +257,32 @@ PROVIDER_SCOPED_CREDENTIAL_PARAMS: Final[Mapping[str, frozenset[str]]] = Mapping ) +def provider_for_generic_call(litellm_params: Mapping[str, object]) -> str | None: + """ + The provider the router hands a deployment's generic SDK call, or None when it cannot be resolved. + + A model that carries its own provider prefix keeps that prefix even where get_llm_provider + would resolve it to a sibling provider (azure_ai/ on an Azure OpenAI host + resolves to azure): the SDK call still receives the prefixed model, and an explicit provider + that contradicts the prefix makes get_llm_provider re-prefix it into a deployment name that + does not exist upstream. + """ + declared: Final = litellm_params.get("custom_llm_provider") + if isinstance(declared, str) and declared: + return declared + model: Final = litellm_params.get("model") + if not isinstance(model, str) or not model: + return None + prefix: Final = model.split("/", 1)[0] + if "/" in model and prefix in litellm.provider_list: + return prefix + try: + _, inferred, _, _ = get_llm_provider(model=model) + except BadRequestError: + return None + return inferred + + def warn_on_provider_credential_mismatch(model_name: str, litellm_params: Mapping[str, object]) -> str | None: """ Warn when a deployment carries one provider's credentials but resolves to another. diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index 9e7f457f631..ef29f7d8fd3 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -12,6 +12,7 @@ from typing_extensions import TypedDict from litellm import verbose_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker if TYPE_CHECKING: @@ -36,10 +37,19 @@ _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0 class CooldownCache: - def __init__(self, cache: DualCache, default_cooldown_time: float): + def __init__( + self, + cache: DualCache, + default_cooldown_time: float, + redis_read_interval_seconds: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, + ): self.cache = cache self.default_cooldown_time = default_cooldown_time self.in_memory_cache = InMemoryCache() + self._cooldown_store = DualCache( + in_memory_cache=self.in_memory_cache, + default_redis_batch_cache_expiry=redis_read_interval_seconds, + ) # Initialize the masker with custom settings for exception strings self.exception_masker = SensitiveDataMasker( visible_prefix=50, # Show first 50 characters @@ -48,6 +58,21 @@ class CooldownCache: mask_short_values=False, # Truncate long messages only; keep short ones readable ) + @property + def cooldown_store(self) -> DualCache: + """ + The cache cooldown entries live in, with the router's Redis attached on first use. + + It is kept separate from the router-wide cache so that a key missing from memory is + re-read from Redis every `redis_read_interval_seconds` rather than on the router + cache's much longer batch interval, which is what lets a sibling replica see a + cooldown another replica wrote, and so that unrelated router keys cannot evict a + cooldown from the in-memory tier before it expires. Redis is attached lazily because + the router builds its cooldown cache before it wires up the shared Redis client. + """ + self._cooldown_store.attach_redis_cache(self.cache.redis_cache) + return self._cooldown_store + def _common_add_cooldown_logic( self, model_id: str, original_exception, exception_status, cooldown_time: float ) -> tuple[str, CooldownCacheValue]: @@ -93,7 +118,7 @@ class CooldownCache: ) # Set the cache with a TTL equal to the cooldown time - self.cache.set_cache( + self.cooldown_store.set_cache( value=cooldown_data, key=cooldown_key, ttl=_cooldown_time, @@ -122,13 +147,13 @@ class CooldownCache: cooldown_cache_value: Final = CooldownCacheValue(**result) # pyright: ignore[reportUnknownArgumentType] - result comes from an untyped cache read, not from our own code remaining: Final = (cooldown_cache_value["timestamp"] + cooldown_cache_value["cooldown_time"]) - current_time if remaining <= 0: - self.cache.in_memory_cache.delete_cache(key) + self.in_memory_cache.delete_cache(key) return None - current_expiry: Final = self.cache.in_memory_cache.ttl_dict.get(key) + current_expiry: Final = self.in_memory_cache.ttl_dict.get(key) if current_expiry is not None and current_expiry > current_time + remaining + 5: corrected_ttl: Final = min(remaining, _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS) - self.cache.in_memory_cache.delete_cache(key) - self.cache.in_memory_cache.set_cache(key, result, ttl=corrected_ttl) + self.in_memory_cache.delete_cache(key) + self.in_memory_cache.set_cache(key, result, ttl=corrected_ttl) return cooldown_cache_value async def async_get_active_cooldowns( @@ -137,12 +162,7 @@ class CooldownCache: # Generate the keys for the deployments keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] - # Retrieve the values for the keys using mget - ## more likely to be none if no models ratelimited. So just check redis every 1s - ## each redis call adds ~100ms latency. - - ## check in memory cache first - results: Final = await self.cache.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) + results: Final = await self.cooldown_store.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) active_cooldowns: Final[list[tuple[str, CooldownCacheValue]]] = [] if results is None or all(v is None for v in results): @@ -164,7 +184,7 @@ class CooldownCache: # Generate the keys for the deployments keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] # Retrieve the values for the keys using mget - results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] + results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] active_cooldowns: Final = [] current_time: Final = time.time() @@ -184,7 +204,7 @@ class CooldownCache: keys: Final = [f"deployment:{model_id}:cooldown" for model_id in model_ids] # Retrieve the values for the keys using mget - results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] + results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] min_cooldown_time: float | None = None # Process the results diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 4534fa114b3..f722b6fd20c 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache from litellm.constants import ( DEFAULT_COOLDOWN_TIME_SECONDS, DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS, @@ -558,9 +559,12 @@ def should_cooldown_based_on_allowed_fails_policy( When *allowed_fails_override* / *cooldown_time_override* are supplied they take precedence over the router-level values (used by deployment-level overrides). + The counter lives in the router's shared ``DualCache`` (Redis when configured), so + every worker process increments the same key and the threshold applies fleet-wide. + When *cache_key_suffix* is supplied the fail counter is keyed as - ``{deployment}:{cache_key_suffix}`` so that different exception types are - tracked independently per deployment. + ``deployment:{deployment}:allowed_fails:{cache_key_suffix}`` so that different + exception types are tracked independently per deployment. Returns: - True if fails exceed the allowed limit (should cooldown) @@ -584,16 +588,25 @@ def should_cooldown_based_on_allowed_fails_policy( else (litellm_router_instance.cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS) ) - cache_key: Final = f"{deployment}:{cache_key_suffix}" if cache_key_suffix else deployment - current_fails: Final = litellm_router_instance.failed_calls.get_cache(key=cache_key) or 0 - updated_fails: Final = current_fails + 1 + base_key: Final = f"deployment:{deployment}:allowed_fails" + cache_key: Final = f"{base_key}:{cache_key_suffix}" if cache_key_suffix else base_key + updated_fails: Final = _increment_allowed_fails( + cache=litellm_router_instance.cache, cache_key=cache_key, ttl=cooldown_time + ) + return updated_fails > allowed_fails - if updated_fails > allowed_fails: - return True - else: - litellm_router_instance.failed_calls.set_cache(key=cache_key, value=updated_fails, ttl=cooldown_time) - return False +def _increment_allowed_fails(cache: DualCache, cache_key: str, ttl: float) -> int: + """ + Return the fleet-wide fail count. ``DualCache.increment_cache`` bumps the in-memory tier + before Redis and re-raises a Redis error, so a Redis outage degrades to this worker's own count. + """ + try: + return cache.increment_cache(key=cache_key, value=1, ttl=ttl) + except Exception as e: # noqa: BLE001 # a Redis outage must not stop failing deployments from cooling down + verbose_router_logger.warning("allowed_fails counter fell back to this worker's in-memory count: %s", e) + local_fails: Final = cache.get_cache(key=cache_key, local_only=True) + return local_fails if isinstance(local_fails, int) else 0 def _is_allowed_fails_set_on_router( diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index 22d816e13e9..c8ca7105392 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -12,6 +12,7 @@ from typing_extensions import TypedDict from litellm import verbose_logger from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -27,6 +28,16 @@ class DeploymentHealthStateValue(TypedDict): reason: str +def _read_shared_health_snapshot(cache: DualCache, key: str) -> object: + redis_cache: Final = cache.redis_cache + if redis_cache is None: + return None + try: + return redis_cache.get_cache(key) + except RedisCircuitBreakerOpenError: + return None + + class DeploymentHealthCache: """ Cache for deployment health states produced by background health checks. @@ -50,13 +61,12 @@ class DeploymentHealthCache: coexist on the one shared entry without erasing each other's results. The snapshot is read from Redis when available, since a pod-local read would only ever see this writer's own previous merge. When the Redis - read comes back empty (a miss, or a swallowed connection error), the - pod-local copy of the last merge is used so peers are not erased. + read comes back empty (a miss, a swallowed connection error, or a read + refused by the open circuit breaker), the pod-local copy of the last + merge is used so peers are not erased. """ try: - redis_raw: Final = ( - self.cache.redis_cache.get_cache(self.CACHE_KEY) if self.cache.redis_cache is not None else None - ) + redis_raw: Final = _read_shared_health_snapshot(self.cache, self.CACHE_KEY) raw: Final = redis_raw if isinstance(redis_raw, dict) else self.cache.get_cache(key=self.CACHE_KEY) existing: Final = raw if isinstance(raw, dict) else {} expiry_seconds: Final = self.staleness_threshold * 1.5 diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 6f3ea8eb78a..c7eb46046ef 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -349,7 +349,9 @@ class DeploymentAffinityCheck(CustomLogger): first write instead of the last. Re-claiming with the stored value refreshes its TTL, the same keepalive the complexity router's model pin documents: an active session must not lose its pin mid-conversation just because it outlives the - original write, so `session_affinity_ttl_seconds` bounds idle time, not total + original write, so the affinity TTL (the Router's + `deployment_affinity_ttl_seconds`, or a pre-routing hook's per-request + `session_affinity_ttl_seconds` override) bounds idle time, not total session length. On Redis one Lua script does the get-or-set-or-refresh atomically (same registration seam the rate limiters use) and the in-memory tier is synchronized to the winner; without Redis, and whenever Redis is diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index b623e31ce06..cdd70e6baf2 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -37,17 +37,21 @@ Safe to enable globally: """ import time +from collections.abc import Iterator, Mapping from typing import TYPE_CHECKING, Final, Optional, Protocol, cast import httpx from litellm._logging import verbose_router_logger from litellm.exceptions import ( - BadRequestError, RateLimitError, ServiceUnavailableError, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_content_of_block, + strip_encrypted_reasoning_from_messages, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues @@ -138,15 +142,48 @@ class EncryptedContentAffinityCheck(CustomLogger): # If no encoded ID, check if encrypted_content itself is wrapped encrypted_content = item.get("encrypted_content") if encrypted_content and isinstance(encrypted_content, str): - ( - model_id, - _, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content) + model_id = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content) if model_id: return model_id return None + @staticmethod + def _anthropic_content_blocks(messages: object) -> Iterator[Mapping[str, object]]: + if not isinstance(messages, list): + return iter(()) + return ( + cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance + for message in cast(list[object], messages) # cast-ok: narrowed by isinstance + if isinstance(message, Mapping) + for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance + if isinstance(content, list) + for block in cast(list[object], content) # cast-ok: narrowed by isinstance + if isinstance(block, Mapping) + ) + + @staticmethod + def _model_id_from_wrapped_encrypted_content(encrypted_content: str) -> str | None: + model_id, _ = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content) + return model_id or None + + @staticmethod + def _extract_model_id_from_anthropic_messages(messages: object) -> str | None: + return next( + ( + model_id + for block in EncryptedContentAffinityCheck._anthropic_content_blocks(messages) + if (encrypted_content := encrypted_content_of_block(block)) is not None + if ( + model_id := EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content( + encrypted_content + ) + ) + is not None + ), + None, + ) + @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: for deployment in healthy_deployments: @@ -158,6 +195,23 @@ class EncryptedContentAffinityCheck(CustomLogger): return deployment return None + @staticmethod + def _request_team_id(request_kwargs: Mapping[str, object]) -> str | None: + containers: Final = (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + team_ids: Final = (c.get("user_api_key_team_id") for c in containers if isinstance(c, Mapping)) + return next((tid for tid in team_ids if isinstance(tid, str)), None) + + def _routed_group_candidate_model_ids(self, request_kwargs: Mapping[str, object], model: str) -> frozenset[str]: + """ + Deployment ids that could serve this turn's routed ``model``, as the router + resolves a route (model_group_alias / routing group / model_name / team / + pattern). Delegates to the router so the full precedence is not re-derived here + and no deployment ids are written into request kwargs bound for the provider. + """ + if self.router is None: + return frozenset() + return self.router.get_candidate_model_ids_for_route(model=model, team_id=self._request_team_id(request_kwargs)) + @staticmethod def _encryption_boundary_key( litellm_params: object, @@ -223,12 +277,17 @@ class EncryptedContentAffinityCheck(CustomLogger): parent_otel_span: Span | None = None, ) -> list[dict]: """ - If the request ``input`` contains litellm-encoded item IDs, decode the - embedded ``model_id`` and pin the request to that deployment. Raises - ``RateLimitError`` / ``ServiceUnavailableError`` / ``BadRequestError`` - when the originating deployment is unavailable and no encryption-boundary - peer exists, rather than dispatching a doomed request to a non-peer - deployment. The 429/503 split mirrors the originating cooldown's status: + If the request ``input`` contains litellm-encoded item IDs, or its Anthropic + ``messages`` replay a bridge-tagged thinking block, decode the embedded + ``model_id`` and pin the request to that deployment. Raises + ``RateLimitError`` / ``ServiceUnavailableError`` when the originating + deployment is a member of the routed model group but currently unavailable + and no encryption-boundary peer exists, rather than dispatching a doomed + request to a non-peer deployment. When the origin is not a member of the + routed group (an auto-router tier change, a model switch with no peer, a + removed deployment, or an unknown/forged marker), the encrypted reasoning is + stripped and the request dispatches with its readable history instead. The + 429/503 split mirrors the originating cooldown's status: a 429-induced cooldown surfaces as 429 (with ``Retry-After`` set to the remaining cooldown window) so OpenAI-compatible clients back off and retry after the deployment is eligible again. @@ -249,12 +308,15 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True request_input: Final = request_kwargs.get("input") - model_id: Final = self._extract_model_id_from_input(request_input) + anthropic_messages: Final = messages or request_kwargs.get("messages") + model_id: Final = self._extract_model_id_from_input( + request_input + ) or self._extract_model_id_from_anthropic_messages(anthropic_messages) if not model_id: return typed_healthy_deployments verbose_router_logger.debug( - "EncryptedContentAffinityCheck: decoded model_id=%s from input item IDs", + "EncryptedContentAffinityCheck: decoded model_id=%s from the request's encrypted content markers", model_id, ) @@ -285,12 +347,35 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["_encrypted_content_affinity_pinned"] = True return boundary_matches - # Dispatching to a non-peer would guarantee an upstream - # `invalid_encrypted_content` 400, so fail fast with a clearer error. + # The origin cannot serve this turn's routed group and no peer shares the boundary, so its + # encrypted reasoning can never decrypt here. Strip it, keep the readable history, and dispatch + # to the routed group instead of failing. Membership is tested by deployment id against the set + # the router actually resolved for this route, not by model-group name, so an alias, a + # provider-qualified spelling, a team-public name, or a pattern route of the same group is not + # mistaken for a tier change. An unknown origin (a removed deployment, or a forged marker) is + # treated the same as a cross-group one, which also denies an authenticated caller a + # deployment-id existence oracle: a real cross-group id and a nonexistent id both strip and + # dispatch rather than returning distinguishable responses. Only a genuine same-group member + # that is currently unavailable falls through to the fail-fast, preserving the cooldown contract. + routed_group_model_ids: Final = ( + self._routed_group_candidate_model_ids(request_kwargs, model) if originating is not None else frozenset() + ) + if str(model_id) not in routed_group_model_ids: + verbose_router_logger.debug( + "EncryptedContentAffinityCheck: model_id=%s is not a candidate for the routed group %s; " + "forwarding without its encrypted reasoning", + model_id, + model, + ) + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + strip_encrypted_reasoning_from_messages(anthropic_messages) + return typed_healthy_deployments + + # The origin is a member of the routed group but currently unavailable (cooled down); fail fast + # rather than dispatching to a non-peer, which would guarantee an upstream 400. raise await self._unavailable_origin_error( model=model, model_id=model_id, - originating=originating, parent_otel_span=parent_otel_span, ) @@ -298,25 +383,11 @@ class EncryptedContentAffinityCheck(CustomLogger): self, model: str, model_id: str, - originating: Deployment | None, parent_otel_span: Span | None, ) -> Exception: # Public error messages intentionally omit the originating ``model_id`` so # an authenticated caller forging encrypted-content markers cannot use the # error surface to enumerate which deployment IDs exist on this router. - if originating is None: - return BadRequestError( - message=( - "The deployment that produced this encrypted_content is no " - "longer configured on this router, and no deployment on the " - "same encryption boundary is available. Re-issue the request " - "without the stale encrypted_content items, or restore the " - "originating deployment." - ), - model=model, - llm_provider="", - ) - cooldown: Final = await self._get_origin_cooldown(model_id=model_id, parent_otel_span=parent_otel_span) if cooldown is not None and str(cooldown.get("status_code")) == "429": diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 01d42627001..fbd3e18e357 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -21,6 +21,7 @@ import litellm from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.router import RouterCacheEnum, RouterErrors from litellm.utils import get_utc_datetime @@ -466,7 +467,7 @@ async def async_io_token_pre_call_check( request_kwargs: Final = get_io_token_rate_limit_request_kwargs() _model: Final = (deployment.get("litellm_params") or {}).get("model") or "" - estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) + estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) dt: Final = get_utc_datetime() diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 0788c8db710..0589e290b47 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -14,6 +14,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -61,7 +62,7 @@ class PromptCachingDeploymentCheck(CustomLogger): if request_kwargs is not None and request_kwargs.get("_target_order") is not None: return healthy_deployments - if messages is not None and is_prompt_caching_valid_prompt( + if messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)( messages=messages, model=model, min_token_count=_get_min_token_count_for_deployments(healthy_deployments), @@ -90,6 +91,7 @@ class PromptCachingDeploymentCheck(CustomLogger): enable_prompt_caching=( request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None ), + request_kwargs=request_kwargs, ) model_id_dict: Final = await prompt_cache.async_get_model_id( @@ -138,7 +140,7 @@ class PromptCachingDeploymentCheck(CustomLogger): return ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider - if is_prompt_caching_valid_prompt( + if await offload_token_count(is_prompt_caching_valid_prompt)( model=model, messages=cast(list[AllMessageValues], messages), ): diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py index 9185d901a28..7b145c15a07 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -10,8 +10,8 @@ opt-in. none is opt-out everywhere except the azure gpt-5 family, whose config r UnsupportedParamsError without an explicit true. xhigh is gated on the request path by the openai and azure gpt-5 configs. max is not gated there at -all: every entry carrying supports_max_reasoning_effort is Claude-family, and -anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort +all: outside the gpt-6-astra rows every entry carrying supports_max_reasoning_effort is Claude-family, +and anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort path maps any level to a thinking budget. Making max opt-in is a deliberate trade, then, since an explicit flag is the only signal that the tier is a real one rather than litellm rounding the level to a budget, and a missing flag costs advisory metadata rather than a rejected request. diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index c599667ab17..674bd8847f7 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -247,8 +247,7 @@ def rust_chat_completions_accepts( return False if stream: return False - request_override: Final = litellm_params.get("rust") if litellm_params is not None else None - if not rust_enabled(request_override=request_override if isinstance(request_override, bool) else None): + if not rust_enabled(): return False if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 515ab6edef1..5582027bb5d 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -1,13 +1,11 @@ from __future__ import annotations import os -import warnings from typing import Final DEFAULT_RUST_ENABLED: Final = False _TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"}) _GLOBAL_ENV_NAME: Final = "LITELLM_RUST" -_LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR" class _RustConfiguration: @@ -26,49 +24,24 @@ def _parse_env_bool(value: str | None) -> bool | None: def resolve_rust_enabled( *, - request_override: bool | None, process_override: bool | None, environment_override: bool | None, - legacy_environment_override: bool | None = None, release_default: bool = DEFAULT_RUST_ENABLED, ) -> bool: - if request_override is not None: - return request_override if process_override is not None: return process_override if environment_override is not None: return environment_override - if legacy_environment_override is not None: - return legacy_environment_override return release_default -def rust_enabled(*, request_override: bool | None = None) -> bool: - if request_override is not None: - return request_override - process_override: Final = _CONFIGURATION.override - if process_override is not None: - return process_override - global_override: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)) - legacy_override: Final = None if global_override is not None else _parse_env_bool(os.getenv(_LEGACY_OCR_ENV_NAME)) - if legacy_override is not None: - warnings.warn( - f"{_LEGACY_OCR_ENV_NAME} is deprecated; use {_GLOBAL_ENV_NAME} instead", - DeprecationWarning, - stacklevel=2, - ) +def rust_enabled() -> bool: return resolve_rust_enabled( - request_override=None, - process_override=None, - environment_override=global_override, - legacy_environment_override=legacy_override, + process_override=_CONFIGURATION.override, + environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)), ) -def rust_ocr_enabled(*, request_override: bool | None = None) -> bool: - return rust_enabled(request_override=request_override) - - def reset_rust_configuration() -> None: _CONFIGURATION.override = None diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 86038438f57..89eab71ccba 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,16 +2,57 @@ from __future__ import annotations -from collections.abc import Awaitable +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables import httpx -from litellm.rust_bridge import configuration as _configuration +import litellm +from litellm.constants import request_timeout +from litellm.llms.azure_ai.ocr.common_utils import is_azure_cohere_parse_model +from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse +from litellm.rust_bridge.bindings import NativeBinding, native_exception_types from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ProviderConfigManager -rust_ocr_enabled = _configuration.rust_ocr_enabled -rust = _configuration.rust +_RUST_OCR_PROVIDERS: Final = frozenset({"mistral", "azure_ai", "vertex_ai"}) +_RUST_OCR_CONFIG_FIELDS: Final = frozenset( + { + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_scope", + "azure_authority_host", + "azure_credential", + "azure_federated_token_file", + "vertex_credentials", + "vertex_ai_credentials", + "vertex_project", + "vertex_ai_project", + "vertex_location", + "vertex_ai_location", + } +) +_RUST_OCR_SECRET_FIELDS: Final = frozenset( + {"azure_ad_token", "client_secret", "azure_federated_token_file", "vertex_credentials", "vertex_ai_credentials"} +) + + +@dataclass(frozen=True, slots=True) +class LiteLLMOcrRequest: + model: str + document: Mapping[str, object] + api_key: str | None + api_base: str | None + timeout: float | httpx.Timeout | None + custom_llm_provider: str | None + extra_headers: dict[str, object] | None + kwargs: Mapping[str, object] + input_sources: Mapping[str, str] | None = None class RustOcr(Protocol): @@ -24,6 +65,7 @@ class RustOcr(Protocol): custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: raise NotImplementedError @@ -39,54 +81,308 @@ class RustAocr(Protocol): custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> Awaitable[dict[str, object]]: raise NotImplementedError -class _Unset: - pass +class _OCRLogging(Protocol): + def update_from_kwargs( + self, + *, + kwargs: dict[str, object], + model: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + custom_llm_provider: str | None, + ) -> None: ... + + def pre_call( + self, + *, + input: str, + api_key: str | None, + additional_args: dict[str, object], + ) -> None: ... -_UNSET: Final[_Unset] = _Unset() +def _as_ocr(value: object) -> RustOcr | None: + return cast(RustOcr, value) if callable(value) else None -_rust_ocr_impl: RustOcr | None = None -_rust_aocr_impl: RustAocr | None = None +def _as_aocr(value: object) -> RustAocr | None: + return cast(RustAocr, value) if callable(value) else None -def set_rust_ocr( - *, - ocr: RustOcr | None | _Unset = _UNSET, - aocr: RustAocr | None | _Unset = _UNSET, -) -> None: - global _rust_ocr_impl, _rust_aocr_impl - if not isinstance(ocr, _Unset): - _rust_ocr_impl = ocr - if not isinstance(aocr, _Unset): - _rust_aocr_impl = aocr +_OCR: Final = NativeBinding("ocr", validate=_as_ocr) +_AOCR: Final = NativeBinding("aocr", validate=_as_aocr) def load_rust_ocr() -> RustOcr | None: - if _rust_ocr_impl is not None: - return _rust_ocr_impl - from litellm.rust_bridge import get_native_bridge - - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustOcr, native_bridge.ocr) + return _OCR.load() def load_rust_aocr() -> RustAocr | None: - if _rust_aocr_impl is not None: - return _rust_aocr_impl - from litellm.rust_bridge import get_native_bridge + return _AOCR.load() - native_bridge: Final = get_native_bridge() - if native_bridge is None: + +def provider(request: LiteLLMOcrRequest) -> str | None: + if request.custom_llm_provider is not None: + return request.custom_llm_provider + prefix: Final = request.model.partition("/")[0] + if prefix in _RUST_OCR_PROVIDERS: + return prefix + if request.model.startswith("mistral-ocr"): + return "mistral" + return None + + +def supported(request: LiteLLMOcrRequest) -> bool: + request_provider: Final = provider(request) + if request_provider not in _RUST_OCR_PROVIDERS: + return False + if request_provider == "azure_ai": + return ( + not is_azure_cohere_parse_model(request.model) + and not callable(request.kwargs.get("azure_ad_token_provider")) + and request.kwargs.get("azure_username") is None + and request.kwargs.get("azure_password") is None + ) + return True + + +def _optional_params(request: LiteLLMOcrRequest, resolve_secret: Callable[[str], str | None]) -> Mapping[str, object]: + optional_params: Final = MappingProxyType( + { + name: value + for name, value in request.kwargs.items() + if (name not in GenericLiteLLMParams.model_fields or name in _RUST_OCR_CONFIG_FIELDS) + and name not in ("litellm_logging_obj", "aocr", "litellm_call_id", "proxy_server_request") + } + ) + request_provider: Final = provider(request) + if request_provider == "azure_ai" and litellm.enable_azure_ad_token_refresh is True: + return MappingProxyType({**optional_params, "enable_azure_ad_token_refresh": True}) + if request_provider != "vertex_ai": + return optional_params + project: Final = ( + request.kwargs.get("vertex_project") + or request.kwargs.get("vertex_ai_project") + or litellm.vertex_project + or resolve_secret("VERTEXAI_PROJECT") + ) + location: Final = ( + request.kwargs.get("vertex_location") + or request.kwargs.get("vertex_ai_location") + or litellm.vertex_location + or resolve_secret("VERTEXAI_LOCATION") + or resolve_secret("VERTEX_LOCATION") + ) + credentials: Final = ( + request.kwargs.get("vertex_credentials") + or request.kwargs.get("vertex_ai_credentials") + or resolve_secret("VERTEXAI_CREDENTIALS") + ) + vertex_params: Final = MappingProxyType( + { + name: value + for name, value in ( + ("vertex_project", project), + ("vertex_location", location), + ("vertex_credentials", credentials), + ) + if value is not None + } + ) + return MappingProxyType({**optional_params, **vertex_params}) + + +def _input_sources(request: LiteLLMOcrRequest, optional_params: Mapping[str, object]) -> Mapping[str, str]: + proxy_request_value: Final = request.kwargs.get("proxy_server_request") + if not isinstance(proxy_request_value, Mapping): + return MappingProxyType({}) + proxy_request: Final = cast( # cast-ok: runtime Mapping check narrows metadata with unknown key and value types + Mapping[object, object], proxy_request_value + ) + credential_fields_value: Final = proxy_request.get("credential_fields", ()) + credential_fields: Final = ( + frozenset(name for name in credential_fields_value if isinstance(name, str)) + if isinstance(credential_fields_value, (list, tuple, set, frozenset)) + else frozenset() + ) + request_fields_value: Final = proxy_request.get("body_fields") + request_fields: Sequence[object] + if isinstance(request_fields_value, Sequence) and not isinstance(request_fields_value, (str, bytes)): + request_fields = cast( # cast-ok: runtime Sequence check excludes scalar strings and bytes + Sequence[object], request_fields_value + ) + else: + body_value: Final = proxy_request.get("body") + request_fields = ( + tuple(cast(Mapping[object, object], body_value)) # cast-ok: runtime Mapping check establishes iterable keys + if isinstance(body_value, Mapping) + else () + ) + names: Final = frozenset(optional_params) | frozenset({"api_key", "api_base", "extra_headers"}) + request_sources: Final = MappingProxyType( + {name: "request" for name in names if name in request_fields or name in credential_fields} + ) + if litellm.enable_azure_ad_token_refresh is True and "enable_azure_ad_token_refresh" in optional_params: + return MappingProxyType({**request_sources, "enable_azure_ad_token_refresh": "deployment"}) + return request_sources + + +def _marshal( + request: LiteLLMOcrRequest, + resolve_secret: Callable[[str], str | None], + convert_file_document: Callable[[dict[str, object]], dict[str, str]], +) -> LiteLLMOcrRequest: + if not isinstance(request.document, dict): + raise TypeError(f"document must be a dict with 'type' and URL/file field, got {type(request.document)}") + document: Final = ( + convert_file_document(request.document) if request.document.get("type") == "file" else request.document + ) + request_provider: Final = provider(request) + api_key: Final = ( + request.api_key or resolve_secret("MISTRAL_API_KEY") if request_provider == "mistral" else request.api_key + ) + optional_params: Final = _optional_params(request, resolve_secret) + input_sources: Final = _input_sources(request, optional_params) + logged_optional_params: Final = MappingProxyType( + {name: "****" if name in _RUST_OCR_SECRET_FIELDS else value for name, value in optional_params.items()} + ) + logged_kwargs: Final = MappingProxyType( + { + name: "****" if name in _RUST_OCR_SECRET_FIELDS else value + for name, value in request.kwargs.items() + if name != "proxy_server_request" + } + ) + logging_obj: Final = cast( # cast-ok: client decorator injects the logging object through untyped kwargs + _OCRLogging, request.kwargs["litellm_logging_obj"] + ) + logging_obj.update_from_kwargs( + kwargs=dict(logged_kwargs), # mutable-ok: legacy logging mutates its kwargs copy + model=request.model, + optional_params=dict(logged_optional_params), # mutable-ok: legacy logging requires concrete dict params + litellm_params={ # mutable-ok: legacy logging requires a concrete params dict + "litellm_call_id": request.kwargs.get("litellm_call_id"), + "api_base": request.api_base, + }, + custom_llm_provider=request_provider, + ) + logging_obj.pre_call( + input="OCR document processing", + api_key=api_key, + additional_args={ # mutable-ok: pre_call mutates the additional_args dict + "complete_input_dict": { # mutable-ok: callbacks consume a JSON-serializable request dict + "model": request.model, + "document": document, + **logged_optional_params, + }, + "api_base": request.api_base or "", + "headers": request.extra_headers or {}, # mutable-ok: logging callbacks consume a concrete headers dict + }, + ) + return LiteLLMOcrRequest( + model=request.model, + document=document, + api_key=api_key, + api_base=request.api_base, + timeout=request.timeout if request.timeout is not None else request_timeout, + custom_llm_provider=request.custom_llm_provider, + extra_headers=request.extra_headers, + kwargs=optional_params, + input_sources=input_sources, + ) + + +def _map_error(error: Exception, request: LiteLLMOcrRequest) -> Exception: + exception_types: Final = native_exception_types() + if exception_types is None or not isinstance(error, exception_types[1]): + return error + request_provider: Final = provider(request) + if request_provider is None: + return error + provider_config: Final = ProviderConfigManager.get_provider_ocr_config( + model=request.model.removeprefix(f"{request_provider}/"), provider=litellm.LlmProviders(request_provider) + ) + if provider_config is None: + return error + error_args: Final = cast( # cast-ok: BaseException.args exposes Any while native errors carry scalar args + tuple[object, ...], error.args + ) + status: Final = error_args[0] if error_args and isinstance(error_args[0], int) else 500 + message: Final = str(error_args[1]) if len(error_args) > 1 else str(error) + error_factory: Final = cast( # cast-ok: legacy provider error factories have untyped callable parameters + Callable[..., Exception], provider_config.get_error_class + ) + return error_factory( + error_message=message, + status_code=status or 500, + headers={}, # mutable-ok: provider error factories require a concrete headers dict + ) + + +def _response(response: Mapping[str, object]) -> OCRResponse: + provider_native_response: Final = response.get(PROVIDER_NATIVE_RESPONSE_KEY) + normalized: Final = OCRResponse.model_validate( + MappingProxyType({key: value for key, value in response.items() if key != PROVIDER_NATIVE_RESPONSE_KEY}) + ) + if isinstance(provider_native_response, Mapping): + normalized.set_provider_native_response(provider_native_response) + return normalized + + +def run( + request: LiteLLMOcrRequest, + resolve_secret: Callable[[str], str | None], + convert_file_document: Callable[[dict[str, object]], dict[str, str]], +) -> OCRResponse | None: + if load_rust_ocr() is None: return None - return cast(RustAocr, getattr(native_bridge, "aocr", None)) + marshalled: Final = _marshal(request, resolve_secret, convert_file_document) + try: + response: Final = ocr( + model=marshalled.model, + document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict + api_key=marshalled.api_key, + api_base=marshalled.api_base, + custom_llm_provider=marshalled.custom_llm_provider, + extra_headers=marshalled.extra_headers, + optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict + input_sources=marshalled.input_sources, + timeout=marshalled.timeout, + ) + except Exception as error: + raise _map_error(error, request) from error + return _response(response) if response is not None else None + + +async def arun( + request: LiteLLMOcrRequest, + resolve_secret: Callable[[str], str | None], + convert_file_document: Callable[[dict[str, object]], dict[str, str]], +) -> OCRResponse | None: + if load_rust_aocr() is None: + return None + marshalled: Final = _marshal(request, resolve_secret, convert_file_document) + try: + response: Final = await aocr( + model=marshalled.model, + document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict + api_key=marshalled.api_key, + api_base=marshalled.api_base, + custom_llm_provider=marshalled.custom_llm_provider, + extra_headers=marshalled.extra_headers, + optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict + input_sources=marshalled.input_sources, + timeout=marshalled.timeout, + ) + except Exception as error: + raise _map_error(error, request) from error + return _response(response) if response is not None else None def ocr( @@ -99,6 +395,7 @@ def ocr( extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, + input_sources: Mapping[str, str] | None = None, ) -> dict[str, object] | None: rust_ocr: Final = load_rust_ocr() if rust_ocr is None: @@ -111,6 +408,7 @@ def ocr( custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, optional_params=optional_params, + input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict timeout_seconds=_timeout_to_seconds(timeout), ) @@ -125,6 +423,7 @@ async def aocr( extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, + input_sources: Mapping[str, str] | None = None, ) -> dict[str, object] | None: rust_aocr: Final = load_rust_aocr() if rust_aocr is None: @@ -137,5 +436,6 @@ async def aocr( custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, optional_params=optional_params, + input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict timeout_seconds=_timeout_to_seconds(timeout), ) diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py new file mode 100644 index 00000000000..ee755b317d8 --- /dev/null +++ b/litellm/rust_bridge/token_counter.py @@ -0,0 +1,79 @@ +"""Thin Python wrapper for the native Rust input token counter.""" + +from __future__ import annotations + +from collections.abc import Awaitable +from dataclasses import dataclass +from functools import lru_cache +from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables + +from pydantic import TypeAdapter + +import litellm +from litellm._logging import verbose_logger +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.configuration import rust_enabled +from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt +from litellm.utils import claude_json_str +from litellm.utils import uses_anthropic_tokenizer as _python_uses_anthropic_tokenizer + + +class RustTokenCounter(Protocol): + def acount_request(self, body: bytes) -> Awaitable[object]: + raise NotImplementedError + + +class RustTokenCounterFactory(Protocol): + def __call__(self, tokenizer_json: str) -> RustTokenCounter: + raise NotImplementedError + + +@dataclass(frozen=True, slots=True) +class InputTokenCount: + model: str | None + input_tokens: int + + +_INPUT_TOKEN_COUNT: Final = TypeAdapter(InputTokenCount) + + +def _as_factory(value: object) -> RustTokenCounterFactory | None: + return ( + cast( # cast-ok: native extension protocol is runtime-defined + RustTokenCounterFactory, value + ) + if callable(value) + else None + ) + + +TOKEN_COUNTER: Final = NativeBinding("TokenCounter", validate=_as_factory) + + +def uses_anthropic_tokenizer(model: str) -> bool: + if litellm.disable_token_counter is True or litellm.disable_hf_tokenizer_download is True: + return False + return _python_uses_anthropic_tokenizer(model) + + +@lru_cache(maxsize=4) +def _anthropic_counter(factory: RustTokenCounterFactory) -> RustTokenCounter: + return factory(claude_json_str) + + +async def count_anthropic_input_tokens(body: bytes) -> InputTokenCount | None: + if not rust_enabled(): + return None + factory: Final = TOKEN_COUNTER.load() + if factory is None: + return None + try: + attempt: Final = await aattempt( + native_call=lambda: _anthropic_counter(factory).acount_request(body), + adapt=_INPUT_TOKEN_COUNT.validate_python, + context=BridgeErrorContext(route="token_counter", provider="anthropic", model=""), + ) + except (RuntimeError, ValueError) as error: + verbose_logger.debug("Rust token counter failed, counting in Python: %s", error) + return None + return attempt.value if isinstance(attempt, RustHandled) else None diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index 3d71f6f8a50..6c81786accd 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -56,7 +56,6 @@ _STATE: Final = _RustTranscriptionState() def configure_rust_transcription( - enabled: bool = True, *, transcription: RustTranscription | None | _Unset = _UNSET, atranscription: RustAtranscription | None | _Unset = _UNSET, diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index e86c8e7c919..d75375a01cc 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -47,6 +47,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): aws_web_identity_token: str | None = None, aws_sts_endpoint: str | None = None, replica_regions: list[str] | None = None, + kms_key_id: str | None = None, **kwargs, ): BaseSecretManager.__init__(self, **kwargs) @@ -61,6 +62,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): self.aws_web_identity_token = aws_web_identity_token self.aws_sts_endpoint = aws_sts_endpoint self.replica_regions: list[str] = replica_regions or [] + self.kms_key_id = kms_key_id @classmethod def validate_environment(cls): @@ -106,7 +108,8 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): # Remove None values aws_kwargs = {k: v for k, v in aws_kwargs.items() if v is not None} - litellm.secret_manager_client = cls(**aws_kwargs) + kms_key_id: Final = key_management_settings.kms_key_id if key_management_settings is not None else None + litellm.secret_manager_client = cls(kms_key_id=kms_key_id, **aws_kwargs) litellm._key_management_system = KeyManagementSystem.AWS_SECRET_MANAGER except Exception as e: @@ -275,6 +278,9 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): if description: data["Description"] = description + if self.kms_key_id: + data["KmsKeyId"] = self.kms_key_id + # ✅ Normalize tags to AWS format if tags: if isinstance(tags, dict): diff --git a/litellm/skills/main.py b/litellm/skills/main.py index 002419dbad4..71fd78f11a3 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -182,10 +182,14 @@ def create_skill( if extra_body: create_request.update(extra_body) - # Route to LiteLLM DB if custom_llm_provider="litellm_proxy" + # Route to LiteLLM DB if custom_llm_provider="litellm_proxy". description/instructions + # arrive as top-level kwargs from the REST form endpoint, or nested in extra_body from + # the SDK convention used by other providers' create_request above. if custom_llm_provider == LlmProviders.LITELLM_PROXY.value: return _get_litellm_skills_handler().create_skill_handler( display_title=display_title, + description=kwargs.get("description") or (extra_body.get("description") if extra_body else None), + instructions=kwargs.get("instructions") or (extra_body.get("instructions") if extra_body else None), files=files, metadata=_get_skill_request_metadata(kwargs, extra_body), user_id=kwargs.get("user_id"), diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 2cb42ce3fac..dbaaab62d86 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -1,9 +1,9 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal from pydantic import BaseModel, PrivateAttr, StrictInt -from typing_extensions import Required, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -172,6 +172,7 @@ class AugmentedAgentCard(AgentCard): class AgentObjectPermission(TypedDict, total=False): mcp_servers: list[str] | None mcp_access_groups: list[str] | None + mcp_toolsets: ReadOnly[Sequence[str] | None] mcp_tool_permissions: dict[str, list[str]] | None models: list[str] | None agents: list[str] | None diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index ef28181eba5..02dee40f2a3 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -778,6 +778,9 @@ class ContentFilterConfigModel(BaseModel): ) +MCP_SECURITY_ON_VIOLATION: Final = frozenset({"block", "alert"}) + + class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails api_key: str | None = Field(default=None, description="API key for the guardrail service") api_base: str | None = Field(default=None, description="Base URL for the guardrail service API") @@ -886,9 +889,13 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up default=None, description="For /v1/realtime sessions: automatically close the session after this many guardrail violations.", ) - on_violation: Literal["warn", "end_session"] | None = Field( + on_violation: Literal["warn", "end_session", "block", "alert"] | None = Field( default=None, - description="For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.", + description=( + "For /v1/realtime sessions: 'warn' speaks the violation message and continues; " + "'end_session' speaks the message and closes the connection. " + "For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning." + ), ) realtime_violation_message: str | None = Field( default=None, @@ -1093,6 +1100,15 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o except (TypeError, ValueError) as e: raise ValueError(f"timeout must be numeric, got {v!r}") from e + @model_validator(mode="after") + def validate_on_violation_for_guardrail(self) -> "LitellmParams": + if ( + self.on_violation in MCP_SECURITY_ON_VIOLATION + and self.guardrail != SupportedGuardrailIntegrations.MCP_SECURITY.value + ): + raise ValueError(f"on_violation={self.on_violation!r} is only supported by guardrail='mcp_security'") + return self + def __init__(self, **kwargs) -> None: default_on: Final = kwargs.pop("default_on", None) if default_on is not None: diff --git a/litellm/types/integrations/datadog_llm_obs.py b/litellm/types/integrations/datadog_llm_obs.py index 17cf5831c96..f4fcabf53ed 100644 --- a/litellm/types/integrations/datadog_llm_obs.py +++ b/litellm/types/integrations/datadog_llm_obs.py @@ -87,6 +87,7 @@ class LLMMetrics(TypedDict, total=False): cache_write_input_tokens: ReadOnly[float] non_cached_input_tokens: ReadOnly[float] reasoning_output_tokens: ReadOnly[float] + tool_output_tokens: ReadOnly[float] class LLMObsPayload(TypedDict, total=False): diff --git a/litellm/types/integrations/newrelic.py b/litellm/types/integrations/newrelic.py index 36e4d02c2a8..b5905ad0b93 100644 --- a/litellm/types/integrations/newrelic.py +++ b/litellm/types/integrations/newrelic.py @@ -27,9 +27,9 @@ NEWRELIC_METRIC_ENDPOINT_BY_REGION: Final[Mapping[str, str]] = MappingProxyType( NEWRELIC_DEFAULT_REGION: Final = "us" #: Metric API caps a payload at 2000 data points / 1MB compressed; each queued -#: record expands to at most 6 metrics, so cap the per-flush record count well -#: below that. -NEWRELIC_METRICS_MAX_BATCH_SIZE: Final = 250 +#: record expands to at most 6 bucket metrics plus 2 team budget gauges (8), so cap +#: the per-flush record count well below 2000 / 8. +NEWRELIC_METRICS_MAX_BATCH_SIZE: Final = 200 #: Hard cap on records retained across failed flushes (5xx/network requeue). #: Beyond this the oldest records are dropped. @@ -48,6 +48,8 @@ NEWRELIC_METRIC_PROMPT_TOKENS: Final = "litellm.tokens.prompt" NEWRELIC_METRIC_COMPLETION_TOKENS: Final = "litellm.tokens.completion" NEWRELIC_METRIC_TOTAL_TOKENS: Final = "litellm.tokens.total" NEWRELIC_METRIC_REQUEST_DURATION_MS: Final = "litellm.request.duration_ms" +NEWRELIC_METRIC_TEAM_MAX_BUDGET: Final = "litellm.team.max_budget" +NEWRELIC_METRIC_TEAM_REMAINING_BUDGET: Final = "litellm.team.remaining_budget" class NewRelicSummaryValue(TypedDict): @@ -66,6 +68,13 @@ class NewRelicCountMetric(TypedDict): attributes: ReadOnly[Mapping[str, str]] +class NewRelicGaugeMetric(TypedDict): + name: ReadOnly[str] + type: ReadOnly[Literal["gauge"]] + value: ReadOnly[float] + attributes: ReadOnly[Mapping[str, str]] + + class NewRelicSummaryMetric(TypedDict): name: ReadOnly[str] type: ReadOnly[Literal["summary"]] @@ -73,7 +82,7 @@ class NewRelicSummaryMetric(TypedDict): attributes: ReadOnly[Mapping[str, str]] -NewRelicMetric = NewRelicCountMetric | NewRelicSummaryMetric +NewRelicMetric = NewRelicCountMetric | NewRelicGaugeMetric | NewRelicSummaryMetric #: ``interval.ms`` has a dot in it, so the functional TypedDict form is required. @@ -108,6 +117,8 @@ class NewRelicMetricRecord: completion_tokens: int total_tokens: int duration_ms: float + team_max_budget: float | None = None + team_spend: float | None = None @property def bucket_key(self) -> tuple[str, str, str, str, str, str]: diff --git a/litellm/types/integrations/pointfive.py b/litellm/types/integrations/pointfive.py new file mode 100644 index 00000000000..9ff17393ecd --- /dev/null +++ b/litellm/types/integrations/pointfive.py @@ -0,0 +1,45 @@ +from dataclasses import dataclass +from typing import Final + +from pydantic import Field + +from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams + +RETRYABLE_UPLOAD_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504}) + +DEFAULT_API_URL: Final = "https://api.pointfive.co/api/v1/ingestion" + + +class PointFiveInitParams(StandardCustomLoggerInitParams): + """ + Params for initializing a PointFive logger on litellm. + + Defaults trade freshness for fewer, larger uploads: every flush becomes one object, so + the interval is minutes rather than seconds. ``batch_size`` also bounds how much a busy + proxy holds in memory between flushes, so it stays modest. ``max_batch_bytes`` bounds + how much a single object may hold, which matters most when message logging is left on, + since an unredacted payload is orders of magnitude larger than a redacted one. + """ + + api_key: str | None = None + api_url: str | None = None + batch_size: int = Field(default=1_000, gt=0) + flush_interval: int = Field(default=300, gt=0) + max_batch_bytes: int = Field(default=8 * 1024 * 1024, gt=0) + max_upload_retries: int = Field(default=3, ge=1) + + +@dataclass(frozen=True, slots=True) +class PointFiveUploadTarget: + """A single-use presigned destination for one batch, issued by the PointFive API.""" + + upload_url: str + object_key: str + + +@dataclass(frozen=True, slots=True) +class PointFiveUploadFailure: + """Why a batch could not be uploaded, and whether a later attempt could still succeed.""" + + detail: str + retryable: bool diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 8498b6f6d00..a024581f600 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -154,6 +154,22 @@ LATENCY_BUCKETS: Final = ( float("inf"), ) +UNKNOWN_INPUT_SEQUENCE_LENGTH: Final = "unknown" +INPUT_SEQUENCE_LENGTH_BUCKETS: Final = ( + (1_000, "0-1k"), + (4_000, "1k-4k"), + (16_000, "4k-16k"), + (64_000, "16k-64k"), + (float("inf"), "64k+"), +) + + +def get_input_sequence_length_bucket(prompt_tokens: object) -> str: + if not isinstance(prompt_tokens, int) or isinstance(prompt_tokens, bool) or prompt_tokens < 0: + return UNKNOWN_INPUT_SEQUENCE_LENGTH + return next(label for upper, label in INPUT_SEQUENCE_LENGTH_BUCKETS if prompt_tokens < upper) + + # Batch jobs can run for minutes to hours; buckets span 1 min → 24 h. BATCH_DURATION_BUCKETS: Final = ( 60.0, @@ -205,6 +221,7 @@ class UserAPIKeyLabelNames(Enum): MCP_TOOL_NAME = "mcp_tool_name" MCP_SERVER_NAME = "mcp_server_name" SERVICE_TIER = "service_tier" + INPUT_SEQUENCE_LENGTH = "input_sequence_length" DEFINED_PROMETHEUS_METRICS = Literal[ @@ -857,6 +874,13 @@ class PrometheusMetricLabels: "litellm_images_generated_metric", } ) + _input_sequence_length_metrics: ClassVar[frozenset[str]] = frozenset( + { + "litellm_llm_api_latency_metric", + "litellm_llm_api_time_to_first_token_metric", + "litellm_request_total_latency_metric", + } + ) # Managed batch metrics _batch_user_labels = [ UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, @@ -955,14 +979,23 @@ class PrometheusMetricLabels: custom_labels.append(label) if label_name in PrometheusMetricLabels._org_label_metrics: - for label in [ + for label in ( UserAPIKeyLabelNames.ORG_ID.value, UserAPIKeyLabelNames.ORG_ALIAS.value, - ]: + ): if label not in default_labels and label not in custom_labels: custom_labels.append(label) - return default_labels + custom_labels + input_sequence_length_labels: Final = ( + (UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value,) + if ( + label_name in PrometheusMetricLabels._input_sequence_length_metrics + and litellm.prometheus_emit_input_sequence_length_label is True + and UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value not in custom_labels + ) + else () + ) + return [*default_labels, *custom_labels, *input_sequence_length_labels] _USER_API_KEY_LABEL_VALUE_INIT_ALIASES: Final[Mapping[str, str]] = MappingProxyType( @@ -1015,6 +1048,7 @@ class UserAPIKeyLabelValues: mcp_tool_name: str | None = None mcp_server_name: str | None = None service_tier: str | None = None + input_sequence_length: str | None = None # Added for test compatibility. def __init__(self, **kwargs: Any) -> None: diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index b3462203c4b..365d59a179b 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -520,8 +520,15 @@ ContentBlockContentBlockDict = ToolUseBlock | TextBlock | ChatCompletionThinking ContentBlockStart = ContentBlockStartToolUse | ContentBlockStartText +class AnthropicStopDetails(TypedDict, total=False): + type: ReadOnly[Literal["refusal"]] + category: ReadOnly[str | None] + explanation: ReadOnly[str | None] + + class MessageDelta(TypedDict, total=False): stop_reason: str | None + stop_details: 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): diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 4fe1dafc73b..038a23a3ca2 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -5,6 +5,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.types.llms.anthropic import ( AnthropicResponseContentBlockText, AnthropicResponseContentBlockToolUse, + AnthropicStopDetails, ContextManagementResponse, ServerToolUsage, ) @@ -78,16 +79,6 @@ class AnthropicUsage(TypedDict, total=False): server_tool_use: NotRequired[ReadOnly[ServerToolUsage]] -class AnthropicStopDetails(TypedDict, total=False): - """ - Safeguard verdict accompanying a `stop_reason: "refusal"` response: - https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback - """ - - category: ReadOnly[str | None] - explanation: ReadOnly[str | None] - - class AnthropicMessagesResponse(TypedDict, total=False): """ Anthropic Messages API Response: https://docs.anthropic.com/en/api/messages diff --git a/litellm/types/llms/anthropic_skills.py b/litellm/types/llms/anthropic_skills.py index 51eefe7154f..4b27f9b17ef 100644 --- a/litellm/types/llms/anthropic_skills.py +++ b/litellm/types/llms/anthropic_skills.py @@ -57,6 +57,15 @@ class Skill(BaseModel): updated_at: str """ISO 8601 timestamp of when the skill was last updated""" + description: str | None = None + """Description of the skill. Populated for the LiteLLM-hosted registry + (custom_llm_provider="litellm_proxy"); Anthropic's list endpoint does not + return a description, so this is None there.""" + + search_score: float | None = None + """Semantic similarity to the ``query`` passed to ``GET /v1/skills``. None + unless a query was given.""" + class ListSkillsResponse(BaseModel): """Response from listing skills""" diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index bed0ba3dc08..76756ac35bb 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,7 +1,7 @@ import json from collections.abc import Sequence from enum import Enum -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias from typing_extensions import ReadOnly, Required, TypedDict, override @@ -557,7 +557,7 @@ class AmazonTitanMultimodalEmbeddingResponse(TypedDict): message: str # Specifies any errors that occur during generation. -# TwelveLabs Marengo Embed 2.7 types +# TwelveLabs Marengo Embed types TWELVELABS_EMBEDDING_INPUT_TYPES = Literal["text", "image", "video", "audio"] TWELVELABS_EMBEDDING_OPTIONS = Literal["visual-text", "visual-image", "audio"] @@ -591,6 +591,113 @@ class TwelveLabsMarengoEmbeddingResponse(TypedDict): endSec: float +TWELVELABS_MARENGO_3_INPUT_TYPES: TypeAlias = Literal["text", "image", "video", "audio", "text_image", "multi_input"] +TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS: TypeAlias = Literal["visual", "audio", "transcription"] +TWELVELABS_MARENGO_3_EMBEDDING_TYPES: TypeAlias = Literal["separate_embedding", "fused_embedding"] +TWELVELABS_MARENGO_3_EMBEDDING_SCOPES: TypeAlias = Literal["clip", "asset"] + + +class TwelveLabsMarengo3FixedSegmentationConfig(TypedDict): + durationSec: ReadOnly[int] + + +class TwelveLabsMarengo3FixedSegmentation(TypedDict): + method: ReadOnly[Literal["fixed"]] + fixed: ReadOnly[TwelveLabsMarengo3FixedSegmentationConfig] + + +class TwelveLabsMarengo3DynamicSegmentationConfig(TypedDict): + minDurationSec: ReadOnly[int] + + +class TwelveLabsMarengo3DynamicSegmentation(TypedDict): + method: ReadOnly[Literal["dynamic"]] + dynamic: ReadOnly[TwelveLabsMarengo3DynamicSegmentationConfig] + + +TwelveLabsMarengo3Segmentation: TypeAlias = TwelveLabsMarengo3FixedSegmentation | TwelveLabsMarengo3DynamicSegmentation + + +class TwelveLabsMarengo3TextInput(TypedDict): + inputText: ReadOnly[str] + + +class TwelveLabsMarengo3ImageInput(TypedDict): + mediaSource: ReadOnly[TwelveLabsMediaSource] + + +class TwelveLabsMarengo3TimedMediaOptions(TypedDict, total=False): + startSec: ReadOnly[float] + endSec: ReadOnly[float] + segmentation: ReadOnly[TwelveLabsMarengo3Segmentation] + embeddingOption: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS]] + embeddingType: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_TYPES]] + embeddingScope: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_SCOPES]] + + +class TwelveLabsMarengo3TimedMediaInput(TwelveLabsMarengo3TimedMediaOptions): + mediaSource: Required[ReadOnly[TwelveLabsMediaSource]] + + +class TwelveLabsMarengo3TextImageInput(TypedDict): + inputText: ReadOnly[str] + mediaSource: ReadOnly[TwelveLabsMediaSource] + + +class TwelveLabsMarengo3NamedMediaSource(TwelveLabsMediaSource): + name: Required[ReadOnly[str]] + mediaType: Required[ReadOnly[Literal["image"]]] + + +class TwelveLabsMarengo3MultiInput(TypedDict, total=False): + inputText: ReadOnly[str] + mediaSources: Required[ReadOnly[Sequence[TwelveLabsMarengo3NamedMediaSource]]] + + +class TwelveLabsMarengo3RequestBase(TypedDict, total=False): + inferenceId: ReadOnly[str] + + +class TwelveLabsMarengo3TextRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["text"]] + text: ReadOnly[TwelveLabsMarengo3TextInput] + + +class TwelveLabsMarengo3ImageRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["image"]] + image: ReadOnly[TwelveLabsMarengo3ImageInput] + + +class TwelveLabsMarengo3VideoRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["video"]] + video: ReadOnly[TwelveLabsMarengo3TimedMediaInput] + + +class TwelveLabsMarengo3AudioRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["audio"]] + audio: ReadOnly[TwelveLabsMarengo3TimedMediaInput] + + +class TwelveLabsMarengo3TextImageRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["text_image"]] + text_image: ReadOnly[TwelveLabsMarengo3TextImageInput] + + +class TwelveLabsMarengo3MultiInputRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["multi_input"]] + multi_input: ReadOnly[TwelveLabsMarengo3MultiInput] + + +TwelveLabsMarengo3EmbeddingRequest: TypeAlias = ( + TwelveLabsMarengo3TextRequest + | TwelveLabsMarengo3ImageRequest + | TwelveLabsMarengo3VideoRequest + | TwelveLabsMarengo3AudioRequest + | TwelveLabsMarengo3TextImageRequest + | TwelveLabsMarengo3MultiInputRequest +) + + class TwelveLabsS3OutputDataConfig(TypedDict): s3Uri: str @@ -601,7 +708,7 @@ class TwelveLabsOutputDataConfig(TypedDict): class TwelveLabsAsyncInvokeRequest(TypedDict): modelId: str - modelInput: TwelveLabsMarengoEmbeddingRequest + modelInput: ReadOnly[TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest] outputDataConfig: TwelveLabsOutputDataConfig @@ -1000,6 +1107,11 @@ class BedrockTag(TypedDict): value: str +class AwsSessionTag(TypedDict): + Key: str # writable-ok: boto3's STS stubs type assume_role Tags as writable TagTypeDef, which rejects ReadOnly + Value: str # writable-ok: boto3's STS stubs type assume_role Tags as writable TagTypeDef, which rejects ReadOnly + + class BedrockCreateBatchRequest(TypedDict, total=False): """ Request structure for creating a Bedrock batch inference job. diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index e87a684aab8..a9c027bd2de 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -2,6 +2,7 @@ from typing import Any, Literal from pydantic import BaseModel from typing_extensions import ( + ReadOnly, Required, TypedDict, ) @@ -57,6 +58,14 @@ class DatabricksMessage(TypedDict, total=False): role: Required[str] content: Required[AllDatabricksContentValues] tool_calls: list[DatabricksTool] | None + reasoning_content: ReadOnly[str | None] + reasoning: ReadOnly[str | None] + + +class DatabricksDelta(TypedDict, total=False): + role: ReadOnly[str] + content: ReadOnly[AllDatabricksContentValues | None] + reasoning_content: ReadOnly[str | None] class DatabricksChoice(TypedDict, total=False): diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b6da9490e01..b7c4371f32f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1564,6 +1564,9 @@ class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject): response: ResponsesAPIResponse +ResponsesTerminalEvent: TypeAlias = ResponseCompletedEvent | ResponseIncompleteEvent | ResponseFailedEvent + + class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.RESPONSE_PART_ADDED] item_id: str diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 50c3515cf01..9f29f27e41d 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -195,6 +195,10 @@ class AutoRouterBenchmarkTotals(BaseModel): avg_session_seconds: float avg_tokens_per_session: float spend: float = Field(description="What the routed traffic actually cost") + classifier_cost: float | None = Field( + description="Recorded LLM classifier cost already included in spend; null when any session turns predate " + "subtotal recording, and zero for an empty window" + ) saved_spend: float = Field( description="Signed dollars saved versus each router's savings baseline (derived from its hardest " "tier, or the configured override), from the same per-request savings record the usage tab reads" @@ -222,6 +226,30 @@ class AutoRouterBenchmarkGroup(AutoRouterBenchmarkTotals): ) +class AutoRouterSessionResponse(BaseModel): + """One auto-routed session as its own key sees it: what the last turn ran on, and what the session cost + against the router's savings baseline (the priciest model in its hardest tier).""" + + session_id: str + router_name: str = Field(description="The auto-router alias the session's requests were sent to") + router_type: str = Field(description="complexity, adaptive or quality") + turns: int = Field(description="Auto-routed turns the rollup has recorded for this session so far") + last_model: str = Field(description="The deployment model the most recent turn was routed to") + spend: float = Field(description="What the session's routed traffic actually cost, classifier calls included") + saved_spend: float = Field(description="Estimated savings against the baseline, net of classifier cost") + baseline_spend: float = Field(description="spend plus saved_spend: the estimated single-model cost") + baseline_model: str | None = Field( + description="The savings baseline most of this session's turns were priced against, recorded turn by " + "turn, so it still names the counterfactual after the router is reconfigured or removed. None when no " + "turn recorded one: rows from before the baseline was recorded, and adaptive and quality routers, " + "which derive no baseline and so report no savings" + ) + baseline_models: Mapping[str, int] = Field( + description="Turns priced against each baseline model; more than one entry means the router's " + "baseline changed mid-session and baseline_spend mixes both" + ) + + class AutoRouterBenchmarksResponse(BaseModel): """Benchmarks for the auto-router dashboard, aggregated from the per-session rollup.""" diff --git a/litellm/types/management_endpoints/cache_settings_endpoints.py b/litellm/types/management_endpoints/cache_settings_endpoints.py index 32b32991449..cb05f3a50ac 100644 --- a/litellm/types/management_endpoints/cache_settings_endpoints.py +++ b/litellm/types/management_endpoints/cache_settings_endpoints.py @@ -237,4 +237,49 @@ CACHE_SETTINGS_FIELDS: Final[list[CacheSettingsField]] = [ ui_field_name="SSL Check Hostname", redis_type=None, ), + CacheSettingsField( + field_name="aws_iam_auth", + field_type="Boolean", + field_value=None, + field_description="Enable AWS ElastiCache IAM authentication", + field_default=False, + ui_field_name="AWS IAM Authentication", + redis_type=None, + ), + CacheSettingsField( + field_name="aws_iam_user_name", + field_type="String", + field_value=None, + field_description="AWS ElastiCache IAM user name", + field_default=None, + ui_field_name="AWS IAM User Name", + redis_type=None, + ), + CacheSettingsField( + field_name="aws_iam_cache_name", + field_type="String", + field_value=None, + field_description="AWS ElastiCache cache name", + field_default=None, + ui_field_name="AWS IAM Cache Name", + redis_type=None, + ), + CacheSettingsField( + field_name="aws_iam_region", + field_type="String", + field_value=None, + field_description="AWS region for ElastiCache IAM authentication", + field_default=None, + ui_field_name="AWS IAM Region", + redis_type=None, + ), + CacheSettingsField( + field_name="aws_iam_serverless", + field_type="Boolean", + field_value=None, + field_description="The ElastiCache cache is serverless rather than a self-designed cluster", + field_default=False, + ui_field_name="AWS IAM Serverless Cache", + redis_type=None, + ), ] diff --git a/litellm/types/management_endpoints/coordination_redis_endpoints.py b/litellm/types/management_endpoints/coordination_redis_endpoints.py index 30033346ed7..d70a921a22f 100644 --- a/litellm/types/management_endpoints/coordination_redis_endpoints.py +++ b/litellm/types/management_endpoints/coordination_redis_endpoints.py @@ -102,4 +102,41 @@ COORDINATION_REDIS_SETTINGS_FIELDS: Final[list[CoordinationRedisSettingsField]] ui_field_name="Service Name", section="sentinel", ), + CoordinationRedisSettingsField( + field_name="aws_iam_auth", + field_type="Boolean", + field_description="Enable AWS ElastiCache IAM authentication", + field_default=False, + ui_field_name="AWS IAM Authentication", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="aws_iam_user_name", + field_type="String", + field_description="AWS ElastiCache IAM user name", + ui_field_name="AWS IAM User Name", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="aws_iam_cache_name", + field_type="String", + field_description="AWS ElastiCache cache name", + ui_field_name="AWS IAM Cache Name", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="aws_iam_region", + field_type="String", + field_description="AWS region for ElastiCache IAM authentication", + ui_field_name="AWS IAM Region", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="aws_iam_serverless", + field_type="Boolean", + field_description="The ElastiCache cache is serverless rather than a self-designed cluster", + field_default=False, + ui_field_name="AWS IAM Serverless Cache", + section="connection", + ), ] diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 9bf3acc601c..58f9dadc1ce 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,7 +1,8 @@ from datetime import datetime from typing import Any, Final, Literal -from pydantic import BaseModel, ConfigDict, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from typing_extensions import Self from litellm.types.mcp import ( DEFAULT_SUBJECT_TOKEN_TYPE, @@ -38,6 +39,26 @@ class MCPOAuthMetadata(BaseModel): usable in memory but must never be persisted as configuration.""" +class MCPOAuthIdentityBinding(BaseModel): + """Per-server policy binding stored per-user OAuth credentials to the authenticated LiteLLM caller. + + When enabled for an interactive oauth2 server, the token relay validates the upstream OIDC + ``id_token`` (signature via the pinned issuer's JWKS, issuer, audience, expiry, nonce) and compares its + principal claim to the LiteLLM caller's trusted identity before the token is returned, stored, + or cached. ``audit`` logs mismatches without changing behavior; ``enforce`` fails closed with + 403 ``oauth_principal_mismatch`` and disables the direct ``oauth-user-credential`` POST, which + would otherwise bypass validation with an arbitrary opaque token. + """ + + mode: Literal["disabled", "audit", "enforce"] = "disabled" + issuer: str + jwks_url: str | None = None + audiences: list[str] = Field(min_length=1) # mutable-ok: public Pydantic schema requires list values + principal_claim: str = "email" + caller_field: Literal["user_email", "user_id"] = "user_email" + require_email_verified: bool = True + + class MCPServer(BaseModel): server_id: str name: str @@ -158,6 +179,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 @@ -172,6 +194,7 @@ class MCPServer(BaseModel): # response (supports dot-notation for nested fields, e.g. "team.enterprise_id"). # Tokens that fail validation are rejected before storage. token_validation: dict[str, Any] | None = None + oauth_identity_binding: MCPOAuthIdentityBinding | None = None # Optional TTL override (seconds) for the Redis per-user token cache, capped # at the token's expires_in minus the expiry buffer so a cached entry never # outlives the token. Defaults to the token's expires_in minus the expiry @@ -224,6 +247,14 @@ class MCPServer(BaseModel): """ return self.oauth2_flow == "client_credentials" + @model_validator(mode="after") + def validate_identity_binding_mode(self) -> Self: + binding: Final = self.oauth_identity_binding + if binding is not None and binding.mode != "disabled": + if not self.needs_user_oauth_token or self.delegate_auth_to_upstream: + raise ValueError("oauth_identity_binding requires gateway-managed per-user OAuth2 credentials") + return self + @property def needs_user_oauth_token(self) -> bool: """True if this is an OAuth2 server that relies on per-user tokens (no client_credentials).""" @@ -241,6 +272,32 @@ 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.""" + if self.auth_type == MCPAuth.oauth2: + return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + if self.auth_type not in ( + None, + MCPAuth.none, + MCPAuth.api_key, + MCPAuth.bearer_token, + MCPAuth.basic, + MCPAuth.authorization, + MCPAuth.token, + MCPAuth.aws_sigv4, + ): + return False + return not any( + header.lower() in ("authorization", "x-api-key", "api-key", "apikey") + for header in (self.extra_headers or ()) + ) + @property def is_true_passthrough(self) -> bool: """True for the transparent-proxy mode: LiteLLM performs no admission auth and forwards the diff --git a/litellm/types/object_permission.py b/litellm/types/object_permission.py index 1b391a3a1ef..68661aed891 100644 --- a/litellm/types/object_permission.py +++ b/litellm/types/object_permission.py @@ -8,7 +8,7 @@ can adopt the type without violating the SDK-must-not-import-from-proxy layering rule. """ -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class ObjectPermissionDict(TypedDict, total=False): @@ -23,3 +23,4 @@ class ObjectPermissionDict(TypedDict, total=False): models: list[str] | None search_tools: list[str] | None mcp_tool_search_enabled: bool | None + skills: ReadOnly[list[str] | None] diff --git a/litellm/types/proxy/claude_code_endpoints.py b/litellm/types/proxy/claude_code_endpoints.py index 2ee1bbbbb98..dcb5561cfeb 100644 --- a/litellm/types/proxy/claude_code_endpoints.py +++ b/litellm/types/proxy/claude_code_endpoints.py @@ -25,10 +25,12 @@ class PluginSpec(BaseModel): source: dict[str, str] = Field( ..., description=( - "Git source reference. Supported formats:\n" + "Plugin source reference. Supported formats:\n" "- GitHub: {'source': 'github', 'repo': 'org/repo'}\n" "- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n" - "- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}" + "- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}\n" + "- Zip archive on any https host (e.g. S3): " + "{'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''}" ), ) version: str | None = Field("1.0.0", description="Semantic version") @@ -46,7 +48,7 @@ class RegisterPluginRequest(PluginSpec): Request body for registering a plugin in the marketplace. LiteLLM acts as a registry/discovery layer. Plugins are hosted on - GitHub/GitLab/Bitbucket and referenced by their git source. + GitHub/GitLab/Bitbucket or as a zip archive on any https host and referenced by their source. """ name: str = Field( @@ -76,7 +78,7 @@ class PluginResponse(BaseModel): name: str = Field(..., description="Plugin name") version: str | None = Field(None, description="Plugin version") description: str | None = Field(None, description="Plugin description") - source: dict[str, str] = Field(..., description="Git source reference") + source: dict[str, str] = Field(..., description="Plugin source reference") enabled: bool = Field(..., description="Whether plugin is enabled") diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py index d517e508596..0df0f0c80ef 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -26,6 +26,10 @@ class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): "forwards the request uncompressed instead of blocking it." ), ) + ccr_retrieval: bool = Field( + default=True, + description="Inject the Headroom retrieval tool for hashes declared by the compression service.", + ) @staticmethod def ui_friendly_name() -> str: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py b/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py index 94f8161f44e..29f1b4bdcd6 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py @@ -16,6 +16,10 @@ class PromptSecurityGuardrailConfigModel(GuardrailConfigModel): default=True, description="Whether file sanitization timeouts allow the original file through instead of blocking the request.", ) + block_on_file_modify: bool = Field( + default=True, + description="Whether a file sanitization `modify` verdict blocks the request instead of replacing the file content.", + ) @staticmethod def ui_friendly_name() -> str: diff --git a/litellm/types/router.py b/litellm/types/router.py index 2dad22751de..fc09c40fe08 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -4,7 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc import datetime import enum -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints @@ -12,13 +12,16 @@ import httpx from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable +from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import normalize_drop_params if TYPE_CHECKING: from litellm.router import Router from .completion import CompletionRequest from .embedding import EmbeddingRequest +from .llms.bedrock import AwsSessionTag from .llms.openai import OpenAIFileObject from .search import SearchProvider from .utils import ( @@ -286,6 +289,7 @@ class CredentialLiteLLMParams(BaseModel): aws_web_identity_token: str | None = None aws_sts_endpoint: str | None = None aws_external_id: str | None = None + aws_session_tags: Sequence[AwsSessionTag] | None = None aws_bedrock_runtime_endpoint: str | None = None aws_bedrock_project_id: str | None = None s3_bucket_name: str | None = None @@ -307,7 +311,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): """ custom_llm_provider: str | None = None - rust: bool | None = None tpm: int | None = None rpm: int | None = None itpm: int | None = None @@ -315,6 +318,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): timeout: float | str | httpx.Timeout | None = None # if str, pass in as os.environ/ stream_timeout: float | str | None = None # timeout when making stream=True calls, if str, pass in as os.environ/ max_retries: int | None = None + drop_params: bool | str | None = None organization: str | None = None # for openai orgs configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None litellm_credential_name: str | None = None @@ -405,6 +409,18 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): return filtered return data + @field_validator("drop_params", mode="before") + @classmethod + def coerce_drop_params(cls, value: object) -> bool | str | None: + normalized: Final = normalize_drop_params(value) + if normalized is not None: + return normalized + if isinstance(value, str): + return value + if value is not None: + verbose_logger.warning("drop_params=%r is not a flag value, treating it as unset", value) + return None + def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) @@ -511,6 +527,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): input_cost_per_second: float | None output_cost_per_second: float | None output_cost_per_second_480p: ReadOnly[float | None] + output_cost_per_second_720p: ReadOnly[float | None] output_cost_per_second_1080p: float | None output_cost_per_second_4k: ReadOnly[float | None] num_retries: int | None @@ -597,6 +614,24 @@ class Deployment(BaseModel): setattr(self, key, value) +@dataclass(frozen=True, slots=True) +class DeploymentModelListingInfo: + """What the deployments behind a model name contribute to its OpenAI-compatible listing entry. + + ``cost_map_keys`` are the names those deployments' underlying models are known by in + ``litellm.model_cost`` (``base_model`` when set, else ``litellm_params.model``), which + is what a request actually reaches; the public model name they are listed under is an + arbitrary alias and often absent from the cost map. Keys are deduplicated in config + order, so the ordinary group -- several interchangeable deployments of one model -- + carries exactly one. The token limits are the widest explicitly set in any + deployment's ``model_info``, which outrank anything the cost map says. + """ + + cost_map_keys: tuple[str, ...] + max_input_tokens: int | None + max_output_tokens: int | None + + class RouterErrors(enum.Enum): """ Enum for router specific errors with common codes diff --git a/litellm/types/secret_managers/main.py b/litellm/types/secret_managers/main.py index 599e5746dfb..148e680a236 100644 --- a/litellm/types/secret_managers/main.py +++ b/litellm/types/secret_managers/main.py @@ -45,6 +45,9 @@ class KeyManagementSettings(LiteLLMPydanticObjectBase): tags: dict[str, str] | None = None """Optional tags to attach when creating secrets (e.g. {"Environment": "Prod", "Owner": "AI-Platform"}).""" + kms_key_id: str | None = None + """Optional customer-managed KMS key (ID, alias or ARN) used to encrypt secrets created in AWS Secrets Manager.""" + custom_secret_manager: str | None = None """ Path to custom secret manager class (e.g. "my_secret_manager.InMemorySecretManager") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 61c2fc8c5a5..e3ea37dc0c8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -35,6 +35,7 @@ from pydantic import ( BaseModel, ConfigDict, Field, + JsonValue, PrivateAttr, SkipValidation, field_serializer, @@ -168,6 +169,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 @@ -270,7 +273,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token_above_272k_tokens_flex: float | None input_cost_per_token_above_512k_tokens: float | None # MiniMax-M3: prompts >512K priced at 2x input input_cost_per_character_above_128k_tokens: float | None # only for vertex ai models - input_cost_per_query: float | None # only for rerank models + input_cost_per_query: float | None # per-request pricing: rerank, search, and Bedrock Marengo embeddings input_cost_per_image: float | None # only for vertex ai models input_cost_per_image_token: float | None # for gpt-image-1 and similar models input_cost_per_video_token: float | None # for gemini omni models with video input @@ -316,6 +319,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): float | None ) # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) output_cost_per_second_480p: ReadOnly[float | None] + output_cost_per_second_720p: ReadOnly[float | None] output_cost_per_second_4k: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models ocr_cost_per_credit: float | None # for OCR models priced by credit @@ -335,6 +339,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): "image_generation", "chat", "audio_transcription", + "audio_speech", "responses", "ocr", "realtime", @@ -577,6 +582,7 @@ CallTypesLiteral = Literal[ "search", "asearch", "_arealtime", + "_aresponses_websocket", "create_batch", "acreate_batch", "create_file", @@ -954,6 +960,14 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { CallTypes.agenerate_content_stream, CallTypes.generate_content_stream, ], + "/v1beta/models/{model}:generateContent": ( + CallTypes.agenerate_content, + CallTypes.generate_content, + ), + "/v1beta/models/{model}:streamGenerateContent": ( + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream, + ), # MCP (Model Context Protocol) "/mcp/call_tool": [CallTypes.call_mcp_tool], # A2A (Agent-to-Agent) @@ -961,12 +975,12 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { "/a2a/{agent_id}/message/send": [CallTypes.asend_message, CallTypes.send_message], # Passthrough endpoints "/llm_passthrough": [ - CallTypes.llm_passthrough_route, CallTypes.allm_passthrough_route, + CallTypes.llm_passthrough_route, ], "/v1/llm_passthrough": [ - CallTypes.llm_passthrough_route, CallTypes.allm_passthrough_route, + CallTypes.llm_passthrough_route, ], "/v1/messages": [CallTypes.anthropic_messages], # OCR @@ -1622,6 +1636,17 @@ class Choices(SafeAttributeModel, OpenAIObject): setattr(self, key, value) +def text_tokens_without_nested_reasoning( + completion_tokens: int, + text_tokens: int, + reasoning_tokens: int, + other_modality_tokens: int, +) -> int: + reported_total: Final = text_tokens + reasoning_tokens + other_modality_tokens + nested_reasoning_tokens: Final = min(reasoning_tokens, text_tokens, max(reported_total - completion_tokens, 0)) + return text_tokens - nested_reasoning_tokens + + class CompletionTokensDetailsWrapper(CompletionTokensDetails): # wrapper for older openai versions text_tokens: int | None = None """Text tokens generated by the model.""" @@ -1671,6 +1696,9 @@ class PromptTokensDetailsWrapper( audio_length_seconds: float | None = None """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second.""" + query_count: int | None = None + """Number of billable requests sent to the model. Used for embeddings priced per request, such as Bedrock Marengo.""" + cache_write_tokens: int | None = None """Number of cache write (creation) tokens sent to the model. OpenAI naming (prompt_tokens_details.cache_write_tokens); this is the canonical field.""" @@ -1712,6 +1740,8 @@ class PromptTokensDetailsWrapper( del self.video_length_seconds if self.audio_length_seconds is None: del self.audio_length_seconds + if self.query_count is None: + del self.query_count if self.web_search_requests is None: del self.web_search_requests if self.google_maps_grounding_requests is None: @@ -2890,6 +2920,7 @@ RoutingDecisionCause = Literal[ # same tier served instead. The displaced group rides in signals. Reported even on a kept # session pin, since the pinned model did not serve the request. "health_failover", + "health_default_fallback", "session_affinity_pin", "session_affinity_escalation", # classification_mode 'user_turn': the request is an agent loop's continuation turn (no new @@ -3050,6 +3081,7 @@ class StandardLoggingPayloadErrorInformation(TypedDict, total=False): llm_provider: str | None traceback: str | None error_message: str | None + error_provider_request_id: ReadOnly[str | None] # error_rate_limit_category: # For 429 / rate-limit errors, the source of the rate limit. One of the # string values defined by `litellm.exceptions.RateLimitErrorCategory` @@ -3343,7 +3375,12 @@ class StandardAuditLogPayload(TypedDict): updated_values: str | None -class StandardLoggingPayload(TypedDict): +class ClassifierAudit(TypedDict, total=False): + classifier_input: ReadOnly[Mapping[str, JsonValue]] + originating_request_masked: ReadOnly[Mapping[str, JsonValue]] + + +class StandardLoggingPayload(ClassifierAudit): id: str trace_id: str # Trace multiple LLM calls belonging to same overall request (e.g. fallbacks/retries) session_id: str # End-user/conversation session id (litellm_session_id), independent of trace_id @@ -3493,6 +3530,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_second: float | None = None output_cost_per_second_1080p: float | None = None output_cost_per_second_480p: float | None = None + output_cost_per_second_720p: float | None = None output_cost_per_second_4k: float | None = None input_cost_per_pixel: float | None = None output_cost_per_pixel: float | None = None diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 8bb0235ea2a..6d2ca308798 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -5,7 +5,7 @@ from enum import Enum from typing import Any, Literal from pydantic import BaseModel -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class SupportedVectorStoreIntegrations(str, Enum): @@ -96,6 +96,17 @@ class VectorStoreSearchResponse(TypedDict, total=False): data: list[VectorStoreSearchResult] | None +VectorStoreSearchFailureMode = Literal["annotate", "error"] + + +class VectorStoreSearchFailure(TypedDict): + """A configured vector store whose search failed, as reported back to the API caller""" + + vector_store_id: ReadOnly[str] + custom_llm_provider: ReadOnly[str | None] + error: ReadOnly[str] + + class VectorStoreSearchOptionalRequestParams(TypedDict, total=False): """TypedDict for Optional parameters supported by the vector store search API.""" diff --git a/litellm/utils.py b/litellm/utils.py index 52c1859b525..bb1bce66d9b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -66,6 +66,7 @@ from litellm.constants import ( DEFAULT_EMBEDDING_PARAM_VALUES, DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_TRIM_RATIO, FUNCTION_DEFINITION_TOKEN_COUNT, @@ -80,6 +81,7 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, ) @@ -277,7 +279,7 @@ except (ImportError, AttributeError, TypeError): # Convert to str (if necessary) claude_json_str = json.dumps(json_data) import importlib.metadata -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args from litellm import utils as litellm_utils @@ -672,11 +674,19 @@ def load_credentials_from_list(kwargs: dict): CredentialAccessor: Final = getattr(sys.modules[__name__], "CredentialAccessor") credential_name: Final = kwargs.get("litellm_credential_name") - if credential_name and litellm.credential_list: - credential_accessor: Final[Mapping[str, object]] = CredentialAccessor.get_credential_values(credential_name) - for key, value in credential_accessor.items(): - if key not in kwargs: - kwargs[key] = value + if not credential_name: + return + credential: Final = CredentialAccessor.find_credential(credential_name) + if credential is None: + verbose_logger.warning( + "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", + credential_name, + len(litellm.credential_list), + ) + return + for key, value in credential.credential_values.items(): + if key not in kwargs: + kwargs[key] = value def get_dynamic_callbacks( @@ -1186,6 +1196,47 @@ def function_setup( raise e +def _dispatch_success_logging( + logging_obj: LiteLLMLoggingObject, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + is_completion_with_fallbacks: bool, + is_litellm_internal_call: bool, +) -> None: + if not is_litellm_internal_call: + if getattr(logging_obj, "_defer_async_logging", False): + + def _enqueue_deferred_logging() -> None: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) + ) + + logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging + else: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) + ) + + logging_obj.handle_sync_success_callbacks_for_async_calls( + result=result, + start_time=start_time, + end_time=end_time, + ) + + async def _client_async_logging_helper( logging_obj: LiteLLMLoggingObject, result, @@ -1653,6 +1704,16 @@ def client(original_function): kwargs=kwargs, ) + _update_response_metadata: Final = getattr(sys.modules[__name__], "update_response_metadata") + _update_response_metadata( + result=result, + logging_obj=logging_obj, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) + # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") # Copy the current context to propagate it to the background thread @@ -1667,15 +1728,6 @@ def client(original_function): end_time, ) # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") - update_response_metadata( - result=result, - logging_obj=logging_obj, - model=model, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) return result except Exception as e: call_type = original_function.__name__ @@ -1834,6 +1886,9 @@ def client(original_function): elif _caching_handler_response.embedding_all_elements_cache_hit is True: return _caching_handler_response.final_embedding_cached_response + if _llm_caching_handler.preset_cache_key is not None: + logging_obj.litellm_params["preset_cache_key"] = _llm_caching_handler.preset_cache_key + # CHECK MAX TOKENS if ( kwargs.get("max_tokens", None) is not None @@ -1932,48 +1987,20 @@ def client(original_function): args=args, ) - # LOG SUCCESS - handle streaming success logging in the _next_ object - # Internal sub-calls (e.g. emulated file-search steps) share the - # parent's logging obj; skip async logging here so only the outer call bills once. - # NOTE: streaming requests return early (before this point) via - # CustomStreamWrapper, so this block is non-streaming only. - if not _is_litellm_internal_call: - if getattr(logging_obj, "_defer_async_logging", False): - - def _enqueue_deferred_logging() -> None: - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) - ) - - logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging - else: - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) - ) - - logging_obj.handle_sync_success_callbacks_for_async_calls( - result=result, - start_time=start_time, - end_time=end_time, - ) # REBUILD EMBEDDING CACHING if ( isinstance(result, EmbeddingResponse) and _caching_handler_response is not None and _caching_handler_response.final_embedding_cached_response is not None ): + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=_is_litellm_internal_call, + ) return _llm_caching_handler._combine_cached_embedding_response_with_api_result( _caching_handler_response=_caching_handler_response, embedding_response=result, @@ -1989,6 +2016,14 @@ def client(original_function): start_time=start_time, end_time=end_time, ) + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=_is_litellm_internal_call, + ) return result except Exception as e: @@ -2188,13 +2223,17 @@ def _return_openai_tokenizer(model: str) -> SelectTokenizerResponse: return {"type": "openai_tokenizer", "tokenizer": _get_default_encoding()} +def uses_anthropic_tokenizer(model: str) -> bool: + return model in litellm.anthropic_models and "claude-3" not in model + + def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None: if model in litellm.cohere_models and "command-r" in model: # cohere cohere_tokenizer: Final = Tokenizer.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer") return {"type": "huggingface_tokenizer", "tokenizer": cohere_tokenizer} # anthropic - elif model in litellm.anthropic_models and "claude-3" not in model: + elif uses_anthropic_tokenizer(model): claude_tokenizer: Final = Tokenizer.from_str(claude_json_str) return {"type": "huggingface_tokenizer", "tokenizer": claude_tokenizer} # llama2 @@ -2284,15 +2323,7 @@ def create_pretrained_tokenizer(identifier: str, revision="main", auth_token: st dict: A dictionary with the tokenizer and its type. """ - try: - tokenizer = Tokenizer.from_pretrained( - identifier, - revision=revision, - auth_token=auth_token, - ) - except Exception as e: - verbose_logger.error("Error creating pretrained tokenizer: %s. Defaulting to version without 'auth_token'.", e) - tokenizer = Tokenizer.from_pretrained(identifier, revision=revision) + tokenizer: Final = Tokenizer.from_pretrained(identifier, revision=revision, token=auth_token) return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} @@ -2805,6 +2836,13 @@ def supports_reasoning(model: str, custom_llm_provider: str | None = None) -> bo return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_reasoning") +def supports_none_reasoning_effort(model: str, custom_llm_provider: str | None = None) -> bool: + """ + Check if the given model accepts reasoning effort "none" and return a boolean value. + """ + return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_none_reasoning_effort") + + def supports_native_structured_output(model: str, custom_llm_provider: str | None = None) -> bool: """ Check if the given model supports native structured outputs and return a boolean value. @@ -2959,24 +2997,33 @@ def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, obje return None +def is_generalized_model_info(model_info: ModelInfo) -> bool: + """Whether ``model_info`` came from a fallback-generalization capability rule. + + Detected as the resolved key missing ``litellm.model_cost`` while matching a + capability rule. A rule-derived entry carries no pricing and only a conservative + family-baseline context window, so callers holding a second candidate name should + prefer an exact cost-map entry from that name over this one. + """ + key: Final = cast("Mapping[str, object]", model_info).get("key") # cast-ok: partial dicts may omit "key" + if not isinstance(key, str): + return False + return key not in litellm.model_cost and match_capability_generalizations(key) is not None + + def _get_builtin_model_info_for_registration(model: str) -> ModelInfo | None: """Resolve ``model`` to its built-in cost-map entry for registration merging. Returns ``None`` when the lookup raises or when it resolved via a - fallback-generalization capability rule, detected as the resolved key missing - ``litellm.model_cost`` while matching a capability rule. A rule-derived entry - carries no pricing, so treating it as a hit would skip the built-in - cache-pricing inheritance for prefix-mangled keys. + fallback-generalization capability rule. A rule-derived entry carries no + pricing, so treating it as a hit would skip the built-in cache-pricing + inheritance for prefix-mangled keys. """ try: info: Final = get_model_info(model=model) except Exception: return None - if info["key"] in litellm.model_cost: - return info - if match_capability_generalizations(info["key"]) is None: - return info - return None + return None if is_generalized_model_info(info) else info _runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload @@ -3063,7 +3110,7 @@ def register_model( # Convert stringified numbers to appropriate numeric types loaded_model_cost = model_cost elif isinstance(model_cost, str): - loaded_model_cost = litellm.get_model_cost_map(url=model_cost) + loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1) if persist_across_reloads: _registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost @@ -3087,7 +3134,10 @@ def register_model( existing_model = cast(dict, builtin_model_info) model_cost_key = existing_model["key"] else: - existing_model = {} + # An exact entry ends the lookup ladder before the capability rules are + # consulted, so seed from them: otherwise registering an unmapped model + # shadows the very defaults it would have resolved to unregistered. + existing_model = dict(match_capability_generalizations(_key_str) or {}) # mutable-ok: merge target model_cost_key = key builtin_entry = _resolve_builtin_model_cost_entry(key=_key_str, provider=provider) if builtin_entry is not None: @@ -3224,7 +3274,7 @@ def get_optional_params_transcription( passed_params.pop("OPENAI_TRANSCRIPTION_PARAMS") custom_llm_provider = passed_params.pop("custom_llm_provider") - drop_params = passed_params.pop("drop_params") + drop_params = normalize_drop_params(passed_params.pop("drop_params")) special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs") for k, v in special_params.items(): passed_params[k] = v @@ -3332,7 +3382,7 @@ def get_optional_params_image_gen( model = passed_params.pop("model", None) custom_llm_provider = passed_params.pop("custom_llm_provider") provider_config = passed_params.pop("provider_config", None) - drop_params = passed_params.pop("drop_params", None) + drop_params = normalize_drop_params(passed_params.pop("drop_params", None)) additional_drop_params = passed_params.pop("additional_drop_params", None) special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs") for k, v in special_params.items(): @@ -3396,7 +3446,7 @@ def get_optional_params_image_gen( non_default_params=non_default_params, optional_params=optional_params, model=model or "", - drop_params=drop_params if drop_params is not None else False, + drop_params=litellm.drop_params is True or drop_params is True, ) elif ( custom_llm_provider == "openai" @@ -3460,7 +3510,7 @@ def get_optional_params_embeddings( custom_llm_provider = passed_params.pop("custom_llm_provider", None) special_params: Final = passed_params.pop("kwargs") - drop_params = passed_params.pop("drop_params", None) + drop_params = normalize_drop_params(passed_params.pop("drop_params", None)) additional_drop_params = passed_params.pop("additional_drop_params", None) allowed_openai_params = passed_params.pop("allowed_openai_params", None) or [] # Remove function objects from passed_params to avoid JSON serialization errors @@ -3616,7 +3666,7 @@ def get_optional_params_embeddings( elif "cohere.embed" in model: object = litellm.BedrockCohereEmbeddingConfig() elif "twelvelabs" in model or "marengo" in model: - object = litellm.TwelveLabsMarengoEmbeddingConfig() + object = litellm.TwelveLabsMarengoEmbeddingConfig(model=model) elif "nova" in model.lower(): object = litellm.AmazonNovaEmbeddingConfig() else: # unmapped model @@ -4187,6 +4237,7 @@ def get_optional_params( base_model: str | None = None, **kwargs, ): + drop_params = normalize_drop_params(drop_params) # rebind-ok: config and DB deployments pass "true" as a string passed_params: Final = locals().copy() special_params: Final = passed_params.pop("kwargs") # Remove base_model from passed_params so it doesn't interfere with @@ -4264,20 +4315,20 @@ def get_optional_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "anthropic_text": optional_params = litellm.AnthropicTextConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) optional_params = litellm.AnthropicTextConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere": @@ -4286,14 +4337,14 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "triton": optional_params = litellm.TritonConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=drop_params if drop_params is not None else False, + drop_params=bool(drop_params), ) elif custom_llm_provider == "maritalk": @@ -4301,35 +4352,35 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "replicate": optional_params = litellm.ReplicateConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "predibase": optional_params = litellm.PredibaseConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "huggingface": optional_params = litellm.HuggingFaceChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "together_ai": optional_params = litellm.TogetherAIChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "vertex_ai" and ( model in litellm.vertex_chat_models @@ -4343,7 +4394,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "gemini": @@ -4351,21 +4402,21 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "vertex_ai_beta" or (custom_llm_provider == "vertex_ai" and "gemini" in model): optional_params = litellm.VertexGeminiConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif litellm.VertexAIAnthropicConfig.is_supported_model(model=model, custom_llm_provider=custom_llm_provider): optional_params = litellm.VertexAIAnthropicConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "vertex_ai": if model in litellm.vertex_mistral_models: @@ -4374,35 +4425,35 @@ def get_optional_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) else: optional_params = litellm.MistralConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif model in litellm.vertex_ai_ai21_models: optional_params = litellm.VertexAIAi21Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif provider_config is not None: optional_params = provider_config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) else: # use generic openai-like param mapping optional_params = litellm.VertexAILlama3Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "sagemaker": @@ -4411,7 +4462,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "bedrock": BedrockModelInfo: Final = getattr(sys.modules[__name__], "BedrockModelInfo") @@ -4422,14 +4473,14 @@ def get_optional_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif bedrock_route == "openai": optional_params = litellm.AmazonBedrockOpenAIConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif "anthropic" in bedrock_base_model and bedrock_route == "invoke": if bedrock_base_model in litellm.AmazonAnthropicConfig.get_legacy_anthropic_model_names(): @@ -4437,21 +4488,21 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) else: optional_params = litellm.AmazonAnthropicClaudeConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif provider_config is not None: optional_params = provider_config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) if bedrock_route == "claude_platform": optional_params = BedrockModelInfo.map_claude_platform_auth_params( @@ -4462,28 +4513,28 @@ def get_optional_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "ollama": optional_params = litellm.OllamaConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "ollama_chat": optional_params = litellm.OllamaChatConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "nlp_cloud": optional_params = litellm.NLPCloudConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "petals": @@ -4491,35 +4542,35 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "deepinfra": optional_params = litellm.DeepInfraConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "perplexity" and provider_config is not None: optional_params = provider_config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "mistral" or custom_llm_provider == "codestral": optional_params = litellm.MistralConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "text-completion-codestral": optional_params = litellm.CodestralTextCompletionConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "text-completion-inception": @@ -4527,7 +4578,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "databricks": @@ -4535,21 +4586,21 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "nvidia_nim": optional_params = litellm.NvidiaNimConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "cerebras": optional_params = litellm.CerebrasConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "xai": optional_params = litellm.XAIChatConfig().map_openai_params( @@ -4562,77 +4613,77 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "fireworks_ai": optional_params = litellm.FireworksAIConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "volcengine": optional_params = litellm.VolcEngineConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "hosted_vllm": optional_params = litellm.HostedVLLMChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "vllm": optional_params = litellm.VLLMConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "groq": optional_params = litellm.GroqChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "bedrock_mantle": optional_params = litellm.BedrockMantleChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "deepseek": optional_params = litellm.DeepSeekChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "tencent": optional_params = litellm.TencentChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "openrouter": optional_params = litellm.OpenrouterConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "watsonx": optional_params = litellm.IBMWatsonXChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) # WatsonX-text param check for param in passed_params: @@ -4645,21 +4696,21 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "openai": optional_params = litellm.OpenAIConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "nebius": optional_params = litellm.NebiusConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif custom_llm_provider == "azure": _azure_detection_model: Final = base_model or model @@ -4668,14 +4719,14 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=_azure_detection_model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=_azure_detection_model): optional_params = litellm.AzureOpenAIGPT5Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=_azure_detection_model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) else: verbose_logger.debug( @@ -4694,21 +4745,21 @@ def get_optional_params( optional_params=optional_params, model=_azure_detection_model, api_version=api_version, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) elif provider_config is not None: optional_params = provider_config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) else: # assume passing in params for openai-like api optional_params = litellm.OpenAILikeChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), + drop_params=bool(drop_params), ) # if user passed in non-default kwargs for specific providers/models, pass them along optional_params = add_provider_specific_params_to_optional_params( @@ -4889,7 +4940,7 @@ def _get_deployment_order(deployment: dict | Any) -> int | None: return order -def _get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list: +def get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list: if target_order is not None: return [d for d in healthy_deployments if _get_deployment_order(d) == target_order] @@ -4908,7 +4959,7 @@ def _get_order_filtered_deployments(healthy_deployments: list[dict], target_orde return healthy_deployments -def _get_excluded_filtered_deployments( +def get_excluded_filtered_deployments( healthy_deployments: list[dict], excluded_deployment_ids: Iterable[str] | None = None, ) -> list: @@ -4919,10 +4970,12 @@ def _get_excluded_filtered_deployments( across the remaining deployments in the same model group after one of them has failed. - If the filter would leave no deployments, an empty list is returned so the - caller raises its usual no-deployments error and the weighted-failover - helper falls through to the cross-group fallback path. Returning the - original unfiltered list here would re-include the just-failed deployment. + If the filter would leave no deployments, an empty list is returned and the + caller decides what that means. Weighted failover lets it raise the usual + no-deployments error and fall through to the cross-group fallback path; the + retry skip in `async_get_healthy_deployments` deliberately falls back to the + unfiltered list, so a request every deployment refused still comes back with + the provider's own error rather than a no-deployments one. """ if not excluded_deployment_ids: return healthy_deployments @@ -5210,13 +5263,16 @@ def _strip_openai_finetune_model_name(model_name: str) -> str: input: ft:gpt-3.5-turbo:my-org:custom_suffix:id output: ft:gpt-3.5-turbo + input: ft:gpt-4o-2024-08-06:my-org::id (OpenAI leaves the suffix empty when none was set) + output: ft:gpt-4o-2024-08-06 + Args: model_name (str): The full model name Returns: str: The stripped model name """ - return re.sub(r"(:[^:]+){3}$", "", model_name) + return re.sub(r"(:[^:]*){3}$", "", model_name) def _strip_model_name(model: str, custom_llm_provider: str | None) -> str: @@ -5891,6 +5947,7 @@ def _get_model_info_helper( output_cost_per_second=_model_info.get("output_cost_per_second", None), output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None), output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None), + output_cost_per_second_720p=_model_info.get("output_cost_per_second_720p", None), output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), @@ -5946,6 +6003,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: @@ -6019,7 +6078,7 @@ def get_model_info( input_cost_per_character_above_128k_tokens: Optional[ float ] # only for vertex ai models - input_cost_per_query: Optional[float] # only for rerank models + input_cost_per_query: Optional[float] # per-request pricing: rerank, search, and Bedrock Marengo embeddings input_cost_per_image: Optional[float] # only for vertex ai models input_cost_per_audio_token: Optional[float] input_cost_per_audio_per_second: Optional[float] # only for vertex ai models @@ -6965,7 +7024,26 @@ class TextCompletionStreamWrapper: raise StopAsyncIteration -def mock_completion_streaming_obj(model_response, mock_response, model, n: int | None = None): +def mock_stream_usage_chunk(model_response: ModelResponseStream, model: str, prompt_tokens: int) -> ModelResponseStream: + return ModelResponseStream( + id=model_response.id, + choices=[], # mutable-ok: ModelResponseStream only treats a list as explicit choices, a tuple gets a default choice + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + total_tokens=prompt_tokens + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + ), + ) + + +def mock_completion_streaming_obj( + model_response: ModelResponseStream, + mock_response: str | MockException | ModelResponseStream, + model: str, + n: int | None = None, + prompt_tokens: int | None = None, +) -> Iterator[ModelResponseStream]: if isinstance(mock_response, litellm.MockException): raise mock_response if isinstance(mock_response, ModelResponseStream): @@ -6985,14 +7063,17 @@ def mock_completion_streaming_obj(model_response, mock_response, model, n: int | _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if prompt_tokens is not None: + yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens) async def async_mock_completion_streaming_obj( - model_response, + model_response: ModelResponseStream, mock_response: str | MockException | ModelResponseStream, - model, + model: str, n: int | None = None, -): + prompt_tokens: int | None = None, +) -> AsyncIterator[ModelResponseStream]: if isinstance(mock_response, litellm.MockException): raise mock_response if isinstance(mock_response, ModelResponseStream): @@ -7012,6 +7093,8 @@ async def async_mock_completion_streaming_obj( _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if prompt_tokens is not None: + yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens) ########## Reading Config File ############################ @@ -8683,6 +8766,8 @@ class ProviderConfigManager: return litellm.OpenRouterResponsesAPIConfig() elif litellm.LlmProviders.HOSTED_VLLM == provider: return litellm.HostedVLLMResponsesAPIConfig() + elif litellm.LlmProviders.FIREWORKS_AI == provider: + return litellm.FireworksAIResponsesAPIConfig() elif litellm.LlmProviders.BEDROCK_MANTLE == provider: # Both decisions are data-driven from the model's price-map entry, with # no model-name logic. Capability (can it serve Responses?) comes from @@ -8831,6 +8916,12 @@ class ProviderConfigManager: ) return AzurePassthroughConfig() + elif LlmProviders.AZURE_AI == provider: + from litellm.llms.azure_ai.passthrough.transformation import ( + AzureAIPassthroughConfig, + ) + + return AzureAIPassthroughConfig() elif LlmProviders.GIGACHAT == provider: from litellm.llms.gigachat.passthrough.transformation import ( GigaChatPassthroughConfig, @@ -9215,6 +9306,10 @@ class ProviderConfigManager: from litellm.llms.openai.image_edit import get_openai_image_edit_config return get_openai_image_edit_config(model=model) + elif LlmProviders.HOSTED_VLLM == provider: + from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config + + return get_hosted_vllm_image_edit_config(model=model) elif LlmProviders.AZURE == provider: from litellm.llms.azure.image_edit.transformation import ( AzureImageEditConfig, @@ -9305,11 +9400,9 @@ class ProviderConfigManager: ReductoParseV3Config, ) - if model == "parse-v3": - return ReductoParseV3Config() if model == "parse-legacy": return ReductoParseLegacyConfig() - return None + return ReductoParseV3Config() MistralOCRConfig: Final = litellm_utils.MistralOCRConfig PROVIDER_TO_CONFIG_MAP: Final = { @@ -9436,9 +9529,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 ( @@ -9446,6 +9542,12 @@ class ProviderConfigManager: ) return MinimaxTextToSpeechConfig() + elif litellm.LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.audio_speech.transformation import ( + MistralTextToSpeechConfig, + ) + + return MistralTextToSpeechConfig() elif litellm.LlmProviders.AWS_POLLY == provider: from litellm.llms.aws_polly.text_to_speech.transformation import ( AWSPollyTextToSpeechConfig, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f607c1e298e..aadd9bd3028 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -364,7 +364,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -380,6 +381,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -399,6 +401,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -416,6 +419,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -435,6 +439,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -452,6 +457,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -471,6 +477,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -488,6 +495,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -507,6 +515,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -537,7 +546,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -551,7 +561,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -650,7 +661,10 @@ }, "twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, "litellm_provider": "bedrock", "max_input_tokens": 77, "max_tokens": 77, @@ -662,7 +676,7 @@ }, "us.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -677,7 +691,7 @@ }, "eu.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -690,6 +704,48 @@ "supports_embedding_image_input": true, "supports_image_input": true }, + "twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "us.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "eu.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, "twelvelabs.pegasus-1-2-v1:0": { "input_cost_per_video_per_second": 0.00049, "output_cost_per_token": 7.5e-06, @@ -711,6 +767,14 @@ "mode": "chat", "supports_video_input": true }, + "global.twelvelabs.pegasus-1-2-v1:0": { + "input_cost_per_video_per_second": 0.00049, + "output_cost_per_token": 7.5e-06, + "litellm_provider": "bedrock", + "mode": "chat", + "supports_video_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "amazon.titan-text-express-v1": { "input_cost_per_token": 1.3e-06, "litellm_provider": "bedrock", @@ -2831,7 +2895,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.7e-08, @@ -2843,7 +2908,8 @@ "output_cost_per_token": 1.48e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "apac.amazon.nova-pro-v1:0": { "input_cost_per_token": 8.4e-07, @@ -2857,7 +2923,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.anthropic.claude-3-5-sonnet-20240620-v1:0": { "deprecation_date": "2026-07-30", @@ -3485,6 +3552,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, @@ -3532,6 +3648,79 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure_ai/codex-mini": { + "cache_read_input_token_cost": 3.75e-07, + "deprecation_date": "2026-11-15", + "input_cost_per_token": 1.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "responses", + "output_cost_per_token": 6e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/whisper": { + "deprecation_date": "2026-12-15", + "input_cost_per_second": 0.0001, + "litellm_provider": "azure_ai", + "mode": "audio_transcription", + "output_cost_per_second": 0.0001, + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" + }, "azure_ai/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, @@ -3935,13 +4124,29 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/model_router": { + "deprecation_date": "2027-05-20", "input_cost_per_token": 1.4e-07, "output_cost_per_token": 0, "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "chat", "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, + "azure_ai/model-router": { + "deprecation_date": "2027-05-20", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 0, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", + "comment": "Catalog-name twin of azure_ai/model_router: the flat $0.14 per M input tokens is the router's own fee, the routed model is priced on top of it" + }, "azure/eu/gpt-4o-2024-08-06": { "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.375e-06, @@ -7189,7 +7394,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 +7660,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, @@ -7881,7 +8086,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -7925,7 +8131,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -7969,7 +8176,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -10253,6 +10461,18 @@ "/v1/ocr" ] }, + "azure_ai/cohere-command-a": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 8182, + "max_tokens": 8182, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/", + "supports_function_calling": true, + "supports_tool_choice": true + }, "azure_ai/doc-intelligence/prebuilt-read": { "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.0015, @@ -10604,6 +10824,41 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/grok-4-20-reasoning": { + "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-04-06", + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_reasoning": true + }, + "azure_ai/grok-4-20-non-reasoning": { + "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-04-06", + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 262000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure_ai/grok-4-fast-non-reasoning": { "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, @@ -12087,7 +12342,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -12266,7 +12522,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "input_cost_per_token": 4.2e-08, @@ -12278,7 +12535,8 @@ "output_cost_per_token": 1.68e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "input_cost_per_token": 9.6e-07, @@ -12292,7 +12550,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -12859,6 +13118,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "cerebras/qwen-3.8-27b": { + "input_cost_per_token": 9.9e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.49e-06, + "source": "https://api.cerebras.ai/public/v1/models/qwen-3.8-27b", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "chatdolphin": { "input_cost_per_token": 5e-07, "litellm_provider": "nlp_cloud", @@ -13743,7 +14018,8 @@ "max_output_tokens": 3072, "max_tokens": 3072, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/meta/llama-2-7b-chat-int8": { "input_cost_per_token": 1.923e-06, @@ -13752,7 +14028,8 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.1": { "input_cost_per_token": 1.923e-06, @@ -13761,7 +14038,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@hf/thebloke/codellama-7b-instruct-awq": { "input_cost_per_token": 1.923e-06, @@ -13770,7 +14048,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/openai/gpt-oss-120b": { "input_cost_per_token": 3.5e-07, @@ -13780,6 +14059,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13790,7 +14070,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-3b-instruct": { "input_cost_per_token": 5.09e-08, @@ -13799,7 +14080,8 @@ "max_output_tokens": 80000, "max_tokens": 80000, "mode": "chat", - "output_cost_per_token": 3.35e-07 + "output_cost_per_token": 3.35e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-guard-3-8b": { "input_cost_per_token": 4.84e-07, @@ -13808,7 +14090,8 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-08 + "output_cost_per_token": 3e-08, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { "input_cost_per_token": 0.0, @@ -13817,7 +14100,8 @@ "max_output_tokens": 15000, "max_tokens": 15000, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 1.9e-07, @@ -13828,6 +14112,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13839,6 +14124,7 @@ "max_tokens": 80000, "mode": "chat", "output_cost_per_token": 4.881e-06, + "rpm": 300, "supports_reasoning": true }, "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { @@ -13848,7 +14134,8 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 2.87e-07 + "output_cost_per_token": 2.87e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-1b-instruct": { "input_cost_per_token": 2.7e-08, @@ -13857,7 +14144,8 @@ "max_output_tokens": 60000, "max_tokens": 60000, "mode": "chat", - "output_cost_per_token": 2.01e-07 + "output_cost_per_token": 2.01e-07, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.6": { "cache_read_input_token_cost": 1.6e-07, @@ -13868,6 +14156,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13879,6 +14168,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13889,7 +14179,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { "input_cost_per_token": 2.93e-07, @@ -13899,6 +14190,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 2.253e-06, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { @@ -13909,6 +14201,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1.12e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { @@ -13918,7 +14211,8 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1e-06 + "output_cost_per_token": 1e-06, + "rpm": 300 }, "cloudflare/@cf/zai-org/glm-5.2": { "cache_read_input_token_cost": 2.6e-07, @@ -13929,6 +14223,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4.4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -13940,6 +14235,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13950,7 +14246,8 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.55e-07 + "output_cost_per_token": 5.55e-07, + "rpm": 300 }, "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { "input_cost_per_token": 5.09e-08, @@ -13960,6 +14257,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 3.35e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13970,7 +14268,8 @@ "max_output_tokens": 3500, "max_tokens": 3500, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/google/gemma-4-26b-a4b-it": { "input_cost_per_token": 1e-07, @@ -13980,6 +14279,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -13991,6 +14291,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5.55e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { @@ -14001,6 +14302,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.76e-07, + "rpm": 300, "supports_vision": true }, "cloudflare/@cf/openai/gpt-oss-20b": { @@ -14011,6 +14313,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14022,6 +14325,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 8.5e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwq-32b": { @@ -14032,6 +14336,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 1e-06, + "rpm": 300, "supports_reasoning": true }, "codestral/codestral-2405": { @@ -14158,6 +14463,28 @@ "output_vector_size": 1536, "supports_embedding_image_input": true }, + "us.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, + "global.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "cohere", @@ -20729,7 +21056,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.amazon.nova-micro-v1:0": { "input_cost_per_token": 4.6e-08, @@ -20741,7 +21069,8 @@ "output_cost_per_token": 1.84e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "eu.amazon.nova-pro-v1:0": { "input_cost_per_token": 1.05e-06, @@ -20756,24 +21085,25 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.anthropic.claude-3-5-haiku-20241022-v1:0": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 8e-07, "litellm_provider": "bedrock", "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 4e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 2.5e-08, - "cache_creation_input_token_cost": 3.125e-07, + "cache_read_input_token_cost": 8e-08, + "cache_creation_input_token_cost": 1e-06, "prompt_cache_min_tokens": 2048 }, "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -23780,9 +24110,9 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2e-06, @@ -23805,12 +24135,12 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_url_context": true, + "supports_url_context": false, "supports_vision": true, "supports_web_search": true, "search_context_cost_per_query": { @@ -27602,6 +27932,70 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-5.5": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.5", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-luna": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-luna", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-sol": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-sol", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-terra": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-terra", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "chatgpt/gpt-5.4": { "litellm_provider": "chatgpt", "max_input_tokens": 1050000, @@ -28208,6 +28602,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -29032,7 +29427,12 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + } }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29049,9 +29449,9 @@ "output_cost_per_token_priority": 1e-06, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29150,9 +29550,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29177,9 +29577,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29290,9 +29690,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29317,9 +29717,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29401,6 +29801,66 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2.5-flare": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-flare-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, "low/1024-x-1024/gpt-image-1.5": { "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, @@ -29851,7 +30311,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": false, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": false }, "gpt-5.1-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, @@ -29896,7 +30356,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": false, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": false }, "gpt-5.1-chat-latest": { "cache_read_input_token_cost": 1.25e-07, @@ -29988,7 +30448,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": false }, "gpt-5.2-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, @@ -30034,7 +30494,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": false }, "gpt-5.2-chat-latest": { "cache_read_input_token_cost": 1.75e-07, @@ -30742,6 +31202,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], "source": "https://developers.openai.com/api/docs/models/chat-latest", "supported_endpoints": [ "/v1/chat/completions", @@ -30759,6 +31222,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -30880,8 +31344,6 @@ "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-pro": { - "cache_read_input_token_cost": 3e-06, - "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, @@ -30930,8 +31392,6 @@ "supports_low_reasoning_effort": false }, "gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 3e-06, - "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, @@ -31026,7 +31486,7 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 @@ -31078,14 +31538,12 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-pro": { - "cache_read_input_token_cost": 3e-06, - "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, @@ -31130,13 +31588,11 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-pro-2026-03-05": { - "cache_read_input_token_cost": 3e-06, - "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "input_cost_per_token_flex": 1.5e-05, @@ -31181,7 +31637,7 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "input_cost_per_token_above_272k_tokens_flex": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, @@ -32628,6 +33084,7 @@ "supports_tool_choice": true }, "groq/gemma-7b-it": { + "deprecation_date": "2024-12-18", "input_cost_per_token": 5e-08, "litellm_provider": "groq", "max_input_tokens": 8192, @@ -33261,13 +33718,14 @@ "supports_tool_choice": true }, "jina-reranker-v2-base-multilingual": { - "input_cost_per_token": 1.8e-08, + "input_cost_per_token": 5e-08, "litellm_provider": "jina_ai", "max_input_tokens": 1024, "max_output_tokens": 1024, "max_tokens": 1024, "mode": "rerank", - "output_cost_per_token": 1.8e-08 + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" }, "jp.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.125e-06, @@ -33441,6 +33899,20 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "inception/mercury-2.5": { + "input_cost_per_token": 2e-07, + "litellm_provider": "inception", + "max_input_tokens": 260000, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://docs.inceptionlabs.ai/get-started/models", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "text-completion-inception/mercury-edit-2": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, @@ -34898,9 +35370,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-latest": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" @@ -37176,6 +37648,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37216,6 +37689,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37427,6 +37901,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37467,6 +37942,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -38673,19 +39149,21 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-chat-v3.1": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 2.5e-07, "input_cost_per_token_cache_hit": 2e-08, "litellm_provider": "openrouter", "max_input_tokens": 163840, "max_output_tokens": 163840, "max_tokens": 163840, "mode": "chat", - "output_cost_per_token": 8e-07, + "output_cost_per_token": 9.5e-07, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.3e-07, + "source": "https://openrouter.ai/deepseek/deepseek-chat-v3.1" }, "openrouter/deepseek/deepseek-v3.2": { "input_cost_per_token": 2.69e-07, @@ -38748,36 +39226,56 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 8.59908e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.719816e-06, "source": "https://openrouter.ai/deepseek/deepseek-v4-pro", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 7.1659e-08 + }, + "openrouter/deepseek/deepseek-v4.1-flash": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 3e-09, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "source": "https://openrouter.ai/deepseek/deepseek-v4.1-flash", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": false, + "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.7948e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.73844e-06, "source": "https://openrouter.ai/deepseek/deepseek-v4-pro-0813", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.9316e-08 }, "openrouter/google/gemini-2.0-flash-001": { "deprecation_date": "2026-06-01", @@ -39517,6 +40015,67 @@ "supports_tool_choice": true, "supports_vision": true }, + "openrouter/openai/gpt-5.6-sol": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "default_reasoning_effort": "medium", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "reasoning_effort_levels": [ + "none", + "low", + "medium", + "high", + "xhigh", + "max" + ], + "source": "https://openrouter.ai/openai/gpt-5.6-sol", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/openai/gpt-5.6-sol-pro": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2e-07, + "cache_creation_input_token_cost": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-5.6-sol-pro", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/openai/gpt-oss-120b": { "input_cost_per_token": 3.7e-08, "litellm_provider": "openrouter", @@ -39645,13 +40204,13 @@ "supports_tool_choice": true }, "openrouter/qwen/qwen3-235b-a22b-2507": { - "input_cost_per_token": 8.75e-08, + "input_cost_per_token": 2.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 3.5e-07, + "output_cost_per_token": 8.8e-07, "source": "https://openrouter.ai/qwen/qwen3-235b-a22b-2507", "supports_function_calling": true, "supports_tool_choice": true @@ -39684,7 +40243,7 @@ "supports_vision": true }, "openrouter/qwen/qwen3.5-35b-a3b": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 3.125e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, @@ -39695,7 +40254,8 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.5625e-07 }, "openrouter/qwen/qwen3.5-27b": { "input_cost_per_token": 1.95e-07, @@ -39712,13 +40272,13 @@ "supports_vision": true }, "openrouter/qwen/qwen3.5-122b-a10b": { - "input_cost_per_token": 2.9e-07, + "input_cost_per_token": 2.6e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 2.4e-06, + "output_cost_per_token": 2.08e-06, "source": "https://openrouter.ai/qwen/qwen3.5-122b-a10b", "supports_function_calling": true, "supports_reasoning": true, @@ -39805,18 +40365,19 @@ "supports_web_search": true }, "openrouter/z-ai/glm-4.6": { - "input_cost_per_token": 5.5e-07, + "input_cost_per_token": 4.3e-07, "litellm_provider": "openrouter", "max_input_tokens": 202800, "max_output_tokens": 131000, "max_tokens": 131000, "mode": "chat", - "output_cost_per_token": 2.2e-06, + "output_cost_per_token": 1.75e-06, "source": "https://openrouter.ai/z-ai/glm-4.6", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8e-08 }, "openrouter/z-ai/glm-4.6:exacto": { "input_cost_per_token": 4.5e-07, @@ -41473,6 +42034,28 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "rerank-v4.0-fast": { + "input_cost_per_query": 0.002, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, + "rerank-v4.0-pro": { + "input_cost_per_query": 0.0025, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -42619,7 +43202,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.2e-06, + "max_input_tokens": 131072, + "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo": { "litellm_provider": "together_ai", @@ -42627,7 +43214,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 3e-07, + "max_input_tokens": 32768, + "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": { "deprecation_date": "2026-07-10", @@ -42839,7 +43430,11 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2e-07, + "max_input_tokens": 32768, + "source": "https://api.together.xyz/v1/models" }, "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": { "deprecation_date": "2026-04-02", @@ -42847,7 +43442,11 @@ "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, + "max_input_tokens": 32768, + "source": "https://api.together.xyz/v1/models" }, "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": { "deprecation_date": "2026-04-16", @@ -43433,7 +44032,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, @@ -43445,7 +44045,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -43474,7 +44075,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, @@ -45653,8 +46255,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -45678,8 +46280,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -47141,6 +47743,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", @@ -47482,27 +48177,29 @@ "supports_tool_choice": true }, "vertex_ai/mistral-small-2503": { - "input_cost_per_token": 1e-06, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-mistral_models", "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-06, + "output_cost_per_token": 3e-07, "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "vertex_ai/mistral-small-2503@001": { - "input_cost_per_token": 1e-06, + "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-mistral_models", "max_input_tokens": 32000, "max_output_tokens": 8191, "max_tokens": 8191, "mode": "chat", - "output_cost_per_token": 3e-06, + "output_cost_per_token": 3e-07, "supports_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "vertex_ai/mistral-ocr-2505": { "litellm_provider": "vertex_ai", @@ -47552,23 +48249,24 @@ "supports_reasoning": true }, "vertex_ai/openai/gpt-oss-20b-maas": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 7e-08, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 3e-07, - "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas", - "supports_reasoning": true + "output_cost_per_token": 2.5e-07, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_reasoning": true, + "cache_read_input_token_cost": 7e-09 }, "vertex_ai/xai/grok-4.1-fast-non-reasoning": { "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -47582,9 +48280,9 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -47597,14 +48295,17 @@ }, "vertex_ai/xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -47613,14 +48314,17 @@ }, "vertex_ai/xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -47628,6 +48332,44 @@ "supports_vision": true, "supports_web_search": true }, + "vertex_ai/xai/grok-4.3": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "vertex_ai/xai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -47886,6 +48628,16 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-multilingual-2": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-3-large": { "input_cost_per_token": 1.8e-07, "litellm_provider": "voyage", @@ -47991,6 +48743,7 @@ "output_cost_per_token": 0.0 }, "wandb/openai/gpt-oss-120b": { + "supports_reasoning": true, "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, @@ -48001,6 +48754,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/openai/gpt-oss-20b": { + "supports_reasoning": true, "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, @@ -48011,6 +48765,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/zai-org/GLM-4.5": { + "supports_reasoning": true, "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, @@ -48039,6 +48794,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3-235B-A22B-Thinking-2507": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, @@ -48095,6 +48851,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/deepseek-ai/DeepSeek-V3.1": { + "supports_reasoning": true, "max_tokens": 128000, "max_input_tokens": 161000, "max_output_tokens": 128000, @@ -48105,6 +48862,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/deepseek-ai/DeepSeek-R1-0528": { + "supports_reasoning": true, "max_tokens": 161000, "max_input_tokens": 161000, "max_output_tokens": 161000, @@ -51959,7 +52717,8 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 8e-07, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/openai/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -51998,7 +52757,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-01" }, "scaleway/hcompany/holo2-30b-a3b": { "input_cost_per_token": 3e-07, @@ -52009,7 +52769,8 @@ "mode": "chat", "output_cost_per_token": 7e-07, "supports_reasoning": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-09" }, "scaleway/mistralai/mistral-medium-3.5-128b": { "input_cost_per_token": 1.5e-06, @@ -52032,7 +52793,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 2e-06, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/voxtral-small-24b-2507": { "input_cost_per_audio_token": 1.5e-07, @@ -52043,7 +52805,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.5e-07, - "supports_audio_input": true + "supports_audio_input": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/mistral-small-3.2-24b-instruct-2506": { "input_cost_per_token": 1.5e-07, @@ -52065,7 +52828,8 @@ "mode": "chat", "output_cost_per_token": 2e-07, "supports_vision": true, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/BAAI/bge-multilingual-gemma2": { "input_cost_per_token": 1e-07, @@ -54515,7 +55279,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-20b": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 7e-08, "output_cost_per_token": 3e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, @@ -54549,8 +55313,8 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 65536, @@ -54596,11 +55360,14 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-terra": { @@ -54635,11 +55402,14 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-cyber": { @@ -54664,12 +55434,49 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, @@ -54702,11 +55509,14 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "us.openai.gpt-5.6-sol": { @@ -54731,8 +55541,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "global.openai.gpt-5.6-sol": { @@ -54757,8 +55570,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "us.openai.gpt-5.6-terra": { @@ -54783,8 +55599,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "global.openai.gpt-5.6-terra": { @@ -54809,8 +55628,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "us.openai.gpt-5.6-luna": { @@ -54835,8 +55657,11 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, "global.openai.gpt-5.6-luna": { @@ -54861,10 +55686,115 @@ "text" ], "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, "supports_tool_choice": true, "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "us.openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "global.openai.gpt-6-astra": { + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "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, + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, @@ -54894,11 +55824,13 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.4": { @@ -54930,11 +55862,13 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_xhigh_reasoning_effort": true, "supports_web_search": true }, "bedrock_mantle/google.gemma-4-31b": { @@ -55666,17 +56600,17 @@ "supports_reasoning": true, "source": "https://serverless.tensormesh.ai/v1/models/openrouter" }, - "deepseek-v4-flash": { + "deepseek-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -55690,19 +56624,45 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, - "deepseek-v4-flash-vision-exp": { + "deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "deepseek-v4-flash-vision-exp": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "deepseek", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -55744,17 +56704,17 @@ "supports_tool_choice": true, "supports_vision": false }, - "deepseek/deepseek-v4-flash": { + "deepseek/deepseek-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -55768,19 +56728,45 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, - "deepseek/deepseek-v4-flash-vision-exp": { + "deepseek/deepseek-v4-flash": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 1.4e-08, - "input_cost_per_token": 4.4e-07, - "input_cost_per_token_cache_hit": 1.4e-08, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", - "output_cost_per_token": 1.32e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "deepseek/deepseek-v4-flash-vision-exp": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "deepseek", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ "/v1/chat/completions" @@ -56136,6 +57122,23 @@ ], "supports_audio_input": true }, + "gpt-live-1": { + "input_cost_per_second": 0.0008333333333333334, + "litellm_provider": "openai", + "mode": "realtime", + "source": "https://developers.openai.com/api/docs/models/gpt-live-1", + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true + }, "gpt-realtime-translate": { "input_cost_per_second": 0.0005666666666666667, "litellm_provider": "openai", @@ -56374,9 +57377,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-2603": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" @@ -56459,6 +57462,14 @@ "model_info": { "supports_mid_conversation_system": true } + }, + { + "name": "wandb-reasoning-baseline", + "pattern": "^wandb/", + "description": "Any Weights & Biases Inference model id, anchored to the wandb/ namespace so only that provider's ids match. W&B's serverless catalog is reasoning-first and grows faster than this registry names it, so an id the map has not described yet is treated as reasoning-capable and keeps the caller's reasoning_effort instead of dropping it or raising UnsupportedParamsError. Rules lose to exact entries, so a mapped non-reasoning model such as wandb/meta-llama/Llama-3.1-8B-Instruct is unaffected. Carries no mode and no pricing, so cost stays on the standard unpriced behavior and the deployment does not read as catalog-mapped to the router's reasoning-effort resolver.", + "model_info": { + "supports_reasoning": true + } } ] }, @@ -56530,8 +57541,8 @@ "rpm": 10 }, "vertex_ai/gemini-3.5-transcribe-preview": { - "input_cost_per_audio_token": 2.5e-06, - "input_cost_per_token": 2.5e-06, + "input_cost_per_audio_token": 2e-06, + "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai", "mode": "audio_transcription", "output_cost_per_token": 1.2e-05, @@ -56566,6 +57577,27 @@ ], "supports_audio_input": true }, + "vertex_ai/gemini-3.5-live-translate-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 2.1e-05, + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "perplexity/pplx-embed-context-v1-0.6b": { "input_cost_per_token": 8e-09, "litellm_provider": "perplexity", @@ -57816,6 +58848,7 @@ "supports_vision": false }, "wandb/deepseek-ai/DeepSeek-V4-Flash": { + "supports_reasoning": true, "max_tokens": 1048576, "max_input_tokens": 1048576, "input_cost_per_token": 1.4e-07, @@ -57828,6 +58861,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/deepseek-ai/DeepSeek-V4-Flash-0731": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 1.3e-07, @@ -57840,6 +58874,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/deepseek-ai/DeepSeek-V4-Pro": { + "supports_reasoning": true, "max_tokens": 1048576, "max_input_tokens": 1048576, "input_cost_per_token": 1.15e-06, @@ -57852,6 +58887,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/google/gemma-4-31B-it": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 1e-07, @@ -57892,6 +58928,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/MiniMaxAI/MiniMax-M3": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 2.3e-07, @@ -57904,6 +58941,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/moonshotai/Kimi-K2.7-Code": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 7.1e-07, @@ -57916,6 +58954,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/moonshotai/Kimi-K2.6": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 6.5e-07, @@ -57928,6 +58967,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 1e-07, @@ -57940,6 +58980,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 7.5e-07, @@ -57962,6 +59003,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.8-27B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 4e-07, @@ -57974,6 +59016,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.6-35B-A3B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 2.5e-07, @@ -57984,6 +59027,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.6-27B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 6e-07, @@ -57996,6 +59040,7 @@ "source": "https://wandb.ai/site/pricing/tokens/" }, "wandb/Qwen/Qwen3.5-35B-A3B": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 2.5e-07, @@ -58015,7 +59060,28 @@ "supports_vision": false, "source": "https://wandb.ai/site/pricing/tokens/" }, + "wandb/deepseek-ai/DeepSeek-V4-Pro-0813": { + "litellm_provider": "wandb", + "mode": "chat", + "supports_reasoning": true, + "input_cost_per_token": 0.00000131, + "output_cost_per_token": 0.00000396, + "cache_read_input_token_cost": 0.000000044, + "supports_prompt_caching": true, + "source": "https://wandb.ai/site/pricing/tokens/" + }, + "wandb/ibm-granite/granite-4.2-8b": { + "litellm_provider": "wandb", + "mode": "chat", + "supports_reasoning": true, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.00000015, + "cache_read_input_token_cost": 0.00000005, + "supports_prompt_caching": true, + "source": "https://wandb.ai/site/pricing/tokens/" + }, "wandb/zai-org/GLM-5.2": { + "supports_reasoning": true, "max_tokens": 262144, "max_input_tokens": 262144, "input_cost_per_token": 7.6e-07, @@ -59177,6 +60243,7 @@ ] }, "xai/grok-imagine-image-quality": { + "deprecation_date": "2026-11-02", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -59193,6 +60260,7 @@ ] }, "xai/grok-imagine-image-quality-20260403": { + "deprecation_date": "2026-11-02", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -59209,6 +60277,7 @@ ] }, "xai/grok-imagine-image-quality-latest": { + "deprecation_date": "2026-11-02", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -59257,6 +60326,77 @@ "image" ] }, + "xai/grok-imagine-video": { + "input_cost_per_image": 0.002, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.05, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_720p": 0.07, + "source": "https://docs.x.ai/docs/models/grok-imagine-video", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-2026-05-30": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-preview": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", @@ -59657,6 +60797,113 @@ "output_cost_per_token": 4.7e-07, "source": "https://docs.together.ai/docs/serverless-models" }, + "together_ai/moonshotai/Kimi-K2.6": { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 4.5e-06, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/moonshotai/Kimi-K2.5-fp4": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.8e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/MiniMaxAI/MiniMax-M2.7": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 196608, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/zai-org/GLM-5": { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3.2e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 202752, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/zai-org/GLM-5.1": { + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 202752, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/deepseek-ai/DeepSeek-R1-0528": { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 7e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 163840, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/Qwen/Qwen3-Coder-Next-FP8": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/Qwen/Qwen3-VL-32B-Instruct": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.5e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/Qwen/Qwen3-VL-8B-Instruct": { + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 6.8e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/mistralai/Ministral-3-14B-Instruct-2512": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { + "input_cost_per_token": 6e-08, + "output_cost_per_token": 2.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/mistralai/Mistral-7B-Instruct-v0.3": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, + "together_ai/Qwen/QwQ-32B": { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "mode": "chat", + "source": "https://api.together.xyz/v1/models" + }, "cerebras/gemma-4-31b": { "input_cost_per_token": 9.9e-07, "litellm_provider": "cerebras", @@ -60160,10 +61407,12 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true, "input_cost_per_token": 2.64e-06, "input_cost_per_token_above_272k_tokens": 5.28e-06, @@ -60193,10 +61442,12 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true, "input_cost_per_token": 2.64e-07, "input_cost_per_token_above_272k_tokens": 5.28e-07, @@ -60225,10 +61476,12 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true, "input_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 3.3e-07, @@ -60387,10 +61640,12 @@ "text" ], "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_xhigh_reasoning_effort": true, "supports_vision": true, "input_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 3.3e-07, @@ -60524,6 +61779,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60534,6 +61790,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60589,6 +61846,31 @@ "supports_web_search": false, "output_cost_per_image": 0.08 }, + "gemini/lyria-3.5": { + "input_cost_per_token": 0, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 0, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "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, + "output_cost_per_image": 0.08 + }, "perplexity/anthropic/claude-fable-5": { "litellm_provider": "perplexity", "mode": "responses", @@ -60900,6 +62182,7 @@ "mode": "responses", "supports_web_search": true, "supports_function_calling": true, + "supports_reasoning": true, "input_cost_per_token": 1.15e-08, "output_cost_per_token": 1.7e-07, "cache_read_input_token_cost": 1.15e-09, @@ -60910,6 +62193,7 @@ "mode": "responses", "supports_web_search": true, "supports_function_calling": true, + "supports_reasoning": true, "input_cost_per_token": 2.5e-07, "output_cost_per_token": 2.5e-06, "cache_read_input_token_cost": 2.5e-07, @@ -61273,6 +62557,28 @@ "cache_read_input_token_cost": 2e-08, "supports_prompt_caching": true }, + "openrouter/openai/gpt-5.6-luna-pro": { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 2e-08, + "cache_creation_input_token_cost": 2.5e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-5.6-luna-pro", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/openai/gpt-5.6-terra": { "input_cost_per_token": 2e-06, "output_cost_per_token": 1.2e-05, @@ -61292,6 +62598,28 @@ "cache_read_input_token_cost": 2e-07, "supports_prompt_caching": true }, + "openrouter/openai/gpt-5.6-terra-pro": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1.2e-05, + "cache_read_input_token_cost": 2e-07, + "cache_creation_input_token_cost": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-5.6-terra-pro", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/openai/o3": { "input_cost_per_token": 2e-06, "output_cost_per_token": 8e-06, @@ -61525,6 +62853,28 @@ "supports_pdf_input": true, "supports_prompt_caching": true }, + "openrouter/openai/gpt-6-astra-pro": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_creation_input_token_cost": 1.25e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-6-astra-pro", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/qwen/qwen3.8-flash": { "input_cost_per_token": 1.5e-07, "output_cost_per_token": 4.7e-07, @@ -61544,9 +62894,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 2.5e-07, - "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 3e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -61661,6 +63011,25 @@ "supports_vision": true, "supports_prompt_caching": true }, + "openrouter/qwen/qwen3.8-max-0902": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost": 2.5e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "source": "https://openrouter.ai/qwen/qwen3.8-max-0902", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": false, + "supports_prompt_caching": true + }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 6.5e-08, "output_cost_per_token": 1.8e-07, @@ -61732,9 +63101,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 3e-06, - "output_cost_per_token": 1.5e-05, - "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 2.1e-06, + "output_cost_per_token": 1.053e-05, + "cache_read_input_token_cost": 2.35e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -61831,9 +63200,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.2": { - "input_cost_per_token": 9.66e-07, - "output_cost_per_token": 3.036e-06, - "cache_read_input_token_cost": 1.932e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2e-06, + "cache_read_input_token_cost": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 131072, @@ -61864,9 +63233,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 6.6e-07, - "output_cost_per_token": 3.4e-06, - "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 7.1e-07, + "output_cost_per_token": 3.5e-06, + "cache_read_input_token_cost": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 235929, @@ -62114,10 +63483,28 @@ "supports_vision": true, "supports_pdf_input": true }, + "openrouter/openai/gpt-chat-latest": { + "input_cost_per_token": 5e-06, + "output_cost_per_token": 3e-05, + "cache_read_input_token_cost": 5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://openrouter.ai/openai/gpt-chat-latest", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": false, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_prompt_caching": true + }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.778e-08, - "output_cost_per_token": 1.7556e-07, - "cache_read_input_token_cost": 1.7556e-08, + "input_cost_per_token": 8.54e-08, + "output_cost_per_token": 1.708e-07, + "cache_read_input_token_cost": 1.708e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -62150,8 +63537,8 @@ "supports_prompt_caching": true }, "openrouter/google/gemma-4-26b-a4b-it": { - "input_cost_per_token": 7e-08, - "output_cost_per_token": 3.4e-07, + "input_cost_per_token": 4.2e-08, + "output_cost_per_token": 2.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -62833,7 +64220,7 @@ "supports_vision": false }, "openrouter/qwen/qwen3-next-80b-a3b-instruct": { - "input_cost_per_token": 1e-07, + "input_cost_per_token": 9e-08, "output_cost_per_token": 1.1e-06, "cache_read_input_token_cost": 7e-08, "litellm_provider": "openrouter", @@ -62959,8 +64346,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-30b-a3b-instruct-2507": { - "input_cost_per_token": 4.815e-08, - "output_cost_per_token": 1.9305e-07, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32000, @@ -63158,8 +64545,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-14b": { - "input_cost_per_token": 1.2e-07, - "output_cost_per_token": 2.4e-07, + "input_cost_per_token": 2.275e-07, + "output_cost_per_token": 9.1e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 16384, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index a51149bf958..7ed1e7e568b 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -478,6 +478,10 @@ "type": "number", "minimum": 0 }, + "output_cost_per_second_720p": { + "type": "number", + "minimum": 0 + }, "output_cost_per_token": { "type": "number", "minimum": 0, @@ -623,6 +627,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 +861,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.", diff --git a/osv-scanner.toml b/osv-scanner.toml index 5b0339bdcd0..3e070fc8cf7 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -1,6 +1,6 @@ [[IgnoredVulns]] id = "GHSA-w8v5-vhqr-4h9v" -ignoreUntil = 2026-09-09 +ignoreUntil = 2026-10-01 reason = "diskcache has no fixed release published; remove this entry once one exists" [[IgnoredVulns]] diff --git a/pyproject.toml b/pyproject.toml index b889a3a0e60..d33d693f794 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.101.0" +version = "1.102.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -67,8 +67,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.94", - "litellm-enterprise==0.1.65", + "litellm-proxy-extras==0.4.96", + "litellm-enterprise==0.1.66", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", @@ -114,7 +114,6 @@ caching = ["diskcache>=5.6.3,<6.0"] mcp = ["mcp>=1.28.1,<2.0"] # Driver for the MongoDB Atlas vector store; Atlas Vector Search has no HTTP query API. # The floor is 4.9 because that is the release AsyncMongoClient landed in. -mongodb = ["pymongo>=4.9,<5.0"] # SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels # bundle the native libxmlsec1/libxml2 libraries, so no system packages are # required. Kept out of the base `proxy` extra so it stays optional. @@ -328,7 +327,7 @@ members = ["enterprise", "litellm-proxy-extras"] profile = "black" [tool.commitizen] -version = "1.101.0" +version = "1.102.0" version_files = [ "pyproject.toml:^version", ] @@ -340,6 +339,7 @@ markers = [ "asyncio: mark test as an asyncio test", "limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')", "no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests", + "requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension", ] filterwarnings = [ # Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests) @@ -388,7 +388,7 @@ mutate_only_covered_lines = true # registered hash algorithm". Nothing to do with mutation coverage, and one # erroring test is enough to end the stats phase before any mutant runs. pytest_add_cli_args = [ - "-p", "no:retry", + "-p", "no:pytest-retry", "-p", "no:rerunfailures", "-p", "no:xdist", "--ignore=tests/test_litellm/proxy/management_endpoints/test_saml_sso.py", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fd7b30bc314..d63a69de76f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2956 + "limit": 2918 }, "ANN002": { "limit": 71 @@ -9,10 +9,10 @@ "limit": 806 }, "ANN201": { - "limit": 1979 + "limit": 1965 }, "ANN202": { - "limit": 831 + "limit": 829 }, "ANN204": { "limit": 683 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2916 + "limit": 2914 }, "C401": { "limit": 8 @@ -189,7 +189,7 @@ "limit": 0 }, "S110": { - "limit": 217 + "limit": 207 }, "S112": { "limit": 22 diff --git a/schema.prisma b/schema.prisma index 1c43668f227..7d521d54791 100644 --- a/schema.prisma +++ b/schema.prisma @@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + skills String[] @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -343,6 +344,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? @@ -491,7 +493,7 @@ model LiteLLM_JWTKeyMapping { updated_at DateTime @default(now()) @updatedAt updated_by String? - litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) @@unique([jwt_claim_name, jwt_claim_value]) @@index([jwt_claim_name, jwt_claim_value, is_active]) @@ -1035,6 +1037,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") updated_at DateTime @updatedAt @@ -1508,7 +1511,10 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") + baseline_models Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index adc4c0664be..485e118efd2 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -34,7 +34,7 @@ import subprocess import sys from pathlib import Path from types import MappingProxyType -from typing import NamedTuple +from typing import Final, NamedTuple if sys.version_info >= (3, 11): import tomllib @@ -42,7 +42,6 @@ else: import tomli as tomllib REPO_ROOT = Path(__file__).resolve().parent.parent -DEFAULT_BASE = "origin/litellm_internal_staging" DEFAULT_BUDGETS: tuple[str, ...] = ( "ruff-strict-budget.json", "type-discipline-budget.json", @@ -182,12 +181,15 @@ def regressions_for( def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("budgets", nargs="*", help="budget files to check") args = parser.parse_args() + from default_branch import resolve_base_ref + + base_ref: Final = resolve_base_ref(args.base, REPO_ROOT) budgets = args.budgets or list(DEFAULT_BUDGETS) - ref = _merge_base(args.base) + ref = _merge_base(base_ref) if not _ref_is_commit(ref): print( f"FAIL: base ref {ref!r} does not resolve to a commit, so the ratchet has nothing " @@ -204,14 +206,14 @@ def main() -> int: if base is None and head is None: continue if base is None: - print(f"skip {rel}: new file (no base at {args.base} to ratchet against)") + print(f"skip {rel}: new file (no base at {base_ref} to ratchet against)") continue checked.append(rel) regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel))) if regressions: print( - f"FAIL: budget limit(s) loosened vs base {args.base} (merge-base {ref[:12]}):" + f"FAIL: budget limit(s) loosened vs base {base_ref} (merge-base {ref[:12]}):" ) for reg in regressions: print(f" {reg.budget} {reg.rule}: {reg.detail}") @@ -223,7 +225,7 @@ def main() -> int: return 1 suffix = f" ({', '.join(checked)})" if checked else "" - print(f"OK: no budget limit increased vs base {args.base}{suffix}") + print(f"OK: no budget limit increased vs base {base_ref}{suffix}") return 0 diff --git a/scripts/default_branch.py b/scripts/default_branch.py new file mode 100644 index 00000000000..fb1852fd21c --- /dev/null +++ b/scripts/default_branch.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import argparse +import os +import subprocess +from pathlib import Path +from typing import Final + + +def _git(repo_root: Path, *args: str) -> str: + try: + result: Final = subprocess.run( + ["git", *args], + cwd=repo_root, + env={**os.environ, "GIT_TERMINAL_PROMPT": "0"}, + check=True, + capture_output=True, + text=True, + timeout=60, + ) + except (OSError, subprocess.SubprocessError) as exc: + raise SystemExit( + "Cannot verify the base branch against origin. Check remote access, " + "or supply an explicit base ref (--base / BASE_REF). " + f"Git operation failed: {exc}" + ) from exc + return result.stdout.strip() + + +def default_branch(repo_root: Path) -> str: + output: Final = _git(repo_root, "ls-remote", "--symref", "origin", "HEAD") + branches: Final = tuple( + line.removeprefix("ref: refs/heads/").removesuffix("\tHEAD") + for line in output.splitlines() + if line.startswith("ref: refs/heads/") and line.endswith("\tHEAD") + ) + if len(branches) != 1: + raise SystemExit("Origin did not advertise a default branch. Supply an explicit base ref (--base / BASE_REF).") + _git(repo_root, "check-ref-format", f"refs/heads/{branches[0]}") + return branches[0] + + +def resolve_base_ref(base_ref: str | None, repo_root: Path) -> str: + if base_ref: + return base_ref + branch: Final = default_branch(repo_root) + _git(repo_root, "fetch", "--quiet", "origin", f"+refs/heads/{branch}:refs/remotes/origin/{branch}") + return f"origin/{branch}" + + +def main() -> None: + parser: Final = argparse.ArgumentParser(description="Resolve the live default branch of origin.") + parser.add_argument("--base", help="Explicit comparison ref; skips default-branch discovery") + parser.add_argument("--repo-root", type=Path, default=Path.cwd()) + parser.add_argument("--branch", action="store_true", help="Print only the default branch name, without fetching") + args: Final = parser.parse_args() + print(default_branch(args.repo_root) if args.branch else resolve_base_ref(args.base, args.repo_root)) + + +if __name__ == "__main__": + main() diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index d38a0eee3de..1abd415d237 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -7,11 +7,15 @@ # - anything staged -> scope is the staged files; changed-but-unstaged files # whose checks were skipped are called out # - nothing staged -> scope is the working tree's diff against the merge base -# with origin/litellm_internal_staging, untracked files included +# with origin's current default branch, untracked files included # The per-area checks: # - litellm/ Python -> `make lint` (test-linting.yml's lint job) # - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) # + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) +# - tests/ Python, ruff-tests.toml, test-quality-budget.json, scripts/check_test_quality.py, +# scripts/test_quality_gate.py +# -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's +# test-tree ruff and test-quality budget steps) # - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) # - proxy/types -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (check-ui-api-types.yml) # @@ -29,8 +33,8 @@ set -eu # at a time instead of thrashing the machine. The wrapper exports # LITELLM_GATE_SLOT_HELD, so this re-exec happens exactly once and everything this # script spawns (make lint, the budget gates) skips its own acquisition. +script_dir=$(python3 -c 'import os, sys; print(os.path.dirname(os.path.realpath(sys.argv[1])))' "$0") if [ -z "${LITELLM_GATE_SLOT_HELD:-}" ]; then - script_dir=$(python3 -c 'import os, sys; print(os.path.dirname(os.path.realpath(sys.argv[1])))' "$0") exec python3 "$script_dir/gate_slot_lock.py" "$0" "$@" fi @@ -61,20 +65,24 @@ untracked=$(git ls-files --others --exclude-standard) if [ -n "$staged" ]; then scope=$staged else - git fetch --quiet origin litellm_internal_staging 2>/dev/null || true - merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || { - echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2 - echo " Fix: git fetch origin litellm_internal_staging" >&2 + base_ref=$(python3 "$script_dir/default_branch.py" --base "${BASE_REF:-}") || { + echo "check: FAIL" + exit 1 + } + export BASE_REF="$base_ref" + merge_base=$(git merge-base "$base_ref" HEAD 2>/dev/null) || { + echo "check: cannot resolve the merge base with $base_ref." >&2 + echo " Fix: fetch the base ref and provide BASE_REF=" >&2 echo "check: FAIL" exit 1 } scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u) if [ -z "$scope" ]; then - echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)" + echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs $base_ref)" echo "check: PASS" exit 0 fi - echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:" + echo "check: nothing staged; scoping to the working tree's diff against the merge base with $base_ref:" printf '%s\n' "$scope" | sed 's/^/ /' fi @@ -88,15 +96,14 @@ existing_files() { litellm_py_pattern='^litellm/.*\.py$' e2e_py_pattern='^tests/e2e/.*\.py$' +test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|test-quality-budget\.json|scripts/(check_test_quality|test_quality_gate)\.py)$' spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$' ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' -# CI's lint job (test-linting.yml) only inspects litellm/, so a tests-only or -# scripts-only commit can't turn it red; scope the trigger there to skip the slow -# make lint when it couldn't catch anything. litellm_py_files=$(scope_match "$litellm_py_pattern") e2e_py_files=$(scope_match "$e2e_py_pattern") +test_tree_files=$(scope_match "$test_tree_pattern") # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files) # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types @@ -136,6 +143,7 @@ if [ -n "$staged" ]; then } warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files" warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files" + warn_skipped "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_pattern" "$test_tree_files" warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed" warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files" fi @@ -288,6 +296,15 @@ if [ -n "$spec_files" ]; then set +m fi +if [ -n "$test_tree_files" ] && [ -z "$litellm_py_files" ]; then + echo "check: linting the test tree (ruff check --config ruff-tests.toml tests)" + uv run --no-sync ruff check --config ruff-tests.toml tests \ + || { echo "✗ Test-tree ruff failed. Fix the errors above, then re-run make check." >&2; status=1; } + echo "check: checking the test-quality budget (make lint-test-quality)" + make lint-test-quality \ + || { echo "✗ Test-quality budget failed. Fix the errors above, then re-run make check." >&2; status=1; } +fi + if [ -n "${python_pid:-}" ]; then wait "$python_pid" || status=1 cat "$python_log"; rm -f "$python_log" @@ -313,10 +330,12 @@ summary_item() { echo "check: summary" summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope" summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope" +summary_item "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_files" \ + "no tests/ Python files or test-tree lint inputs in scope" summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope" summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope" -if [ -z "$litellm_py_files$e2e_py_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then +if [ -z "$litellm_py_files$e2e_py_files$test_tree_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2 printf '%s\n' "$scope" | sed 's/^/ /' >&2 echo " A pass here is a no-op, not a lint verdict." >&2 diff --git a/scripts/ruff_strict_gate.py b/scripts/ruff_strict_gate.py index bf070beeb0f..8da10dd76f0 100644 --- a/scripts/ruff_strict_gate.py +++ b/scripts/ruff_strict_gate.py @@ -24,7 +24,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent STRICT_CONFIG = REPO_ROOT / "ruff-strict.toml" BUDGET_PATH = REPO_ROOT / "ruff-strict-budget.json" TARGET = "litellm" -DEFAULT_BASE = "origin/litellm_internal_staging" _HUNK = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@") @@ -193,7 +192,7 @@ def ratcheted_budget(budget: dict, current: dict, base: dict) -> dict: } -def cmd_update(base_ref: str = DEFAULT_BASE) -> None: +def cmd_update(base_ref: str) -> None: """Ratchet each rule's limit down by the violations this branch fixed. The working-tree count is compared against a ruff pass over a detached @@ -212,13 +211,15 @@ def cmd_update(base_ref: str = DEFAULT_BASE) -> None: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("--update", action="store_true") args = parser.parse_args() + from default_branch import resolve_base_ref from gate_slot_lock import held_slot + base_ref: Final = resolve_base_ref(args.base, REPO_ROOT) with held_slot(): - cmd_update(args.base) if args.update else cmd_check(args.base) + cmd_update(base_ref) if args.update else cmd_check(base_ref) if __name__ == "__main__": diff --git a/scripts/test_quality_gate.py b/scripts/test_quality_gate.py index 7d34b194f1c..e486324c741 100644 --- a/scripts/test_quality_gate.py +++ b/scripts/test_quality_gate.py @@ -10,17 +10,12 @@ base. Every rule is seeded at exactly its count on the day the gate landed, so the suite's existing debt is grandfathered and any net-new violation trips the gate -immediately. ``--update`` ratchets a limit down by the violations this branch -fixed relative to its branch point (the merge-base), so the ceilings only ever -fall. Base counts are measured with the *current* checker, so a rule introduced -on this branch is counted at the base too and ratchets like every other one. - -Only ever falling is not the same as always falling, so the gate enforces the -second half: a branch that clears violations and leaves the ceiling above its -new count fails, naming the rules and telling the author to run -``make lint-budget-update``. Without that, a removed violation could come back -later under a ceiling nobody lowered. Drift already in the base is never -blamed, so this fires only on the branch that did the clearing. +immediately. ``--update`` ratchets a limit down by the violations fixed relative +to ``--base``, so the ceilings only ever fall. Base counts are measured with the +*current* checker, so a rule introduced on this branch is counted at the base too +and ratchets like every other one. The ratchet runs as a scheduled automation +against the repository's default branch, not on PR branches, so concurrent PRs never +race to edit the same limit. The deliberate difference from its sibling: this gate has no headroom anywhere. Type discipline seeded LIT010/LIT011 at 1.5x to leave room for an in-flight @@ -34,20 +29,21 @@ import argparse import json import re import shutil +import signal import subprocess import sys import tempfile from collections import Counter from collections.abc import Mapping, Sequence from pathlib import Path -from types import MappingProxyType +from types import FrameType, MappingProxyType from typing import Final, NamedTuple REPO_ROOT: Final = Path(__file__).resolve().parent.parent CHECKER: Final = REPO_ROOT / "scripts" / "check_test_quality.py" BUDGET_PATH: Final = REPO_ROOT / "test-quality-budget.json" TARGET: Final = "tests" -DEFAULT_BASE: Final = "origin/litellm_internal_staging" +TERMINATION_SIGNALS: Final = (signal.SIGTERM, signal.SIGHUP) _HUNK: Final = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@", re.MULTILINE) _FILE_HEADER: Final = re.compile(r"^\+\+\+ b/(.+)$", re.MULTILINE) @@ -116,22 +112,33 @@ def count_by_rule(violations: Sequence[Violation]) -> Mapping[str, int]: return MappingProxyType(dict(Counter(v.code for v in violations))) -def base_counts(ref: str) -> Mapping[str, int]: +def _exit_on_termination(signum: int, _frame: FrameType | None) -> None: + raise SystemExit(128 + signum) + + +def _install_termination_handlers() -> None: + for termination in TERMINATION_SIGNALS: + if signal.getsignal(termination) == signal.SIG_DFL: + signal.signal(termination, _exit_on_termination) + + +def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER) -> Mapping[str, int]: """Rule counts at `ref`, measured with the *current* rule logic rather than whatever the checker looked like at that commit.""" + _install_termination_handlers() parent: Final = Path(tempfile.mkdtemp(prefix="tq_base_")) worktree: Final = parent / "wt" try: - _run(["git", "worktree", "add", "--detach", str(worktree), ref]) + _run(["git", "worktree", "add", "--detach", str(worktree), ref], cwd=repo_root) (worktree / "scripts").mkdir(parents=True, exist_ok=True) - checker: Final = worktree / "scripts" / "check_test_quality.py" - shutil.copy(CHECKER, checker) - return count_by_rule(_check(worktree, checker)) + base_checker: Final = worktree / "scripts" / "check_test_quality.py" + shutil.copy(checker, base_checker) + return count_by_rule(_check(worktree, base_checker)) finally: # Teardown must never raise, or it masks the real error when the body failed. subprocess.run( ["git", "worktree", "remove", "--force", str(worktree)], - cwd=REPO_ROOT, capture_output=True, text=True, + cwd=repo_root, capture_output=True, text=True, ) shutil.rmtree(parent, ignore_errors=True) @@ -144,21 +151,6 @@ def over_ceiling(head: Mapping[str, int], budget: Mapping[str, Mapping[str, int] ) -def unratcheted( - head: Mapping[str, int], - base: Mapping[str, int], - budget: Mapping[str, Mapping[str, int]], -) -> tuple[Breach, ...]: - """Rules this branch cleared without lowering the ceiling behind them. Requires - both `head < base`, so drift already in the base is never blamed on this change, - and `head < limit`, so a ceiling already at the count is left alone.""" - return tuple(sorted( - Breach(rule, head.get(rule, 0), spec["limit"], head.get(rule, 0) - base.get(rule, 0)) - for rule, spec in budget.items() - if head.get(rule, 0) < base.get(rule, 0) and head.get(rule, 0) < spec["limit"] - )) - - def evaluate( head: Mapping[str, int], base: Mapping[str, int], @@ -198,38 +190,15 @@ def introduced( return tuple(v for v in violations if v.line in changed.get(v.file, frozenset())) -def touches_measured_tree(base_point: str) -> bool: - """Whether this branch changed anything that can move a count. A branch that - touches neither the test tree nor the checker cannot have cleared a violation, - so the base scan is skipped and the gate stays cheap on the common change.""" - changed: Final = _run( - ["git", "diff", "--name-only", base_point, "--", TARGET, str(CHECKER.relative_to(REPO_ROOT))] - ) - return bool(changed.strip()) - - def cmd_check(base: str) -> None: budget: Final = json.loads(BUDGET_PATH.read_text()) head: Final = head_violations() head_counts: Final = count_by_rule(head) - base_point: Final = resolve_base_point(base) - if not over_ceiling(head_counts, budget) and not touches_measured_tree(base_point): + if not over_ceiling(head_counts, budget): print(f"OK: every TQ rule is within its test-suite ceiling (base {base})") return + base_point: Final = resolve_base_point(base) base_at_point: Final = base_counts(base_point) - stale: Final = unratcheted(head_counts, base_at_point, budget) - if stale: - print(f"FAIL: TQ-rule limits were left above the count this branch reached (base {base}):") - for breach in stale: - print( - f" {breach.rule}: this branch cleared {-breach.added} down to {breach.total}, " - f"but the limit is still {breach.cap}" - ) - print( - "Run `make lint-budget-update` and commit the lowered limits, so the " - "violations you cleared cannot come back under a ceiling nobody moved." - ) - raise SystemExit(1) breaches: Final = evaluate(head_counts, base_at_point, budget) if not breaches: print(f"OK: every TQ rule is within its test-suite ceiling (base {base})") @@ -270,7 +239,7 @@ def ratcheted_budget( }) -def cmd_update(base_ref: str = DEFAULT_BASE) -> None: +def cmd_update(base_ref: str) -> None: """Ratchet each rule's limit down by the violations this branch fixed.""" budget: Final = json.loads(BUDGET_PATH.read_text()) base_point: Final = resolve_base_point(base_ref) @@ -294,19 +263,20 @@ def cmd_seed() -> None: def main() -> None: parser: Final = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("--update", action="store_true") parser.add_argument("--seed", action="store_true") args: Final = parser.parse_args() + from default_branch import resolve_base_ref from gate_slot_lock import held_slot with held_slot(): if args.seed: cmd_seed() elif args.update: - cmd_update(args.base) + cmd_update(resolve_base_ref(args.base, REPO_ROOT)) else: - cmd_check(args.base) + cmd_check(resolve_base_ref(args.base, REPO_ROOT)) if __name__ == "__main__": diff --git a/scripts/type_check_gate.py b/scripts/type_check_gate.py index 763835e6d2e..78f74ec65a1 100644 --- a/scripts/type_check_gate.py +++ b/scripts/type_check_gate.py @@ -71,7 +71,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent BUDGET_PATH = REPO_ROOT / "basedpyright-code-budget.json" PYRIGHT_CONFIG = REPO_ROOT / "pyrightconfig.json" UV_LOCK = REPO_ROOT / "uv.lock" -DEFAULT_BASE = "origin/litellm_internal_staging" CACHE_FILE_PREFIX = "basedpyright-base-" CACHE_KEEP_ENTRIES = 8 ARTIFACT_NAME_PREFIX = "basedpyright-counts-" @@ -578,7 +577,7 @@ def ratcheted_budget( } -def cmd_update(current: Mapping[str, int], base_ref: str = DEFAULT_BASE) -> None: +def cmd_update(current: Mapping[str, int], base_ref: str) -> None: """Ratchet each rule's limit down by the errors this branch fixed. `current` is the working-tree count; the reference count comes @@ -666,12 +665,14 @@ def cmd_check(head: Mapping[str, int], base_ref: str) -> None: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("--update", action="store_true") parser.add_argument("--emit-counts-dir", type=Path) args = parser.parse_args() + from default_branch import resolve_base_ref from gate_slot_lock import held_slot + base_ref: Final = None if args.emit_counts_dir is not None else resolve_base_ref(args.base, REPO_ROOT) with held_slot(): ensure_typecheck_env() head = count_basedpyright(run_basedpyright()) @@ -679,10 +680,8 @@ def main() -> None: cmd_emit_counts( head, args.emit_counts_dir, _run(["git", "rev-parse", "HEAD"]).strip() ) - elif args.update: - cmd_update(head, args.base) - else: - cmd_check(head, args.base) + elif base_ref is not None: + cmd_update(head, base_ref) if args.update else cmd_check(head, base_ref) if __name__ == "__main__": diff --git a/scripts/type_discipline_gate.py b/scripts/type_discipline_gate.py index 5f6474f20bc..40e61cf7265 100644 --- a/scripts/type_discipline_gate.py +++ b/scripts/type_discipline_gate.py @@ -44,7 +44,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent CHECKER = REPO_ROOT / "scripts" / "check_type_discipline.py" BUDGET_PATH = REPO_ROOT / "type-discipline-budget.json" TARGET = "litellm" -DEFAULT_BASE = "origin/litellm_internal_staging" _HUNK = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@") _LINE = re.compile(r"^(?P.+?):(?P\d+): (?PLIT\d+) ") @@ -239,7 +238,7 @@ def _base_budget_rules(base_point: str) -> frozenset: return frozenset(json.loads(proc.stdout)) -def cmd_update(base_ref: str = DEFAULT_BASE) -> None: +def cmd_update(base_ref: str) -> None: """Ratchet each rule's limit down by the violations this branch fixed. The working-tree count is compared against a checker pass over a detached @@ -264,13 +263,15 @@ def cmd_update(base_ref: str = DEFAULT_BASE) -> None: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("--update", action="store_true") args = parser.parse_args() + from default_branch import resolve_base_ref from gate_slot_lock import held_slot + base_ref: Final = resolve_base_ref(args.base, REPO_ROOT) with held_slot(): - cmd_update(args.base) if args.update else cmd_check(args.base) + cmd_update(base_ref) if args.update else cmd_check(base_ref) if __name__ == "__main__": diff --git a/terraform/litellm/README.md b/terraform/litellm/README.md index d4b40741052..7cc6e4c08ba 100644 --- a/terraform/litellm/README.md +++ b/terraform/litellm/README.md @@ -183,6 +183,7 @@ only where the underlying cloud forces it. | Extra secret-backed env | `gateway_extra_secrets`, `backend_extra_secrets` (ARNs) | `gateway_extra_secrets`, `backend_extra_secrets` (resource IDs) | | Uvicorn `--workers` on gateway | `gateway_num_workers` | `gateway_num_workers` | | OpenTelemetry v2 (opt-in) | `otel_endpoint`, `otel_exporter`, `otel_environment_name`, `otel_capture_message_content`, `otel_headers_secret_arn` | `otel_endpoint`, `otel_exporter`, `otel_environment_name`, `otel_capture_message_content`, `otel_headers_secret` | +| Collector sidecar (opt-in) | `collector_enabled`, `collector_port`, `collector_cpu`, `collector_memory`, `collector_buffer_size`, `collector_on_unavailable`, `collector_drain_timeout_seconds` | same names; `collector_cpu` / `collector_memory` take Cloud Run strings | Each module stamps its own stack-identity tag (`litellm:stack` on AWS, `litellm-stack` on GCP — GCP label keys forbid colons) plus diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 389027bf5ca..6fcbdc2f500 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -242,6 +242,149 @@ this with `litellm_license`. To tune the export cadence, set `LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` / `backend_extra_env` +### Prometheus metrics sidecar + +`gateway_metrics_port` adds a `metrics` sidecar +(`python -m litellm.proxy.prometheus_metrics_server`) to the gateway task that +aggregates the workers' samples over a shared task volume, so a scrape never +runs on an inference worker. The ALB never routes to that port and the tasks +security group only opens it to `gateway_metrics_scrape_cidrs`. Needs +`gateway_image` v1.101.0 or newer. See +[Prometheus metrics](https://docs.litellm.ai/docs/proxy/prometheus) for the +metrics themselves. + +```hcl +gateway_metrics_port = 4001 +gateway_metrics_scrape_cidrs = ["10.0.0.0/16"] +``` + +### In-container connection pool + +Each of the `gateway_num_workers` uvicorn workers opens its own Prisma pool +straight to Postgres, so one task holds `workers x connection_limit` +connections and the fleet's footprint against the database ceiling grows with +every task. `gateway_connection_pool_enabled` runs a PgBouncer (transaction +mode, loopback) inside the gateway container that all workers share, capping +the task at `gateway_pool_max_db_connections` upstream connections however +many workers it runs. `gateway_pool_max_client_conn` bounds the worker-side +connections the pooler accepts. The module sets +`LITELLM_PGBOUNCER_ENABLED`, `LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS` and +`LITELLM_PGBOUNCER_MAX_CLIENT_CONN` on the gateway container only; the backend +and the migration task keep the direct connection. + +```hcl +gateway_num_workers = 4 +gateway_connection_pool_enabled = true +gateway_pool_max_db_connections = 20 +gateway_pool_max_client_conn = 1000 +``` + +The pool works with the module-created Aurora as well as an existing database +via `database_url`. Against Aurora it authenticates with the same rotating IAM +tokens the workers used to (see [Aurora + IAM auth](#aurora--iam-auth)): the +pooler mints a token from the task role, renews it before it expires and hands +the workers a loopback URL with a static password instead + +The componentized `gateway_image` starts through `python -m gateway.launch`, +which reads these variables, starts the pooler once per task and hands the +workers its loopback URL; the classic `litellm` image honours them the same +way. + +### Scaling the gateway on requests and tokens + +By default the gateway service target-tracks CPU (`gateway_cpu_target`) and +memory (`gateway_memory_target`). Two more targets add workload signals next +to them. Application Auto Scaling evaluates every attached policy and follows +the one asking for the most tasks, so the resource policies keep working as a +floor while requests or tokens drive scale-out + +Both targets are per task per second, the way load is usually quoted (1k +rps, 75M tok/s). CloudWatch is the limit on how fast they react: target +tracking evaluates every metric, predefined or custom, aggregated over +60-second periods and has no period setting, so ECS reacts on a roughly +one-minute cadence whatever unit the variable is written in. The Kubernetes +charts get a faster signal because the Prometheus `rate()` window and scrape +interval are theirs to shorten + +`gateway_target_requests_per_second` adds an `ALBRequestCountPerTarget` +policy on the gateway target group. The ALB publishes that metric as requests +per minute per registered task, so the policy's target value is 60 times the +variable: 90 rps becomes a target of 5,400 per minute. No agent or sidecar is +needed + +`gateway_target_tokens_per_second` adds a metric-math policy over a +CloudWatch metric of the gateway's `litellm_total_tokens_metric_total` +counter and the service's `RunningTaskCount` from Container Insights. Nothing +native to ECS carries token throughput, so you publish that metric yourself +with the CloudWatch agent's Prometheus scraper pointed at the metrics sidecar +above. The agent emits the delta of a counter between scrapes, so `Sum` over +the 60-second period is the tokens served in that minute; the expression +divides by 60 (`tokens_per_second`) and then by the task count +(`tokens_per_second_per_task`). Tokens are counted when a response completes, +so long streams show up late in this signal. `gateway_tokens_metric` tells the +policy where the agent publishes: the namespace, the metric name (defaults to +the counter name) and the dimensions from your `metric_declaration` + +```hcl +gateway_metrics_port = 4001 +gateway_target_requests_per_second = 90 +gateway_target_tokens_per_second = 6000000 +gateway_tokens_metric = { + namespace = "LiteLLM/Prometheus" + dimensions = { ClusterName = "acme-litellm-prod", TaskDefinitionFamily = "acme-litellm-prod-gateway" } +} +``` + +Worked example for the request policy: 1,000 rps across 10 tasks is 100 rps +per task (the ALB reports it as 6,000 per minute per target) against a target +of 90 (5,400), so target tracking sizes the service to +`ceil(10 * 100 / 90) = 12` tasks. The token policy does the same arithmetic: +ten tasks handle 4,200,000,000 tokens in a minute, `tokens / 60` is +70,000,000 tokens per second and `tokens_per_second / running_tasks` is +7,000,000 against a target of 6,000,000, so the service grows to +`ceil(10 * 7000000 / 6000000) = 12`. Container Insights must be enabled on the +cluster for `RunningTaskCount` to exist + +### Collector sidecar + +`collector_enabled = true` adds a second container to the gateway task +that runs `python -m litellm.proxy.collector` from the gateway image, and sets +`LITELLM_COLLECTOR_ENABLED=true` on the gateway so its uvicorn workers +ship spend events (SpendLogs writes, key/team/user spend updates, budget +alerts) to the sidecar instead of running that pipeline in the request +path. This is the Terraform counterpart of helm's `gateway.collector`. +The default (`false`) leaves the task definition exactly as before. + +Fargate tasks share one network namespace, so the sidecar listens on +loopback TCP (`tcp://127.0.0.1:${collector_port}`, default 4010) instead +of the Unix socket helm uses; the proxy rejects any non-loopback address. +The sidecar gets the same database, Redis, master-key, license, proxy +config, and `gateway_extra_env` / `gateway_extra_secrets` values as the +gateway container, runs with `LITELLM_JOB_ROLE=collector`, and is +non-essential with an ECS restart policy, so a sidecar crash restarts it in +place while the gateway falls back to in-process spend tracking. With +`gateway_connection_pool_enabled` it also gets the `LITELLM_PGBOUNCER_*` env, +so with a password-authenticated database (`create_database = false`) its +Prisma client goes through the task-local PgBouncer instead of opening a +second pool straight to the database. Under IAM token auth (the module-managed +Aurora cluster) the collector keeps its own direct connection on purpose: the +pooler's auth file only holds the token the gateway container minted, which +the sidecar cannot present, so it mints its own. + +```hcl +collector_enabled = true +# collector_cpu = 512 # carved out of gateway_cpu +# collector_memory = 2048 # MiB, carved out of gateway_memory +# collector_buffer_size = 1000 +# collector_on_unavailable = "fallback" # or "drop" +# collector_drain_timeout_seconds = 10 +``` + +Both sidecar reservations must leave room for the gateway container inside +`gateway_cpu` / `gateway_memory` (the plan fails otherwise). Service +autoscaling keeps tracking the whole task's CPU and memory, sidecar +included + ## Tenant deployment Every resource the stack creates is named `${tenant}-litellm-${env}` (or diff --git a/terraform/litellm/aws/autoscaling.tf b/terraform/litellm/aws/autoscaling.tf index 71b6c24fac7..5197311d1b3 100644 --- a/terraform/litellm/aws/autoscaling.tf +++ b/terraform/litellm/aws/autoscaling.tf @@ -52,6 +52,105 @@ resource "aws_appautoscaling_policy" "gateway_memory" { } } +resource "aws_appautoscaling_policy" "gateway_requests" { + count = var.gateway_autoscaling_enabled && var.gateway_target_requests_per_second > 0 ? 1 : 0 + name = "${local.name}-gateway-requests" + policy_type = "TargetTrackingScaling" + service_namespace = aws_appautoscaling_target.gateway[0].service_namespace + resource_id = aws_appautoscaling_target.gateway[0].resource_id + scalable_dimension = aws_appautoscaling_target.gateway[0].scalable_dimension + + target_tracking_scaling_policy_configuration { + predefined_metric_specification { + predefined_metric_type = "ALBRequestCountPerTarget" + resource_label = "${aws_lb.this.arn_suffix}/${aws_lb_target_group.gateway.arn_suffix}" + } + # ALBRequestCountPerTarget is a per-minute count + target_value = var.gateway_target_requests_per_second * 60 + } +} + +resource "aws_appautoscaling_policy" "gateway_tokens" { + count = var.gateway_autoscaling_enabled && var.gateway_target_tokens_per_second > 0 ? 1 : 0 + name = "${local.name}-gateway-tokens" + policy_type = "TargetTrackingScaling" + service_namespace = aws_appautoscaling_target.gateway[0].service_namespace + resource_id = aws_appautoscaling_target.gateway[0].resource_id + scalable_dimension = aws_appautoscaling_target.gateway[0].scalable_dimension + + lifecycle { + precondition { + condition = var.gateway_tokens_metric != null + error_message = "gateway_tokens_metric is required when gateway_target_tokens_per_second > 0." + } + } + + target_tracking_scaling_policy_configuration { + target_value = var.gateway_target_tokens_per_second + + # target tracking has no period setting and always aggregates over 60s + customized_metric_specification { + metrics { + id = "tokens" + return_data = false + + metric_stat { + stat = "Sum" + + metric { + namespace = var.gateway_tokens_metric.namespace + metric_name = var.gateway_tokens_metric.name + + dynamic "dimensions" { + for_each = var.gateway_tokens_metric.dimensions + content { + name = dimensions.key + value = dimensions.value + } + } + } + } + } + + metrics { + id = "running_tasks" + return_data = false + + metric_stat { + stat = "Average" + + metric { + namespace = "ECS/ContainerInsights" + metric_name = "RunningTaskCount" + + dimensions { + name = "ClusterName" + value = aws_ecs_cluster.this.name + } + dimensions { + name = "ServiceName" + value = aws_ecs_service.gateway.name + } + } + } + } + + metrics { + id = "tokens_per_second" + expression = "tokens / 60" + return_data = false + } + + metrics { + id = "tokens_per_second_per_task" + expression = "tokens_per_second / running_tasks" + label = "Tokens per second per gateway task" + return_data = true + } + } + } +} + # ---------- Backend ---------- resource "aws_appautoscaling_target" "backend" { count = var.backend_autoscaling_enabled ? 1 : 0 diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index 2fb1ca08205..2b235c2bad5 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -212,9 +212,54 @@ locals { # pull the config from S3 first, so the command goes through `sh -c`; # otherwise we keep the image's ENTRYPOINT and only override `command`. gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}" + + gateway_pool_env = var.gateway_connection_pool_enabled ? [ + { name = "LITELLM_PGBOUNCER_ENABLED", value = "true" }, + { name = "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", value = tostring(var.gateway_pool_max_db_connections) }, + { name = "LITELLM_PGBOUNCER_MAX_CLIENT_CONN", value = tostring(var.gateway_pool_max_client_conn) }, + ] : [] + + metrics_enabled = var.gateway_metrics_port != null + metrics_multiproc_dir = "/tmp/litellm_prometheus_multiproc" + metrics_volume = "prometheus-multiproc" + metrics_env = local.metrics_enabled ? [{ name = "PROMETHEUS_MULTIPROC_DIR", value = local.metrics_multiproc_dir }] : [] + metrics_mount_points = local.metrics_enabled ? [{ sourceVolume = local.metrics_volume, containerPath = local.metrics_multiproc_dir }] : [] + metrics_health_cmd = "import socket; socket.create_connection(('127.0.0.1', ${coalesce(var.gateway_metrics_port, 0)}), timeout=2).close()" + + gateway_metrics_container = local.metrics_enabled ? [ + { + name = "metrics" + image = var.gateway_image + essential = false + entryPoint = ["python", "-m", "litellm.proxy.prometheus_metrics_server"] + command = ["--port", tostring(var.gateway_metrics_port)] + + portMappings = [{ containerPort = var.gateway_metrics_port, protocol = "tcp" }] + environment = local.metrics_env + mountPoints = local.metrics_mount_points + + healthCheck = { + command = ["CMD", "python", "-c", local.metrics_health_cmd] + interval = 30 + timeout = 5 + retries = 3 + startPeriod = 30 + } + + logConfiguration = { + logDriver = "awslogs" + options = { + awslogs-group = aws_cloudwatch_log_group.gateway.name + awslogs-region = var.region + awslogs-stream-prefix = "metrics" + } + } + } + ] : [] + backend_uvicorn_args = "--host 0.0.0.0 --port 4001" - gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac" + gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run python -m gateway.launch ${local.gateway_uvicorn_args};; *) exec python -m gateway.launch ${local.gateway_uvicorn_args};; esac" backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac" gateway_proxy_overrides = local.proxy_config_enabled ? { @@ -233,6 +278,62 @@ locals { "${local.proxy_config_fetch_cmd} && ${local.backend_launch_cmd}" ] } : {} + + collector_address = "tcp://127.0.0.1:${var.collector_port}" + collector_env = var.collector_enabled ? [ + { name = "LITELLM_COLLECTOR_ENABLED", value = "true" }, + { name = "LITELLM_COLLECTOR_ADDRESS", value = local.collector_address }, + { name = "LITELLM_COLLECTOR_BUFFER_SIZE", value = tostring(var.collector_buffer_size) }, + { name = "LITELLM_COLLECTOR_ON_UNAVAILABLE", value = var.collector_on_unavailable }, + { name = "LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS", value = tostring(var.collector_drain_timeout_seconds) }, + ] : [] + + gateway_environment = concat( + local.shared_env, + local.gateway_otel_env, + local.billing_metrics_env, + local.gateway_extra_env_list, + local.proxy_config_env, + local.metrics_env, + local.gateway_pool_env, + local.collector_env, + ) + + collector_launch_cmd = "exec python -m litellm.proxy.collector" + collector_command = [ + local.proxy_config_enabled ? "${local.proxy_config_fetch_cmd} && ${local.collector_launch_cmd}" : local.collector_launch_cmd + ] + + collector_container = var.collector_enabled ? [{ + name = "collector" + image = var.gateway_image + essential = false + cpu = var.collector_cpu + memory = var.collector_memory + + restartPolicy = { enabled = true } + + entryPoint = ["sh", "-c"] + command = local.collector_command + environment = concat( + local.shared_env, + local.gateway_extra_env_list, + local.proxy_config_env, + local.gateway_pool_env, + local.collector_env, + [{ name = "LITELLM_JOB_ROLE", value = "collector" }], + ) + secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list) + + logConfiguration = { + logDriver = "awslogs" + options = { + awslogs-group = aws_cloudwatch_log_group.gateway.name + awslogs-region = var.region + awslogs-stream-prefix = "collector" + } + } + }] : [] } # ---------- Gateway ---------- @@ -259,6 +360,21 @@ resource "aws_ecs_task_definition" "gateway" { ) error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." } + + precondition { + condition = !var.gateway_connection_pool_enabled || local.database_enabled + error_message = "gateway_connection_pool_enabled needs a database: set create_database = true or pass database_url." + } + + precondition { + condition = !var.collector_enabled || (var.collector_cpu < var.gateway_cpu && var.collector_memory < var.gateway_memory) + error_message = "collector_cpu and collector_memory are carved out of gateway_cpu / gateway_memory and must leave room for the gateway container." + } + + precondition { + condition = !var.collector_enabled || var.gateway_metrics_port == null || var.collector_port != var.gateway_metrics_port + error_message = "collector_port and gateway_metrics_port must differ: both sidecars bind loopback in the same task." + } } family = "${local.name}-gateway" @@ -269,7 +385,7 @@ resource "aws_ecs_task_definition" "gateway" { execution_role_arn = aws_iam_role.task_execution.arn task_role_arn = aws_iam_role.task.arn - container_definitions = jsonencode([ + container_definitions = jsonencode(concat([ merge( { name = "gateway" @@ -277,14 +393,9 @@ resource "aws_ecs_task_definition" "gateway" { essential = true portMappings = [{ containerPort = 4000, protocol = "tcp" }] - environment = concat( - local.shared_env, - local.gateway_otel_env, - local.billing_metrics_env, - local.gateway_extra_env_list, - local.proxy_config_env, - ) - secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list) + environment = local.gateway_environment + secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list) + mountPoints = local.metrics_mount_points # Container-level healthCheck intentionally omitted — the wolfi # runtime image doesn't ship curl/wget. The ALB target group polls @@ -301,7 +412,14 @@ resource "aws_ecs_task_definition" "gateway" { }, local.gateway_proxy_overrides, ) - ]) + ], local.gateway_metrics_container, local.collector_container)) + + dynamic "volume" { + for_each = local.metrics_enabled ? [1] : [] + content { + name = local.metrics_volume + } + } tags = local.tags } diff --git a/terraform/litellm/aws/examples/default/main.tf b/terraform/litellm/aws/examples/default/main.tf index 2eeaf6adb50..115003f7ddd 100644 --- a/terraform/litellm/aws/examples/default/main.tf +++ b/terraform/litellm/aws/examples/default/main.tf @@ -48,4 +48,7 @@ module "litellm" { backend_extra_env = var.backend_extra_env gateway_extra_secrets = var.gateway_extra_secrets backend_extra_secrets = var.backend_extra_secrets + + gateway_metrics_port = var.gateway_metrics_port + gateway_metrics_scrape_cidrs = var.gateway_metrics_scrape_cidrs } diff --git a/terraform/litellm/aws/examples/default/terraform.tfvars.example b/terraform/litellm/aws/examples/default/terraform.tfvars.example index 59301ea6aa5..880ecf56555 100644 --- a/terraform/litellm/aws/examples/default/terraform.tfvars.example +++ b/terraform/litellm/aws/examples/default/terraform.tfvars.example @@ -102,6 +102,13 @@ env = "stage" # } # } +# ---------- Prometheus metrics sidecar ---------- +# Serve /metrics from a sidecar in the gateway task instead of the inference +# workers. The port is not behind the ALB and has no auth: open it only to +# your Prometheus subnets. +# gateway_metrics_port = 4001 +# gateway_metrics_scrape_cidrs = ["10.0.0.0/16"] + # ---------- Extra env / secrets ---------- # Plain-text env vars (non-sensitive). Land directly in the ECS task def. # gateway_extra_env = { diff --git a/terraform/litellm/aws/examples/default/variables.tf b/terraform/litellm/aws/examples/default/variables.tf index d8ab56b13af..f8140266fca 100644 --- a/terraform/litellm/aws/examples/default/variables.tf +++ b/terraform/litellm/aws/examples/default/variables.tf @@ -158,3 +158,15 @@ variable "backend_extra_secrets" { type = map(string) default = {} } + +variable "gateway_metrics_port" { + description = "Port for the Prometheus metrics sidecar in the gateway task. Null keeps /metrics on the gateway port only." + type = number + default = null +} + +variable "gateway_metrics_scrape_cidrs" { + description = "CIDRs allowed to scrape gateway_metrics_port." + type = list(string) + default = [] +} diff --git a/terraform/litellm/aws/network.tf b/terraform/litellm/aws/network.tf index 4563eefbba5..c54949b4a59 100644 --- a/terraform/litellm/aws/network.tf +++ b/terraform/litellm/aws/network.tf @@ -156,6 +156,17 @@ resource "aws_security_group" "tasks" { security_groups = [aws_security_group.alb.id] } + dynamic "ingress" { + for_each = local.metrics_enabled && length(var.gateway_metrics_scrape_cidrs) > 0 ? [1] : [] + content { + description = "Prometheus scrapers to the gateway metrics sidecar" + from_port = var.gateway_metrics_port + to_port = var.gateway_metrics_port + protocol = "tcp" + cidr_blocks = var.gateway_metrics_scrape_cidrs + } + } + egress { description = "All egress (LLM providers, RDS, Redis)" from_port = 0 diff --git a/terraform/litellm/aws/tests/collector.tftest.hcl b/terraform/litellm/aws/tests/collector.tftest.hcl new file mode 100644 index 00000000000..1465130232f --- /dev/null +++ b/terraform/litellm/aws/tests/collector.tftest.hcl @@ -0,0 +1,144 @@ +# Plan-only coverage for the opt-in collector sidecar in the gateway task. +# The rendered container_definitions JSON is unknown at plan time (it embeds +# Aurora/ElastiCache endpoints and secret ARNs), so the assertions target the +# locals it is built from. Run from terraform/litellm/aws with `terraform test`. + +mock_provider "aws" { + mock_data "aws_iam_policy_document" { + defaults = { + json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}" + } + } +} +mock_provider "random" {} + +variables { + region = "us-east-1" + tenant = "acme" + env = "test" + allow_plaintext_alb = true + azs = ["us-east-1a", "us-east-1b"] +} + +run "disabled_by_default_leaves_the_task_untouched" { + command = plan + + assert { + condition = length(local.collector_container) == 0 + error_message = "The gateway task must stay single-container unless collector_enabled is set." + } + + assert { + condition = !anytrue([for e in local.gateway_environment : startswith(e.name, "LITELLM_COLLECTOR_")]) + error_message = "No LITELLM_COLLECTOR_* env may reach the gateway while the sidecar is disabled." + } +} + +run "enabled_adds_a_sidecar_that_shares_the_gateway_transport" { + command = plan + + variables { + collector_enabled = true + collector_port = 4321 + collector_buffer_size = 250 + collector_on_unavailable = "drop" + gateway_extra_env = { OPENAI_API_BASE = "https://example.invalid" } + gateway_extra_secrets = { OPENAI_API_KEY = "arn:aws:secretsmanager:us-east-1:111122223333:secret:openai-AbCdEf" } + } + + assert { + condition = length(local.collector_container) == 1 && local.collector_container[0].name == "collector" + error_message = "Enabling the sidecar must add exactly one collector container." + } + + assert { + condition = alltrue([ + for env in [local.gateway_environment, local.collector_container[0].environment] : ( + { for e in env : e.name => e.value }["LITELLM_COLLECTOR_ENABLED"] == "true" && + { for e in env : e.name => e.value }["LITELLM_COLLECTOR_ADDRESS"] == "tcp://127.0.0.1:4321" && + { for e in env : e.name => e.value }["LITELLM_COLLECTOR_BUFFER_SIZE"] == "250" && + { for e in env : e.name => e.value }["LITELLM_COLLECTOR_ON_UNAVAILABLE"] == "drop" && + { for e in env : e.name => e.value }["LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS"] == "10" + ) + ]) + error_message = "Gateway and sidecar must agree on the loopback address and the collector knobs." + } + + assert { + condition = ( + local.collector_container[0].image == var.gateway_image && + local.collector_container[0].entryPoint == ["sh", "-c"] && + local.collector_container[0].command == ["exec python -m litellm.proxy.collector"] && + local.collector_container[0].essential == false && + local.collector_container[0].restartPolicy.enabled == true && + { for e in local.collector_container[0].environment : e.name => e.value }["LITELLM_JOB_ROLE"] == "collector" + ) + error_message = "The sidecar must run litellm.proxy.collector from the gateway image as a restartable, non-essential collector." + } + + assert { + condition = ( + { for e in local.collector_container[0].environment : e.name => e.value }["OPENAI_API_BASE"] == "https://example.invalid" && + contains([for e in local.collector_container[0].environment : e.name], "DATABASE_HOST") && + contains([for e in local.collector_container[0].environment : e.name], "REDIS_HOST") && + contains([for s in local.collector_container[0].secrets : s.name], "LITELLM_MASTER_KEY") && + contains([for s in local.collector_container[0].secrets : s.name], "OPENAI_API_KEY") + ) + error_message = "The sidecar must receive the gateway's database, Redis, and shared secrets plus gateway_extra_env / gateway_extra_secrets." + } + + assert { + condition = !contains(keys(local.collector_container[0]), "portMappings") + error_message = "The sidecar must not expose a port to the task's load balancer." + } + + assert { + condition = local.collector_container[0].cpu == 512 && local.collector_container[0].memory == 2048 + error_message = "The sidecar defaults must mirror helm's collector resources (500m / 2Gi)." + } +} + +run "proxy_config_is_fetched_by_the_sidecar_too" { + command = plan + + variables { + collector_enabled = true + proxy_config = { model_list = [] } + } + + assert { + condition = ( + startswith(local.collector_container[0].command[0], local.proxy_config_fetch_cmd) && + endswith(local.collector_container[0].command[0], "exec python -m litellm.proxy.collector") && + contains([for e in local.collector_container[0].environment : e.name], "CONFIG_FILE_PATH") + ) + error_message = "The sidecar must pull the proxy config from S3 before starting, like the gateway does." + } +} + +run "sidecar_must_leave_room_for_the_gateway" { + command = plan + + variables { + collector_enabled = true + collector_cpu = 1024 + } + + expect_failures = [ + aws_ecs_task_definition.gateway, + ] +} + +run "sidecars_must_not_share_a_loopback_port" { + command = plan + + variables { + collector_enabled = true + collector_port = 4001 + gateway_metrics_port = 4001 + } + + expect_failures = [ + aws_ecs_task_definition.gateway, + ] +} diff --git a/terraform/litellm/aws/tests/connection_pool.tftest.hcl b/terraform/litellm/aws/tests/connection_pool.tftest.hcl new file mode 100644 index 00000000000..7408b1179ed --- /dev/null +++ b/terraform/litellm/aws/tests/connection_pool.tftest.hcl @@ -0,0 +1,164 @@ +# Plan-only coverage for the in-container PgBouncer knobs on the gateway task. +# `mock_provider` keeps this offline: no AWS credentials, no API calls, no +# resources. Run from terraform/litellm/aws with `terraform test`. + +mock_provider "aws" { + mock_data "aws_iam_policy_document" { + defaults = { + json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}" + } + } +} +mock_provider "random" {} + +variables { + region = "us-east-1" + tenant = "acme" + env = "test" + azs = ["us-east-1a", "us-east-1b"] + allow_plaintext_alb = true +} + +run "pool_off_by_default" { + command = plan + + assert { + condition = length(local.gateway_pool_env) == 0 + error_message = "The gateway must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } +} + +run "pool_enabled_renders_the_three_vars_with_configured_sizes" { + command = plan + + variables { + create_database = false + database_url = "postgresql://litellm:pw@db.internal:5432/litellm" + gateway_num_workers = 4 + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + length(local.gateway_pool_env) == 3, + local.gateway_pool_env[0].name == "LITELLM_PGBOUNCER_ENABLED" && local.gateway_pool_env[0].value == "true", + local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "8", + local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "250", + ]) + error_message = "The pool env must carry the enabled flag and the configured sizes as strings." + } +} + +run "collector_sidecar_gets_the_same_pool_env_as_the_gateway" { + command = plan + + variables { + create_database = false + database_url = "postgresql://litellm:pw@db.internal:5432/litellm" + collector_enabled = true + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + for env in [local.gateway_environment, local.collector_container[0].environment] : ( + { for e in env : e.name => e.value }["LITELLM_PGBOUNCER_ENABLED"] == "true" && + { for e in env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"] == "8" && + { for e in env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_CLIENT_CONN"] == "250" + ) + ]) + error_message = "The collector sidecar must carry the same three LITELLM_PGBOUNCER_* vars as the gateway so its Prisma connects to the task-local pool." + } +} + +run "collector_sidecar_gets_no_pool_env_when_the_pool_is_off" { + command = plan + + variables { + collector_enabled = true + } + + assert { + condition = !anytrue([for e in local.collector_container[0].environment : startswith(e.name, "LITELLM_PGBOUNCER_")]) + error_message = "The collector sidecar must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } +} + +run "pool_enabled_uses_the_module_default_sizes" { + command = plan + + variables { + create_database = false + database_url = "postgresql://litellm:pw@db.internal:5432/litellm" + gateway_connection_pool_enabled = true + } + + assert { + condition = alltrue([ + length(local.gateway_pool_env) == 3, + local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "20", + local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "1000", + ]) + error_message = "The pool env must fall back to the module defaults of 20 upstream and 1000 client connections." + } +} + +run "gateway_starts_through_the_pool_aware_launcher" { + command = plan + + variables { + gateway_num_workers = 4 + } + + assert { + condition = alltrue([ + strcontains(local.gateway_launch_cmd, "exec python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"), + strcontains(local.gateway_launch_cmd, "exec ddtrace-run python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"), + !strcontains(local.gateway_launch_cmd, "uvicorn gateway.main:app"), + local.gateway_proxy_overrides.command[0] == local.gateway_launch_cmd, + ]) + error_message = "The gateway must start through gateway.launch (with and without ddtrace) so the pooler starts once before uvicorn forks the workers." + } +} + +run "pool_with_module_created_iam_aurora_plans_with_both_the_pool_and_iam_auth" { + command = plan + + variables { + gateway_connection_pool_enabled = true + } + + assert { + condition = alltrue([ + length(local.gateway_pool_env) == 3, + contains(local.managed_db_env, { name = "IAM_TOKEN_DB_AUTH", value = "true" }), + ]) + error_message = "With the module-created Aurora the gateway must get the pool env alongside IAM token auth." + } +} + +run "pool_without_any_database_fails_at_plan" { + command = plan + + variables { + create_database = false + gateway_connection_pool_enabled = true + } + + expect_failures = [ + aws_ecs_task_definition.gateway, + ] +} + +run "module_created_iam_aurora_without_the_pool_still_plans" { + command = plan + + assert { + condition = contains(local.managed_db_env, { name = "IAM_TOKEN_DB_AUTH", value = "true" }) + error_message = "Without the pool the module-created Aurora must keep IAM token auth." + } +} diff --git a/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl b/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl new file mode 100644 index 00000000000..de839d7e5be --- /dev/null +++ b/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl @@ -0,0 +1,108 @@ +# Plan-only coverage for the Prometheus metrics sidecar wiring. Offline via +# mock_provider, same as byo_infrastructure.tftest.hcl. + +mock_provider "aws" { + mock_data "aws_iam_policy_document" { + defaults = { + json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}" + } + } +} +mock_provider "random" {} + +variables { + region = "us-east-1" + tenant = "acme" + env = "test" + allow_plaintext_alb = true + azs = ["us-east-1a", "us-east-1b"] +} + +run "defaults_change_nothing" { + command = plan + + assert { + condition = alltrue([ + length(local.gateway_metrics_container) == 0, + length(local.metrics_env) == 0, + length(local.metrics_mount_points) == 0, + length([for r in aws_security_group.tasks.ingress : r if r.description == "Prometheus scrapers to the gateway metrics sidecar"]) == 0, + ]) + error_message = "The metrics sidecar, its env, its volume, and its security-group rule must all be absent by default." + } +} + +run "metrics_port_adds_a_sidecar_volume_and_scrape_rule" { + command = plan + + variables { + gateway_metrics_port = 9464 + gateway_metrics_scrape_cidrs = ["10.20.0.0/16"] + } + + assert { + condition = length(local.metrics_env) == 1 && local.metrics_env[0].name == "PROMETHEUS_MULTIPROC_DIR" && local.metrics_env[0].value == "/tmp/litellm_prometheus_multiproc" + error_message = "The gateway workers must write multiprocess samples to the shared dir." + } + + assert { + condition = length(local.metrics_mount_points) == 1 && local.metrics_mount_points[0].sourceVolume == "prometheus-multiproc" && local.metrics_mount_points[0].containerPath == "/tmp/litellm_prometheus_multiproc" + error_message = "Gateway and sidecar must mount the same task volume at the multiproc dir." + } + + assert { + condition = alltrue([ + length(local.gateway_metrics_container) == 1, + local.gateway_metrics_container[0].name == "metrics", + local.gateway_metrics_container[0].essential == false, + join(" ", local.gateway_metrics_container[0].entryPoint) == "python -m litellm.proxy.prometheus_metrics_server", + join(" ", local.gateway_metrics_container[0].command) == "--port 9464", + one(local.gateway_metrics_container[0].portMappings).containerPort == 9464, + one(local.gateway_metrics_container[0].environment).value == "/tmp/litellm_prometheus_multiproc", + one(local.gateway_metrics_container[0].mountPoints).sourceVolume == "prometheus-multiproc", + strcontains(local.gateway_metrics_container[0].healthCheck.command[3], "9464"), + ]) + error_message = "The metrics sidecar must run prometheus_metrics_server on the configured port, share the multiproc volume, and health-check that port." + } + + assert { + condition = length(aws_ecs_task_definition.gateway.volume) == 1 && one(aws_ecs_task_definition.gateway.volume).name == "prometheus-multiproc" + error_message = "The gateway task must declare the multiproc volume." + } + + assert { + condition = length([ + for r in aws_security_group.tasks.ingress : r + if r.from_port == 9464 && r.to_port == 9464 && r.protocol == "tcp" && r.cidr_blocks == tolist(["10.20.0.0/16"]) + ]) == 1 + error_message = "The scrape CIDRs must be allowed to reach the metrics port on the tasks security group." + } + + assert { + condition = aws_lb_target_group.gateway.port == 4000 && one(aws_ecs_service.gateway.load_balancer).container_port == 4000 + error_message = "The ALB must keep targeting the gateway port only; the metrics port is never load balanced." + } +} + +run "metrics_port_without_scrape_cidrs_opens_nothing" { + command = plan + + variables { + gateway_metrics_port = 9464 + } + + assert { + condition = length(local.gateway_metrics_container) == 1 && length([for r in aws_security_group.tasks.ingress : r if r.from_port == 9464]) == 0 + error_message = "Without scrape CIDRs the sidecar runs but the metrics port stays closed to everything but the ALB group." + } +} + +run "metrics_port_may_not_reuse_the_gateway_port" { + command = plan + + variables { + gateway_metrics_port = 4000 + } + + expect_failures = [var.gateway_metrics_port] +} diff --git a/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl b/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl new file mode 100644 index 00000000000..94e281faf94 --- /dev/null +++ b/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl @@ -0,0 +1,194 @@ +# Plan-only coverage for the gateway request and token autoscaling policies. +# Offline via mock_provider, same as byo_infrastructure.tftest.hcl. + +mock_provider "aws" { + mock_data "aws_iam_policy_document" { + defaults = { + json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}" + } + } +} +mock_provider "random" {} + +variables { + region = "us-east-1" + tenant = "acme" + env = "test" + allow_plaintext_alb = true + azs = ["us-east-1a", "us-east-1b"] +} + +run "defaults_scale_on_cpu_and_memory_only" { + command = plan + + assert { + condition = alltrue([ + length(aws_appautoscaling_policy.gateway_cpu) == 1, + length(aws_appautoscaling_policy.gateway_memory) == 1, + length(aws_appautoscaling_policy.gateway_requests) == 0, + length(aws_appautoscaling_policy.gateway_tokens) == 0, + ]) + error_message = "Request and token policies must be absent by default while the CPU and memory policies stay." + } +} + +run "requests_per_second_adds_an_alb_request_count_policy" { + command = plan + + variables { + gateway_target_requests_per_second = 90 + } + + assert { + condition = length(aws_appautoscaling_policy.gateway_requests) == 1 && length(aws_appautoscaling_policy.gateway_tokens) == 0 + error_message = "A request target alone must add exactly the request policy." + } + + assert { + condition = alltrue([ + aws_appautoscaling_policy.gateway_requests[0].name == "acme-litellm-test-gateway-requests", + aws_appautoscaling_policy.gateway_requests[0].policy_type == "TargetTrackingScaling", + aws_appautoscaling_policy.gateway_requests[0].service_namespace == "ecs", + aws_appautoscaling_policy.gateway_requests[0].resource_id == "service/acme-litellm-test/acme-litellm-test-gateway", + aws_appautoscaling_policy.gateway_requests[0].scalable_dimension == "ecs:service:DesiredCount", + ]) + error_message = "The request policy must be a target-tracking policy on the gateway service's desired count." + } + + assert { + condition = alltrue([ + one(aws_appautoscaling_policy.gateway_requests[0].target_tracking_scaling_policy_configuration).target_value == 5400, + one(one(aws_appautoscaling_policy.gateway_requests[0].target_tracking_scaling_policy_configuration).predefined_metric_specification).predefined_metric_type == "ALBRequestCountPerTarget", + length(one(aws_appautoscaling_policy.gateway_requests[0].target_tracking_scaling_policy_configuration).customized_metric_specification) == 0, + ]) + error_message = "The request policy must track ALBRequestCountPerTarget at 60 times the configured requests per second per task." + } +} + +run "tokens_per_second_adds_a_metric_math_policy" { + command = plan + + variables { + gateway_target_tokens_per_second = 6000000 + gateway_tokens_metric = { + namespace = "LiteLLM/Prometheus" + dimensions = { ClusterName = "acme-litellm-test", TaskDefinitionFamily = "acme-litellm-test-gateway" } + } + } + + assert { + condition = length(aws_appautoscaling_policy.gateway_tokens) == 1 && length(aws_appautoscaling_policy.gateway_requests) == 0 + error_message = "A token target alone must add exactly the token policy." + } + + assert { + condition = alltrue([ + aws_appautoscaling_policy.gateway_tokens[0].name == "acme-litellm-test-gateway-tokens", + aws_appautoscaling_policy.gateway_tokens[0].resource_id == "service/acme-litellm-test/acme-litellm-test-gateway", + one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).target_value == 6000000, + length(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).predefined_metric_specification) == 0, + ]) + error_message = "The token policy must track a customized metric at the configured tokens per second per task." + } + + assert { + condition = alltrue([ + length(one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics) == 4, + { for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens"].id == "tokens", + { for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens"].return_data == false, + one({ for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens"].metric_stat).stat == "Sum", + one(one({ for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens"].metric_stat).metric).namespace == "LiteLLM/Prometheus", + one(one({ for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens"].metric_stat).metric).metric_name == "litellm_total_tokens_metric_total", + { for d in one(one({ for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens"].metric_stat).metric).dimensions : d.name => d.value } == { ClusterName = "acme-litellm-test", TaskDefinitionFamily = "acme-litellm-test-gateway" }, + ]) + error_message = "The first metric must sum the published token counter deltas under the configured namespace and dimensions." + } + + assert { + condition = alltrue([ + { for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["running_tasks"].id == "running_tasks", + { for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["running_tasks"].return_data == false, + one({ for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["running_tasks"].metric_stat).stat == "Average", + one(one({ for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["running_tasks"].metric_stat).metric).namespace == "ECS/ContainerInsights", + one(one({ for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["running_tasks"].metric_stat).metric).metric_name == "RunningTaskCount", + { for d in one(one({ for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["running_tasks"].metric_stat).metric).dimensions : d.name => d.value } == { ClusterName = "acme-litellm-test", ServiceName = "acme-litellm-test-gateway" }, + ]) + error_message = "The second metric must read the gateway service's Container Insights RunningTaskCount." + } + + assert { + condition = alltrue([ + { for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens_per_second"].expression == "tokens / 60", + { for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens_per_second"].return_data == false, + ]) + error_message = "The 60s period Sum must be divided by 60 to yield tokens per second." + } + + assert { + condition = alltrue([ + { for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens_per_second_per_task"].expression == "tokens_per_second / running_tasks", + { for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens_per_second_per_task"].return_data == true, + length([for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m if m.return_data]) == 1, + ]) + error_message = "Only the per-task tokens per second may return data to the scaling policy." + } +} + +run "tokens_per_second_needs_the_metric_location" { + command = plan + + variables { + gateway_target_tokens_per_second = 6000000 + } + + expect_failures = [ + aws_appautoscaling_policy.gateway_tokens, + ] +} + +run "requests_and_tokens_scale_next_to_cpu_and_memory" { + command = plan + + variables { + gateway_target_requests_per_second = 90 + gateway_target_tokens_per_second = 6000000 + gateway_tokens_metric = { namespace = "LiteLLM/Prometheus" } + } + + assert { + condition = alltrue([ + length(aws_appautoscaling_policy.gateway_cpu) == 1, + length(aws_appautoscaling_policy.gateway_memory) == 1, + length(aws_appautoscaling_policy.gateway_requests) == 1, + length(aws_appautoscaling_policy.gateway_tokens) == 1, + one(aws_appautoscaling_policy.gateway_cpu[0].target_tracking_scaling_policy_configuration).target_value == 70, + one(aws_appautoscaling_policy.gateway_memory[0].target_tracking_scaling_policy_configuration).target_value == 80, + ]) + error_message = "Workload policies must coexist with the CPU and memory policies at their default targets." + } + + assert { + condition = length(one(one({ for m in one(one(aws_appautoscaling_policy.gateway_tokens[0].target_tracking_scaling_policy_configuration).customized_metric_specification).metrics : m.id => m }["tokens"].metric_stat).metric).dimensions) == 0 + error_message = "Omitting dimensions must query the token metric without any." + } +} + +run "workload_targets_are_ignored_when_autoscaling_is_off" { + command = plan + + variables { + gateway_autoscaling_enabled = false + gateway_target_requests_per_second = 90 + gateway_target_tokens_per_second = 6000000 + gateway_tokens_metric = { namespace = "LiteLLM/Prometheus" } + } + + assert { + condition = alltrue([ + length(aws_appautoscaling_target.gateway) == 0, + length(aws_appautoscaling_policy.gateway_requests) == 0, + length(aws_appautoscaling_policy.gateway_tokens) == 0, + ]) + error_message = "Disabling gateway autoscaling must drop the workload policies with the target." + } +} diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index 522138953d6..580a0cc657a 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -200,6 +200,44 @@ variable "gateway_num_workers" { } } +variable "gateway_connection_pool_enabled" { + description = <<-EOT + Run an in-container PgBouncer (transaction mode, loopback) in each gateway + task, shared by every uvicorn worker. Without it each of the + `gateway_num_workers` workers opens its own Prisma pool straight to + Postgres, so a task's footprint against the database connection ceiling is + workers x connection_limit and grows with every task. Sets + LITELLM_PGBOUNCER_ENABLED / LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / + LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. Works with + the module-created Aurora too: the pooler mints the IAM token itself and + renews it before it expires. + EOT + type = bool + default = false +} + +variable "gateway_pool_max_db_connections" { + description = "Upstream Postgres connections one gateway task may hold when gateway_connection_pool_enabled is set, regardless of gateway_num_workers. 20 suits 4 workers; a 5000-connection database then fits roughly 200 tasks." + type = number + default = 20 + + validation { + condition = var.gateway_pool_max_db_connections >= 1 + error_message = "gateway_pool_max_db_connections must be >= 1." + } +} + +variable "gateway_pool_max_client_conn" { + description = "Client connections the in-container PgBouncer accepts from the gateway workers when gateway_connection_pool_enabled is set." + type = number + default = 1000 + + validation { + condition = var.gateway_pool_max_client_conn >= 1 + error_message = "gateway_pool_max_client_conn must be >= 1." + } +} + variable "backend_cpu" { description = "Fargate CPU units for the backend task (1024 = 1 vCPU)." type = number @@ -272,6 +310,47 @@ variable "gateway_memory_target" { default = 80 } +variable "gateway_target_requests_per_second" { + description = <<-EOT + Requests per second one gateway task should serve. Adds an + ALBRequestCountPerTarget target-tracking policy next to the CPU/memory + ones (Application Auto Scaling follows whichever asks for more tasks). + CloudWatch publishes that metric as a 1-minute count, so the policy + targets 60x this value and ECS reacts on a ~1 minute cadence. 0 skips + the policy. + EOT + type = number + default = 0 +} + +variable "gateway_target_tokens_per_second" { + description = <<-EOT + Tokens per second one gateway task should serve. Adds a target-tracking + policy on gateway_tokens_metric summed over each 60s period, divided by + 60 and by the service's Container Insights RunningTaskCount. Tokens are + counted when a response completes, so the signal trails long streams. + 0 skips the policy. + EOT + type = number + default = 0 +} + +variable "gateway_tokens_metric" { + description = <<-EOT + CloudWatch metric carrying the gateway's litellm_total_tokens_metric_total + counter, as published by the CloudWatch agent's Prometheus scraper (it + emits the delta between scrapes, so Sum over a period is the tokens + served in it). Required when gateway_target_tokens_per_second > 0. + dimensions must match the metric_declaration the agent publishes with. + EOT + type = object({ + namespace = string + name = optional(string, "litellm_total_tokens_metric_total") + dimensions = optional(map(string), {}) + }) + default = null +} + variable "backend_autoscaling_enabled" { description = "Toggle Application Auto Scaling target-tracking on the backend service." type = bool @@ -549,6 +628,44 @@ variable "proxy_config" { default = {} } +# ---------- Prometheus metrics sidecar ---------- + +variable "gateway_metrics_port" { + description = <<-EOT + Serve Prometheus /metrics from a `metrics` sidecar container in the + gateway task on this port (1-65535, not 4000), so a scrape never runs on + an inference worker. The sidecar runs the gateway image with + `python -m litellm.proxy.prometheus_metrics_server` and aggregates the + workers' PROMETHEUS_MULTIPROC_DIR samples over a task volume. Null (the + default) leaves /metrics on the gateway port only. The sidecar port has + no virtual-key auth and is not routed through the ALB; open it to your + scrapers with gateway_metrics_scrape_cidrs. Needs gateway_image v1.101.0 + or newer. + EOT + type = number + default = null + + validation { + condition = var.gateway_metrics_port == null || (var.gateway_metrics_port >= 1 && var.gateway_metrics_port <= 65535 && var.gateway_metrics_port != 4000) + error_message = "gateway_metrics_port must be between 1 and 65535 and must not be 4000 (the gateway port)." + } +} + +variable "gateway_metrics_scrape_cidrs" { + description = <<-EOT + CIDR blocks allowed to reach gateway_metrics_port on the gateway tasks + (your Prometheus or collector subnets). Empty by default, so only the + ALB can reach the tasks. Ignored when gateway_metrics_port is null. + EOT + type = list(string) + default = [] + + validation { + condition = alltrue([for c in var.gateway_metrics_scrape_cidrs : can(cidrnetmask(c))]) + error_message = "gateway_metrics_scrape_cidrs must contain valid IPv4 CIDR blocks." + } +} + variable "log_retention_days" { description = "CloudWatch log retention for the three services." type = number @@ -691,3 +808,73 @@ variable "billing_metrics_ca_cert_pem" { default = "" sensitive = true } + +# ---------- Collector sidecar ---------- +# +# Opt-in offload of spend tracking from the gateway's uvicorn workers to a +# `python -m litellm.proxy.collector` sidecar in the same Fargate task (helm's +# `gateway.collector`). Fargate awsvpc tasks share one network namespace, +# so the sidecar listens on loopback TCP. Disabled (the default) adds nothing +# to the task definition. + +variable "collector_enabled" { + description = "Run the collector sidecar next to the gateway container and have the gateway ship spend events to it (sets LITELLM_COLLECTOR_ENABLED=true on both). Autoscaling still targets the whole task's CPU/memory, sidecar included." + type = bool + default = false +} + +variable "collector_port" { + description = "Loopback TCP port the sidecar listens on (LITELLM_COLLECTOR_ADDRESS=tcp://127.0.0.1:)." + type = number + default = 4010 + + validation { + condition = var.collector_port >= 1024 && var.collector_port <= 65535 && var.collector_port != 4000 + error_message = "collector_port must be in 1024-65535 and not 4000." + } +} + +variable "collector_cpu" { + description = "CPU units reserved for the sidecar container, carved out of gateway_cpu. Matches helm's collector.resources.requests.cpu (500m)." + type = number + default = 512 +} + +variable "collector_memory" { + description = "Hard memory limit (MiB) for the sidecar container, carved out of gateway_memory. Matches helm's collector.resources.limits.memory (2Gi)." + type = number + default = 2048 +} + +variable "collector_buffer_size" { + description = "Per-worker in-memory queue of spend events waiting to be shipped to the sidecar (LITELLM_COLLECTOR_BUFFER_SIZE)." + type = number + default = 1000 + + validation { + condition = var.collector_buffer_size >= 1 + error_message = "collector_buffer_size must be >= 1." + } +} + +variable "collector_on_unavailable" { + description = "What the gateway does with spend events when the sidecar is unreachable or the buffer is full (LITELLM_COLLECTOR_ON_UNAVAILABLE): `fallback` runs the pipeline in-process, `drop` discards them." + type = string + default = "fallback" + + validation { + condition = contains(["fallback", "drop"], var.collector_on_unavailable) + error_message = "collector_on_unavailable must be one of: fallback, drop." + } +} + +variable "collector_drain_timeout_seconds" { + description = "Seconds a gateway worker waits on shutdown for its buffered spend events to reach the sidecar (LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS)." + type = number + default = 10 + + validation { + condition = var.collector_drain_timeout_seconds > 0 + error_message = "collector_drain_timeout_seconds must be > 0." + } +} diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index c93e5f6b303..4b2f576adc1 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -238,6 +238,129 @@ this with `litellm_license`. To tune the export cadence, set Behavior matches the AWS stack 1:1; the variable names are identical +### Prometheus metrics sidecar + +`gateway_metrics_port` adds a `metrics` sidecar +(`python -m litellm.proxy.prometheus_metrics_server`) to the gateway Cloud Run +service that aggregates the workers' samples over an in-memory volume shared +with the gateway container, so the collector's scrape never runs on an +inference worker. Cloud Run only routes traffic to the gateway container, so +the load balancer keeps hitting port 4000 (including the gateway's own +authenticated `/metrics`, which stays as it was) and the sidecar port is +reachable on localhost inside the instance only. To get the series out, the +stack also adds Google's +[Managed Service for Prometheus sidecar](https://cloud.google.com/stackdriver/docs/managed-prometheus/cloudrun-sidecar) +(`gateway_metrics_collector_image`) with a `RunMonitoring` config stored in +Secret Manager that scrapes `localhost:/metrics` every 30s and writes to +Cloud Monitoring as `prometheus.googleapis.com/...` metrics. Enabling it grants +the runtime service account `roles/monitoring.metricWriter` and +`roles/logging.logWriter` on the project. Needs `gateway_image` v1.101.0 or +newer. See [Prometheus metrics](https://docs.litellm.ai/docs/proxy/prometheus) +for the metrics themselves + +```hcl +gateway_metrics_port = 4001 +``` + +The collector scrapes from inside the instance, so scrapes on an instance with +no in-flight requests can fail when CPU is throttled between requests. Keep +`gateway_min_instances` at 1 or more and, if you see gaps, enable +instance-based billing on the gateway service. Unlike the AWS stack there is +no `gateway_metrics_scrape_cidrs`: nothing outside the instance can reach the +sidecar port, so there is no network rule to open + +### Autoscaling + +Cloud Run scales the gateway on request concurrency (plus its built-in CPU +target), not on a metric you attach. Each instance takes up to +`gateway_max_instance_request_concurrency` requests at once (default 80) +and Cloud Run adds instances between `gateway_min_instances` and +`gateway_max_instances` when the in-flight count fills up. That is the +request-rate signal for this stack: lower the concurrency for LLM streams +that hold a worker for tens of seconds, since a stream counts as one request +for as long as it is open + +There is no tokens-per-second path here. Cloud Run's autoscaler has no +custom-metric input, so the `litellm_total_tokens_metric_total` counter the +proxy exposes cannot drive it. If you need token-based scaling on GCP, run +the gateway on GKE with the Helm chart's `targetTokensPerSecond` (see +"Dependencies only" below) rather than wiring the counter into Cloud +Monitoring, which the autoscaler would ignore + +### In-container connection pool + +Each of the `gateway_num_workers` uvicorn workers opens its own Prisma pool +straight to Cloud SQL, so one instance holds `workers x connection_limit` +connections and the fleet's footprint against the database ceiling grows with +every instance Cloud Run adds. `gateway_connection_pool_enabled` runs a +PgBouncer (transaction mode, loopback) inside the gateway container that all +workers share, capping the instance at `gateway_pool_max_db_connections` +upstream connections however many workers it runs. +`gateway_pool_max_client_conn` bounds the worker-side connections the pooler +accepts. The module sets `LITELLM_PGBOUNCER_ENABLED`, +`LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS` and `LITELLM_PGBOUNCER_MAX_CLIENT_CONN` +on the gateway service only; the backend service and the migrations job keep +the direct connection + +```hcl +gateway_num_workers = 4 +gateway_connection_pool_enabled = true +gateway_pool_max_db_connections = 20 +gateway_pool_max_client_conn = 1000 +``` + +The pooler holds one static database password for the life of the instance. +This stack authenticates to Cloud SQL with the Secret Manager password (see +[Database authentication](#database-authentication)), so nothing else is +needed; a Cloud SQL Auth Proxy sidecar with IAM auth would not work with the +pool + +The gateway container starts through `python -m gateway.launch` (the +componentized image's own entrypoint) rather than `uvicorn` directly. The +launcher reads these variables, starts the pooler once per instance before +uvicorn forks the workers and hands them its loopback `DATABASE_URL`. It also +honours `KEEPALIVE_TIMEOUT` from `gateway_extra_env` the way the image does + +### Collector sidecar + +`collector_enabled = true` adds a `spend-collector` container to the gateway +Cloud Run service that runs `python -m litellm.proxy.collector` from the gateway +image, and sets `LITELLM_COLLECTOR_ENABLED=true` on the gateway so its +uvicorn workers ship spend events (SpendLogs writes, key/team/user spend +updates, budget alerts) to the sidecar instead of running that pipeline in +the request path. This is the Terraform counterpart of helm's +`gateway.collector`. The default (`false`) leaves the service exactly as +before. It is independent of the metrics sidecars above, whose GMP scraper +already owns the `collector` container name. + +Containers in one Cloud Run instance share localhost, so the sidecar listens +on loopback TCP (`tcp://127.0.0.1:${collector_port}`, default 4010) +instead of the Unix socket helm uses; the proxy rejects any non-loopback +address. The sidecar runs the same Redis CA + `DATABASE_URL` bootstrap as +the gateway container, gets the same database, Redis, master-key, license, +proxy config, and `gateway_extra_env` / `gateway_extra_secrets` values, and +runs with `LITELLM_JOB_ROLE=collector`. With `gateway_connection_pool_enabled` +it also gets the `LITELLM_PGBOUNCER_*` env, so its Prisma client goes through +the instance-local PgBouncer instead of opening a second pool straight to the +database. When it is unreachable the gateway falls back to in-process spend +tracking. + +```hcl +collector_enabled = true +# collector_cpu = "1000m" # added on top of gateway_cpu +# collector_memory = "2Gi" # added on top of gateway_memory +# collector_buffer_size = 1000 +# collector_on_unavailable = "fallback" # or "drop" +# collector_drain_timeout_seconds = 10 +``` + +Cloud Run allocates CPU per instance while requests are in flight, and the +sidecar shares that allocation. Spend events are shipped right after each +response, so this works with request-based billing, but keep +`gateway_min_instances >= 1` if spend must keep draining while an instance +is otherwise idle. Variable names match the AWS stack; only the resource +units differ (Cloud Run strings vs Fargate units) + ## Tenant deployment Every resource the stack creates is named `${tenant}-litellm-${env}` (or diff --git a/terraform/litellm/gcp/cloudrun.tf b/terraform/litellm/gcp/cloudrun.tf index 84ae8b9247f..d0b32a367d6 100644 --- a/terraform/litellm/gcp/cloudrun.tf +++ b/terraform/litellm/gcp/cloudrun.tf @@ -138,10 +138,16 @@ locals { "export DATABASE_URL_READ_REPLICA=\"postgresql://$${DATABASE_USER}:$${DATABASE_PASSWORD}@$${DATABASE_HOST_READ_REPLICA}:$${DATABASE_PORT_READ_REPLICA}/$${DATABASE_NAME}\"", ] + gateway_pool_env = var.gateway_connection_pool_enabled ? [ + { name = "LITELLM_PGBOUNCER_ENABLED", value = "true" }, + { name = "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", value = tostring(var.gateway_pool_max_db_connections) }, + { name = "LITELLM_PGBOUNCER_MAX_CLIENT_CONN", value = tostring(var.gateway_pool_max_client_conn) }, + ] : [] + gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}" backend_uvicorn_args = "--host 0.0.0.0 --port 4001" - gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac" + gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run python -m gateway.launch ${local.gateway_uvicorn_args};; *) exec python -m gateway.launch ${local.gateway_uvicorn_args};; esac" backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac" gateway_args = join(" && ", concat( @@ -150,12 +156,55 @@ locals { [local.gateway_launch_cmd], )) + metrics_enabled = var.create_runtime && var.gateway_metrics_port != null + metrics_multiproc_dir = "/tmp/litellm_prometheus_multiproc" + metrics_volume = "prometheus-multiproc" + metrics_env_kv = local.metrics_enabled ? [{ name = "PROMETHEUS_MULTIPROC_DIR", value = local.metrics_multiproc_dir }] : [] + metrics_config_volume = "gmp-config" + + metrics_run_monitoring_yaml = local.metrics_enabled ? yamlencode({ + apiVersion = "monitoring.googleapis.com/v1beta" + kind = "RunMonitoring" + metadata = { name = "${local.name}-gateway" } + spec = { + endpoints = [{ port = var.gateway_metrics_port, path = "/metrics", interval = "30s" }] + } + }) : "" + backend_args = join(" && ", concat( local.redis_ca_fragment, local.database_url_fragment, [local.backend_launch_cmd], )) + collector_address = "tcp://127.0.0.1:${var.collector_port}" + collector_env_kv = var.collector_enabled ? [ + { name = "LITELLM_COLLECTOR_ENABLED", value = "true" }, + { name = "LITELLM_COLLECTOR_ADDRESS", value = local.collector_address }, + { name = "LITELLM_COLLECTOR_BUFFER_SIZE", value = tostring(var.collector_buffer_size) }, + { name = "LITELLM_COLLECTOR_ON_UNAVAILABLE", value = var.collector_on_unavailable }, + { name = "LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS", value = tostring(var.collector_drain_timeout_seconds) }, + ] : [] + + gateway_env_kv = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env, local.metrics_env_kv, local.gateway_pool_env, local.collector_env_kv) + gateway_env_secrets = concat(local.shared_env_secrets, local.otel_env_secrets, local.billing_metrics_env_secrets, local.gateway_extra_secret_kv) + + collector_env_kv_all = concat( + local.shared_env_kv, + local.gateway_extra_env_kv, + local.proxy_config_env, + local.gateway_pool_env, + local.collector_env_kv, + [{ name = "LITELLM_JOB_ROLE", value = "collector" }], + ) + collector_env_secrets = concat(local.shared_env_secrets, local.gateway_extra_secret_kv) + + collector_args = join(" && ", concat( + local.redis_ca_fragment, + local.database_url_fragment, + ["exec python -m litellm.proxy.collector"], + )) + # Env shipped to the migrations Job. The migrations image runs run.py # which assembles DATABASE_URL from these discrete vars itself, so we # only need writer-side DB env (no read replica, no proxy_config, no @@ -182,6 +231,13 @@ resource "google_cloud_run_v2_service" "gateway" { labels = local.labels deletion_protection = false + lifecycle { + precondition { + condition = !var.collector_enabled || var.gateway_metrics_port == null || var.collector_port != var.gateway_metrics_port + error_message = "collector_port and gateway_metrics_port must differ: both sidecars bind loopback in the same instance." + } + } + template { service_account = google_service_account.runtime.email max_instance_request_concurrency = var.gateway_max_instance_request_concurrency @@ -197,6 +253,7 @@ resource "google_cloud_run_v2_service" "gateway" { } containers { + name = "gateway" image = local.gateway_image command = ["sh", "-c"] args = [local.gateway_args] @@ -213,7 +270,7 @@ resource "google_cloud_run_v2_service" "gateway" { } dynamic "env" { - for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env) + for_each = local.gateway_env_kv content { name = env.value.name value = env.value.value @@ -221,7 +278,7 @@ resource "google_cloud_run_v2_service" "gateway" { } dynamic "env" { - for_each = concat(local.shared_env_secrets, local.otel_env_secrets, local.billing_metrics_env_secrets, local.gateway_extra_secret_kv) + for_each = local.gateway_env_secrets content { name = env.value.name value_source { @@ -241,6 +298,14 @@ resource "google_cloud_run_v2_service" "gateway" { } } + dynamic "volume_mounts" { + for_each = local.metrics_enabled ? [1] : [] + content { + name = local.metrics_volume + mount_path = local.metrics_multiproc_dir + } + } + startup_probe { http_get { path = "/health/readiness" @@ -262,6 +327,117 @@ resource "google_cloud_run_v2_service" "gateway" { } } + dynamic "containers" { + for_each = local.metrics_enabled ? [1] : [] + content { + name = "metrics" + image = local.gateway_image + command = ["python", "-m", "litellm.proxy.prometheus_metrics_server"] + args = ["--port", tostring(var.gateway_metrics_port)] + + dynamic "env" { + for_each = local.metrics_env_kv + content { + name = env.value.name + value = env.value.value + } + } + + volume_mounts { + name = local.metrics_volume + mount_path = local.metrics_multiproc_dir + } + + startup_probe { + http_get { + path = "/health" + port = var.gateway_metrics_port + } + period_seconds = 5 + timeout_seconds = 3 + failure_threshold = 12 + } + + liveness_probe { + http_get { + path = "/health" + port = var.gateway_metrics_port + } + period_seconds = 30 + timeout_seconds = 5 + } + } + } + + dynamic "containers" { + for_each = local.metrics_enabled ? [1] : [] + content { + name = "collector" + image = var.gateway_metrics_collector_image + depends_on = ["metrics"] + + volume_mounts { + name = local.metrics_config_volume + mount_path = "/etc/rungmp" + } + + liveness_probe { + http_get { + path = "/liveness" + port = 13133 + } + period_seconds = 30 + timeout_seconds = 30 + } + } + } + + dynamic "containers" { + for_each = var.collector_enabled ? [1] : [] + content { + name = "spend-collector" + image = local.gateway_image + command = ["sh", "-c"] + args = [local.collector_args] + + resources { + limits = { + cpu = var.collector_cpu + memory = var.collector_memory + } + } + + dynamic "env" { + for_each = local.collector_env_kv_all + content { + name = env.value.name + value = env.value.value + } + } + + dynamic "env" { + for_each = local.collector_env_secrets + content { + name = env.value.name + value_source { + secret_key_ref { + secret = env.value.secret + version = env.value.version + } + } + } + } + + dynamic "volume_mounts" { + for_each = local.proxy_config_enabled ? [1] : [] + content { + name = local.proxy_config_volume + mount_path = local.proxy_config_mount_path + } + } + } + } + dynamic "volumes" { for_each = local.proxy_config_enabled ? [1] : [] content { @@ -272,6 +448,31 @@ resource "google_cloud_run_v2_service" "gateway" { } } } + + dynamic "volumes" { + for_each = local.metrics_enabled ? [1] : [] + content { + name = local.metrics_volume + empty_dir { + medium = "MEMORY" + size_limit = "256Mi" + } + } + } + + dynamic "volumes" { + for_each = local.metrics_enabled ? [1] : [] + content { + name = local.metrics_config_volume + secret { + secret = google_secret_manager_secret.metrics_run_monitoring[0].secret_id + items { + version = "latest" + path = "config.yaml" + } + } + } + } } depends_on = [ @@ -283,6 +484,8 @@ resource "google_cloud_run_v2_service" "gateway" { google_secret_manager_secret_iam_member.billing_metrics_client_cert, google_secret_manager_secret_iam_member.billing_metrics_client_key, google_secret_manager_secret_iam_member.billing_metrics_ca_cert, + google_secret_manager_secret_iam_member.metrics_run_monitoring, + google_project_iam_member.runtime_metric_writer, google_storage_bucket_iam_member.proxy_config_runtime, google_sql_user.app, # Don't go live until the schema is migrated; otherwise the proxy boots, @@ -332,7 +535,7 @@ resource "google_cloud_run_v2_service" "backend" { } dynamic "env" { - for_each = concat(local.shared_env_kv, local.backend_default_env_kv, local.backend_otel_env_kv, local.billing_metrics_env_kv, local.backend_extra_env_kv, local.proxy_config_env) + for_each = concat(local.shared_env_kv, local.backend_default_env_kv, local.backend_otel_env_kv, local.billing_metrics_env_kv, local.backend_extra_env_kv, local.proxy_config_env, local.metrics_env_kv) content { name = env.value.name value = env.value.value diff --git a/terraform/litellm/gcp/examples/default/main.tf b/terraform/litellm/gcp/examples/default/main.tf index f44b2a9a001..782d4ee4b65 100644 --- a/terraform/litellm/gcp/examples/default/main.tf +++ b/terraform/litellm/gcp/examples/default/main.tf @@ -53,4 +53,6 @@ module "litellm" { backend_extra_env = var.backend_extra_env gateway_extra_secrets = var.gateway_extra_secrets backend_extra_secrets = var.backend_extra_secrets + + gateway_metrics_port = var.gateway_metrics_port } diff --git a/terraform/litellm/gcp/examples/default/terraform.tfvars.example b/terraform/litellm/gcp/examples/default/terraform.tfvars.example index c35206503bb..ec6206734e2 100644 --- a/terraform/litellm/gcp/examples/default/terraform.tfvars.example +++ b/terraform/litellm/gcp/examples/default/terraform.tfvars.example @@ -107,3 +107,9 @@ env = "stage" # main.tf (otel_endpoint, otel_exporter, otel_environment_name, # otel_capture_message_content, otel_headers_secret). Full docs in # ../../variables.tf. + +# ---------- Prometheus metrics sidecar ---------- +# Serve /metrics from a sidecar in the gateway service instead of the inference +# workers. Scraped inside the instance by the Managed Service for Prometheus +# sidecar and written to Cloud Monitoring; see ../../README.md. +# gateway_metrics_port = 4001 diff --git a/terraform/litellm/gcp/examples/default/variables.tf b/terraform/litellm/gcp/examples/default/variables.tf index 88b57ce27eb..08b78346df3 100644 --- a/terraform/litellm/gcp/examples/default/variables.tf +++ b/terraform/litellm/gcp/examples/default/variables.tf @@ -142,3 +142,9 @@ variable "backend_extra_secrets" { type = map(string) default = {} } + +variable "gateway_metrics_port" { + description = "Port for the Prometheus metrics sidecar in the gateway service. Null keeps /metrics on the gateway port only." + type = number + default = null +} diff --git a/terraform/litellm/gcp/iam.tf b/terraform/litellm/gcp/iam.tf index 509e6d48ffd..5c2bc38bc12 100644 --- a/terraform/litellm/gcp/iam.tf +++ b/terraform/litellm/gcp/iam.tf @@ -116,3 +116,27 @@ resource "google_secret_manager_secret_iam_member" "billing_metrics_ca_cert" { role = "roles/secretmanager.secretAccessor" member = "serviceAccount:${google_service_account.runtime.email}" } + +resource "google_secret_manager_secret_iam_member" "metrics_run_monitoring" { + count = local.metrics_enabled ? 1 : 0 + + secret_id = google_secret_manager_secret.metrics_run_monitoring[0].id + role = "roles/secretmanager.secretAccessor" + member = "serviceAccount:${google_service_account.runtime.email}" +} + +resource "google_project_iam_member" "runtime_metric_writer" { + count = local.metrics_enabled ? 1 : 0 + + project = var.project_id + role = "roles/monitoring.metricWriter" + member = "serviceAccount:${google_service_account.runtime.email}" +} + +resource "google_project_iam_member" "runtime_log_writer" { + count = local.metrics_enabled ? 1 : 0 + + project = var.project_id + role = "roles/logging.logWriter" + member = "serviceAccount:${google_service_account.runtime.email}" +} diff --git a/terraform/litellm/gcp/secrets.tf b/terraform/litellm/gcp/secrets.tf index 6ec77139996..b3b656721ab 100644 --- a/terraform/litellm/gcp/secrets.tf +++ b/terraform/litellm/gcp/secrets.tf @@ -118,3 +118,20 @@ resource "google_secret_manager_secret_version" "billing_metrics_ca_cert" { secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id secret_data = var.billing_metrics_ca_cert_pem } + +resource "google_secret_manager_secret" "metrics_run_monitoring" { + count = local.metrics_enabled ? 1 : 0 + + secret_id = "${local.name}-gateway-run-monitoring" + labels = local.labels + replication { + auto {} + } +} + +resource "google_secret_manager_secret_version" "metrics_run_monitoring" { + count = local.metrics_enabled ? 1 : 0 + + secret = google_secret_manager_secret.metrics_run_monitoring[0].id + secret_data = local.metrics_run_monitoring_yaml +} diff --git a/terraform/litellm/gcp/tests/collector.tftest.hcl b/terraform/litellm/gcp/tests/collector.tftest.hcl new file mode 100644 index 00000000000..7a5ea781acb --- /dev/null +++ b/terraform/litellm/gcp/tests/collector.tftest.hcl @@ -0,0 +1,165 @@ +# Plan-only coverage for the opt-in collector sidecar on the gateway Cloud +# Run service. `mock_provider` keeps this offline: no GCP credentials, no API +# calls. Run from terraform/litellm/gcp with `terraform test`. + +mock_provider "google" {} +mock_provider "google-beta" {} +mock_provider "random" {} + +variables { + project_id = "acme-test" + region = "us-central1" + tenant = "acme" + env = "test" + allow_plaintext_lb = true +} + +run "disabled_by_default_leaves_the_service_untouched" { + command = plan + + assert { + condition = [for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name] == ["gateway"] + error_message = "The gateway service must stay single-container unless collector_enabled is set." + } + + assert { + condition = !anytrue([ + for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : startswith(e.name, "LITELLM_COLLECTOR_") + ]) + error_message = "No LITELLM_COLLECTOR_* env may reach the gateway while the sidecar is disabled." + } +} + +run "enabled_adds_a_sidecar_that_shares_the_gateway_transport" { + command = plan + + variables { + collector_enabled = true + collector_port = 4321 + collector_buffer_size = 250 + collector_on_unavailable = "drop" + collector_cpu = "500m" + collector_memory = "1Gi" + gateway_extra_env = { OPENAI_API_BASE = "https://example.invalid" } + gateway_extra_secrets = { OPENAI_API_KEY = "projects/acme-test/secrets/openai-api-key" } + } + + assert { + condition = [for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name] == ["gateway", "spend-collector"] + error_message = "Enabling the sidecar must append a spend-collector container after the gateway container." + } + + assert { + condition = alltrue([ + for c in google_cloud_run_v2_service.gateway[0].template[0].containers : ( + { for e in c.env : e.name => e.value }["LITELLM_COLLECTOR_ENABLED"] == "true" && + { for e in c.env : e.name => e.value }["LITELLM_COLLECTOR_ADDRESS"] == "tcp://127.0.0.1:4321" && + { for e in c.env : e.name => e.value }["LITELLM_COLLECTOR_BUFFER_SIZE"] == "250" && + { for e in c.env : e.name => e.value }["LITELLM_COLLECTOR_ON_UNAVAILABLE"] == "drop" && + { for e in c.env : e.name => e.value }["LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS"] == "10" + ) + ]) + error_message = "Gateway and sidecar must agree on the loopback address and the collector knobs." + } + + assert { + condition = ( + google_cloud_run_v2_service.gateway[0].template[0].containers[1].image == local.gateway_image && + google_cloud_run_v2_service.gateway[0].template[0].containers[1].command == tolist(["sh", "-c"]) && + endswith(google_cloud_run_v2_service.gateway[0].template[0].containers[1].args[0], " && exec python -m litellm.proxy.collector") && + strcontains(google_cloud_run_v2_service.gateway[0].template[0].containers[1].args[0], "export DATABASE_URL=") && + strcontains(google_cloud_run_v2_service.gateway[0].template[0].containers[1].args[0], "REDIS_SSL_CA_CERTS") && + { for e in google_cloud_run_v2_service.gateway[0].template[0].containers[1].env : e.name => e.value }["LITELLM_JOB_ROLE"] == "collector" + ) + error_message = "The sidecar must run litellm.proxy.collector from the gateway image with the same Redis CA + DATABASE_URL bootstrap as the gateway." + } + + assert { + condition = ( + length(google_cloud_run_v2_service.gateway[0].template[0].containers[1].ports) == 0 && + google_cloud_run_v2_service.gateway[0].template[0].containers[1].resources[0].limits.cpu == "500m" && + google_cloud_run_v2_service.gateway[0].template[0].containers[1].resources[0].limits.memory == "1Gi" + ) + error_message = "The sidecar must not claim the ingress port and must carry its own resource limits." + } + + assert { + condition = ( + { for e in google_cloud_run_v2_service.gateway[0].template[0].containers[1].env : e.name => e.value }["OPENAI_API_BASE"] == "https://example.invalid" && + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[1].env : e.name], "DATABASE_HOST") && + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[1].env : e.name], "REDIS_HOST") && + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[1].env : e.name if length(e.value_source) > 0], "LITELLM_MASTER_KEY") && + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[1].env : e.name if length(e.value_source) > 0], "DATABASE_PASSWORD") && + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[1].env : e.name if length(e.value_source) > 0], "OPENAI_API_KEY") + ) + error_message = "The sidecar must receive the gateway's database, Redis, and Secret Manager env plus gateway_extra_env / gateway_extra_secrets." + } +} + +run "coexists_with_the_metrics_sidecars" { + command = plan + + variables { + collector_enabled = true + gateway_metrics_port = 4001 + } + + assert { + condition = [for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name] == ["gateway", "metrics", "collector", "spend-collector"] + error_message = "The spend collector must keep its own container name next to the GMP metrics collector." + } + + assert { + condition = ( + { for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name => e.value }["PROMETHEUS_MULTIPROC_DIR"] == local.metrics_multiproc_dir && + { for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name => e.value }["LITELLM_COLLECTOR_ENABLED"] == "true" + ) + error_message = "The gateway container must keep both the metrics and the collector env when both sidecars are on." + } +} + +run "sidecars_must_not_share_a_loopback_port" { + command = plan + + variables { + collector_enabled = true + collector_port = 4001 + gateway_metrics_port = 4001 + } + + expect_failures = [ + google_cloud_run_v2_service.gateway, + ] +} + +run "collector_cannot_take_the_metrics_sidecar_health_port" { + command = plan + + variables { + collector_enabled = true + collector_port = 13133 + } + + expect_failures = [ + var.collector_port, + ] +} + +run "proxy_config_is_mounted_into_the_sidecar_too" { + command = plan + + variables { + collector_enabled = true + proxy_config = { model_list = [] } + } + + assert { + condition = alltrue([ + for c in google_cloud_run_v2_service.gateway[0].template[0].containers : ( + [for m in c.volume_mounts : m.name] == [local.proxy_config_volume] && + contains([for e in c.env : e.name], "CONFIG_FILE_PATH") + ) + ]) + error_message = "Both containers must mount the proxy-config GCS volume and point CONFIG_FILE_PATH at it." + } +} diff --git a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl new file mode 100644 index 00000000000..999e4f0ff95 --- /dev/null +++ b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl @@ -0,0 +1,174 @@ +# Plan-only coverage for the in-container PgBouncer knobs on the gateway +# service. `mock_provider` keeps this offline: no GCP credentials, no API +# calls, no resources. Run from terraform/litellm/gcp with `terraform test`. + +mock_provider "google" { + mock_resource "google_redis_instance" { + defaults = { + host = "10.0.0.4" + port = 6379 + server_ca_certs = [{ + cert = "-----BEGIN CERTIFICATE-----\nmock\n-----END CERTIFICATE-----" + }] + } + } +} + +mock_provider "google-beta" {} +mock_provider "random" {} + +variables { + project_id = "test-project" + tenant = "tenant" + env = "test" + allow_plaintext_lb = true + image_registry = "us-central1-docker.pkg.dev/test-project/litellm" +} + +run "pool_off_by_default" { + command = plan + + assert { + condition = length(local.gateway_pool_env) == 0 + error_message = "The gateway must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } + + assert { + condition = !anytrue([ + for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_") + ]) + error_message = "The gateway service must carry no LITELLM_PGBOUNCER_* env by default." + } +} + +run "pool_enabled_renders_the_three_vars_with_configured_sizes" { + command = plan + + variables { + gateway_num_workers = 4 + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + length(local.gateway_pool_env) == 3, + local.gateway_pool_env[0].name == "LITELLM_PGBOUNCER_ENABLED" && local.gateway_pool_env[0].value == "true", + local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "8", + local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "250", + ]) + error_message = "The pool env must carry the enabled flag and the configured sizes as strings." + } + + assert { + condition = alltrue([ + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_ENABLED"), + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"), + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_MAX_CLIENT_CONN"), + ]) + error_message = "The gateway service must receive all three LITELLM_PGBOUNCER_* env vars." + } + + assert { + condition = !anytrue(concat( + [for e in google_cloud_run_v2_service.backend[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")], + [for e in google_cloud_run_v2_job.migrations[0].template[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")], + )) + error_message = "The backend service and the migrations job must keep their direct database connection." + } +} + +run "collector_sidecar_gets_the_same_pool_env_as_the_gateway" { + command = plan + + variables { + collector_enabled = true + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + for c in google_cloud_run_v2_service.gateway[0].template[0].containers : ( + { for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_ENABLED"] == "true" && + { for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"] == "8" && + { for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_CLIENT_CONN"] == "250" + ) if c.name == "spend-collector" + ]) && length([for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name if c.name == "spend-collector"]) == 1 + error_message = "The spend-collector sidecar must carry the same three LITELLM_PGBOUNCER_* vars as the gateway so its Prisma connects to the instance-local pool." + } +} + +run "collector_sidecar_gets_no_pool_env_when_the_pool_is_off" { + command = plan + + variables { + collector_enabled = true + } + + assert { + condition = !anytrue(flatten([ + for c in google_cloud_run_v2_service.gateway[0].template[0].containers : [ + for e in c.env : startswith(e.name, "LITELLM_PGBOUNCER_") + ] if c.name == "spend-collector" + ])) && length([for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name if c.name == "spend-collector"]) == 1 + error_message = "The spend-collector sidecar must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } +} + +run "pool_enabled_uses_the_module_default_sizes" { + command = plan + + variables { + gateway_connection_pool_enabled = true + } + + assert { + condition = alltrue([ + length(local.gateway_pool_env) == 3, + local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "20", + local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "1000", + ]) + error_message = "The pool env must fall back to the module defaults of 20 upstream and 1000 client connections." + } +} + +run "gateway_starts_through_the_pool_aware_launcher" { + command = plan + + variables { + gateway_num_workers = 4 + } + + assert { + condition = alltrue([ + strcontains(local.gateway_launch_cmd, "exec python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"), + strcontains(local.gateway_launch_cmd, "exec ddtrace-run python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"), + !strcontains(local.gateway_launch_cmd, "uvicorn gateway.main:app"), + endswith(google_cloud_run_v2_service.gateway[0].template[0].containers[0].args[0], local.gateway_launch_cmd), + ]) + error_message = "The gateway must start through gateway.launch (with and without ddtrace) so the pooler starts once before uvicorn forks the workers." + } + + assert { + condition = strcontains(local.backend_launch_cmd, "uvicorn backend.main:app") + error_message = "The backend has no workers to share a pooler and keeps starting uvicorn directly." + } +} + +run "pool_sizes_below_one_fail_at_plan" { + command = plan + + variables { + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 0 + gateway_pool_max_client_conn = 0 + } + + expect_failures = [ + var.gateway_pool_max_db_connections, + var.gateway_pool_max_client_conn, + ] +} diff --git a/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl b/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl new file mode 100644 index 00000000000..7de62065d23 --- /dev/null +++ b/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl @@ -0,0 +1,200 @@ +mock_provider "google" { + mock_resource "google_redis_instance" { + defaults = { + host = "10.0.0.4" + port = 6379 + server_ca_certs = [{ + cert = "-----BEGIN CERTIFICATE-----\nmock\n-----END CERTIFICATE-----" + }] + } + } +} + +mock_provider "google-beta" {} +mock_provider "random" {} + +variables { + project_id = "test-project" + tenant = "tenant" + env = "test" + allow_plaintext_lb = true + image_registry = "us-central1-docker.pkg.dev/test-project/litellm" +} + +run "metrics_sidecar_off_by_default" { + command = plan + + assert { + condition = length(google_cloud_run_v2_service.gateway[0].template[0].containers) == 1 + error_message = "The gateway service must run only the gateway container when gateway_metrics_port is null." + } + + assert { + condition = length([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e if e.name == "PROMETHEUS_MULTIPROC_DIR"]) == 0 + error_message = "PROMETHEUS_MULTIPROC_DIR must not be set when the metrics sidecar is off." + } + + assert { + condition = length(google_cloud_run_v2_service.gateway[0].template[0].volumes) == 0 + error_message = "No shared multiproc or collector config volume must exist when the metrics sidecar is off." + } + + assert { + condition = alltrue([ + length(google_secret_manager_secret.metrics_run_monitoring) == 0, + length(google_secret_manager_secret_version.metrics_run_monitoring) == 0, + length(google_secret_manager_secret_iam_member.metrics_run_monitoring) == 0, + length(google_project_iam_member.runtime_metric_writer) == 0, + length(google_project_iam_member.runtime_log_writer) == 0, + ]) + error_message = "No RunMonitoring secret or monitoring IAM must be created when the metrics sidecar is off." + } +} + +run "metrics_sidecar_enabled" { + command = plan + + variables { + gateway_metrics_port = 4001 + } + + assert { + condition = join(",", [for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name]) == "gateway,metrics,collector" + error_message = "gateway_metrics_port must add the metrics and collector sidecars after the gateway container." + } + + assert { + condition = alltrue([ + google_cloud_run_v2_service.gateway[0].template[0].containers[1].image == local.gateway_image, + join(" ", google_cloud_run_v2_service.gateway[0].template[0].containers[1].command) == "python -m litellm.proxy.prometheus_metrics_server", + join(" ", google_cloud_run_v2_service.gateway[0].template[0].containers[1].args) == "--port 4001", + ]) + error_message = "The metrics sidecar must run the gateway image's prometheus_metrics_server on the configured port." + } + + assert { + condition = alltrue([ + length([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e if e.name == "PROMETHEUS_MULTIPROC_DIR" && e.value == "/tmp/litellm_prometheus_multiproc"]) == 1, + length([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[1].env : e if e.name == "PROMETHEUS_MULTIPROC_DIR" && e.value == "/tmp/litellm_prometheus_multiproc"]) == 1, + ]) + error_message = "Gateway and metrics containers must share PROMETHEUS_MULTIPROC_DIR." + } + + assert { + condition = alltrue([ + length([for m in google_cloud_run_v2_service.gateway[0].template[0].containers[0].volume_mounts : m if m.name == "prometheus-multiproc" && m.mount_path == "/tmp/litellm_prometheus_multiproc"]) == 1, + length([for m in google_cloud_run_v2_service.gateway[0].template[0].containers[1].volume_mounts : m if m.name == "prometheus-multiproc" && m.mount_path == "/tmp/litellm_prometheus_multiproc"]) == 1, + length([for v in google_cloud_run_v2_service.gateway[0].template[0].volumes : v if v.name == "prometheus-multiproc" && length(v.empty_dir) == 1 && v.empty_dir[0].medium == "MEMORY"]) == 1, + ]) + error_message = "Gateway and metrics containers must mount the same in-memory empty_dir at the multiproc dir." + } + + assert { + condition = alltrue([ + google_cloud_run_v2_service.gateway[0].template[0].containers[1].startup_probe[0].http_get[0].path == "/health", + google_cloud_run_v2_service.gateway[0].template[0].containers[1].startup_probe[0].http_get[0].port == 4001, + google_cloud_run_v2_service.gateway[0].template[0].containers[1].liveness_probe[0].http_get[0].path == "/health", + google_cloud_run_v2_service.gateway[0].template[0].containers[1].liveness_probe[0].http_get[0].port == 4001, + ]) + error_message = "The metrics sidecar must be probed on /health at the configured port." + } + + assert { + condition = length(google_cloud_run_v2_service.gateway[0].template[0].containers[1].ports) == 0 && length(google_cloud_run_v2_service.gateway[0].template[0].containers[2].ports) == 0 + error_message = "Only the gateway container may declare a port; Cloud Run routes ingress to exactly one container." + } + + assert { + condition = alltrue([ + google_cloud_run_v2_service.gateway[0].template[0].containers[0].ports[0].container_port == 4000, + google_cloud_run_v2_service.gateway[0].template[0].containers[0].startup_probe[0].http_get[0].port == 4000, + google_cloud_run_v2_service.gateway[0].template[0].containers[0].liveness_probe[0].http_get[0].port == 4000, + google_compute_region_network_endpoint_group.gateway[0].cloud_run[0].service == "${local.name}-gateway", + ]) + error_message = "The gateway must stay on port 4000 and remain the load balancer's Cloud Run target." + } + + assert { + condition = alltrue([ + google_cloud_run_v2_service.gateway[0].template[0].containers[2].image == var.gateway_metrics_collector_image, + join(",", google_cloud_run_v2_service.gateway[0].template[0].containers[2].depends_on) == "metrics", + length([for m in google_cloud_run_v2_service.gateway[0].template[0].containers[2].volume_mounts : m if m.name == "gmp-config" && m.mount_path == "/etc/rungmp"]) == 1, + google_cloud_run_v2_service.gateway[0].template[0].containers[2].liveness_probe[0].http_get[0].port == 13133, + ]) + error_message = "The collector sidecar must start after the metrics server and read its RunMonitoring config from /etc/rungmp." + } + + assert { + condition = alltrue([ + length([for v in google_cloud_run_v2_service.gateway[0].template[0].volumes : v if v.name == "gmp-config" && length(v.secret) == 1 && v.secret[0].items[0].path == "config.yaml"]) == 1, + google_secret_manager_secret.metrics_run_monitoring[0].secret_id == "${local.name}-gateway-run-monitoring", + google_secret_manager_secret_iam_member.metrics_run_monitoring[0].role == "roles/secretmanager.secretAccessor", + ]) + error_message = "The RunMonitoring config must be mounted from a Secret Manager secret readable by the runtime SA." + } + + assert { + condition = alltrue([ + yamldecode(google_secret_manager_secret_version.metrics_run_monitoring[0].secret_data).kind == "RunMonitoring", + yamldecode(google_secret_manager_secret_version.metrics_run_monitoring[0].secret_data).spec.endpoints[0].port == 4001, + yamldecode(google_secret_manager_secret_version.metrics_run_monitoring[0].secret_data).spec.endpoints[0].path == "/metrics", + ]) + error_message = "The RunMonitoring config must scrape /metrics on the configured metrics port." + } + + assert { + condition = alltrue([ + google_project_iam_member.runtime_metric_writer[0].role == "roles/monitoring.metricWriter", + google_project_iam_member.runtime_log_writer[0].role == "roles/logging.logWriter", + google_project_iam_member.runtime_metric_writer[0].project == "test-project", + ]) + error_message = "The runtime SA must be able to write metrics and logs for the collector sidecar." + } +} + +run "metrics_sidecar_ignored_in_deps_only" { + command = plan + + variables { + create_runtime = false + gateway_metrics_port = 4001 + } + + assert { + condition = alltrue([ + length(google_secret_manager_secret.metrics_run_monitoring) == 0, + length(google_project_iam_member.runtime_metric_writer) == 0, + ]) + error_message = "Dependencies-only mode must not create metrics sidecar resources." + } +} + +run "metrics_port_rejects_gateway_port" { + command = plan + + variables { + gateway_metrics_port = 4000 + } + + expect_failures = [var.gateway_metrics_port] +} + +run "metrics_port_rejects_collector_health_port" { + command = plan + + variables { + gateway_metrics_port = 13133 + } + + expect_failures = [var.gateway_metrics_port] +} + +run "metrics_port_rejects_fractional_port" { + command = plan + + variables { + gateway_metrics_port = 4000.5 + } + + expect_failures = [var.gateway_metrics_port] +} diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index 9c68ed3db76..412f919ab89 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -206,6 +206,45 @@ variable "gateway_num_workers" { } } +variable "gateway_connection_pool_enabled" { + description = <<-EOT + Run an in-container PgBouncer (transaction mode, loopback) in each gateway + instance, shared by every uvicorn worker. Without it each of the + `gateway_num_workers` workers opens its own Prisma pool straight to + Cloud SQL, so an instance's footprint against the database connection + ceiling is workers x connection_limit and grows with every instance. Sets + LITELLM_PGBOUNCER_ENABLED / LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / + LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. The + module's Cloud SQL authenticates with the static password in Secret + Manager, which is what the pooler needs. Mirrors the AWS stack's + gateway_connection_pool_enabled. + EOT + type = bool + default = false +} + +variable "gateway_pool_max_db_connections" { + description = "Upstream Cloud SQL connections one gateway instance may hold when gateway_connection_pool_enabled is set, regardless of gateway_num_workers. 20 suits 4 workers." + type = number + default = 20 + + validation { + condition = var.gateway_pool_max_db_connections >= 1 + error_message = "gateway_pool_max_db_connections must be >= 1." + } +} + +variable "gateway_pool_max_client_conn" { + description = "Client connections the in-container PgBouncer accepts from the gateway workers when gateway_connection_pool_enabled is set." + type = number + default = 1000 + + validation { + condition = var.gateway_pool_max_client_conn >= 1 + error_message = "gateway_pool_max_client_conn must be >= 1." + } +} + # Cloud Run autoscales out of the box (request-rate driven). The min/max # bounds mirror the HPA replica bounds in helm/litellm/values.yaml so each # stack scales over the same range. Cloud Run has no direct CPU-utilization @@ -517,6 +556,44 @@ variable "otel_capture_message_content" { } } +# ---------- Prometheus metrics sidecar ---------- + +variable "gateway_metrics_port" { + description = <<-EOT + Serve Prometheus /metrics from a `metrics` sidecar container in the + gateway Cloud Run service on this port (a whole number 1-65535, not 4000 + or 13133), so the collector's scrape never runs on an inference worker. + The sidecar runs the gateway image with + `python -m litellm.proxy.prometheus_metrics_server` and aggregates the + workers' PROMETHEUS_MULTIPROC_DIR samples over an in-memory volume shared + with the gateway container. Cloud Run only routes ingress to the gateway + container, so the sidecar port is reachable on localhost inside the + instance only; a Managed Service for Prometheus collector sidecar + (gateway_metrics_collector_image) scrapes it and writes the series to + Cloud Monitoring. The load balancer keeps serving the authenticated + /metrics on the gateway port as before. Null (the default) leaves /metrics + on the gateway port only. Needs gateway_image v1.101.0 or newer. + EOT + type = number + default = null + + validation { + condition = var.gateway_metrics_port == null || (var.gateway_metrics_port >= 1 && var.gateway_metrics_port <= 65535 && floor(var.gateway_metrics_port) == var.gateway_metrics_port && !contains([4000, 13133], var.gateway_metrics_port)) + error_message = "gateway_metrics_port must be a whole number between 1 and 65535 and must not be 4000 (the gateway port) or 13133 (the collector health port)." + } +} + +variable "gateway_metrics_collector_image" { + description = <<-EOT + Managed Service for Prometheus sidecar image that scrapes + localhost:/metrics and writes to Cloud Monitoring. + Override only to pin a different release or pull through your own + Artifact Registry. Ignored when gateway_metrics_port is null. + EOT + type = string + default = "us-docker.pkg.dev/cloud-ops-agents-artifacts/cloud-run-gmp-sidecar/cloud-run-gmp-sidecar:1.9.2" +} + # ---------- Enterprise billing metrics ---------- # # License-gated request metering. Opt-in and gated entirely on @@ -579,3 +656,73 @@ variable "billing_metrics_ca_cert_pem" { default = "" sensitive = true } + +# ---------- Collector sidecar ---------- +# +# Opt-in offload of spend tracking from the gateway's uvicorn workers to a +# `python -m litellm.proxy.collector` sidecar container in the same Cloud Run +# instance (helm's `gateway.collector`, mirrors the AWS stack). Containers +# in one instance share localhost, so the sidecar listens on loopback TCP. +# Disabled (the default) adds nothing to the service. + +variable "collector_enabled" { + description = "Run the collector sidecar next to the gateway container and have the gateway ship spend events to it (sets LITELLM_COLLECTOR_ENABLED=true on both). The sidecar shares the instance's request-based CPU allocation, so pair it with a non-zero gateway_min_instances if spend must keep flowing between requests." + type = bool + default = false +} + +variable "collector_port" { + description = "Loopback TCP port the sidecar listens on (LITELLM_COLLECTOR_ADDRESS=tcp://127.0.0.1:)." + type = number + default = 4010 + + validation { + condition = var.collector_port >= 1024 && var.collector_port <= 65535 && !contains([4000, 13133], var.collector_port) + error_message = "collector_port must be in 1024-65535 and not 4000 (the gateway port) or 13133 (the metrics sidecar health port)." + } +} + +variable "collector_cpu" { + description = "Cloud Run CPU limit for the sidecar container, on top of gateway_cpu. Matches helm's collector.resources.limits.cpu." + type = string + default = "1000m" +} + +variable "collector_memory" { + description = "Cloud Run memory limit for the sidecar container, on top of gateway_memory. Matches helm's collector.resources.limits.memory." + type = string + default = "2Gi" +} + +variable "collector_buffer_size" { + description = "Per-worker in-memory queue of spend events waiting to be shipped to the sidecar (LITELLM_COLLECTOR_BUFFER_SIZE)." + type = number + default = 1000 + + validation { + condition = var.collector_buffer_size >= 1 + error_message = "collector_buffer_size must be >= 1." + } +} + +variable "collector_on_unavailable" { + description = "What the gateway does with spend events when the sidecar is unreachable or the buffer is full (LITELLM_COLLECTOR_ON_UNAVAILABLE): `fallback` runs the pipeline in-process, `drop` discards them." + type = string + default = "fallback" + + validation { + condition = contains(["fallback", "drop"], var.collector_on_unavailable) + error_message = "collector_on_unavailable must be one of: fallback, drop." + } +} + +variable "collector_drain_timeout_seconds" { + description = "Seconds a gateway worker waits on shutdown for its buffered spend events to reach the sidecar (LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS)." + type = number + default = 10 + + validation { + condition = var.collector_drain_timeout_seconds > 0 + error_message = "collector_drain_timeout_seconds must be > 0." + } +} diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 842bfb4bdb1..e5e0a164a83 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -16,6 +16,7 @@ longer signal it. ### Added +- **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement - **jwt_key_mapping**: New `litellm_jwt_key_mapping` resource for the proxy's JWT to virtual key mappings, so JWT clients identified by a claim (`client_id`, `azp`, `sub`) map to virtual keys and inherit their models, budgets and rate limits. Supports `description` and `is_active`, rotating the mapped key in place, and forces replacement when the claim name or value changes - **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it - **user**: New `litellm_user` resource and `litellm_user` / `litellm_users` data sources for managing internal users @@ -38,12 +39,14 @@ longer signal it. - **team**: Read now decodes the `team_info` envelope `/team/info` actually returns, so team attributes refresh from the proxy instead of always falling back to the prior state - **key**: Read now unwraps the `info` envelope `/key/info` actually returns; previously reads mapped nothing back into state, so drift on a key was never detected +- **key**: Read now picks up `model_rpm_limit`, `model_tpm_limit`, `guardrails`, `tags`, `enforced_params`, `allowed_passthrough_routes`, `rpm_limit_type`, `tpm_limit_type` and `prompts` from `info.metadata`, where the proxy actually stores them; previously they stayed empty in state, so a matching config showed a permanent phantom diff on them and out-of-band changes to them were never detected - **key**: Updates no longer send an empty `budget_duration`, which the proxy rejects with a 400; any update to a key without a configured `budget_duration` previously failed outright - **key**: A config-supplied `key` value (write-only) is now forwarded to `/key/generate`; previously it was silently dropped and the proxy generated a random key instead - **security**: The `litellm_key` data source and `litellm_key_block` resource normalize raw `sk-` keys to their SHA-256 token hash before building request URLs and resource IDs, so plaintext keys no longer land in reverse-proxy access logs, Terraform plan output, or state IDs ### Changed +- **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying ## [0.4.0] - 2026-08-06 diff --git a/terraform/provider/README.md b/terraform/provider/README.md index 0a6d15c7844..b392fd6279d 100644 --- a/terraform/provider/README.md +++ b/terraform/provider/README.md @@ -103,9 +103,12 @@ resource "litellm_key" "example_key" { permissions = { can_create_keys = "true" } - model_max_budget = { - "gpt-4" = 50.0 - } + model_max_budget = jsonencode({ + "gpt-4" = { + budget_limit = 50.0 + time_period = "30d" + } + }) model_rpm_limit = { "claude-3.5-sonnet" = 30 } diff --git a/terraform/provider/RELEASING.md b/terraform/provider/RELEASING.md index 59f4c5f066c..5b72621c291 100644 --- a/terraform/provider/RELEASING.md +++ b/terraform/provider/RELEASING.md @@ -79,7 +79,7 @@ Before publishing to the Terraform Registry: ## What a change needs -1. **Land it in `BerriAI/litellm`.** Open a PR against `litellm_internal_staging` with the source change and a `CHANGELOG.md` entry under `[Unreleased]`. CI runs `gofmt`, `go vet`, build, tests and the endpoint-drift audit. A change that breaks existing configurations or state must say so in the changelog: the version number cannot signal it any more +1. **Land it in `BerriAI/litellm`.** Open a PR against the repository's current default branch with the source change and a `CHANGELOG.md` entry under `[Unreleased]`. CI runs `gofmt`, `go vet`, build, tests and the endpoint-drift audit. A change that breaks existing configurations or state must say so in the changelog: the version number cannot signal it any more 2. **Wait for the next LiteLLM release.** The nightly dev release carries it within a day; it reaches a stable version on the next stable cut 3. **Verify** (optional): the version appears at https://registry.terraform.io/providers/BerriAI/litellm and https://github.com/BerriAI/terraform-provider-litellm/releases. If the tag is on the mirror but there is no release, the goreleaser run failed: https://github.com/BerriAI/terraform-provider-litellm/actions diff --git a/terraform/provider/docs/resources/key.md b/terraform/provider/docs/resources/key.md index 5094b77cbec..0ef0688830f 100644 --- a/terraform/provider/docs/resources/key.md +++ b/terraform/provider/docs/resources/key.md @@ -30,9 +30,12 @@ resource "litellm_key" "example" { permissions = { "can_create_keys" = "true" } - model_max_budget = { - "gpt-4" = 50.0 - } + model_max_budget = jsonencode({ + "gpt-4" = { + budget_limit = 50.0 + time_period = "30d" + } + }) model_rpm_limit = { "gpt-3.5-turbo" = 30 } @@ -73,7 +76,7 @@ The following arguments are supported: * `key_alias` - (Optional) Alias for this key. This provides a human-readable identifier for the key. -* `duration` - (Optional) Duration for which this key is valid. This sets an expiration time for the key. +* `duration` - (Optional) How long the key stays valid, e.g. "30d" or "12h". The proxy stores this as an absolute `expires` timestamp. Changing the value resets the expiry to the time of the update plus the new duration; removing it from the configuration leaves the current expiry in place. * `aliases` - (Optional) Map of model aliases. This allows you to create custom names for models when using this key. @@ -81,7 +84,7 @@ The following arguments are supported: * `permissions` - (Optional) Permissions associated with this key. This defines what actions are allowed with this key. -* `model_max_budget` - (Optional) Maximum budget per model. This allows setting different budget limits for each model. +* `model_max_budget` - (Optional) JSON string of per-model budget config, e.g. `jsonencode({"gpt-4" = {budget_limit = 50.0, time_period = "30d"}})`. Each model maps to an object with `budget_limit` (or `max_budget`), `time_period` (or `budget_duration`), `tpm_limit` and `rpm_limit`. * `model_rpm_limit` - (Optional) Requests per minute limit per model. This allows setting different RPM limits for each model. diff --git a/terraform/provider/docs/resources/team.md b/terraform/provider/docs/resources/team.md index 821d8c1dee3..19575907269 100644 --- a/terraform/provider/docs/resources/team.md +++ b/terraform/provider/docs/resources/team.md @@ -14,6 +14,16 @@ resource "litellm_team" "engineering" { } ``` +### Team with a Custom ID + +```hcl +resource "litellm_team" "platform" { + team_id = "platform-team" + team_alias = "platform" + models = ["gpt-4-proxy"] +} +``` + ### Team with Comprehensive Configuration ```hcl @@ -92,6 +102,8 @@ resource "litellm_team" "model_dependent_team" { The following arguments are supported: +* `team_id` - (Optional) A stable, human-readable ID for the team (for example `platform-team`). If omitted, the provider generates a random UUID. Changing this forces a new team to be created. + * `team_alias` - (Required) A human-readable identifier for the team. * `organization_id` - (Optional) The ID of the organization this team belongs to. @@ -152,7 +164,7 @@ The following arguments are supported: In addition to the arguments above, the following attributes are exported: -* `id` - The unique identifier for the team. +* `id` - The unique identifier for the team, equal to `team_id`. ## Import @@ -162,7 +174,7 @@ Teams can be imported using the team ID: terraform import litellm_team.engineering ``` -Note: The team ID is generated when the team is created and is different from the `team_alias`. +Note: Unless `team_id` is set, the team ID is generated when the team is created and is different from the `team_alias`. ## Note on Team Members diff --git a/terraform/provider/litellm/client.go b/terraform/provider/litellm/client.go index 0f825d85d31..e68b8a3a80b 100644 --- a/terraform/provider/litellm/client.go +++ b/terraform/provider/litellm/client.go @@ -4,6 +4,7 @@ import ( "bytes" "crypto/tls" "encoding/json" + "errors" "fmt" "io" "log" @@ -19,6 +20,20 @@ type Client struct { InsecureSkipVerify bool } +type apiError struct { + StatusCode int + Body string +} + +func (e *apiError) Error() string { + return fmt.Sprintf("API request failed with status code %d: %s", e.StatusCode, e.Body) +} + +func isNotFound(err error) bool { + var apiErr *apiError + return errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound +} + func NewClient(apiBase, apiKey string, insecureSkipVerify bool) *Client { tr := &http.Transport{ TLSClientConfig: &tls.Config{InsecureSkipVerify: insecureSkipVerify}, @@ -57,6 +72,9 @@ func (c *Client) CreateKey(key *Key) (*Key, error) { func (c *Client) GetKey(keyID string) (*Key, error) { resp, err := c.sendRequest("GET", fmt.Sprintf("/key/info?key=%s", keyID), nil) + if isNotFound(err) { + return nil, nil + } if err != nil { return nil, err } @@ -69,32 +87,71 @@ func (c *Client) GetKey(keyID string) (*Key, error) { info["key"] = k } } + hoistKeyFieldsStoredInMetadata(info) return c.parseKeyResponse(info) } return c.parseKeyResponse(resp) } +var keyFieldsStoredInMetadata = []string{ + "model_rpm_limit", + "model_tpm_limit", + "guardrails", + "tags", + "enforced_params", + "allowed_passthrough_routes", + "rpm_limit_type", + "tpm_limit_type", + "prompts", +} + +func hoistKeyFieldsStoredInMetadata(info map[string]interface{}) { + metadata, ok := info["metadata"].(map[string]interface{}) + if !ok { + return + } + for _, field := range keyFieldsStoredInMetadata { + if existing, present := info[field]; present && existing != nil { + continue + } + if v, present := metadata[field]; present { + info[field] = v + } + } +} + func (c *Client) UpdateKey(key *Key) (*Key, error) { // Create a new map with only the fields that can be updated updateData := map[string]interface{}{ "key": key.Key, "team_id": key.TeamID, - "metadata": key.Metadata, "key_alias": key.KeyAlias, "aliases": key.Aliases, "permissions": key.Permissions, "model_max_budget": key.ModelMaxBudget, - "model_rpm_limit": key.ModelRPMLimit, - "model_tpm_limit": key.ModelTPMLimit, "blocked": key.Blocked, } + // The proxy keeps the stored metadata only when the field is absent, so nil means omit. + if key.Metadata != nil { + updateData["metadata"] = key.Metadata + } + if key.ModelRPMLimit != nil { + updateData["model_rpm_limit"] = key.ModelRPMLimit + } + if key.ModelTPMLimit != nil { + updateData["model_tpm_limit"] = key.ModelTPMLimit + } + // The proxy rejects an empty-string budget_duration with a 400, so only // send it when set. if key.BudgetDuration != "" { updateData["budget_duration"] = key.BudgetDuration } + if key.Duration != "" { + updateData["duration"] = key.Duration + } // Only add pointer fields if they are explicitly set if key.MaxBudget != nil { @@ -366,7 +423,7 @@ func (c *Client) sendRequest(method, path string, body interface{}) (map[string] log.Printf("Response body: %s", c.redactSensitiveData(string(bodyBytes))) if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("API request failed with status code %d: %s", resp.StatusCode, string(bodyBytes)) + return nil, &apiError{StatusCode: resp.StatusCode, Body: string(bodyBytes)} } var result map[string]interface{} diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index 0d8674f2d4c..018d01f75a8 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -2,7 +2,9 @@ package litellm import ( "context" + "encoding/json" "fmt" + "log" "github.com/hashicorp/go-cty/cty" "github.com/hashicorp/terraform-plugin-sdk/v2/diag" @@ -10,7 +12,7 @@ import ( ) func resourceKey() *schema.Resource { - return &schema.Resource{ + r := &schema.Resource{ CreateContext: resourceKeyCreate, ReadContext: resourceKeyRead, UpdateContext: resourceKeyUpdate, @@ -18,6 +20,7 @@ func resourceKey() *schema.Resource { Importer: &schema.ResourceImporter{ StateContext: schema.ImportStatePassthroughContext, }, + SchemaVersion: 1, Schema: map[string]*schema.Schema{ "key": { Type: schema.TypeString, @@ -86,8 +89,9 @@ func resourceKey() *schema.Resource { Optional: true, }, "duration": { - Type: schema.TypeString, - Optional: true, + Type: schema.TypeString, + Optional: true, + Description: "How long the key stays valid, e.g. \"30d\" or \"12h\". Changing it resets the expiry to the time of the update plus the new duration; removing it leaves the current expiry in place", }, "aliases": { Type: schema.TypeMap, @@ -105,9 +109,11 @@ func resourceKey() *schema.Resource { Elem: &schema.Schema{Type: schema.TypeString}, }, "model_max_budget": { - Type: schema.TypeMap, - Optional: true, - Elem: &schema.Schema{Type: schema.TypeFloat, Computed: true}, + Type: schema.TypeString, + Optional: true, + ValidateFunc: validateKeyModelMaxBudget, + DiffSuppressFunc: budgetSuppressEquivalentJSON, + Description: "JSON string of per-model budget config (e.g. '{\"gpt-4o-mini\": {\"budget_limit\": 50, \"time_period\": \"30d\"}}')", }, "model_rpm_limit": { Type: schema.TypeMap, @@ -182,6 +188,79 @@ func resourceKey() *schema.Resource { }, }, } + r.StateUpgraders = []schema.StateUpgrader{{ + Version: 0, + Type: resourceKeyV0Type(r.Schema), + Upgrade: resourceKeyStateUpgradeV0, + }} + return r +} + +// Schema version 0 typed model_max_budget as map(number), which the proxy +// rejects; version 1 stores the per-model BudgetConfig objects as a JSON string. +func resourceKeyV0Type(current map[string]*schema.Schema) cty.Type { + v0 := make(map[string]*schema.Schema, len(current)) + for k, v := range current { + v0[k] = v + } + v0["model_max_budget"] = &schema.Schema{ + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeFloat}, + } + return (&schema.Resource{Schema: v0}).CoreConfigSchema().ImpliedType() +} + +func resourceKeyStateUpgradeV0(_ context.Context, rawState map[string]interface{}, _ interface{}) (map[string]interface{}, error) { + delete(rawState, "model_max_budget") + return rawState, nil +} + +var keyModelBudgetFields = map[string]bool{ + "budget_limit": true, + "max_budget": true, + "time_period": true, + "budget_duration": true, + "tpm_limit": true, + "rpm_limit": true, +} + +func validateKeyModelMaxBudget(v interface{}, k string) ([]string, []error) { + var parsed map[string]json.RawMessage + if err := json.Unmarshal([]byte(v.(string)), &parsed); err != nil || parsed == nil { + return nil, []error{fmt.Errorf("%q must be a JSON object keyed by model name, got %s", k, v)} + } + for model, cfg := range parsed { + var budget map[string]json.RawMessage + if err := json.Unmarshal(cfg, &budget); err != nil || len(budget) == 0 { + return nil, []error{fmt.Errorf("%q[%q] must be a budget object such as {\"budget_limit\": 50, \"time_period\": \"30d\"}, got %s", k, model, cfg)} + } + for field := range budget { + if !keyModelBudgetFields[field] { + return nil, []error{fmt.Errorf("%q[%q] has unknown budget field %q; supported fields are budget_limit, max_budget, time_period, budget_duration, tpm_limit, rpm_limit", k, model, field)} + } + } + } + return nil, nil +} + +func parseKeyModelMaxBudget(raw string) map[string]interface{} { + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(raw), &parsed); err != nil || parsed == nil { + return map[string]interface{}{} + } + return parsed +} + +func keyModelMaxBudgetJSON(modelMaxBudget map[string]interface{}) string { + if len(modelMaxBudget) == 0 { + return "" + } + encoded, err := json.Marshal(modelMaxBudget) + if err != nil { + return "" + } + return string(encoded) } func resourceKeyCreate(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics { @@ -219,10 +298,12 @@ func resourceKeyRead(ctx context.Context, d *schema.ResourceData, m interface{}) } if key == nil { + log.Printf("[WARN] Key %s not found, removing from state", d.Id()) d.SetId("") return nil } + key.Metadata = declaredKeyMetadata(key.Metadata, d.Get("metadata").(map[string]interface{})) mapKeyToResourceData(d, key) return nil } @@ -232,15 +313,76 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{ key := &Key{Key: d.Id()} mapResourceDataToKey(d, key) + if !d.HasChange("duration") { + key.Duration = "" + } + key.ModelRPMLimit = changedMap(d, "model_rpm_limit") + key.ModelTPMLimit = changedMap(d, "model_tpm_limit") - _, err := c.UpdateKey(key) + metadata, err := plannedKeyMetadata(c, d) if err != nil { + d.Partial(true) + return diag.FromErr(fmt.Errorf("error updating key: %s", err)) + } + key.Metadata = metadata + + if _, err := c.UpdateKey(key); err != nil { + d.Partial(true) return diag.FromErr(fmt.Errorf("error updating key: %s", err)) } return resourceKeyRead(ctx, d, m) } +func changedMap(d *schema.ResourceData, name string) map[string]interface{} { + if !d.HasChange(name) { + return nil + } + return d.Get(name).(map[string]interface{}) +} + +func plannedKeyMetadata(c *Client, d *schema.ResourceData) (map[string]interface{}, error) { + if !d.HasChange("metadata") { + return nil, nil + } + current, err := c.GetKey(d.Id()) + if err != nil { + return nil, err + } + if current == nil { + return nil, fmt.Errorf("key %s no longer exists", d.Id()) + } + oldDeclared, newDeclared := d.GetChange("metadata") + return mergeKeyMetadata(current.Metadata, oldDeclared.(map[string]interface{}), newDeclared.(map[string]interface{})), nil +} + +func declaredKeyMetadata(server, declared map[string]interface{}) map[string]interface{} { + if server == nil { + return nil + } + result := make(map[string]interface{}, len(declared)) + for k := range declared { + if v, ok := server[k]; ok { + result[k] = v + } + } + return result +} + +func mergeKeyMetadata(server, oldDeclared, newDeclared map[string]interface{}) map[string]interface{} { + result := make(map[string]interface{}, len(server)+len(newDeclared)) + for k, v := range server { + result[k] = v + } + for k := range oldDeclared { + delete(result, k) + } + for k, v := range newDeclared { + result[k] = v + } + return result +} + func resourceKeyDelete(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics { c := m.(*Client) @@ -285,7 +427,7 @@ func mapResourceDataToKey(d *schema.ResourceData, key *Key) { key.Aliases = d.Get("aliases").(map[string]interface{}) key.Config = d.Get("config").(map[string]interface{}) key.Permissions = d.Get("permissions").(map[string]interface{}) - key.ModelMaxBudget = d.Get("model_max_budget").(map[string]interface{}) + key.ModelMaxBudget = parseKeyModelMaxBudget(d.Get("model_max_budget").(string)) key.ModelRPMLimit = d.Get("model_rpm_limit").(map[string]interface{}) key.ModelTPMLimit = d.Get("model_tpm_limit").(map[string]interface{}) key.Guardrails = expandStringList(d.Get("guardrails").([]interface{})) @@ -358,9 +500,7 @@ func mapKeyToResourceData(d *schema.ResourceData, key *Key) { if key.Permissions != nil { d.Set("permissions", key.Permissions) } - if key.ModelMaxBudget != nil { - d.Set("model_max_budget", key.ModelMaxBudget) - } + d.Set("model_max_budget", keyModelMaxBudgetJSON(key.ModelMaxBudget)) if key.ModelRPMLimit != nil { d.Set("model_rpm_limit", key.ModelRPMLimit) } diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go index 91f0061a9ef..66291eadcc5 100644 --- a/terraform/provider/litellm/resource_key_test.go +++ b/terraform/provider/litellm/resource_key_test.go @@ -6,9 +6,11 @@ import ( "io" "net/http" "net/http/httptest" + "reflect" "testing" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" ) func newKeyResourceData(t *testing.T, raw map[string]interface{}) *schema.ResourceData { @@ -193,6 +195,102 @@ func TestCreateKeySendsConfigSuppliedKey(t *testing.T) { } } +// The proxy validates each model_max_budget entry as a BudgetConfig object and +// 500s on a bare number, so the JSON string must reach /key/generate as nested +// objects and the proxy's response must map back to equivalent JSON in state. +func TestCreateKeySendsModelMaxBudgetAsBudgetObjects(t *testing.T) { + var captured map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/key/generate" { + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, &captured) + w.Write([]byte(`{"key": "sk-test", "token_id": "hash-1"}`)) + return + } + w.Write([]byte(`{"key": "hash-1", "info": {"model_max_budget": {"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d", "rpm_limit": 60}}}}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newKeyResourceData(t, map[string]interface{}{ + "model_max_budget": `{"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d"}}`, + }) + + if diags := resourceKeyCreate(context.Background(), d, client); diags.HasError() { + t.Fatalf("create returned error: %v", diags) + } + + budgets, ok := captured["model_max_budget"].(map[string]interface{}) + if !ok { + t.Fatalf("create payload model_max_budget = %v, want object", captured["model_max_budget"]) + } + cfg, ok := budgets["gpt-4o-mini"].(map[string]interface{}) + if !ok { + t.Fatalf("model_max_budget[gpt-4o-mini] = %v, want BudgetConfig object", budgets["gpt-4o-mini"]) + } + if cfg["budget_limit"] != float64(50) || cfg["time_period"] != "30d" { + t.Errorf("BudgetConfig = %v, want budget_limit 50 and time_period 30d", cfg) + } + + var state map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("model_max_budget").(string)), &state); err != nil { + t.Fatalf("state model_max_budget %q is not JSON: %v", d.Get("model_max_budget"), err) + } + if got, _ := state["gpt-4o-mini"].(map[string]interface{}); got["budget_limit"] != float64(50) || got["rpm_limit"] != float64(60) { + t.Errorf("state model_max_budget = %v, want the BudgetConfig read back from /key/info", state) + } +} + +// Schema version 0 stored model_max_budget as map(number); that state cannot +// decode into the version 1 string attribute, so the upgrader must drop it. +func TestKeyStateUpgradeV0DropsMapModelMaxBudget(t *testing.T) { + upgraded, err := resourceKey().StateUpgraders[0].Upgrade(context.Background(), map[string]interface{}{ + "id": "hash-1", + "key_alias": "legacy", + "model_max_budget": map[string]interface{}{"gpt-4o-mini": 50.0}, + }, nil) + if err != nil { + t.Fatalf("upgrade returned error: %v", err) + } + if _, present := upgraded["model_max_budget"]; present { + t.Errorf("upgraded state still carries map model_max_budget: %v", upgraded["model_max_budget"]) + } + if upgraded["key_alias"] != "legacy" { + t.Errorf("upgrade dropped unrelated attribute: %v", upgraded) + } +} + +func TestKeyModelMaxBudgetValidationRequiresBudgetObjects(t *testing.T) { + validate := resourceKey().Schema["model_max_budget"].ValidateFunc + for _, valid := range []string{ + `{}`, + `{"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d"}}`, + `{"gpt-4o-mini": {"max_budget": 50, "rpm_limit": 60}, "gpt-4o": {"budget_duration": "1d", "tpm_limit": 1000}}`, + } { + if _, errs := validate(valid, "model_max_budget"); len(errs) != 0 { + t.Errorf("validate(%s) = %v, want accepted", valid, errs) + } + } + for _, invalid := range []string{ + `null`, + `[]`, + `"gpt-4o-mini"`, + `50`, + `{"gpt-4o-mini": 50}`, + `{"gpt-4o-mini": null}`, + `{"gpt-4o-mini": [50]}`, + `{"gpt-4o-mini": {}}`, + `{"gpt-4o-mini": {"budget_limt": 50}}`, + `{"gpt-4o-mini": {"budget_limit": 50, "max_tokens": 100}}`, + `not json`, + } { + if _, errs := validate(invalid, "model_max_budget"); len(errs) == 0 { + t.Errorf("validate(%s) accepted a value that would send no per-model budget", invalid) + } + } +} + // The proxy 400s on budget_duration: "", so an unset duration must be // omitted from the update payload entirely. func TestUpdateKeyOmitsEmptyBudgetDuration(t *testing.T) { @@ -221,6 +319,47 @@ func TestUpdateKeyOmitsEmptyBudgetDuration(t *testing.T) { } } +func TestResourceKeyUpdateFailureKeepsPriorState(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/key/update" { + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(`{"error":{"message":"Invalid budget_duration 'bad'"}}`)) + return + } + w.Write([]byte(`{"key":"hash-1","info":{"key_alias":"demo","models":["fake-model"]}}`)) + })) + defer srv.Close() + + res := resourceKey() + priorData := newKeyResourceData(t, map[string]interface{}{ + "key_alias": "demo", + "models": []interface{}{"fake-model"}, + }) + priorData.SetId("hash-1") + prior := priorData.State() + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "key_alias": "demo", + "models": []interface{}{"fake-model"}, + "budget_duration": "bad", + }) + diff, err := res.Diff(context.Background(), prior, config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + + newState, diags := res.Apply(context.Background(), prior, diff, NewClient(srv.URL, "test-key", true)) + if !diags.HasError() { + t.Fatal("apply succeeded, want the proxy's 400 surfaced as an error") + } + if got, ok := newState.Attributes["budget_duration"]; ok { + t.Errorf("failed update persisted budget_duration=%q into state, want it absent", got) + } + if newState.Attributes["key_alias"] != "demo" { + t.Errorf("prior key_alias lost from state: %v", newState.Attributes) + } +} + // /key/info nests the key's fields under "info"; GetKey must unwrap that // envelope or reads map nothing back into state. func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) { @@ -254,3 +393,296 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) { t.Errorf("RPMLimit not parsed: %+v", key.RPMLimit) } } + +func TestGetKeyReadsFieldsStoredInMetadata(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "key": "hash-1", + "info": { + "models": ["gpt-4o-mini"], + "metadata": { + "team": "core-infra", + "model_rpm_limit": {"gpt-4o-mini": 7}, + "model_tpm_limit": {"gpt-4o-mini": 10000}, + "guardrails": ["pii-guard"], + "tags": ["prod"], + "enforced_params": ["user"], + "allowed_passthrough_routes": ["/v1/foo"], + "rpm_limit_type": "guaranteed_throughput", + "tpm_limit_type": "dynamic", + "prompts": ["p1"] + } + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + key, err := client.GetKey("hash-1") + if err != nil { + t.Fatalf("GetKey returned error: %v", err) + } + if got, ok := key.ModelRPMLimit["gpt-4o-mini"].(float64); !ok || got != 7 { + t.Errorf("ModelRPMLimit = %v, want gpt-4o-mini=7 read from metadata", key.ModelRPMLimit) + } + if got, ok := key.ModelTPMLimit["gpt-4o-mini"].(float64); !ok || got != 10000 { + t.Errorf("ModelTPMLimit = %v, want gpt-4o-mini=10000 read from metadata", key.ModelTPMLimit) + } + if len(key.Guardrails) != 1 || key.Guardrails[0] != "pii-guard" { + t.Errorf("Guardrails = %v, want [pii-guard]", key.Guardrails) + } + if len(key.Tags) != 1 || key.Tags[0] != "prod" { + t.Errorf("Tags = %v, want [prod]", key.Tags) + } + if len(key.EnforcedParams) != 1 || key.EnforcedParams[0] != "user" { + t.Errorf("EnforcedParams = %v, want [user]", key.EnforcedParams) + } + if len(key.AllowedPassthroughRoutes) != 1 || key.AllowedPassthroughRoutes[0] != "/v1/foo" { + t.Errorf("AllowedPassthroughRoutes = %v, want [/v1/foo]", key.AllowedPassthroughRoutes) + } + if key.RPMLimitType != "guaranteed_throughput" || key.TPMLimitType != "dynamic" { + t.Errorf("limit types = %q/%q, want guaranteed_throughput/dynamic", key.RPMLimitType, key.TPMLimitType) + } + if len(key.Prompts) != 1 || key.Prompts[0] != "p1" { + t.Errorf("Prompts = %v, want [p1]", key.Prompts) + } + if key.Metadata["team"] != "core-infra" { + t.Errorf("Metadata = %v, want team=core-infra preserved", key.Metadata) + } +} + +func TestGetKeyPrefersTopLevelOverMetadataCopy(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "key": "hash-1", + "info": { + "tags": ["top-level"], + "guardrails": null, + "metadata": { + "tags": ["from-metadata"], + "guardrails": ["from-metadata"] + } + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + key, err := client.GetKey("hash-1") + if err != nil { + t.Fatalf("GetKey returned error: %v", err) + } + if len(key.Tags) != 1 || key.Tags[0] != "top-level" { + t.Errorf("Tags = %v, want [top-level]", key.Tags) + } + if len(key.Guardrails) != 1 || key.Guardrails[0] != "from-metadata" { + t.Errorf("Guardrails = %v, want [from-metadata] (null top-level must not shadow)", key.Guardrails) + } +} + +func TestResourceKeyReadDropsMissingKeyFromState(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + w.Write([]byte(`{"error":{"message":"Key not found in database","type":"not_found_error","param":"key","code":"404"}}`)) + })) + defer srv.Close() + + d := newKeyResourceData(t, map[string]interface{}{"key_alias": "stale"}) + d.SetId("deleted-out-of-band") + + diags := resourceKeyRead(context.Background(), d, NewClient(srv.URL, "test-key", true)) + if diags.HasError() { + t.Fatalf("read of a missing key must not error, got: %v", diags) + } + if d.Id() != "" { + t.Errorf("Id = %q, want empty so Terraform plans a recreate", d.Id()) + } +} + +func TestResourceKeyReadStillFailsOnNon404Errors(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":{"message":"db down"}}`)) + })) + defer srv.Close() + + d := newKeyResourceData(t, map[string]interface{}{"key_alias": "live"}) + d.SetId("still-exists") + + diags := resourceKeyRead(context.Background(), d, NewClient(srv.URL, "test-key", true)) + if !diags.HasError() { + t.Fatal("a 500 from /key/info must surface as an error, not be treated as a deleted key") + } + if d.Id() != "still-exists" { + t.Errorf("Id = %q, want unchanged on a transient error", d.Id()) + } +} + +// fakeKeyProxy serves /key/info from stored metadata and applies /key/update +// the way the proxy does: an absent "metadata" keeps the stored map, a +// present one replaces it wholesale. +type fakeKeyProxy struct { + metadata map[string]interface{} + updates []map[string]interface{} +} + +func (p *fakeKeyProxy) handler() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/key/info": + json.NewEncoder(w).Encode(map[string]interface{}{ + "key": "hash-1", + "info": map[string]interface{}{"key_alias": "alias-1", "models": []string{"gpt-4o-mini"}, "metadata": p.metadata}, + }) + case "/key/update": + var body map[string]interface{} + json.NewDecoder(r.Body).Decode(&body) + p.updates = append(p.updates, body) + if m, ok := body["metadata"].(map[string]interface{}); ok { + p.metadata = m + } + json.NewEncoder(w).Encode(map[string]interface{}{"key": "hash-1", "metadata": p.metadata}) + default: + http.NotFound(w, r) + } + } +} + +func applyKeyUpdate(t *testing.T, client *Client, stateAttrs map[string]string, config map[string]interface{}) *terraform.InstanceState { + t.Helper() + r := resourceKey() + state := &terraform.InstanceState{ID: "hash-1", Attributes: stateAttrs} + diff, err := r.Diff(context.Background(), state, terraform.NewResourceConfigRaw(config), client) + if err != nil { + t.Fatalf("Diff returned error: %v", err) + } + if diff == nil { + t.Fatalf("expected a non-empty diff between %v and %v", stateAttrs, config) + } + newState, diags := r.Apply(context.Background(), state, diff, client) + if diags.HasError() { + t.Fatalf("Apply returned error: %v", diags) + } + return newState +} + +func TestKeyUpdateWithoutMetadataChangePreservesServerMetadata(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": float64(5)}}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + newState := applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "max_budget": "10", "metadata.%": "1", "metadata.a": "1"}, + map[string]interface{}{"key_alias": "alias-1", "max_budget": 20, "metadata": map[string]interface{}{"a": "1"}}, + ) + + if len(proxy.updates) != 1 { + t.Fatalf("expected one /key/update call, got %d", len(proxy.updates)) + } + for _, field := range []string{"metadata", "model_rpm_limit", "model_tpm_limit"} { + if _, present := proxy.updates[0][field]; present { + t.Errorf("unchanged %q was sent on /key/update: %v", field, proxy.updates[0][field]) + } + } + if proxy.metadata["server_side"] != "x" { + t.Errorf("server-side metadata lost: %v", proxy.metadata) + } + if got := newState.Attributes["metadata.%"]; got != "1" { + t.Errorf("state metadata should hold only the declared entry, got %v", newState.Attributes) + } + if got := newState.Attributes["metadata.a"]; got != "1" { + t.Errorf("metadata.a = %q, want 1", got) + } +} + +func TestKeyUpdateWithMetadataChangeMergesOverServerMetadata(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "b": "2", "server_side": "x"}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "metadata.%": "2", "metadata.a": "1", "metadata.b": "2"}, + map[string]interface{}{"key_alias": "alias-1", "metadata": map[string]interface{}{"a": "2", "c": "3"}}, + ) + + want := map[string]interface{}{"a": "2", "c": "3", "server_side": "x"} + if !reflect.DeepEqual(proxy.metadata, want) { + t.Errorf("metadata after update = %v, want %v", proxy.metadata, want) + } +} + +func TestKeyUpdateSendsChangedModelLimits(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "model_rpm_limit.%": "1", "model_rpm_limit.gpt-4o-mini": "5"}, + map[string]interface{}{"key_alias": "alias-1", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 7}}, + ) + + got, ok := proxy.updates[0]["model_rpm_limit"].(map[string]interface{}) + if !ok || got["gpt-4o-mini"] != float64(7) { + t.Errorf("changed model_rpm_limit not sent: %v", proxy.updates[0]) + } +} + +func TestKeyReadKeepsOnlyDeclaredMetadata(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x"}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + d := newKeyResourceData(t, map[string]interface{}{"metadata": map[string]interface{}{"a": "1"}}) + d.SetId("hash-1") + if diags := resourceKeyRead(context.Background(), d, client); diags.HasError() { + t.Fatalf("Read returned error: %v", diags) + } + + want := map[string]interface{}{"a": "1"} + if got := d.Get("metadata"); !reflect.DeepEqual(got, want) { + t.Errorf("metadata in state = %v, want %v", got, want) + } +} + +func TestKeyUpdateSendsChangedDuration(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "duration": "30d"}, + map[string]interface{}{"key_alias": "alias-1", "duration": "90d"}, + ) + + if got := proxy.updates[0]["duration"]; got != "90d" { + t.Errorf("update payload duration = %v, want 90d", got) + } +} + +func TestKeyUpdateOmitsUnchangedDuration(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "duration": "30d"}, + map[string]interface{}{"key_alias": "alias-2", "duration": "30d"}, + ) + + if got := proxy.updates[0]["key_alias"]; got != "alias-2" { + t.Fatalf("update payload key_alias = %v, want alias-2", got) + } + if v, present := proxy.updates[0]["duration"]; present { + t.Errorf("update payload unexpectedly contains duration = %v", v) + } +} diff --git a/terraform/provider/litellm/resource_key_utils.go b/terraform/provider/litellm/resource_key_utils.go index d426fec05b2..7046ef7097a 100644 --- a/terraform/provider/litellm/resource_key_utils.go +++ b/terraform/provider/litellm/resource_key_utils.go @@ -65,7 +65,7 @@ func buildKeyData(d *schema.ResourceData) map[string]interface{} { keyData["permissions"] = v.(map[string]interface{}) } if v, ok := d.GetOkExists("model_max_budget"); ok { - keyData["model_max_budget"] = v.(map[string]interface{}) + keyData["model_max_budget"] = parseKeyModelMaxBudget(v.(string)) } if v, ok := d.GetOkExists("model_rpm_limit"); ok { keyData["model_rpm_limit"] = v.(map[string]interface{}) @@ -107,7 +107,7 @@ func setKeyResourceData(d *schema.ResourceData, key *Key) error { "aliases": key.Aliases, "config": key.Config, "permissions": key.Permissions, - "model_max_budget": key.ModelMaxBudget, + "model_max_budget": keyModelMaxBudgetJSON(key.ModelMaxBudget), "model_rpm_limit": key.ModelRPMLimit, "model_tpm_limit": key.ModelTPMLimit, "guardrails": key.Guardrails, diff --git a/terraform/provider/litellm/resource_team.go b/terraform/provider/litellm/resource_team.go index 24c47843cd1..bf7d2508077 100644 --- a/terraform/provider/litellm/resource_team.go +++ b/terraform/provider/litellm/resource_team.go @@ -31,6 +31,13 @@ func ResourceLiteLLMTeam() *schema.Resource { }, Schema: map[string]*schema.Schema{ + "team_id": { + Type: schema.TypeString, + Optional: true, + Computed: true, + ForceNew: true, + Description: "Unique ID for the team. Generated by the provider if not provided", + }, "team_alias": { Type: schema.TypeString, Required: true, @@ -162,7 +169,7 @@ func ResourceLiteLLMTeam() *schema.Resource { func resourceLiteLLMTeamCreate(d *schema.ResourceData, m interface{}) error { client := m.(*Client) - teamID := uuid.New().String() + teamID := resolveTeamID(d) teamData := buildTeamData(d, teamID) // Throughput limit types are only accepted by /team/new, not /team/update. @@ -214,6 +221,7 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { teamResp := infoResp.TeamInfo // Update the state with values from the response or fall back to the data passed in during creation + d.Set("team_id", d.Id()) d.Set("team_alias", GetStringValue(teamResp.TeamAlias, d.Get("team_alias").(string))) d.Set("organization_id", GetStringValue(teamResp.OrganizationID, d.Get("organization_id").(string))) @@ -263,11 +271,11 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { d.Set("team_member_tpm_limit", *teamResp.TeamMemberTPMLimit) } d.Set("team_member_key_duration", GetStringValue(teamResp.TeamMemberKeyDuration, d.Get("team_member_key_duration").(string))) - if teamResp.ModelRPMLimit != nil { - d.Set("model_rpm_limit", teamResp.ModelRPMLimit) + if v := teamModelLimit(teamResp.ModelRPMLimit, teamResp.Metadata, "model_rpm_limit"); v != nil { + d.Set("model_rpm_limit", v) } - if teamResp.ModelTPMLimit != nil { - d.Set("model_tpm_limit", teamResp.ModelTPMLimit) + if v := teamModelLimit(teamResp.ModelTPMLimit, teamResp.Metadata, "model_tpm_limit"); v != nil { + d.Set("model_tpm_limit", v) } if teamResp.AllowedPassthroughRoutes != nil { d.Set("allowed_passthrough_routes", teamResp.AllowedPassthroughRoutes) @@ -354,6 +362,13 @@ func resourceLiteLLMTeamDelete(d *schema.ResourceData, m interface{}) error { return nil } +func resolveTeamID(d *schema.ResourceData) string { + if v, ok := d.GetOk("team_id"); ok { + return v.(string) + } + return uuid.New().String() +} + func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} { teamData := map[string]interface{}{ "team_id": teamID, @@ -364,14 +379,19 @@ func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} "organization_id", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models", "blocked", "team_member_permissions", "model_aliases", "guardrails", "prompts", "team_member_budget", "team_member_budget_duration", "team_member_rpm_limit", - "team_member_tpm_limit", "team_member_key_duration", "model_rpm_limit", - "model_tpm_limit", "allowed_passthrough_routes", + "team_member_tpm_limit", "team_member_key_duration", "allowed_passthrough_routes", } { if v, ok := d.GetOk(key); ok { teamData[key] = v } } + for _, key := range []string{"model_rpm_limit", "model_tpm_limit"} { + if v, ok := d.GetOk(key); ok || d.HasChange(key) { + teamData[key] = v + } + } + if v, ok := d.GetOk("soft_budget"); ok { teamData["soft_budget"] = v } else if d.HasChange("soft_budget") { @@ -404,6 +424,14 @@ func buildTeamMetadata(d *schema.ResourceData) map[string]interface{} { return metadata } +func teamModelLimit(topLevel, metadata map[string]interface{}, key string) map[string]interface{} { + if topLevel != nil { + return topLevel + } + nested, _ := metadata[key].(map[string]interface{}) + return nested +} + func splitTeamMetadata(raw map[string]interface{}) (map[string]string, []string, []string) { metadata := map[string]string{} var tags, alertEmails []string diff --git a/terraform/provider/litellm/resource_team_test.go b/terraform/provider/litellm/resource_team_test.go index 9638378cdfe..35d60401d30 100644 --- a/terraform/provider/litellm/resource_team_test.go +++ b/terraform/provider/litellm/resource_team_test.go @@ -9,6 +9,7 @@ import ( "reflect" "testing" + "github.com/google/uuid" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" ) @@ -85,6 +86,87 @@ func TestTeamCreateSendsSoftBudgetTagsAndAlertEmails(t *testing.T) { } } +func TestTeamCreateSendsConfiguredTeamID(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"platform-team","team_info":{"team_id":"platform-team","team_alias":"platform"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{ + "team_id": "platform-team", + "team_alias": "platform", + }) + + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if got := captured["team_id"]; got != "platform-team" { + t.Fatalf("payload team_id = %v, want platform-team", got) + } + if got := d.Id(); got != "platform-team" { + t.Fatalf("resource id = %q, want platform-team", got) + } + if got := d.Get("team_id"); got != "platform-team" { + t.Fatalf("state team_id = %v, want platform-team", got) + } +} + +func TestTeamCreateGeneratesTeamIDWhenUnset(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"x","team_info":{"team_alias":"eng"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"}) + + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + sent, _ := captured["team_id"].(string) + if _, err := uuid.Parse(sent); err != nil { + t.Fatalf("payload team_id = %q, want a generated UUID: %v", sent, err) + } + if d.Id() != sent || d.Get("team_id") != sent { + t.Fatalf("id = %q, state team_id = %v, want both to equal the sent id %q", d.Id(), d.Get("team_id"), sent) + } +} + +func TestTeamReadSetsTeamIDFromResourceID(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"imported-team","team_info":{"team_id":"imported-team","team_alias":"imported"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{}) + d.SetId("imported-team") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + if got := d.Get("team_id"); got != "imported-team" { + t.Fatalf("team_id = %v, want imported-team", got) + } +} + +func TestTeamIDChangeForcesReplacement(t *testing.T) { + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{ + "team_id": "old-team", + "team_alias": "eng", + }) + priorData.SetId("old-team") + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "team_id": "new-team", + "team_alias": "eng", + }) + diff, err := res.Diff(context.Background(), priorData.State(), config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + if diff == nil || !diff.RequiresNew() { + t.Fatalf("changing team_id must force replacement, diff = %+v", diff) + } +} + func TestTeamReadMapsTeamInfoEnvelope(t *testing.T) { var captured map[string]interface{} srv := newTeamTestServer(t, &captured, teamInfoWithSoftBudget) @@ -250,6 +332,77 @@ func TestTeamReadMapsNewFields(t *testing.T) { } } +func TestTeamReadMapsPerModelLimitsFromMetadata(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{ + "team_id": "team-1", + "team_info": { + "team_id": "team-1", + "team_alias": "eng", + "model_rpm_limit": null, + "model_tpm_limit": null, + "metadata": { + "department": "eng", + "model_rpm_limit": {"gpt-4o-mini": 250}, + "model_tpm_limit": {"gpt-4o-mini": 5000} + } + } + }`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{ + "team_alias": "eng", + "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 100}, + }) + d.SetId("team-1") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read returned error: %v", err) + } + if got := d.Get("model_rpm_limit"); !reflect.DeepEqual(got, map[string]interface{}{"gpt-4o-mini": 250}) { + t.Errorf("model_rpm_limit = %v, want server value 250", got) + } + if got := d.Get("model_tpm_limit"); !reflect.DeepEqual(got, map[string]interface{}{"gpt-4o-mini": 5000}) { + t.Errorf("model_tpm_limit = %v, want server value 5000", got) + } + if got := d.Get("metadata"); !reflect.DeepEqual(got, map[string]interface{}{"department": "eng"}) { + t.Errorf("metadata = %v, want per-model limits kept out of the string map", got) + } +} + +func TestTeamUpdateClearsRemovedPerModelLimits(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{ + "team_alias": "eng", + "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 100}, + "model_tpm_limit": map[string]interface{}{"gpt-4o-mini": 5000}, + }) + priorData.SetId("team-1") + prior := priorData.State() + config := terraform.NewResourceConfigRaw(map[string]interface{}{"team_alias": "eng"}) + diff, err := res.Diff(context.Background(), prior, config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + d, err := schema.InternalMap(res.Schema).Data(prior, diff) + if err != nil { + t.Fatalf("data failed: %v", err) + } + + if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + for _, k := range []string{"model_rpm_limit", "model_tpm_limit"} { + if got, ok := captured[k]; !ok || !reflect.DeepEqual(got, map[string]interface{}{}) { + t.Errorf("payload %s = %v (present=%v), want explicit empty map", k, got, ok) + } + } +} + // rpm_limit_type / tpm_limit_type are accepted by /team/new but not // /team/update, so create must send them and update must not. func TestTeamLimitTypesSentOnCreateOnly(t *testing.T) { diff --git a/tests/batches_tests/test_batches_logging_unit_tests.py b/tests/batches_tests/test_batches_logging_unit_tests.py index 5bde40d90b0..73adb391481 100644 --- a/tests/batches_tests/test_batches_logging_unit_tests.py +++ b/tests/batches_tests/test_batches_logging_unit_tests.py @@ -144,18 +144,20 @@ def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict): @pytest.mark.asyncio async def test_batch_cost_calculator(sample_file_content_dict): """ - mock litellm.completion_cost to return 0.5 + mock batch_cost_calculator to return (0.3, 0.2) per line we know sample_file_content_dict has 2 successful responses - so we expect the cost to be 0.5 * 2 = 1.0 + so we expect the cost to be (0.3 + 0.2) * 2 = 1.0, split 0.6 / 0.4 """ - with patch("litellm.completion_cost", return_value=0.5): + with patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.3, 0.2)): result = _aggregate_batch_cost_usage_models( entries=sample_file_content_dict, custom_llm_provider="openai", ) - assert result.cost == 1.0 # 0.5 * 2 successful responses + assert result.cost == pytest.approx(1.0) # (0.3 + 0.2) * 2 successful responses + assert result.prompt_cost == pytest.approx(0.6) + assert result.completion_cost == pytest.approx(0.4) def test_get_response_from_batch_job_output_file(sample_file_content_dict): @@ -402,6 +404,56 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): assert mock_batch.usage == explicit_usage +@pytest.mark.asyncio +async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown(): + """The poller passes the batch's prompt/completion cost split so the spend row's + cost_breakdown carries real input/output costs; without it the UI's Cost Breakdown + card renders blank for every batch. Regression for the split being dropped.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CallTypes, LiteLLMBatch + + mock_batch = LiteLLMBatch( + id="batch-breakdown-1", + object="batch", + endpoint="/v1/chat/completions", + errors=None, + input_file_id="file-input-1", + completion_window="24h", + status="completed", + output_file_id="file-output-1", + created_at=1234567890, + ) + mock_batch._hidden_params = {} + + logging_obj = Logging( + model="gpt-5-mini", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type=CallTypes.aretrieve_batch.value, + litellm_call_id="test-call-breakdown", + function_id="test-function", + start_time=time.time(), + dynamic_success_callbacks=[], + ) + logging_obj.custom_llm_provider = "openai" + + await logging_obj.async_success_handler( + result=mock_batch, + start_time=time.time(), + end_time=time.time() + 1, + batch_cost=0.10, + batch_usage=litellm.Usage(prompt_tokens=200, completion_tokens=100, total_tokens=300), + batch_models=["gpt-5-mini"], + batch_prompt_cost=0.06, + batch_completion_cost=0.04, + ) + + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown["input_cost"] == 0.06 + assert logging_obj.cost_breakdown["output_cost"] == 0.04 + assert logging_obj.cost_breakdown["total_cost"] == 0.10 + + @pytest.mark.asyncio async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batch(): """ diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py new file mode 100644 index 00000000000..ae2c1d80c8f --- /dev/null +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -0,0 +1,521 @@ +"""Catch workflow jobs that publish check runs under the same name. + +A ruleset's required status check names a check run and GitHub matches it by that +name alone. When two jobs publish the same name the required context stops +mapping to the job that proves it: the commit carries two check runs under one +name and nothing says which one the ruleset required. Both being green hides the +clash completely, so the context quietly stops meaning what the ruleset intended. +One job lands in the same place when its `name:` holds no matrix value, since +every combination it runs then reports under that one name. + +`.github/workflows/auto-close-duplicates.yml` shipped a job id `test` while +`.github/workflows/test-mcp.yml` already published the required `test` context, +and commit ed5761daef4ae17152446d182c860630c38b7268 carried both check runs. +This invariant has to be enforced here because CI cannot enforce it on itself. + +A job publishes its `name:` when it sets one, and otherwise its job id plus the +values of the combination it runs, the way GitHub writes `build (3.12)`. A name +carrying `${{ ... }}` publishes one check run per combination the matrix +produces: `exclude` rows drop combinations before `include` rows fold into the +survivors, and each `include` row's values stay together rather than crossing +with the other rows', so two shard lists that overlap collide even though their +templates read differently. Each expression is evaluated per combination over the +pieces a job name can hold: string literals, `matrix.`, `format()`, `==` and +`!=`, and the ` && || ` idiom, which is how the shards reach their +real ` / Run tests` names rather than staying opaque. + +Whatever the sweep cannot work out is left out of the comparison and reported +instead of guessed, because a guess that lands wrong fails a workflow GitHub +would have published perfectly well. A name still holding an expression once the +combination is filled in is usually one GitHub resolves per job, so it is one of +those: guessing that two jobs sharing such a template clash would fail workflows +over a context this sweep cannot read. The exception is a name whose leftover +expressions all read a `github.` property other than `github.job`, which one run +fills in the same way for every job in it, so those are compared against the +other jobs of their own workflow and stay out of the comparison across files, +where two workflows can run on different events. A matrix that is itself an +expression or that lists values which are not scalars, an `include` or `exclude` +row shaped the same way, a whole `strategy:` that comes from an expression, and a +call this sweep cannot follow, go in the same bucket. The cost is that a real clash hiding behind +one of them goes unseen, which leaves a merge no worse off than before this check +existed, where the opposite direction would block work that was fine. + +A job calling a local reusable workflow publishes one check run per job of the +callee, named ` / ` and chained through however many levels of +local calls it takes, which is why a caller's name never collides with a plain +job that happens to match it. A file under `.github/workflows/` that does not +read as one workflow at all is reported rather than skipped, since skipping it +silently would hide every job it holds. +""" + +import itertools +import operator +import re +import sys +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import yaml +from pydantic import BaseModel, Field, ValidationError + +REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent +WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows" +EXPRESSION: Final = re.compile(r"\$\{\{(?P.*?)\}\}", re.DOTALL) +MATRIX_REF: Final = re.compile(r"^matrix\.(?P[\w-]+)$") +LITERAL: Final = re.compile(r"^'(?P[^']*)'$") +FORMAT_CALL: Final = re.compile(r"^format\((?P.*)\)$", re.DOTALL) +COMPARISON: Final = re.compile(r"^(?P.+?)\s*(?P==|!=)\s*(?P.+)$", re.DOTALL) +RUN_WIDE: Final = re.compile(r"^github\.(?!job\b)[\w.]+$") +GITHUB_PLACEHOLDER: Final = re.compile(r"\{\{|\}\}|\{\d+\}") +NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({}) +NO_CALLERS: Final[frozenset[str]] = frozenset() +SCALAR: Final = (str, int, float) +MATRIX_DIRECTIVES: Final = frozenset({"include", "exclude"}) +LOCAL_CALL_PREFIX: Final = "./" + + +@dataclass(frozen=True, slots=True) +class Unreadable: + reason: str + + +@dataclass(frozen=True, slots=True) +class Opaque: + reason: str + + +@dataclass(frozen=True, slots=True) +class Names: + """The check-run names a job publishes, beside the reasons the rest of them stay unknown.""" + + known: tuple[str, ...] = () + unknown: tuple[str, ...] = () + local: tuple[str, ...] = () + + +class Job(BaseModel): + name: object = None + uses: str | None = None + strategy: object = Field(default_factory=dict) + + +class Workflow(BaseModel): + jobs: Mapping[str, Job] = Field(default_factory=dict) + + +def scalar_text(value: object) -> str: + """A YAML scalar the way GitHub renders it, so `true` never reaches a name as `True`.""" + return str(value).lower() if isinstance(value, bool) else str(value) + + +def parse(source: str) -> tuple[Workflow, object] | Unreadable: + """The workflow plus its raw `on:` value, or why the file does not read as one.""" + try: + parsed: Final = yaml.safe_load(source) + except yaml.YAMLError: + return Unreadable("it does not read as one YAML document") + if not isinstance(parsed, dict): + return Unreadable("its top level is not a mapping of workflow keys") + try: + return Workflow.model_validate(parsed), parsed.get(True, parsed.get("on")) + except ValidationError as error: + return Unreadable(f"{error.error_count()} of its job definitions have a shape GitHub would reject") + + +def events(raw_on: object) -> frozenset[str]: + if isinstance(raw_on, Mapping): + return frozenset(str(key) for key in raw_on) + if isinstance(raw_on, str): + return frozenset({raw_on}) + if isinstance(raw_on, Sequence): + return frozenset(str(event) for event in raw_on) + return frozenset() + + +def publishes_check_runs(raw_on: object) -> bool: + """A `workflow_call`-only workflow posts its check runs through callers, never itself.""" + return events(raw_on) != frozenset({"workflow_call"}) + + +def scalar_list(value: object) -> tuple[str, ...] | Opaque: + """One matrix key's values, or why the combinations it produces cannot be worked out.""" + if not isinstance(value, Sequence) or isinstance(value, str): + return Opaque("a matrix key holds something other than a list of values") + if any(not isinstance(item, SCALAR) for item in value): + return Opaque("a matrix key lists values that are not plain scalars") + return tuple(scalar_text(item) for item in value) + + +def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...] | Opaque: + listed: Final = tuple( + (str(key), scalar_list(values)) for key, values in matrix.items() if str(key) not in MATRIX_DIRECTIVES + ) + opaque: Final = next((values for _, values in listed if isinstance(values, Opaque)), None) + if opaque is not None: + return opaque + return tuple((key, values) for key, values in listed if not isinstance(values, Opaque)) + + +def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...] | Opaque: + """One `include` or `exclude` row, or why the combinations they shape cannot be worked out.""" + rows: Final = matrix.get(directive) + if rows is None: + return () + if not isinstance(rows, Sequence) or isinstance(rows, str): + return Opaque(f"a matrix `{directive}` is itself an expression rather than a list of rows") + mappings: Final = tuple(row for row in rows if isinstance(row, Mapping)) + if len(mappings) != len(rows): + return Opaque(f"a matrix `{directive}` row is not a mapping of values") + if any(not isinstance(value, SCALAR) for row in mappings for value in row.values()): + return Opaque(f"a matrix `{directive}` row holds a value that is not a plain scalar") + return tuple(MappingProxyType({str(key): scalar_text(value) for key, value in row.items()}) for row in mappings) + + +def drops(row: Mapping[str, str], combination: Mapping[str, str]) -> bool: + """GitHub removes a combination that carries every value one `exclude` row names.""" + return all(combination.get(key) == value for key, value in row.items()) + + +def extends(row: Mapping[str, str], combination: Mapping[str, str]) -> bool: + """GitHub folds an `include` row into a combination only where it overwrites no listed value.""" + return all(combination[key] == value for key, value in row.items() if key in combination) + + +def extended(combination: Mapping[str, str], rows: Sequence[Mapping[str, str]]) -> Mapping[str, str]: + additions: Final = {key: value for row in rows if extends(row, combination) for key, value in row.items()} + return MappingProxyType({**combination, **additions}) + + +def crossed_values(listed: Sequence[tuple[str, tuple[str, ...]]]) -> tuple[Mapping[str, str], ...]: + if not listed: + return () + return tuple( + MappingProxyType(dict(zip((key for key, _ in listed), values))) + for values in itertools.product(*(values for _, values in listed)) + ) + + +def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...] | Opaque: + """One mapping per job the matrix produces, `exclude` applied before `include` as GitHub does.""" + if not isinstance(job.strategy, Mapping): + return Opaque("its whole `strategy` comes from an expression") + matrix: Final = job.strategy.get("matrix") + if matrix is None: + return () + if not isinstance(matrix, Mapping): + return Opaque("the matrix itself comes from an expression") + listed: Final = listed_values(matrix) + if isinstance(listed, Opaque): + return listed + rows: Final = directive_rows(matrix, "include") + if isinstance(rows, Opaque): + return rows + dropped: Final = directive_rows(matrix, "exclude") + if isinstance(dropped, Opaque): + return dropped + kept: Final = tuple( + combination for combination in crossed_values(listed) if not any(drops(row, combination) for row in dropped) + ) + standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in kept)) + return (*(extended(combination, rows) for combination in kept), *standalone) + + +def scanned(state: tuple[int, bool], char: str) -> tuple[int, bool]: + depth, quoted = state + if char == "'": + return depth, not quoted + if quoted: + return depth, quoted + return depth + int(char == "(") - int(char == ")"), quoted + + +def split_outside(text: str, token: str) -> tuple[str, ...]: + """`text` cut on every `token` that sits outside quotes and parentheses.""" + states: Final = tuple(itertools.accumulate(text, scanned, initial=(0, False))) + cuts: Final = tuple( + index + for index in range(len(text) - len(token) + 1) + if text.startswith(token, index) and states[index] == (0, False) + ) + starts: Final = (0, *(cut + len(token) for cut in cuts)) + return tuple(text[start:end] for start, end in zip(starts, (*cuts, len(text)))) + + +def formatted(template: str, arguments: Sequence[str]) -> str | None: + """GitHub's `format()` fills `{0}`-style holes and escapes braces, so anything richer resolves to nothing.""" + residue: Final = GITHUB_PLACEHOLDER.sub("", template) + if "{" in residue or "}" in residue: + return None + try: + return template.format(*arguments) + except (IndexError, KeyError, ValueError): + return None + + +def value_of(text: str, values: Mapping[str, str]) -> str | None: + expression: Final = text.strip() + literal: Final = LITERAL.match(expression) + if literal is not None: + return literal.group("text") + reference: Final = MATRIX_REF.match(expression) + if reference is not None: + return values.get(reference.group("key")) + call: Final = FORMAT_CALL.match(expression) + if call is None: + return None + arguments: Final = tuple(value_of(part, values) for part in split_outside(call.group("args"), ",")) + resolved: Final = tuple(argument for argument in arguments if argument is not None) + if not resolved or len(resolved) != len(arguments): + return None + return formatted(resolved[0], resolved[1:]) + + +def holds(condition: str, values: Mapping[str, str]) -> bool | None: + comparison: Final = COMPARISON.match(condition.strip()) + if comparison is None: + return None + left: Final = value_of(comparison.group("left"), values) + right: Final = value_of(comparison.group("right"), values) + if left is None or right is None: + return None + return (left == right) == (comparison.group("operator") == "==") + + +def evaluate(body: str, values: Mapping[str, str]) -> str | None: + """The single string this expression yields, or None when its shape is not understood.""" + branches: Final = tuple(split_outside(alternative, "&&") for alternative in split_outside(body, "||")) + outcomes: Final = tuple(tuple(holds(part, values) for part in branch[:-1]) for branch in branches) + if any(outcome is None for branch in outcomes for outcome in branch): + return None + taken: Final = next((branch[-1] for branch, outcome in zip(branches, outcomes) if all(outcome)), None) + return None if taken is None else value_of(taken, values) + + +def resolved_span(span: re.Match[str], values: Mapping[str, str]) -> str: + substitution: Final = evaluate(span.group("body"), values) + return span.group(0) if substitution is None else substitution + + +def rendered(template: str, values: Mapping[str, str]) -> str: + return EXPRESSION.sub(lambda span: resolved_span(span, values), template) + + +def comparable(name: str) -> bool: + """A name still holding an expression is one GitHub resolves per job, so it is nothing to compare.""" + return EXPRESSION.search(name) is None + + +def run_wide(name: str) -> bool: + """A name whose leftover expressions one workflow run fills in the same way for every job in it.""" + return all(RUN_WIDE.match(span.group("body").strip()) is not None for span in EXPRESSION.finditer(name)) + + +def settled(names: Sequence[str]) -> Names: + unresolved: Final = tuple(name for name in names if not comparable(name)) + return Names( + tuple(name for name in names if comparable(name)), + tuple(f"its name stays `{name}`" for name in unresolved if not run_wide(name)), + tuple(name for name in unresolved if run_wide(name)), + ) + + +def expand(template: str, job: Job) -> Names: + combinations: Final = matrix_combinations(job) + if isinstance(combinations, Opaque): + return Names((), (combinations.reason,)) + over: Final = combinations or (NO_MATRIX,) + return settled(tuple(rendered(template, values) for values in over)) + + +def suffixed(job_id: str, combination: Mapping[str, str]) -> str: + """The name GitHub gives a job with no `name:`, its id plus the combination it runs.""" + return f"{job_id} ({', '.join(combination.values())})" if combination else job_id + + +def published_names(job_id: str, job: Job) -> Names: + if job.name is not None: + return expand(scalar_text(job.name), job) + combinations: Final = matrix_combinations(job) + if isinstance(combinations, Opaque): + return Names((), (combinations.reason,)) + suffixes: Final = tuple(dict.fromkeys(suffixed(job_id, values) for values in combinations)) + return Names(suffixes or (job_id,)) + + +def callee_path(job: Job) -> str | None: + if job.uses is None or not job.uses.startswith(LOCAL_CALL_PREFIX): + return None + return job.uses[len(LOCAL_CALL_PREFIX) :].split("@")[0] + + +def joined(groups: Sequence[Names]) -> Names: + return Names( + tuple(name for group in groups for name in group.known), + tuple(reason for group in groups for reason in group.unknown), + tuple(name for group in groups for name in group.local), + ) + + +def tagged(names: Names) -> tuple[tuple[str, bool], ...]: + """Each name a job publishes beside whether only its own workflow's run settles it.""" + return (*((name, False) for name in names.known), *((name, True) for name in names.local)) + + +def call_blocker(job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str]) -> str | None: + path: Final = callee_path(job) + if path is None: + return "it calls a reusable workflow outside this repository" + if path in callers: + return f"its call to {path} loops back on itself" + return None if path in workflows else f"it calls {path}, which this checkout does not hold" + + +def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str] = NO_CALLERS) -> Names: + prefixes: Final = published_names(job_id, job) + if job.uses is None: + return prefixes + blocker: Final = call_blocker(job, workflows, callers) + if blocker is not None: + return Names((), (*prefixes.unknown, blocker)) + path: Final = callee_path(job) or "" + suffixes: Final = joined( + tuple( + job_names(callee_id, callee_job, workflows, callers | {path}) + for callee_id, callee_job in workflows[path].jobs.items() + ) + ) + composed: Final = tuple( + (f"{prefix} / {suffix}", prefix_local or suffix_local) + for prefix, prefix_local in tagged(prefixes) + for suffix, suffix_local in tagged(suffixes) + ) + return Names( + tuple(name for name, is_local in composed if not is_local), + (*prefixes.unknown, *suffixes.unknown), + tuple(name for name, is_local in composed if is_local), + ) + + +def readable(sources: Mapping[str, str]) -> Mapping[str, tuple[Workflow, object]]: + parsed: Final = {rel: parse(source) for rel, source in sources.items()} + return MappingProxyType({rel: entry for rel, entry in parsed.items() if not isinstance(entry, Unreadable)}) + + +def unreadable(sources: Mapping[str, str]) -> tuple[str, ...]: + parsed: Final = {rel: parse(source) for rel, source in sources.items()} + return tuple( + f"{rel} sits in the workflows directory but {entry.reason}, so none of its jobs were checked." + for rel, entry in sorted(parsed.items()) + if isinstance(entry, Unreadable) + ) + + +def scanned_jobs(sources: Mapping[str, str]) -> Iterator[tuple[str, str, Names]]: + parsed: Final = readable(sources) + workflows: Final = {rel: workflow for rel, (workflow, _) in parsed.items()} + for rel, (workflow, raw_on) in parsed.items(): + if not publishes_check_runs(raw_on): + continue + for job_id, job in workflow.jobs.items(): + yield rel, job_id, job_names(job_id, job, workflows) + + +def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]: + for rel, job_id, names in scanned_jobs(sources): + for name in names.known: + yield name, f"{rel} job `{job_id}`" + + +def blind_spots(sources: Mapping[str, str]) -> tuple[str, ...]: + """Jobs whose published names GitHub decides at run time, which no offline sweep can compare.""" + return tuple( + f"{rel} job `{job_id}` publishes a name this check cannot work out because {reason}." + for rel, job_id, names in scanned_jobs(sources) + for reason in sorted(names.unknown) + ) + + +def owners_by_name(sources: Mapping[str, str]) -> Iterator[tuple[str, tuple[str, ...]]]: + for name, pairs in itertools.groupby(sorted(published(sources)), key=operator.itemgetter(0)): + yield name, tuple(owner for _, owner in pairs) + + +def clash(name: str, owners: Sequence[str]) -> str | None: + """Why one name is ambiguous, whether two jobs carry it or one job repeats it over its matrix.""" + jobs: Final = tuple(dict.fromkeys(owners)) + if len(jobs) > 1: + return ( + f"`{name}` is published by {len(jobs)} jobs: {', '.join(jobs)}. A required status check matching " + f"that name cannot say which job proves it; give one of them a distinct `name:` or job id." + ) + if len(owners) > 1: + return ( + f"`{name}` is published {len(owners)} times by {jobs[0]}, once per matrix combination. A required " + f"status check matching that name cannot say which run proves it; put a matrix value in its `name:`." + ) + return None + + +def local_published(sources: Mapping[str, str]) -> Iterator[tuple[tuple[str, str], str]]: + """Names their own workflow's run settles, keyed by the file whose run settles them.""" + for rel, job_id, names in scanned_jobs(sources): + for name in names.local: + yield (rel, name), f"job `{job_id}`" + + +def local_clash(rel: str, name: str, owners: Sequence[str]) -> str | None: + """Why one workflow's own run lands several of its jobs on one check run.""" + if len(owners) < 2: + return None + jobs: Final = tuple(dict.fromkeys(owners)) + return ( + f"`{name}` is published {len(owners)} times inside {rel}, by {', '.join(jobs)}. One run fills that " + f"expression in the same way throughout, so they all land on one check run; make the names differ." + ) + + +def local_clashes(sources: Mapping[str, str]) -> tuple[str, ...]: + grouped: Final = itertools.groupby(sorted(local_published(sources)), key=operator.itemgetter(0)) + found: Final = tuple(local_clash(rel, name, tuple(owner for _, owner in pairs)) for (rel, name), pairs in grouped) + return tuple(message for message in found if message is not None) + + +def collisions(sources: Mapping[str, str]) -> tuple[str, ...]: + found: Final = tuple(clash(name, owners) for name, owners in owners_by_name(sources)) + return (*(message for message in found if message is not None), *local_clashes(sources)) + + +def workflow_sources() -> Mapping[str, str]: + """Repo-relative posix paths to text, the keys `uses: ./...` resolves against.""" + return {path.relative_to(REPO_ROOT).as_posix(): path.read_text() for path in sorted(WORKFLOWS_DIR.glob("*.y*ml"))} + + +def report(header: str, problems: Sequence[str]) -> None: + if problems: + print(f"ERROR: {header}:\n - " + "\n - ".join(problems), file=sys.stderr) + + +def exit_code(sources: Mapping[str, str]) -> int: + unread: Final = unreadable(sources) + found: Final = collisions(sources) + blind: Final = blind_spots(sources) + if blind: + print("NOTE: names left out of the comparison:\n - " + "\n - ".join(blind)) + report("Some workflows could not be read", unread) + report("Check-run names are not unique", found) + if unread or found: + return 1 + + print(f"Check-run names are unique across {len(sources)} workflows") + return 0 + + +def main() -> int: + return exit_code(workflow_sources()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 0af29f069c6..a11f015743b 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -87,6 +87,7 @@ ignored_function_names = [ "_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py "_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py + "_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name) ] diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py new file mode 100644 index 00000000000..d2b842364a6 --- /dev/null +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -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", ("", "')) +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() diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py new file mode 100644 index 00000000000..0f5ba43bd7a --- /dev/null +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -0,0 +1,880 @@ +from typing import Final + +from check_workflow_job_name_collisions import ( + Unreadable, + blind_spots, + callee_path, + collisions, + exit_code, + parse, + published, + unreadable, + workflow_sources, +) + +REUSABLE_BASE: Final = """on: + workflow_call: +jobs: + run: + name: >- + ${{ matrix.python-version == '3.12' && 'Run tests' + || format('Run tests (Python {0})', matrix.python-version) }} + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12", "3.13"] +""" + +SHARD_CALLER: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.shard }} + uses: ./.github/workflows/base.yml + strategy: + matrix: + include: + - shard: core-utils +""" + + +CORRELATED_ROWS: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.shard }} on ${{ matrix.test-path }} + runs-on: ubuntu-latest + strategy: + matrix: + include: + - shard: core-utils + test-path: tests/core + - shard: proxy + test-path: tests/proxy +""" + +LISTED_PLUS_ROW: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.python-version }} ${{ matrix.label }} + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12", "3.13"] + include: + - label: fast +""" + +NAMELESS_MATRIX: Final = """on: pull_request +jobs: + build: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12", "3.13"] +""" + +EXCLUDED_PAIR: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.os }}-${{ matrix.python-version }} + runs-on: ubuntu-latest + strategy: + matrix: + os: [ubuntu, macos] + python-version: ["3.12", "3.13"] + exclude: + - os: macos + python-version: "3.13" +""" + +EXCLUDED_KEY: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.os }}-${{ matrix.python-version }} + runs-on: ubuntu-latest + strategy: + matrix: + os: [ubuntu, macos] + python-version: ["3.12", "3.13"] + exclude: + - os: macos +""" + +BOOLEAN_MATRIX: Final = """on: pull_request +jobs: + unit: + name: cache ${{ matrix.cached }} + runs-on: ubuntu-latest + strategy: + matrix: + cached: [true, false] +""" + +UNFILLABLE_FORMAT: Final = """on: pull_request +jobs: + unit: + name: ${{ format('{0} {1}', matrix.shard) }} + runs-on: ubuntu-latest + strategy: + matrix: + shard: [core] +""" + + +def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None: + assert collisions(workflow_sources()) == () + + +def test_every_workflow_in_the_repo_parses_into_jobs() -> None: + unparsed: Final = tuple( + rel + for rel, source in workflow_sources().items() + if isinstance(entry := parse(source), Unreadable) or not entry[0].jobs + ) + + assert unparsed == () + + +def test_every_local_reusable_call_in_the_repo_resolves_to_a_workflow() -> None: + sources: Final = workflow_sources() + parsed: Final = tuple(entry for text in sources.values() if not isinstance(entry := parse(text), Unreadable)) + unresolved: Final = tuple( + job.uses + for workflow, _ in parsed + for job in workflow.jobs.values() + if (callee := callee_path(job)) is not None and callee not in sources + ) + + assert unresolved == () + + +def test_two_jobs_falling_back_to_the_same_job_id_collide() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`test` is published by 2 jobs" in found[0] + assert "a.yml job `test`" in found[0] and "b.yml job `test`" in found[0] + + +def test_an_explicit_name_overrides_the_job_id_and_clears_the_collision() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n test:\n name: Sweep tests\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + + +def test_an_explicit_name_matching_another_job_id_collides() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n sweep:\n name: test\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`test` is published by 2 jobs" in found[0] + + +def test_two_callers_of_one_reusable_workflow_collide_on_a_shared_matrix_value() -> None: + base: Final = "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n" + caller: Final = ( + "on: pull_request\n" + "jobs:\n" + " {job}:\n" + " name: ${{{{ matrix.shard }}}}\n" + " uses: ./.github/workflows/base.yml\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: {shard}\n" + ) + sources: Final = { + ".github/workflows/base.yml": base, + "unit.yml": caller.format(job="unit", shard="proxy-auth"), + "proxy-db.yml": caller.format(job="proxy-db", shard="proxy-auth"), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`proxy-auth / Run tests` is published by 2 jobs" in found[0] + + +def test_distinct_matrix_values_through_one_reusable_workflow_do_not_collide() -> None: + base: Final = "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n" + caller: Final = ( + "on: pull_request\n" + "jobs:\n" + " {job}:\n" + " name: ${{{{ matrix.shard }}}}\n" + " uses: ./.github/workflows/base.yml\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: {shard}\n" + ) + sources: Final = { + ".github/workflows/base.yml": base, + "unit.yml": caller.format(job="unit", shard="proxy-auth"), + "proxy-db.yml": caller.format(job="proxy-db", shard="budgets"), + } + + assert collisions(sources) == () + + +def test_a_reusable_caller_does_not_collide_with_a_plain_job_of_the_same_name() -> None: + sources: Final = { + ".github/workflows/base.yml": ( + "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n" + ), + "unit.yml": ( + "on: pull_request\n" + "jobs:\n" + " unit:\n" + " name: ${{ matrix.shard }}\n" + " uses: ./.github/workflows/base.yml\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: proxy-behavior\n" + ), + "postgres.yml": ( + "on: pull_request\n" + "jobs:\n" + " postgres:\n" + " name: ${{ matrix.shard }}\n" + " runs-on: ubuntu-latest\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: proxy-behavior\n" + ), + } + + assert collisions(sources) == () + + +def test_a_workflow_call_only_workflow_publishes_nothing_of_its_own() -> None: + sources: Final = { + "base.yml": "on:\n workflow_call:\njobs:\n run:\n runs-on: ubuntu-latest\n", + "other.yml": "on:\n workflow_call:\njobs:\n run:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + + +def test_a_workflow_call_workflow_that_also_runs_on_pull_request_still_publishes() -> None: + sources: Final = { + "base.yml": "on:\n workflow_call:\n pull_request:\njobs:\n run:\n runs-on: ubuntu-latest\n", + "other.yml": "on: pull_request\njobs:\n run:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`run` is published by 2 jobs" in found[0] + + +def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\n" + "jobs:\n" + " build:\n" + " name: Analyze (${{ matrix.language }})\n" + " runs-on: ubuntu-latest\n" + " strategy:\n" + " matrix:\n" + " language: [python, go]\n" + ), + "b.yml": "on: pull_request\njobs:\n go:\n name: Analyze (go)\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`Analyze (go)` is published by 2 jobs" in found[0] + + +def test_two_workflows_sharing_a_run_wide_template_are_not_called_a_collision() -> None: + template: Final = ( + "on: pull_request\njobs:\n {job}:\n name: ${{{{ github.event_name }}}}-build\n runs-on: ubuntu-latest\n" + ) + sources: Final = { + "a.yml": template.format(job="one"), + "b.yml": template.format(job="two"), + } + + assert collisions(sources) == () + assert blind_spots(sources) == () + + +def test_two_jobs_of_one_workflow_sharing_a_run_wide_template_are_a_collision() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n" + " two:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n" + ), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "is published 2 times inside a.yml, by job `one`, job `two`" in found[0] + assert exit_code(sources) == 1 + + +def test_a_run_wide_template_carrying_a_matrix_value_does_not_collide_inside_one_workflow() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}-${{ matrix.shard }}\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n shard: [core, extras]\n" + ), + } + + assert collisions(sources) == () + assert blind_spots(sources) == () + + +def test_a_run_wide_name_repeated_over_a_matrix_by_one_job_is_a_collision() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n shard: [core, extras]\n" + ), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "is published 2 times inside a.yml, by job `one`" in found[0] + + +def test_a_name_reading_the_job_it_sits_in_stays_out_of_the_comparison() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.job }}-build\n runs-on: ubuntu-latest\n" + " two:\n name: ${{ github.job }}-build\n runs-on: ubuntu-latest\n" + ), + } + + assert collisions(sources) == () + assert len(blind_spots(sources)) == 2 + + +def test_a_run_wide_caller_name_collides_through_the_workflow_it_calls() -> None: + sources: Final = { + ".github/workflows/a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}\n uses: ./.github/workflows/c.yml\n" + " two:\n name: ${{ github.event_name }}\n uses: ./.github/workflows/c.yml\n" + ), + ".github/workflows/c.yml": "on:\n workflow_call:\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "github.event_name }} / build` is published 2 times inside .github/workflows/a.yml" in found[0] + + +def test_a_run_wide_name_inside_a_called_workflow_collides_under_the_caller() -> None: + sources: Final = { + ".github/workflows/a.yml": ("on: pull_request\njobs:\n one:\n uses: ./.github/workflows/c.yml\n"), + ".github/workflows/c.yml": ( + "on:\n workflow_call:\njobs:\n" + " build:\n name: ${{ github.event_name }}\n runs-on: ubuntu-latest\n" + " lint:\n name: ${{ github.event_name }}\n runs-on: ubuntu-latest\n" + ), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`one / ${{ github.event_name }}` is published 2 times inside .github/workflows/a.yml" in found[0] + + +def test_a_name_reading_the_workflow_it_sits_in_is_not_called_a_collision() -> None: + template: Final = ( + "on: pull_request\njobs:\n {job}:\n name: ${{{{ github.workflow }}}} / build\n runs-on: ubuntu-latest\n" + ) + sources: Final = { + "a.yml": template.format(job="one"), + "b.yml": template.format(job="two"), + } + + assert collisions(sources) == () + assert exit_code(sources) == 0 + + +def test_a_format_call_python_accepts_but_github_does_not_publishes_nothing_to_compare() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: ${{ format('{0.real}', matrix.shard) }}\n" + " runs-on: ubuntu-latest\n strategy:\n matrix:\n shard: [core]\n" + ), + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_a_format_call_padding_its_argument_publishes_nothing_to_compare() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: ${{ format('{0:>8}', matrix.shard) }}\n" + " runs-on: ubuntu-latest\n strategy:\n matrix:\n shard: [core]\n" + ), + "b.yml": "on: pull_request\njobs:\n two:\n name: ' core'\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_an_exclude_row_that_is_not_a_mapping_is_reported_rather_than_skipped() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n v: [1, 2]\n exclude:\n - oops\n" + ), + "b.yml": "on: pull_request\njobs:\n other:\n name: build (1)\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_an_exclude_row_holding_a_non_scalar_never_drops_every_combination() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n v: [1, 2]\n exclude:\n - cfg: {k: 1}\n" + ), + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_two_jobs_sharing_a_template_that_reads_per_job_are_not_called_a_collision() -> None: + template: Final = ( + "on: pull_request\njobs:\n {job}:\n name: ${{{{ matrix.shard }}}}\n runs-on: ubuntu-latest\n" + ) + sources: Final = { + "a.yml": template.format(job="one"), + "b.yml": template.format(job="two"), + } + + assert collisions(sources) == () + assert len(blind_spots(sources)) == 2 + + +def test_a_file_that_is_not_a_workflow_is_reported_rather_than_skipped() -> None: + sources: Final = { + "notes.yml": "just a string\n", + "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + found: Final = unreadable(sources) + + assert len(found) == 1 + assert "notes.yml" in found[0] + assert collisions(sources) == () + + +def test_a_workflow_holding_a_job_shape_github_would_reject_is_reported() -> None: + sources: Final = {"a.yml": "on: pull_request\njobs:\n test:\n uses: [not, a, string]\n"} + + found: Final = unreadable(sources) + + assert len(found) == 1 + assert "a.yml" in found[0] + + +def test_a_conditional_name_expands_to_the_branch_each_matrix_value_takes() -> None: + names: Final = frozenset( + name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER}) + ) + + assert names == frozenset({"core-utils / Run tests", "core-utils / Run tests (Python 3.13)"}) + + +def test_a_conditional_name_never_publishes_the_branch_its_condition_rules_out() -> None: + names: Final = frozenset( + name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER}) + ) + + assert "core-utils / Run tests (Python 3.12)" not in names + + +def test_a_conditional_reusable_name_collides_with_a_plain_job_publishing_the_same_name() -> None: + sources: Final = { + ".github/workflows/base.yml": REUSABLE_BASE, + "unit.yml": SHARD_CALLER, + "postgres.yml": ("on: pull_request\njobs:\n legacy:\n name: core-utils / Run tests\n"), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`core-utils / Run tests` is published by 2 jobs" in found[0] + + +def test_a_name_reading_two_matrix_keys_publishes_only_the_pairs_each_include_row_holds() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS})) + + assert names == frozenset({"core-utils on tests/core", "proxy on tests/proxy"}) + + +def test_a_name_reading_two_matrix_keys_never_publishes_a_pair_no_include_row_holds() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS})) + + assert "core-utils on tests/proxy" not in names + assert "proxy on tests/core" not in names + + +def test_an_include_row_carrying_no_listed_key_extends_every_listed_combination() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": LISTED_PLUS_ROW})) + + assert names == frozenset({"3.12 fast", "3.13 fast"}) + + +def test_every_workflow_in_the_repo_resolves_every_expression_in_its_job_names() -> None: + unresolved: Final = tuple(f"{owner}: {name}" for name, owner in published(workflow_sources()) if "${{" in name) + + assert unresolved == () + + +def test_a_matrix_job_with_no_name_publishes_the_id_and_values_github_appends() -> None: + names: Final = frozenset(name for name, _ in published({"a.yml": NAMELESS_MATRIX})) + + assert names == frozenset({"build (3.12)", "build (3.13)"}) + + +def test_a_matrix_job_with_no_name_does_not_collide_with_a_plain_job_carrying_its_id() -> None: + sources: Final = { + "a.yml": NAMELESS_MATRIX, + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + + +def test_a_matrix_job_with_no_name_collides_with_the_suffixed_name_github_writes() -> None: + sources: Final = { + "a.yml": NAMELESS_MATRIX, + "b.yml": "on: pull_request\njobs:\n legacy:\n name: build (3.13)\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`build (3.13)` is published by 2 jobs" in found[0] + + +def test_an_excluded_combination_publishes_no_check_run() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_PAIR})) + + assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13", "macos-3.12"}) + + +def test_an_exclude_row_naming_one_key_drops_every_combination_carrying_it() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_KEY})) + + assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13"}) + + +def test_a_boolean_matrix_value_renders_the_way_github_writes_it() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": BOOLEAN_MATRIX})) + + assert names == frozenset({"cache true", "cache false"}) + + +def test_a_format_call_its_arguments_cannot_fill_publishes_nothing_to_compare() -> None: + sources: Final = {"unit.yml": UNFILLABLE_FORMAT} + + assert frozenset(name for name, _ in published(sources)) == frozenset() + assert "its name stays" in blind_spots(sources)[0] + + +def test_a_call_to_a_workflow_outside_the_repo_is_reported_rather_than_guessed() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n unit:\n uses: BerriAI/other/.github/workflows/base.yml@main\n", + "b.yml": "on: pull_request\njobs:\n unit:\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"unit"}) + assert collisions(sources) == () + assert "outside this repository" in blind_spots(sources)[0] + + +def test_a_chain_of_local_reusable_calls_publishes_every_level_of_the_chain() -> None: + sources: Final = { + ".github/workflows/leaf.yml": ( + "on:\n workflow_call:\njobs:\n run:\n name: Leaf\n runs-on: ubuntu-latest\n" + ), + ".github/workflows/mid.yml": ( + "on:\n workflow_call:\njobs:\n call:\n name: Mid\n uses: ./.github/workflows/leaf.yml\n" + ), + "top.yml": "on: pull_request\njobs:\n top:\n name: Top\n uses: ./.github/workflows/mid.yml\n", + } + + names: Final = frozenset(name for name, _ in published(sources)) + + assert names == frozenset({"Top / Mid / Leaf"}) + + +def test_a_job_name_that_is_not_a_string_still_publishes_the_value_github_renders() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n sweep:\n name: 2024\n runs-on: ubuntu-latest\n", + "b.yml": 'on: pull_request\njobs:\n other:\n name: "2024"\n runs-on: ubuntu-latest\n', + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`2024` is published by 2 jobs" in found[0] + + +def test_the_check_fails_when_a_file_in_the_workflows_directory_cannot_be_read() -> None: + assert exit_code({"notes.yml": "just a string\n"}) == 1 + + +def test_the_check_fails_when_two_jobs_publish_one_check_run_name() -> None: + plain: Final = "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n" + + assert exit_code({"a.yml": plain, "b.yml": plain}) == 1 + + +def test_the_check_passes_when_every_file_reads_and_every_name_is_unique() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n sweep:\n runs-on: ubuntu-latest\n", + } + + assert exit_code(sources) == 0 + + +def test_two_callers_of_one_reusable_workflow_named_from_its_inputs_do_not_collide() -> None: + sources: Final = { + ".github/workflows/callee.yml": ( + "on:\n workflow_call:\njobs:\n run:\n name: ${{ inputs.suite }}\n runs-on: ubuntu-latest\n" + ), + "caller.yml": ( + "on: pull_request\njobs:\n" + " alpha:\n name: A\n uses: ./.github/workflows/callee.yml\n with:\n suite: alpha\n" + " beta:\n name: A\n uses: ./.github/workflows/callee.yml\n with:\n suite: beta\n" + ), + } + + assert collisions(sources) == () + assert len(blind_spots(sources)) == 2 + + +def test_a_matrix_that_is_itself_an_expression_never_collapses_onto_the_bare_job_id() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n strategy:\n" + " matrix: ${{ fromJson(needs.plan.outputs.matrix) }}\n runs-on: ubuntu-latest\n" + ), + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"build"}) + assert collisions(sources) == () + assert "the matrix itself comes from an expression" in blind_spots(sources)[0] + + +def test_a_matrix_listing_objects_never_collapses_onto_the_bare_job_id() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n strategy:\n matrix:\n target:\n" + " - os: ubuntu\n - os: windows\n runs-on: ubuntu-latest\n" + ), + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"build"}) + assert collisions(sources) == () + assert "not plain scalars" in blind_spots(sources)[0] + + +def test_a_call_to_a_workflow_file_the_checkout_does_not_hold_is_reported() -> None: + sources: Final = {"a.yml": "on: pull_request\njobs:\n unit:\n uses: ./.github/workflows/gone.yml\n"} + + assert collisions(sources) == () + assert "which this checkout does not hold" in blind_spots(sources)[0] + + +def test_reusable_workflows_calling_each_other_in_a_loop_are_reported_not_followed() -> None: + sources: Final = { + ".github/workflows/a.yml": ( + "on:\n workflow_call:\njobs:\n call:\n name: A\n uses: ./.github/workflows/b.yml\n" + ), + ".github/workflows/b.yml": ( + "on:\n workflow_call:\njobs:\n call:\n name: B\n uses: ./.github/workflows/a.yml\n" + ), + "top.yml": "on: pull_request\njobs:\n top:\n name: Top\n uses: ./.github/workflows/a.yml\n", + } + + assert collisions(sources) == () + assert any("loops back on itself" in spot for spot in blind_spots(sources)) + + +def test_a_caller_still_publishes_the_callee_jobs_it_can_read() -> None: + sources: Final = { + ".github/workflows/callee.yml": ( + "on:\n workflow_call:\njobs:\n" + " lint:\n name: Lint\n runs-on: ubuntu-latest\n" + " suite:\n name: ${{ inputs.suite }}\n runs-on: ubuntu-latest\n" + ), + "caller.yml": "on: pull_request\njobs:\n call:\n name: A\n uses: ./.github/workflows/callee.yml\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"A / Lint"}) + assert len(blind_spots(sources)) == 1 + + +def test_a_name_the_check_cannot_work_out_is_reported_without_failing_the_check() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n unit:\n uses: BerriAI/other/.github/workflows/base.yml@main\n", + } + + assert blind_spots(sources) != () + assert exit_code(sources) == 0 + + +def test_a_caller_whose_own_name_is_unreadable_publishes_none_of_its_callee_names() -> None: + sources: Final = { + ".github/workflows/callee.yml": ( + "on:\n workflow_call:\njobs:\n lint:\n name: Lint\n runs-on: ubuntu-latest\n" + ), + "caller.yml": ( + "on: pull_request\njobs:\n call:\n name: ${{ matrix.suite }}\n" + " uses: ./.github/workflows/callee.yml\n" + ), + "other.yml": "on: pull_request\njobs:\n plain:\n name: Lint\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"Lint"}) + assert collisions(sources) == () + assert "its name stays" in blind_spots(sources)[0] + + +def test_an_include_row_naming_a_listed_key_extends_only_the_combinations_it_matches() -> None: + sources: Final = { + "unit.yml": ( + "on: pull_request\njobs:\n unit:\n" + " name: ${{ matrix.python-version }} ${{ matrix.label }}\n" + " runs-on: ubuntu-latest\n strategy:\n matrix:\n" + ' python-version: ["3.12", "3.13"]\n' + " include:\n" + ' - python-version: "3.12"\n' + " label: fast\n" + ) + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"3.12 fast"}) + assert len(blind_spots(sources)) == 1 + + +def test_a_job_whose_whole_strategy_is_an_expression_is_reported_rather_than_rejecting_the_file() -> None: + sources: Final = { + "plan.yml": ( + "on: pull_request\njobs:\n plan:\n name: Plan\n runs-on: ubuntu-latest\n" + " fan:\n strategy: ${{ fromJSON(needs.plan.outputs.strategy) }}\n runs-on: ubuntu-latest\n" + ) + } + + assert unreadable(sources) == () + assert frozenset(name for name, _ in published(sources)) == frozenset({"Plan"}) + assert "`strategy` comes from an expression" in blind_spots(sources)[0] + assert exit_code(sources) == 0 + + +def test_one_job_publishing_one_name_for_every_matrix_combination_is_a_collision() -> None: + sources: Final = { + "unit.yml": ( + "on: pull_request\njobs:\n build:\n name: Run tests\n runs-on: ubuntu-latest\n" + ' strategy:\n matrix:\n python-version: ["3.12", "3.13"]\n' + ) + } + + found: Final = collisions(sources) + assert len(found) == 1 + assert "`Run tests` is published 2 times by unit.yml job `build`" in found[0] + assert exit_code(sources) == 1 + + +def test_a_name_carrying_a_matrix_value_publishes_one_name_per_combination_without_colliding() -> None: + sources: Final = { + "unit.yml": ( + "on: pull_request\njobs:\n build:\n name: Run tests ${{ matrix.python-version }}\n" + ' runs-on: ubuntu-latest\n strategy:\n matrix:\n python-version: ["3.12", "3.13"]\n' + ) + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"Run tests 3.12", "Run tests 3.13"}) + assert collisions(sources) == () + assert exit_code(sources) == 0 + + +def test_a_file_that_is_not_valid_yaml_is_reported_rather_than_raising() -> None: + sources: Final = {"broken.yml": "jobs:\n build: [\n"} + + assert unreadable(sources) == ( + "broken.yml sits in the workflows directory but it does not read as one YAML " + "document, so none of its jobs were checked.", + ) + assert exit_code(sources) == 1 + + +def test_a_file_holding_two_yaml_documents_is_reported_rather_than_raising() -> None: + sources: Final = {"two.yml": "on: pull_request\n---\non: push\n"} + + assert len(unreadable(sources)) == 1 + assert exit_code(sources) == 1 + + +def test_an_exclude_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: build\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n python: ['3.11', '3.12']\n" + " exclude: ${{ fromJson(vars.SKIP) }}\n" + ), + } + + found: Final = blind_spots(sources) + + assert collisions(sources) == () + assert len(found) == 1 + assert "a matrix `exclude` is itself an expression" in found[0] + + +def test_an_include_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: build-${{ matrix.python }}\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n python: ['3.11']\n" + " include: ${{ fromJson(vars.EXTRA) }}\n" + ), + } + + found: Final = blind_spots(sources) + + assert collisions(sources) == () + assert len(found) == 1 + assert "a matrix `include` is itself an expression" in found[0] diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index c954c03dfb7..2945daae17c 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -181,13 +181,15 @@ quota_management... | team_multi_window | fallback | spend_counter chat_completions | stream | messages_bridge | embeddings | cache_hit | key_rollup | concurrent_burst | tags | end_user - | per_model | failure | spend_calculate | pagination + | per_model | failure | spend_calculate | pagination | key_attribution assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm | blocks_then_resets | resets_windows_independently | alerts_without_blocking | isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys | routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows - | writes_failure_row | returns_cost | keeps_total + | writes_failure_row | returns_cost | keeps_total | joins_key | reports_alias_and_email + | health_rows_keep_service_account | retrieve_batch_cost_joins_retrieving_key + | poller_batch_cost_joins_creating_key e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages] quota_management.budget.key.blocks_over_limit exercised_on=[chat_completions] ``` diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 3383ce79cf8..53fcdbfe381 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -88,8 +88,31 @@ The suites run against a live proxy, so bring one up first by running the litell Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy. The presidio guardrail tests need a running Presidio analyzer and anonymizer the proxy can reach, addressed by `PRESIDIO_ANALYZER_API_BASE` / `PRESIDIO_ANONYMIZER_API_BASE` +The Presidio spend-log audit test also requires prompt storage on the proxy, because its assertions inspect the detected entities and their scores. Add the following to the proxy config before starting it: + +```yaml +general_settings: + store_prompts_in_spend_logs: true +``` + +Alternatively, start the proxy with `STORE_PROMPTS_IN_SPEND_LOGS=true litellm --config .yml --port 4000`. Setting the variable only on the pytest process does not configure the proxy. With prompt storage disabled, the proxy correctly redacts `guardrail_response`, so that configuration cannot exercise this test's entity-detail assertions. Enable this only on a test stack using synthetic prompts + 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 diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 8a7b68511ec..919c39f21a2 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -120,6 +120,36 @@ create traverse gateway -> gateway -> OpenAI (LIT-5347, PR #36240). The pin: nested managed ids round-trip retrieve. This self-chaining only needs the proxy to reach its own `PROXY_BASE_URL`, which holds both locally and on the e2e stage. +## Cleanup + +Batch teardown cancels active batches before deleting their input files and keys. +Raw file IDs from both `model_param` and `provider_fallback` uploads use the upload +provider when deleted. Model-encoded and managed file IDs route themselves + +File deletion and batch cancellation check their responses and retry transient +failures up to three times. Teardown attempts every registered cleanup before +reporting failures as test errors. Already deleted files and batches that are +terminal are safe to clean up again. Managed batch cancellation polls for up to eleven minutes +before input deletion: the ten-minute provider window plus a propagation margin. +Accepted cancellation may still report validating or in_progress while the provider +updates its state. Raw and model-encoded batches are polled until cancelling or +terminal before input deletion. OpenAI and Azure lifecycle cleanup also deletes +output and error files returned by terminal batches. Bedrock deletion uses a signed S3 DELETE +restricted to the configured storage buckets and managed file prefixes. The low-RPM +test submits with its restricted key and cleans up with the test administrator key + +Managed deletion forwards the deployment's trusted bucket configuration and returns +the requested managed file ID even when stored output metadata carries a provider ID + +Azure input uploads request `expires_after` anchored to `created_at` with +`seconds=1209600`, and the lifecycle tests check the returned expiry. This is a +fallback for interrupted runs: immediate deletion remains the normal cleanup. +Azure's minimum supported native expiry is 14 days, so a three-day expiry cannot +be requested through its Files API + +The Azure entry in `files_settings` must use `api_version: 2025-04-01-preview` +for raw uploads to honor expiry, matching the batch deployment's API version + ## Terminal state + cost write-back (cross-run marker baton) The 24h completion window rules out submit-and-wait inside one run, so diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py new file mode 100644 index 00000000000..9284882ad82 --- /dev/null +++ b/tests/e2e/batches/batch_cleanup.py @@ -0,0 +1,140 @@ +from builtins import ExceptionGroup +from collections.abc import Callable +from itertools import count +from time import monotonic, sleep +from typing import Final, Protocol + +from batch_client import BatchObject, FileDeleteResponse +from capabilities import is_managed_id +from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError +from pydantic import BaseModel + +CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0) +BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"}) +BATCH_PENDING_STATUSES: Final = frozenset({"validating", "in_progress", "finalizing", "cancelling"}) +BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0 +BATCH_CANCEL_POLL_SECONDS: Final = 10.0 + + +class BatchCleanupClient(Protocol): + def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: ... + + def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ... + + def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ... + + +def cleanup_result[R: BaseModel]( + action: Callable[[], Result[R]], *, wait: Callable[[float], None] = sleep +) -> Result[R]: + for delay, result in ((delay, action()) for delay in CLEANUP_DELAYS): + match result: + case NetworkError() | RateLimitedError(): + wait(delay) + case UnknownApiError(status_code=code) if code in {408, 429, 500, 502, 503, 504}: + wait(delay) + case _: + return result + return action() + + +def _require_cleanup_success[R: BaseModel](result: Result[R], operation: str) -> R: + match result: + case Success(data=data): + return data + case UnknownApiError(status_code=code): + raise AssertionError(f"{operation} failed: HTTP {code}") + case _: + raise AssertionError(f"{operation} failed: {result.kind}") + + +def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None = None) -> None: + result: Final = cleanup_result(lambda: client.delete_file(file_id, key=key, provider=provider)) + if isinstance(result, UnknownApiError) and result.status_code == 404: + return + deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}") + assert deleted.deleted is True or ( + deleted.deleted is None and is_managed_id(file_id) and deleted.id == file_id and deleted.object == "file" + ), f"Delete file {file_id} did not confirm deletion" + + +def cleanup_batch( + client: BatchCleanupClient, + batch_id: str, + *, + key: str, + provider: str | None = None, + delete_output_files: bool = False, + wait: Callable[[float], None] = sleep, + clock: Callable[[], float] = monotonic, +) -> None: + needs_terminal_state: Final = is_managed_id(batch_id) + fetched: Final = _require_cleanup_success( + cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), + f"Retrieve batch {batch_id} for cleanup", + ) + if fetched.status in BATCH_TERMINAL_STATUSES: + if delete_output_files: + _cleanup_batch_outputs(client, fetched, key=key, provider=provider) + return + if fetched.status == "cancelling" and not needs_terminal_state: + return + result: Final = ( + Success(status_code=200, data=fetched) + if fetched.status == "cancelling" + else cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider)) + ) + conflicted: Final = isinstance(result, UnknownApiError) and result.status_code in {400, 409} + if not conflicted: + cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}") + assert cancelled.status in BATCH_TERMINAL_STATUSES | BATCH_PENDING_STATUSES, ( + f"Cancel batch {batch_id} left status {cancelled.status}" + ) + if cancelled.status in BATCH_TERMINAL_STATUSES: + if delete_output_files: + _cleanup_batch_outputs(client, cancelled, key=key, provider=provider) + return + if cancelled.status == "cancelling" and not needs_terminal_state: + return + deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS + for current in ( + _require_cleanup_success( + cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), + f"Retrieve batch {batch_id} after cancellation", + ) + for _ in count() + ): + if current.status in BATCH_TERMINAL_STATUSES: + if delete_output_files: + _cleanup_batch_outputs(client, current, key=key, provider=provider) + return + assert current.status in ({"cancelling"} if conflicted else BATCH_PENDING_STATUSES), ( + f"Cancel batch {batch_id} left status {current.status}" + ) + if current.status == "cancelling" and not needs_terminal_state: + return + assert clock() < deadline, ( + f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s" + ) + wait(BATCH_CANCEL_POLL_SECONDS) + + +def _cleanup_batch_outputs(client: BatchCleanupClient, batch: BatchObject, *, key: str, provider: str | None) -> None: + errors: Final = tuple( + error + for file_id in dict.fromkeys((batch.output_file_id, batch.error_file_id)) + if file_id is not None and file_id != batch.input_file_id + if (error := _output_cleanup_error(client, file_id, key=key, provider=provider)) is not None + ) + if errors: + raise ExceptionGroup(f"Batch {batch.id} output cleanup failed", errors) + + +def _output_cleanup_error( + client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None +) -> Exception | None: + try: + cleanup_file(client, file_id, key=key, provider=provider) + except Exception as error: + return error + return None diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 31e49f22450..c9c77e1f12e 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -13,8 +13,9 @@ co-located here because only this suite uses them. from __future__ import annotations from dataclasses import dataclass +from typing import Final, Literal -from pydantic import BaseModel +from pydantic import BaseModel, Field from proxy_client import ProxyClient from e2e_http import ( @@ -27,6 +28,18 @@ from e2e_http import ( from models import LiteLLMParamsBody UPLOAD_FILENAME = "batch_input.jsonl" +AZURE_FILE_EXPIRY_SECONDS: Final = 14 * 24 * 60 * 60 + + +class ExpiringFileUploadForm(FileUploadForm): + expires_after_anchor: Literal["created_at"] = Field(default="created_at", alias="expires_after[anchor]") + expires_after_seconds: int = Field(default=AZURE_FILE_EXPIRY_SECONDS, alias="expires_after[seconds]") + + +def batch_upload_form(provider: str, *, target_model_names: str | None = None) -> FileUploadForm: + if provider == "azure": + return ExpiringFileUploadForm(target_model_names=target_model_names) + return FileUploadForm(target_model_names=target_model_names) class FileObject(BaseModel): @@ -37,6 +50,7 @@ class FileObject(BaseModel): bytes: int | None = None status: str | None = None created_at: int | None = None + expires_at: int | None = None class FileList(BaseModel): @@ -85,7 +99,7 @@ class BatchList(BaseModel): class FileDeleteResponse(BaseModel): id: str object: str | None = None - deleted: bool + deleted: bool | None = None class BatchCreateBody(BaseModel): diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 1bcea0a61ee..17749c2fb87 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -108,6 +108,10 @@ class Capability: def id(self) -> str: return f"{self.provider}-{self.scenario}" + @property + def file_provider(self) -> str | None: + return self.provider if self.scenario in {"model_param", "provider_fallback"} else None + @property def jsonl_model(self) -> str: # Always the provider deployment name. Unified routes via diff --git a/tests/e2e/batches/conftest.py b/tests/e2e/batches/conftest.py index 3b133fab680..91a365b6b92 100644 --- a/tests/e2e/batches/conftest.py +++ b/tests/e2e/batches/conftest.py @@ -13,7 +13,7 @@ the proxy config. from __future__ import annotations import os -from typing import Iterator +from typing import Final, Iterator import pytest @@ -21,6 +21,7 @@ from batch_client import BatchClient, build_client from capabilities import PROVIDERS from e2e_config import MANAGED_FILES_OPT_IN_ENV from e2e_http import NoBody +from lifecycle import ResourceManager from proxy_client import ProxyClient @@ -52,6 +53,13 @@ def client(proxy: ProxyClient) -> BatchClient: return build_client(proxy) +@pytest.fixture +def resources(client: BatchClient) -> Iterator[ResourceManager]: + manager: Final = ResourceManager(client=client.proxy, strict_cleanup=True) + yield manager + manager.teardown() + + @pytest.fixture(scope="session") def batch_deployments(client: BatchClient) -> Iterator[None]: probe = client.proxy.probe("/health/liveliness", params=NoBody()) diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py new file mode 100644 index 00000000000..d0038139dcf --- /dev/null +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -0,0 +1,313 @@ +from builtins import ExceptionGroup +from collections.abc import Callable +from typing import Final +from unittest.mock import Mock, call + +import pytest +from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result +from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form +from capabilities import CAPABILITIES, Capability +from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError +from lifecycle import ResourceManager +from models import KeyGenerateBody + +MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE=" +MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" + + +class ExpectedCalls[T]: + def __init__(self, values: tuple[T, ...]) -> None: + self.values: Final = values + self.recorder: Final = Mock() + + def __call__(self, value: T) -> None: + self.recorder(value) + + def assert_done(self) -> None: + assert tuple(self.recorder.call_args_list) == tuple(call(value) for value in self.values) + + +class CleanupClient: + def __init__( + self, + *, + calls: ExpectedCalls[str], + files: tuple[Result[FileDeleteResponse], ...] = (), + batches: tuple[Result[BatchObject], ...] = (), + cancellations: tuple[Result[BatchObject], ...] = (), + ) -> None: + self.calls: Final = calls + self.file_response: Final[Callable[[], Result[FileDeleteResponse]]] = Mock(side_effect=files) + self.batch_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=batches) + self.cancel_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=cancellations) + + def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: + self.calls(f"delete {provider} {file_id}") + return self.file_response() + + def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: + self.calls(f"retrieve {provider} {batch_id}") + return self.batch_response() + + def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: + self.calls(f"cancel {provider} {batch_id}") + return self.cancel_response() + + def generate_key(self, body: KeyGenerateBody) -> str: + return "test-key" + + def delete_key(self, key: str) -> None: + self.calls(f"delete key {key}") + + def delete_customers(self, user_ids: list[str]) -> None: + self.calls(f"delete customers {user_ids}") + + +def batch(status: str) -> Success[BatchObject]: + return Success(status_code=200, data=BatchObject(id="batch-1", status=status)) + + +def deleted_file(*, deleted: bool = True) -> Success[FileDeleteResponse]: + return Success(status_code=200, data=FileDeleteResponse(id="file-1", deleted=deleted)) + + +class TestFileCleanup: + def test_managed_delete_accepts_the_deleted_file_object(self) -> None: + response: Final = Success( + status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"}) + ) + client: Final = CleanupClient(calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)), files=(response,)) + cleanup_file(client, MANAGED_FILE_ID, key="test-key") + client.calls.assert_done() + + @pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID]) + def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls((f"delete None {file_id}",)), + files=(Success(status_code=200, data=FileDeleteResponse(id=file_id)),), + ) + with pytest.raises(AssertionError, match="did not confirm deletion"): + cleanup_file(client, file_id, key="test-key") + client.calls.assert_done() + + @pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES]) + def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None: + expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None + client: Final = CleanupClient( + calls=ExpectedCalls((f"delete {expected_provider} file-1",)), files=(deleted_file(),) + ) + cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider) + client.calls.assert_done() + + def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls(("delete azure file-1", "delete key test-key")), + files=(UnknownApiError(status_code=403, body="secret response"),), + ) + manager: Final = ResourceManager(client=client, strict_cleanup=True) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="azure")) + with pytest.raises(ExceptionGroup) as caught: + manager.teardown() + client.calls.assert_done() + assert len(caught.value.exceptions) == 1 + assert str(caught.value.exceptions[0]) == "Delete file file-1 failed: HTTP 403" + + def test_success_response_must_confirm_deletion(self) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls(("delete None file-1",)), files=(deleted_file(deleted=False),) + ) + with pytest.raises(AssertionError, match="did not confirm deletion"): + cleanup_file(client, "file-1", key="test-key") + client.calls.assert_done() + + def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls(("delete azure file-1",)), + files=(UnknownApiError(status_code=404, body="missing"),), + ) + cleanup_file(client, "file-1", key="test-key", provider="azure") + client.calls.assert_done() + + def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls(("delete None file-1", "delete key test-key")), + files=(UnknownApiError(status_code=403, body="forbidden"),), + ) + manager: Final = ResourceManager(client=client) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key)) + manager.teardown() + client.calls.assert_done() + + +class TestCleanupRetries: + @pytest.mark.parametrize( + "failure", + [NetworkError(message="offline"), RateLimitedError(), UnknownApiError(status_code=503, body="unavailable")], + ) + def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None: + responses: Final = (failure, deleted_file()) + outcomes: Final = Mock(side_effect=responses) + delays: Final = ExpectedCalls((1.0,)) + result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays) + assert isinstance(result, Success) and result.data.deleted + delays.assert_done() + + def test_persistent_error_has_bounded_retries(self) -> None: + failure: Final = UnknownApiError(status_code=503, body="unavailable") + outcomes: Final = Mock(return_value=failure) + delays: Final = ExpectedCalls(CLEANUP_DELAYS) + result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays) + assert result is failure + delays.assert_done() + assert outcomes.call_count == len(CLEANUP_DELAYS) + 1 + + def test_permanent_error_is_not_retried(self) -> None: + failure: Final = UnknownApiError(status_code=403, body="forbidden") + responses: Final = (failure, deleted_file()) + outcomes: Final = Mock(side_effect=responses) + delays: Final = ExpectedCalls[float](()) + assert cleanup_result(outcomes, wait=delays) is failure + delays.assert_done() + assert outcomes.call_count == 1 + + +class TestBatchCancellation: + def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 3), + batches=(batch("cancelling"), batch("cancelling"), batch("cancelled")), + ) + delays: Final = ExpectedCalls((10.0,)) + cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays) + client.calls.assert_done() + delays.assert_done() + + def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls( + ( + f"retrieve None {MANAGED_BATCH_ID}", + f"retrieve None {MANAGED_BATCH_ID}", + "delete None file-1", + "delete key test-key", + ) + ), + batches=(batch("cancelling"), batch("cancelling")), + files=(deleted_file(),), + ) + times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS) + ticks: Final[Callable[[], float]] = Mock(side_effect=times) + manager: Final = ResourceManager(client=client, strict_cleanup=True) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key)) + manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=ticks)) + with pytest.raises(ExceptionGroup) as caught: + manager.teardown() + assert "cancellation did not finish" in str(caught.value.exceptions[0]) + client.calls.assert_done() + + @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"]) + def test_inactive_batch_needs_no_cancellation(self, status: str) -> None: + client: Final = CleanupClient(calls=ExpectedCalls(("retrieve None batch-1",)), batches=(batch(status),)) + cleanup_batch(client, "batch-1", key="test-key") + client.calls.assert_done() + + def test_active_batch_is_cancelled_through_its_provider(self) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls(("retrieve azure batch-1", "cancel azure batch-1")), + batches=(batch("in_progress"), batch("cancelled")), + cancellations=(batch("cancelling"),), + ) + cleanup_batch(client, "batch-1", key="test-key", provider="azure") + client.calls.assert_done() + + @pytest.mark.parametrize("batch_id", ["batch-1", MANAGED_BATCH_ID]) + @pytest.mark.parametrize("pending_status", ["validating", "in_progress"]) + def test_accepted_cancellation_waits_through_stale_provider_status( + self, batch_id: str, pending_status: str + ) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls( + ( + f"retrieve vertex_ai {batch_id}", + f"cancel vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + "delete vertex_ai file-1", + "delete key test-key", + ) + ), + batches=(batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled")), + cancellations=(batch(pending_status),), + files=(deleted_file(),), + ) + delays: Final = ExpectedCalls((10.0, 10.0)) + manager: Final = ResourceManager(client=client, strict_cleanup=True) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="vertex_ai")) + manager.defer(lambda: cleanup_batch(client, batch_id, key=key, provider="vertex_ai", wait=delays)) + manager.teardown() + client.calls.assert_done() + delays.assert_done() + + @pytest.mark.parametrize("output_delete_fails", [False, True]) + def test_batch_that_completed_before_cleanup_deletes_output_and_error_files( + self, output_delete_fails: bool + ) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")), + batches=( + Success( + status_code=200, + data=BatchObject( + id="batch-1", + status="completed", + input_file_id="file-input", + output_file_id="file-output", + error_file_id="file-error", + ), + ), + ), + files=( + UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(), + deleted_file(), + ), + ) + if output_delete_fails: + with pytest.raises(ExceptionGroup, match="output cleanup failed"): + cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True) + else: + cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True) + client.calls.assert_done() + + @pytest.mark.parametrize("status", ["completed", "in_progress"]) + def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1")), + batches=(batch("in_progress"), batch(status)), + cancellations=(UnknownApiError(status_code=409, body="conflict"),), + ) + if status == "completed": + cleanup_batch(client, "batch-1", key="test-key") + else: + with pytest.raises(AssertionError, match="Cancel batch batch-1 left status in_progress"): + cleanup_batch(client, "batch-1", key="test-key") + client.calls.assert_done() + + +class TestAzureFileExpiry: + def test_azure_form_serializes_native_expiry_for_the_proxy(self) -> None: + form: Final = batch_upload_form("azure", target_model_names="azure-test") + assert form.model_dump(by_alias=True, exclude_none=True) == { + "purpose": "batch", + "target_model_names": "azure-test", + "expires_after[anchor]": "created_at", + "expires_after[seconds]": AZURE_FILE_EXPIRY_SECONDS, + } + + @pytest.mark.parametrize("provider", ["openai", "vertex_ai", "bedrock"]) + def test_other_providers_keep_their_existing_upload_fields(self, provider: str) -> None: + assert batch_upload_form(provider).model_dump(by_alias=True, exclude_none=True) == {"purpose": "batch"} diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index ed7cf656d01..c4b699190b8 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -21,14 +21,16 @@ import os import re import time from datetime import datetime, timedelta, timezone -from typing import Callable import pytest from pydantic import BaseModel -from e2e_config import PROXY_BASE_URL, unique_marker +from e2e_config import MASTER_KEY, PROXY_BASE_URL, unique_marker +from batch_cleanup import cleanup_batch, cleanup_file from batch_client import ( + AZURE_FILE_EXPIRY_SECONDS, + batch_upload_form, UPLOAD_FILENAME, BatchClient, BatchCreateBody, @@ -155,19 +157,19 @@ def upload_for_scenario( if cap.scenario == "encoded": return client.upload_file( content=content, - form=FileUploadForm(purpose="batch"), + form=batch_upload_form(cap.provider), model=cap.model, key=key, ) if cap.scenario == "unified": return client.upload_file( content=content, - form=FileUploadForm(purpose="batch", target_model_names=cap.model), + form=batch_upload_form(cap.provider, target_model_names=cap.model), key=key, ) return client.upload_file( content=content, - form=FileUploadForm(purpose="batch"), + form=batch_upload_form(cap.provider), key=key, provider=cap.provider, ) @@ -188,20 +190,11 @@ def create_for_scenario( def op_provider(cap: Capability) -> str | None: - """provider_fallback ids are raw, so retrieve/cancel/list/delete need the provider + """provider_fallback batch ids are raw, so retrieve/cancel/list need the provider hint; the other scenarios encode it into the id and route automatically.""" return cap.provider if cap.scenario == "provider_fallback" else None -def quietly(action: Callable[[], object]) -> Callable[[], None]: - """Adapt a value-returning call into a best-effort cleanup the teardown can run.""" - - def run() -> None: - action() - - return run - - def assert_file_object(file: FileObject, *, provider: str) -> None: assert file.object == "file", f"file.object={file.object!r}" assert file.purpose == "batch", f"file.purpose={file.purpose!r}" @@ -209,6 +202,10 @@ def assert_file_object(file: FileObject, *, provider: str) -> None: if provider != "bedrock": assert file.bytes > 0, f"file.bytes={file.bytes!r}" assert file.status, "file.status missing" + if provider == "azure": + assert file.expires_at is not None, "Azure batch input has no automatic expiry" + assert file.created_at is not None + assert file.expires_at - file.created_at == AZURE_FILE_EXPIRY_SECONDS assert ( file.created_at is not None and file.created_at > 0 ), "file.created_at missing" @@ -249,7 +246,7 @@ def test_batch_lifecycle( file = unwrap(upload_for_scenario(client, cap, render_jsonl(cap.jsonl_model), key)) resources.defer( - quietly(lambda: client.delete_file(file.id, key=key, provider=provider)) + lambda: cleanup_file(client, file.id, key=key, provider=cap.file_provider) ) assert_file_object(file, provider=cap.provider) assert matches_id_shape( @@ -260,7 +257,9 @@ def test_batch_lifecycle( require_successful_call(created) batch = BatchObject.model_validate_json(created.body) resources.defer( - quietly(lambda: client.cancel_batch(batch.id, key=key, provider=provider)) + lambda: cleanup_batch( + client, batch.id, key=key, provider=provider, delete_output_files=cap.provider in {"openai", "azure"} + ) ) assert batch.id, f"create returned no batch id (body={created.body[:200]})" @@ -339,7 +338,7 @@ def test_batch_key_model_access_denied( denied_upload = client.upload_file( content=render_jsonl(AZURE_BATCH_MODEL), - form=FileUploadForm(purpose="batch"), + form=batch_upload_form("azure"), model=AZURE_BATCH_MODEL, key=key, ) @@ -356,7 +355,7 @@ def test_batch_key_model_access_denied( ) ).id resources.defer( - quietly(lambda: client.delete_file(raw_file, key=key, provider="openai")) + lambda: cleanup_file(client, raw_file, key=key, provider="openai") ) denied_create = client.create_batch( @@ -383,6 +382,7 @@ def test_file_upload_and_delete_outputs( key=key, ) ) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="openai") deleted = unwrap(client.delete_file(file.id, key=key)) @@ -458,12 +458,12 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) _ = client.proxy.poll_logs_for_key(key, min_rows=1) @@ -517,7 +517,7 @@ class TestBatchFileContent: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert file.id downloaded = client.proxy.transport.download( @@ -559,11 +559,11 @@ class TestBatchFileContent: file = unwrap( client.upload_file( content=payload, - form=FileUploadForm(purpose="batch", target_model_names=provider.model), + form=batch_upload_form(provider.name, target_model_names=provider.model), key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider=provider.name) assert is_managed_id(file.id), ( f"{provider.name}: unified upload must return a managed file id, got {file.id!r}" @@ -626,7 +626,7 @@ class TestOpenAIFiles: ) ) resources.defer( - quietly(lambda: client.delete_file(file.id, key=key, provider="openai")) + lambda: cleanup_file(client, file.id, key=key, provider="openai") ) listed = unwrap(client.list_files(key=key)) @@ -690,7 +690,7 @@ class TestOpenAIFiles: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) fetched = unwrap(client.retrieve_file(file.id, key=key)) assert fetched.id == file.id, "retrieve must echo the uploaded file id" @@ -760,7 +760,7 @@ class TestBatchRateLimitErrorMapping: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) @@ -803,7 +803,7 @@ class TestBatchEnqueuedTokenLimit: """ def _upload_batch_file( - self, client: BatchClient, resources: ResourceManager, key: str + self, client: BatchClient, resources: ResourceManager, key: str, *, cleanup_key: str | None = None ) -> FileObject: file = unwrap( client.upload_file( @@ -813,7 +813,7 @@ class TestBatchEnqueuedTokenLimit: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=cleanup_key or key)) return file def _generate_enqueued_key( @@ -850,7 +850,7 @@ class TestBatchEnqueuedTokenLimit: marker="rpm", rpm_limit=BATCH_RL_RPM_LIMIT, ) - file = self._upload_batch_file(client, resources, key) + file = self._upload_batch_file(client, resources, key, cleanup_key=MASTER_KEY) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) @@ -861,7 +861,7 @@ class TestBatchEnqueuedTokenLimit: ) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=MASTER_KEY, delete_output_files=True)) @pytest.mark.covers( "quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted", @@ -904,7 +904,7 @@ class TestBatchEnqueuedTokenLimit: first = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(first) first_batch = BatchObject.model_validate_json(first.body) - resources.defer(quietly(lambda: client.cancel_batch(first_batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, first_batch.id, key=key)) blocked = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) assert blocked.status_code == 429, ( @@ -928,7 +928,7 @@ class TestBatchEnqueuedTokenLimit: ) require_successful_call(retried) retry_batch = BatchObject.model_validate_json(retried.body) - resources.defer(quietly(lambda: client.cancel_batch(retry_batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, retry_batch.id, key=key)) ASSUME_ROLE_RAW_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" @@ -984,13 +984,13 @@ class TestBedrockBatchAssumeRole: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="bedrock") created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert batch.id, f"assume-role create returned no batch id: {created.body[:200]}" assert is_managed_id(batch.id), ( @@ -1044,7 +1044,7 @@ class TestGeminiFiles: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="gemini") assert file.id, "gemini file upload returned no id" @@ -1099,13 +1099,13 @@ class TestHostedVllmBatch: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="hosted_vllm") created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert batch.id, f"hosted_vllm create returned no batch id: {created.body[:200]}" assert batch.status in CREATED_BATCH_STATUSES, ( @@ -1192,7 +1192,7 @@ class TestBatchFailurePaths: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) @@ -1243,12 +1243,12 @@ class TestBatchFailurePaths: file = unwrap( client.upload_file( content=render_jsonl(AZURE_BATCH_RAW_MODEL), - form=FileUploadForm(purpose="batch"), + form=batch_upload_form("azure"), model=AZURE_BATCH_MODEL, key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert decoded_model_from_id(file.id) == AZURE_BATCH_MODEL, ( f"upload did not encode the azure deployment into the file id: {file.id!r}" ) @@ -1258,7 +1258,7 @@ class TestBatchFailurePaths: ) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert decoded_model_from_id(batch.id) == AZURE_BATCH_MODEL, ( "create with a foreign encoded file id must route by the file's embedded model, " @@ -1307,7 +1307,7 @@ class TestBatchSecondHop: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert is_managed_id(file.id), ( f"second-hop unified upload must return a managed file id, got {file.id!r}" ) @@ -1315,7 +1315,7 @@ class TestBatchSecondHop: created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert is_managed_id(batch.id), ( f"second-hop create must return a managed batch id, got {batch.id!r}" diff --git a/tests/e2e/batches/test_managed_files_enforcement_e2e.py b/tests/e2e/batches/test_managed_files_enforcement_e2e.py index 7ad0b16adc3..4f703cf0fdc 100644 --- a/tests/e2e/batches/test_managed_files_enforcement_e2e.py +++ b/tests/e2e/batches/test_managed_files_enforcement_e2e.py @@ -21,6 +21,7 @@ from typing import Iterator import pytest from batch_client import BatchClient, FileObject +from batch_cleanup import cleanup_file from capabilities import batch_model_name, is_managed_id, openai_batch_params from e2e_config import unique_marker from e2e_http import FileUploadForm, Result, UnknownApiError, unwrap @@ -108,7 +109,7 @@ def test_cross_user_managed_id_denied_owner_allowed( key=owner_key, ) ) - resources.defer(lambda: client.delete_file(uploaded.id, key=owner_key)) + resources.defer(lambda: cleanup_file(client, uploaded.id, key=owner_key)) assert is_managed_id(uploaded.id), f"expected a managed unified file id, got {uploaded.id}" denied = client.retrieve_file(uploaded.id, key=other_key) diff --git a/tests/e2e/claude_code/_env.py b/tests/e2e/claude_code/_env.py index ed2da1fe1e2..889d8f848dd 100644 --- a/tests/e2e/claude_code/_env.py +++ b/tests/e2e/claude_code/_env.py @@ -99,5 +99,6 @@ def require_proxy_client( base_url=cfg.base_url, master_key=cfg.api_key, control_plane_base_url=cfg.base_url, + replica_urls=(cfg.base_url,), ) return ProxyClientConfig(client=client, api_key=cfg.api_key) diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index aaad7667936..6e3dce0377e 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -594,6 +594,7 @@ def _build_control_plane_client(proxy_config: ProxyConfig): base_url=proxy_config.base_url, master_key=proxy_config.api_key, control_plane_base_url=proxy_config.base_url, + replica_urls=(proxy_config.base_url,), ) diff --git a/tests/e2e/coverage_registry/collector.py b/tests/e2e/coverage_registry/collector.py index e20f7884f55..c5d09820fe7 100644 --- a/tests/e2e/coverage_registry/collector.py +++ b/tests/e2e/coverage_registry/collector.py @@ -105,6 +105,8 @@ def collect_markers(e2e_dir: Path = E2E_DIR) -> CollectedMarkers: "--continue-on-collection-errors", "-p", "no:cacheprovider", + "-p", + "no:pytest-retry", str(e2e_dir), ], plugins=[sink], diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index c64fd6150af..81832bebf49 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -3,6 +3,7 @@ - {id: guardrail.presidio.pre_call.masks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "PII masking pre-call; data-leak blast radius"} - {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"} - {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"} +- {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"} - {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"} - {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"} - {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"} diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..f853e9ff8d6 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -119,3 +119,11 @@ assertions: [succeeds] source: "server.py:1089" rationale: Smoke; rarely used; same auth model as tools +- id: mcp.list_tools.api_key.toolset_scoped + module: mcp + tier: P0 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "user_api_key_auth_mcp.py:2137" + rationale: "A key granted a toolset lists exactly the toolset's tools: the rest of the server's catalog stays hidden and every stored name resolves" diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index c3d7892aa79..d1227fe7c0c 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -5,6 +5,8 @@ - {id: mgmt.key.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:2462", rationale: "Budget/model changes persist"} - {id: mgmt.key.update.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:2462", rationale: "Non-admin cannot escalate perms"} - {id: mgmt.key.update.happy_path, module: mgmt, tier: P1, surface: ui, assertions: [happy_path], source: "key_management_endpoints.py:2462", rationale: "Key edit through the dashboard"} +- {id: mgmt.key.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "key_management_endpoints.py:2829", rationale: "A partial /key/update changes only the field it names; alias, models, limits, budget window, team and metadata read back unchanged on every gateway replica"} +- {id: mgmt.key.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "key_management_endpoints.py:2829", rationale: "An explicit null on /key/update clears max_budget and budget_duration, and the derived budget_reset_at with it, on every gateway replica"} - {id: mgmt.key.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3122", rationale: "Deletion revokes future calls"} - {id: mgmt.key.delete.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:3122", rationale: "Non-owner cannot delete"} - {id: mgmt.key.info.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3380", rationale: "Info reflects all writes"} @@ -78,3 +80,13 @@ - {id: mgmt.key.jwt.lifecycle, module: mgmt, tier: P0, surface: api, assertions: [lifecycle], source: "management_endpoints/key_management_endpoints.py", rationale: "An IdP-issued admin JWT creates, reads, updates, clears and deletes a key; omitted fields survive updates"} - {id: mgmt.key.jwt.member_denied, module: mgmt, tier: P0, surface: api, assertions: [member_denied], source: "auth/handle_jwt.py", rationale: "A valid member JWT cannot update an admin-managed key and denial leaves it unchanged"} - {id: mgmt.key.jwt.other_team_denied, module: mgmt, tier: P0, surface: api, assertions: [other_team_denied], source: "auth/handle_jwt.py", rationale: "A valid JWT for another existing team cannot read the key"} +- {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} +- {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} +- {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} +- {id: mgmt.mcp_server.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:2665", rationale: "An explicit null clears the stored field (absent keeps, null clears)"} +- {id: mgmt.mcp_server.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:2139", rationale: "A deleted server is gone by id and from the list on every replica"} +- {id: mgmt.mcp_toolset.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3009", rationale: "Toolset tools read back under the exact server_id and tool_name written; a toolset stored under one name and read under another granted nothing"} +- {id: mgmt.mcp_toolset.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:3098", rationale: "Editing the description leaves the tools and name intact"} +- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"} +- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"} +- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"} diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index d0afcaca848..ad0914d455b 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -58,3 +58,8 @@ - {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"} - {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"} - {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"} +- {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, batch create, and a replayed callback log carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"} +- {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"} +- {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"} +- {id: quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"} +- {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"} diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index b50551ec105..6b69677d490 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -7,7 +7,7 @@ - {id: reliability.retry.timeout.succeeds_within_retries, module: reliability, tier: P0, behavior: retry, variant: timeout, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:44", rationale: "Timeout retried per policy"} - {id: reliability.retry.429.succeeds_within_retries, module: reliability, tier: P0, behavior: retry, variant: "429", assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:46", rationale: "429 retried per RateLimitErrorRetries policy"} - {id: reliability.retry.auth.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: auth, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:42", rationale: "Transient auth glitch retry"} -- {id: reliability.retry.context_window.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: context_window, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:51", rationale: "Multi-attempt on context error"} +- {id: reliability.retry.context_window.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: context_window, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:51", fail_before_fix: proven, rationale: "A context-window 400 under BadRequestErrorRetries retries onto a sibling deployment in the same model group, instead of coming straight back as the 400 the deployment that just refused it returned"} - {id: reliability.cooldown.5xx.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "5xx", assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:40", rationale: "Deployment cools after repeated 5xx, recovers after cooldown_time"} - {id: reliability.cooldown.429.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "429", assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:69", rationale: "Cools on 429, avoids hammering exhausted provider"} - {id: reliability.cooldown.auth.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: auth, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:74", rationale: "Cools on 401 auth error"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 09d17b9a0db..691335ffdd5 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -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") diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 061fbc23fde..ce069720c6e 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -49,6 +49,11 @@ class AnthropicHeaders(AuthHeaders): anthropic_version: str = Field(default="2023-06-01", alias="anthropic-version") +class PartialBody(BaseModel): + """A body for a partial-update route (absent = keep, null = clear): a field left + unset is omitted from the wire, and a field set to None is sent as JSON null.""" + + class NoBody(BaseModel): """Empty body/query for routes that take none.""" @@ -266,6 +271,12 @@ def assert_auth_denied(result: StreamingResponse, context: str) -> None: f"{context}: expected 401/403, got {result.status_code}: {result.body[:300]}" ) +def wire_body(json: BaseModel) -> dict[str, object]: + if isinstance(json, PartialBody): + return json.model_dump(by_alias=True, exclude_unset=True) + return json.model_dump(by_alias=True, exclude_none=True) + + def _flat(model: BaseModel) -> dict[str, str]: dumped: dict[str, object] = model.model_dump(by_alias=True, exclude_none=True) return {key: str(value) for key, value in dumped.items()} @@ -322,9 +333,26 @@ def request_with_retry[T: RetryableResponse]( return issue() -def _classify[R: BaseModel]( - resp: requests.Response, response_type: type[R] -) -> Result[R]: +class ClassifiableResponse(Protocol): + """What classifying an outcome reads off a response. requests.Response satisfies + it, and so does a fake, so the classification rules are testable on their own.""" + + @property + def status_code(self) -> int: ... + + @property + def ok(self) -> bool: ... + + @property + def text(self) -> str: ... + + @property + def content(self) -> bytes: ... + + def json(self) -> object: ... + + +def classify[R: BaseModel](resp: ClassifiableResponse, response_type: type[R]) -> Result[R]: if resp.status_code == 401: return UnauthorizedError(body=resp.text) if resp.status_code == 429: @@ -332,7 +360,8 @@ def _classify[R: BaseModel]( if not resp.ok: return UnknownApiError(status_code=resp.status_code, body=resp.text) try: - return Success(status_code=resp.status_code, data=response_type.model_validate(resp.json())) + payload: Final[object] = resp.json() if resp.content else {} + return Success(status_code=resp.status_code, data=response_type.model_validate(payload)) except Exception as exc: # noqa: BLE001 - any parse/validation failure is a value return ValidationError(message=str(exc)) @@ -350,13 +379,13 @@ def post[R: BaseModel]( lambda: requests.post( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def get[R: BaseModel]( @@ -378,7 +407,7 @@ def get[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def get_external[R: BaseModel]( @@ -398,7 +427,7 @@ def get_external[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def post_form_external[R: BaseModel]( @@ -421,7 +450,7 @@ def post_form_external[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def post_json_external( @@ -471,14 +500,14 @@ def delete[R: BaseModel]( lambda: requests.delete( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), params=_params(params), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def patch[R: BaseModel]( @@ -494,13 +523,13 @@ def patch[R: BaseModel]( lambda: requests.patch( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def put[R: BaseModel]( @@ -516,13 +545,13 @@ def put[R: BaseModel]( lambda: requests.put( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def probe( @@ -626,7 +655,7 @@ def send( str(url), headers=_headers(headers), params=_params(params), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), stream=stream, timeout=timeout, ) @@ -676,7 +705,7 @@ def upload[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def stream_binary( @@ -694,7 +723,7 @@ def stream_binary( resp = requests.post( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), stream=True, timeout=timeout, ) diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml new file mode 100644 index 00000000000..229e8514dee --- /dev/null +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -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 diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 1f55a0f9a56..ed112a79b9b 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -7,7 +7,7 @@ from __future__ import annotations import time from collections.abc import Callable from dataclasses import dataclass -from typing import Literal +from typing import Final, Literal from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap @@ -405,6 +405,29 @@ def build_client(proxy: ProxyClient) -> GuardrailsClient: return GuardrailsClient(proxy=proxy) +def poll_until_guardrail_applied( + call: Callable[[], StreamingResponse], + guardrail_name: str, + *, + timeout: float = POLL_TIMEOUT, + interval: float = POLL_INTERVAL, + now: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, +) -> StreamingResponse: + deadline: Final = now() + timeout + if not (result := call()).ok: + return result + while ( + guardrail_name + not in (name.strip() for name in result.headers.get("x-litellm-applied-guardrails", "").split(",")) + and (remaining := deadline - now()) > 0 + ): + sleep(min(interval, remaining)) + if now() >= deadline or not (result := call()).ok: + break + return result + + def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]: """Retry a call that a guardrail should reject until it is, returning the last result. diff --git a/tests/e2e/guardrails/test_guardrails_client.py b/tests/e2e/guardrails/test_guardrails_client.py new file mode 100644 index 00000000000..423c2ede599 --- /dev/null +++ b/tests/e2e/guardrails/test_guardrails_client.py @@ -0,0 +1,66 @@ +from dataclasses import dataclass +from itertools import chain, repeat +from typing import Final + +import pytest + +from e2e_http import StreamingResponse +from guardrails_client import poll_until_guardrail_applied + + +@dataclass +class Clock: + elapsed: float = 0.0 + + def now(self) -> float: + return self.elapsed + + def sleep(self, seconds: float) -> None: + self.elapsed += seconds + + +def _response(applied: str, status: int = 200) -> StreamingResponse: + return StreamingResponse(status_code=status, body="{}", headers={"x-litellm-applied-guardrails": applied}) + + +def test_waits_for_requested_guardrail_after_an_unrelated_global_guardrail() -> None: + clock: Final = Clock() + expected: Final = _response("global-filter, tool-permission") + responses: Final = iter((_response("global-filter"), expected)) + + result: Final = poll_until_guardrail_applied( + lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + ) + + assert result is expected + assert clock.elapsed == 2 + + +@pytest.mark.parametrize("applied", ("", "global-filter", "tool-permission-sibling")) +def test_missing_exact_guardrail_returns_failure_evidence_at_deadline(applied: str) -> None: + clock: Final = Clock() + missing: Final = _response(applied) + responses: Final = iter((missing, missing, missing)) + + result: Final = poll_until_guardrail_applied( + lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + ) + + assert result is missing + assert clock.elapsed == 5 + with pytest.raises(StopIteration): + next(responses) + + +@pytest.mark.parametrize("status", (400, 401, 429, 500)) +def test_http_failure_is_not_hidden_by_a_later_success(status: int) -> None: + clock: Final = Clock() + failed: Final = _response("", status) + responses: Final = iter(chain((failed,), repeat(_response("tool-permission")))) + + result: Final = poll_until_guardrail_applied( + lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + ) + + assert result is failed + assert clock.elapsed == 0 diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py index c6d87473c21..49d698938ce 100644 --- a/tests/e2e/guardrails/test_presidio_masking_e2e.py +++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py @@ -19,6 +19,12 @@ PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip. Each guardrail registers with an explicit presidio_filter_scope so only the configured hook's callback exists (the default "both" registers input masking AND a post_call output masker), and is deleted on teardown. + +The spend-log audit test requires general_settings.store_prompts_in_spend_logs: +true in the proxy config, or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy +process environment before startup. Setting it only on pytest has no effect. +Without this opt-in, redacting guardrail_response is expected proxy behavior; +this suite deliberately requires the detected-entity details to remain visible. """ from __future__ import annotations @@ -26,16 +32,22 @@ from __future__ import annotations import os import time from collections.abc import Callable -from typing import Literal +from typing import Final, Literal import pytest -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from e2e_config import unique_marker -from e2e_http import Result, Success +from e2e_http import Result, StreamingResponse, Success from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody from lifecycle import ResourceManager -from models import AnthropicMessagesResponse, ChatResponse +from models import ( + AnthropicMessagesResponse, + ChatResponse, + GuardrailEntityMatch, + GuardrailRunRecord, + SpendLogRow, +) pytestmark = pytest.mark.e2e @@ -276,3 +288,114 @@ class TestPresidioPostCallMasking: f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}" ) time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) + + +_LOGGED_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"} + +_ENTITY_LIST_ADAPTER: Final = TypeAdapter(list[GuardrailEntityMatch]) + + +def _applied_guardrails(outcome: StreamingResponse) -> str: + return outcome.headers.get("x-litellm-applied-guardrails", "") + + +def _guardrail_records(row: SpendLogRow) -> tuple[GuardrailRunRecord, ...]: + metadata = row.metadata + if metadata is None: + return () + return tuple(metadata.guardrail_information or ()) + + +def _poll_until_guardrail_applied( + client: GuardrailsClient, key: str, guardrail_name: str, prompt: str +) -> StreamingResponse: + deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS + last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128) + while time.monotonic() < deadline: + if last.ok and guardrail_name in _applied_guardrails(last): + return last + time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) + last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128) + return last + + +class TestPresidioSpendLogRecord: + @pytest.mark.covers( + "guardrail.presidio.pre_call.logs_masked_entities", + exercised_on=["chat_completions"], + ) + def test_masking_run_is_recorded_on_the_spend_log( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + name = f"e2e-presidio-log-{unique_marker()}" + _register_presidio(client, resources, name=name, entities=_LOGGED_ENTITIES) + + email = _fake_email() + prompt = _pii_prompt(unique_marker(), email) + + outcome = _poll_until_guardrail_applied(client, scoped_key, name, prompt) + assert outcome.ok, f"the guarded call must be served, got {outcome.status_code}: {outcome.body[:400]}" + assert name in _applied_guardrails(outcome), ( + "the response must carry x-litellm-applied-guardrails naming the guardrail; without it " + f"the 200 only proves the guardrail never attached. Got {_applied_guardrails(outcome)!r}" + ) + + request_id = ChatResponse.model_validate_json(outcome.body).id + assert request_id, f"the served response must carry an id to look the spend log up by: {outcome.body[:400]}" + + rows = client.proxy.poll_logs_for_request_id( + request_id, + predicate=lambda logged: bool(_guardrail_records(logged[0])), + ) + assert rows, f"no spend log row ever appeared for request {request_id}" + + records = _guardrail_records(rows[0]) + assert records, ( + f"the spend log for {request_id} carries no guardrail_information, so the dashboard's " + "guardrail panel would render nothing for a request the guardrail demonstrably ran on" + ) + + record = next( + (entry for entry in records if entry.guardrail_name == name and entry.guardrail_mode == "pre_call"), + None, + ) + assert record is not None, ( + "guardrail_information carries no pre_call record for this guardrail, only " + f"{[(entry.guardrail_name, entry.guardrail_mode) for entry in records]}" + ) + assert record.guardrail_status == "success", ( + f"the recorded status must be success for a run that masked and served, got {record.guardrail_status!r}" + ) + assert record.guardrail_provider == "presidio", ( + f"the record must attribute the run to presidio so the dashboard picks the right " + f"renderer, got {record.guardrail_provider!r}" + ) + + counts = record.masked_entity_count or {} + assert _LOGGED_ENTITIES.keys() <= counts.keys(), ( + f"every entity the guardrail was configured to mask must appear in masked_entity_count, got {counts}" + ) + assert all(counts[entity] >= 1 for entity in _LOGGED_ENTITIES), ( + f"each masked entity must be counted at least once, got {counts}" + ) + + assert not isinstance(record.guardrail_response, str), ( + "This audit test requires general_settings.store_prompts_in_spend_logs: true in the proxy config " + "or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy process environment before startup " + "(not just the pytest environment). Default prompt redaction is valid proxy behavior, " + f"but prevents entity-detail assertions; got guardrail_response={record.guardrail_response!r}" + ) + entities = _ENTITY_LIST_ADAPTER.validate_python(record.guardrail_response) + assert entities, ( + "guardrail_response must carry the detected entities; the dashboard's Detected Entities " + "list and its per-entity scores are rendered from exactly this array" + ) + assert {entity.entity_type for entity in entities} >= _LOGGED_ENTITIES.keys(), ( + f"the detected entities must cover what was masked, got " + f"{sorted(entity.entity_type for entity in entities)}" + ) + for entity in entities: + assert 0.0 < entity.score <= 1.0, f"{entity.entity_type} carries an out-of-range score: {entity.score}" + assert 0 <= entity.start < entity.end, ( + f"{entity.entity_type} carries a degenerate span: {entity.start}-{entity.end}" + ) diff --git a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py index 9ef3650625c..8d1047e53c7 100644 --- a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py @@ -30,6 +30,7 @@ from guardrails_client import ( ToolPermissionParamsBody, ToolPermissionRuleBody, poll_until_blocked, + poll_until_guardrail_applied, ) from lifecycle import ResourceManager from models import ChatResponse, ChatTool, ChatToolFunction @@ -84,8 +85,8 @@ def _register_tool_permission(client: GuardrailsClient, resources: ResourceManag resources.defer(lambda: client.delete_guardrail(guardrail_id)) -def _applied_guardrails(outcome: StreamingResponse) -> str: - return outcome.headers.get("x-litellm-applied-guardrails", "") +def _applied_guardrails(outcome: StreamingResponse) -> tuple[str, ...]: + return tuple(name.strip() for name in outcome.headers.get("x-litellm-applied-guardrails", "").split(",")) def _tool_call_names(response: ChatResponse) -> tuple[str, ...]: @@ -144,14 +145,17 @@ class TestToolPermissionPreCall: name = f"e2e-toolperm-allow-{unique_marker()}" _register_tool_permission(client, resources, name=name) - outcome = client.chat_raw( - scoped_key, - MODEL, - TOOL_PROMPT, - guardrails=[name], - max_tokens=128, - tools=[ALLOWED_TOOL], - tool_choice="required", + outcome = poll_until_guardrail_applied( + lambda: client.chat_raw( + scoped_key, + MODEL, + TOOL_PROMPT, + guardrails=[name], + max_tokens=128, + tools=[ALLOWED_TOOL], + tool_choice="required", + ), + name, ) assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}" diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py index c9a67ebdb8c..eb9704d4dcb 100644 --- a/tests/e2e/lifecycle.py +++ b/tests/e2e/lifecycle.py @@ -8,8 +8,9 @@ ResourceManager; the test registers a cleanup for every resource it creates, and the fixture's teardown releases them all even when the test body raises. """ +from builtins import ExceptionGroup from dataclasses import dataclass, field -from typing import Callable, List, Protocol, runtime_checkable +from typing import Callable, Final, List, Protocol, runtime_checkable from proxy_client import ProxyClient from models import KeyGenerateBody @@ -52,6 +53,7 @@ class ResourceManager: """ client: ResourceClient + strict_cleanup: bool = False _cleanups: List[Callable[[], object]] = field( default_factory=list ) # mutable-ok: append-only teardown registry @@ -82,8 +84,17 @@ class ResourceManager: return customer_id def teardown(self) -> None: - for cleanup in reversed(self._cleanups): - try: - cleanup() - except Exception: - pass # best-effort: a failed cleanup must not block the rest + failures: Final = tuple( + failure for cleanup in reversed(self._cleanups) + if (failure := _run_cleanup(cleanup)) is not None + ) + if failures and self.strict_cleanup: + raise ExceptionGroup("Resource cleanup failed", failures) + + +def _run_cleanup(cleanup: Callable[[], object]) -> Exception | None: + try: + cleanup() + except Exception as exc: + return exc + return None diff --git a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py index 9a45743a0cd..75817340876 100644 --- a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py +++ b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py @@ -16,8 +16,10 @@ into a chat completion chunk. Two customer-visible contracts only hold on that p from __future__ import annotations +from typing import Final, Literal + import pytest -from pydantic import BaseModel +from pydantic import BaseModel, Field from e2e_config import unique_marker from e2e_http import StreamingResponse @@ -51,7 +53,8 @@ class _BridgeChoice(BaseModel): class _BridgeChunk(BaseModel): id: str - choices: list[_BridgeChoice] = [] + object: Literal["chat.completion.chunk"] + choices: list[_BridgeChoice] = Field(default_factory=list) class _WeatherArgs(BaseModel): @@ -103,16 +106,19 @@ class TestResponsesBridgeChatCompletionsStreaming: resources.key(), ChatBody( model=bridged_model, - messages=[ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}")], + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], max_tokens=64, stream=True, ), ) - chunks = _bridge_chunks(result) - ids = {chunk.id for chunk in chunks} + chunks: Final = _bridge_chunks(result) + assert len(chunks) > 1, "the shared-id contract needs more than one streamed chunk" + ids: Final = frozenset(chunk.id for chunk in chunks) assert len(ids) == 1, f"bridged stream used {len(ids)} different chunk ids: {sorted(ids)[:5]}" - assert ids.pop().startswith("chatcmpl-"), f"bridged chunk id is not chat-completion shaped: {chunks[0].id}" + assert chunks[0].id.strip(), "bridged stream emitted an empty chunk id" @pytest.mark.covers( "llm.chat_completions.openai.basic.stream.bridge_streams_sse", @@ -134,9 +140,9 @@ class TestResponsesBridgeChatCompletionsStreaming: chunks = _bridge_chunks(result) content = "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) assert content.strip(), f"bridged stream completed with no content deltas: {result.stream_events[:3]}" - assert any( - choice.finish_reason for chunk in chunks for choice in chunk.choices - ), f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + assert any(choice.finish_reason for chunk in chunks for choice in chunk.choices), ( + f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + ) assert result.stream_done, f"bridged stream did not terminate with [DONE]: {result.stream_events[-2:]}" @pytest.mark.covers( diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index d0f478185c2..368c20cb6aa 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -12,8 +12,12 @@ empty result. External reads go through ``e2e_http``. from __future__ import annotations +import math +import random import time -from dataclasses import dataclass +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import Final import pytest from pydantic import BaseModel, ConfigDict, Field @@ -27,17 +31,33 @@ from e2e_config import ( DD_SITE, POLL_TIMEOUT, ) -from e2e_http import URL, Headers, RateLimitedError, Success, post +from e2e_http import URL, Headers, StreamingResponse, send -#: How many rate-limited responses in a row one search tolerates before the -#: hard fail; each retry sleeps a full search interval, so this rides out a -#: burst from a concurrent consumer of the org-wide search budget. -_RATE_LIMIT_RETRIES = 5 +type SearchCall = Callable[[str, float], StreamingResponse] + + +def _seconds(value: str | None) -> float | None: + if value is None: + return None + try: + seconds: Final = float(value) + except ValueError: + return None + return seconds if math.isfinite(seconds) and seconds >= 0 else None + + +def _rate_limit_delay(headers: Mapping[str, str]) -> float: + delays: Final = tuple( + delay + for name in ("x-ratelimit-reset", "retry-after") + if (delay := _seconds(headers.get(name))) is not None + ) + return max(1.0, max(delays, default=DD_SEARCH_INTERVAL)) class _DdAuthHeaders(Headers): - api_key: str = Field(serialization_alias="DD-API-KEY") - app_key: str = Field(serialization_alias="DD-APPLICATION-KEY") + api_key: str = Field(serialization_alias="DD-API-KEY", repr=False) + app_key: str = Field(serialization_alias="DD-APPLICATION-KEY", repr=False) class _SearchFilter(BaseModel): @@ -88,8 +108,12 @@ class _SearchResponse(BaseModel): @dataclass(frozen=True, slots=True) class DdLogsReader: site: str - api_key: str - app_key: str + api_key: str = field(repr=False) + app_key: str = field(repr=False) + search: SearchCall | None = field(default=None, repr=False) + now: Callable[[], float] = field(default=time.monotonic, repr=False) + sleep: Callable[[float], None] = field(default=time.sleep, repr=False) + jitter: Callable[[], float] = field(default=random.random, repr=False) def events_for_marker(self, marker: str) -> list[DdLogEvent]: """Every ingested event whose attributes carry the marker. DataDog @@ -108,25 +132,28 @@ class DdLogsReader: a single event. A 429 backs off and retries - the search budget is org-wide, so another consumer can empty it under us - while any other failure stays a hard fail.""" - for _ in range(_RATE_LIMIT_RETRIES): - result = post( - URL(f"https://api.{self.site}/api/v2/logs/events/search"), - headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), - json=_SearchRequest(filter=_SearchFilter(query=query)), - response_type=_SearchResponse, - timeout=30.0, - ) - match result: - case Success(data=page): - return [event.attributes for event in page.data] - case RateLimitedError(retry_after_seconds=retry_after): - time.sleep(retry_after if retry_after else DD_SEARCH_INTERVAL) - case failure: - pytest.fail(f"DataDog Logs Search API at api.{self.site} failed: {failure}") + return self._events_for_query(query, self.now() + POLL_TIMEOUT) + + def _events_for_query(self, query: str, deadline: float) -> list[DdLogEvent]: + search: Final = self.search or self._search_page + while (remaining := deadline - self.now()) > 0: + if (result := search(query, min(30.0, remaining))).ok: + return [event.attributes for event in _SearchResponse.model_validate_json(result.body).data] + if result.status_code != 429: + pytest.fail(f"DataDog Logs Search API at api.{self.site} failed with HTTP {result.status_code}") + if (delay := min(_rate_limit_delay(result.headers) + self.jitter(), deadline - self.now())) > 0: + self.sleep(delay) pytest.fail( - f"DataDog Logs Search API at api.{self.site} still rate-limited after " - f"{_RATE_LIMIT_RETRIES} retries {DD_SEARCH_INTERVAL}s apart - the org-wide " - "logs_public_search_api budget (2 requests per 10s) is exhausted by another consumer" + f"DataDog Logs Search API at api.{self.site} remained rate-limited for {POLL_TIMEOUT}s; " + "the org-wide logs_public_search_api budget is exhausted" + ) + + def _search_page(self, query: str, timeout: float) -> StreamingResponse: + return send( + URL(f"https://api.{self.site}/api/v2/logs/events/search"), + headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), + json=_SearchRequest(filter=_SearchFilter(query=query)), + timeout=timeout, ) def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]: @@ -140,33 +167,42 @@ class DdLogsReader: hide from the exactly-one assertion - real-DataDog jitter can surface one call's two events tens of seconds apart. Searches pace at DD_SEARCH_INTERVAL, not POLL_INTERVAL, to respect the search API's - request budget. At the deadline the last result is returned as-is.""" - deadline = time.monotonic() + POLL_TIMEOUT - while time.monotonic() < deadline: - events = self.events_for_query(query) + request budget. Discovery, quota retries, and duplicate detection share + one POLL_TIMEOUT deadline; an incomplete settle window fails closed.""" + deadline: Final = self.now() + POLL_TIMEOUT + while (remaining := deadline - self.now()) > 0: + events = self._events_for_query(query, deadline) if events: - return self._settled_events_for_query(query, events) - time.sleep(DD_SEARCH_INTERVAL) - return self.events_for_query(query) + return self._settled_events_for_query(query, events, deadline) + if (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + return [] - def _settled_events_for_query(self, query: str, events: list[DdLogEvent]) -> list[DdLogEvent]: + def _settled_events_for_query(self, query: str, events: list[DdLogEvent], deadline: float) -> list[DdLogEvent]: """Re-read at every search interval until the settle window closes; a duplicate ends the watch early because more waiting cannot clear it. Keep the last non-empty result: a transient empty search (index lag) must not erase events already confirmed earlier in the settle window. + A successful final search must reach the full settle window before the + shared read-back deadline; otherwise duplicate detection is incomplete. """ - settle_deadline = time.monotonic() + DD_SETTLE_SECONDS + settle_deadline: Final = self.now() + DD_SETTLE_SECONDS last_nonempty = events - while time.monotonic() < settle_deadline: - time.sleep(DD_SEARCH_INTERVAL) - latest = self.events_for_query(query) - if not latest: - continue + if len(events) > 1: + return events + while (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + if self.now() >= deadline: + break + latest = self._events_for_query(query, deadline) if len(latest) > 1: return latest - last_nonempty = latest - return last_nonempty + if latest: + last_nonempty = latest + if self.now() >= settle_deadline: + return last_nonempty + pytest.fail(f"DataDog log delivery could not complete its duplicate-detection window within {POLL_TIMEOUT}s") def build_dd_logs_reader() -> DdLogsReader: diff --git a/tests/e2e/logging/test_datadog_reader.py b/tests/e2e/logging/test_datadog_reader.py new file mode 100644 index 00000000000..910a1cefd42 --- /dev/null +++ b/tests/e2e/logging/test_datadog_reader.py @@ -0,0 +1,223 @@ +import json +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from typing import Final + +import pytest + +from datadog_reader import DdLogsReader +from datadog_reader import _DdAuthHeaders # pyright: ignore[reportPrivateUsage] # verifies private auth-header serialization +from e2e_config import DD_SEARCH_INTERVAL, POLL_TIMEOUT +from e2e_http import StreamingResponse + + +def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> None: + api_key: Final = "test-datadog-api-secret" + app_key: Final = "test-datadog-app-secret" + reader: Final = DdLogsReader(site="datadoghq.com", api_key=api_key, app_key=app_key) + headers: Final = _DdAuthHeaders(api_key=api_key, app_key=app_key) + + for value in (reader, headers): + assert api_key not in repr(value) + assert app_key not in repr(value) + + assert headers.model_dump(by_alias=True) == { + "DD-API-KEY": api_key, + "DD-APPLICATION-KEY": app_key, + } + + +@dataclass +class Clock: + elapsed: float = 0.0 + + def now(self) -> float: + return self.elapsed + + def sleep(self, seconds: float) -> None: + self.elapsed += seconds + + +@dataclass +class Search: + responses: Iterator[StreamingResponse] + calls: tuple[tuple[str, float], ...] = () + + def __call__(self, query: str, timeout: float) -> StreamingResponse: + self.calls += ((query, timeout),) + return next(self.responses) + + +def _page(*event_ids: str) -> StreamingResponse: + return StreamingResponse( + status_code=200, + body=json.dumps({"data": [{"attributes": {"attributes": {"id": event_id}}} for event_id in event_ids]}), + ) + + +def _reader(responses: Sequence[StreamingResponse], clock: Clock) -> tuple[DdLogsReader, Search]: + search: Final = Search(iter(responses)) + return DdLogsReader( + site="us5.datadoghq.com", + api_key="test-api-secret", + app_key="test-app-secret", + search=search, + now=clock.now, + sleep=clock.sleep, + jitter=lambda: 0.25, + ), search + + +def test_429_honors_server_reset_and_preserves_duplicate_events() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "6"}), _page("first", "duplicate")), + clock, + ) + + events: Final = reader.events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 6.25 + assert search.calls == (("test-marker", 30.0), ("test-marker", 30.0)) + + +@pytest.mark.parametrize("reset", ("", "invalid", "nan", "inf", "-1")) +def test_invalid_reset_uses_search_interval(reset: str) -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": reset}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == DD_SEARCH_INTERVAL + 0.25 + + +def test_zero_reset_cannot_create_a_busy_retry_loop() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "0"}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 1.25 + + +def test_retry_after_is_not_shortened_by_an_earlier_reset() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "2", "retry-after": "8"}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 8.25 + + +def test_rate_limit_wait_stops_at_deadline_without_issuing_another_request() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT * 10)}),), clock + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls == (("test-marker", 30.0),) + + +def test_late_retry_cannot_receive_a_fresh_request_timeout() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT - 5)}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert search.calls == (("test-marker", 30.0), ("test-marker", 4.75)) + + +@pytest.mark.parametrize("status", (-1, 401, 403, 500)) +def test_non_quota_failures_are_not_retried_or_treated_as_empty_results(status: int) -> None: + clock: Final = Clock() + reader, search = _reader((StreamingResponse(status_code=status, body=""), _page()), clock) + + with pytest.raises(pytest.fail.Exception, match=f"failed with HTTP {status}"): + reader.events_for_query("test-marker") + + assert search.calls == (("test-marker", 30.0),) + assert clock.elapsed == 0 + + +def test_polling_quota_retries_share_the_original_deadline() -> None: + clock: Final = Clock() + reader, search = _reader( + (_page(), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == 2 + + +def test_empty_polling_does_not_start_a_final_search_after_its_deadline() -> None: + clock: Final = Clock() + attempts: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) + reader, search = _reader((_page(),) * attempts, clock) + + assert reader.poll_events_for_query("test-marker") == [] + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == attempts + + +def test_settlement_quota_retries_keep_the_remaining_readback_budget() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader( + (_page(),) * empty_reads + + (_page("first"), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls[-1] == ("test-marker", DD_SEARCH_INTERVAL) + assert len(search.calls) == empty_reads + 2 + + +def test_settlement_detects_a_duplicate_on_the_final_search() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page("first"), _page(), _page("first", "duplicate")), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 30 + + +def test_settlement_keeps_confirmed_events_through_empty_searches() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page(), _page(), _page()), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first",) + assert clock.elapsed == 30 + + +def test_late_delivery_cannot_pass_without_a_complete_settle_window() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader((_page(),) * empty_reads + (_page("first"), _page("first")), clock) + + with pytest.raises(pytest.fail.Exception, match="duplicate-detection window"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == empty_reads + 2 diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index e588d19d5e8..e17b92a13ed 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -43,6 +43,9 @@ from models import ( KeyResetSpendBody, KeyResetSpendResponse, KeyUpdateBody, + McpServerCreateBody, + McpServerRow, + McpServerUpdateBody, ModelDeleteBody, OrgDeleteBody, OrgInfoParams, @@ -546,6 +549,38 @@ class ManagementClient: ).root ) + def create_mcp_server(self, body: McpServerCreateBody) -> McpServerRow: + return unwrap( + self.proxy.transport.post( + "/v1/mcp/server", + headers=self.proxy.transport.master, + json=body, + response_type=McpServerRow, + ) + ) + + def update_mcp_server(self, body: McpServerUpdateBody) -> McpServerRow: + """PUT /v1/mcp/server, the call behind the dashboard's Save Changes: a partial + update where a field left unset keeps its stored value and None clears it.""" + return unwrap( + self.proxy.transport.put( + "/v1/mcp/server", + headers=self.proxy.transport.master, + json=body, + response_type=McpServerRow, + ) + ) + + def delete_mcp_server(self, server_id: str) -> Result[NoBody]: + """DELETE /v1/mcp/server/{server_id}. Returns the outcome so the act phase can + unwrap it while a deferred teardown can ignore an already-deleted server.""" + return self.proxy.transport.delete( + f"/v1/mcp/server/{server_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + def chat_status(self, key: str, model: str, content: str) -> StreamingResponse: return self.proxy.transport.send( "/chat/completions", diff --git a/tests/e2e/management/test_key_lifecycle_e2e.py b/tests/e2e/management/test_key_lifecycle_e2e.py new file mode 100644 index 00000000000..4c8effc4d24 --- /dev/null +++ b/tests/e2e/management/test_key_lifecycle_e2e.py @@ -0,0 +1,299 @@ +"""Live e2e: one virtual key walked through its whole lifecycle, read back on every +gateway replica. + +Create, read, partial update, clear, enforce, delete: one method per step, and every +step creates its own team and key (both deleted on teardown) so a step reruns or skips +on its own. Writes go through the control plane; read-backs poll every URL in +PROXY_REPLICA_URLS until each replica converges, because a write that is visible on the +gateway that took it and stale on its neighbour is exactly the failure this file exists +to catch. Revocation is the slowest of those: a deleted key stays usable on the other +replicas until their auth cache entry expires, so the delete step polls each of them +rather than asserting once. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +import pytest + +from e2e_config import unique_marker +from e2e_http import Result, StreamingResponse, Success, UnknownApiError, unwrap +from lifecycle import ResourceManager +from management_client import MODEL_ACCESS_DENIED_MARKER, ManagementClient +from models import ( + CLEAR, + ChatBody, + ChatMessage, + KeyGenerateBody, + KeyGenerateResponse, + KeyInfo, + KeyInfoParams, + KeyInfoResponse, + KeyMetadata, + KeyUpdateBody, + LiteLLMParamsBody, + TeamNewBody, +) +from proxy_client import Converged, NotConverged, Poller, await_converged, await_converged_everywhere +from transport import Transport + +pytestmark = pytest.mark.e2e + +BACKING_MODEL: Final = "gpt-4o-mini" +DENIED_MODEL: Final = "gpt-5.5" +MAX_BUDGET: Final = 25.0 +TPM_LIMIT: Final = 313131 +RPM_LIMIT: Final = 323232 +UPDATED_RPM_LIMIT: Final = 424242 +BUDGET_DURATION: Final = "30d" + + +@dataclass(frozen=True, slots=True) +class CreatedKey: + written: KeyGenerateBody + response: KeyGenerateResponse + + @property + def key(self) -> str: + return self.response.key + + +@pytest.fixture(scope="module") +def mock_deployment(client: ManagementClient) -> Iterator[str]: + """A deployment that answers from a canned response, so the enforcement step needs no + provider key. The alias carries a unique marker, like every other model this suite + registers, so concurrent runs never share one model group.""" + model_name: Final = f"e2e-key-lifecycle-{unique_marker()}" + model_id: Final = client.proxy.create_model(model_name, LiteLLMParamsBody(model=BACKING_MODEL, mock_response="ok")) + try: + yield model_name + finally: + client.proxy.delete_model(model_id) + + +def _await[T](client: ManagementClient, poller: Poller[T], converged: Callable[[T], bool], failure: str) -> T: + outcome: Final = await_converged( + poller, + converged=converged, + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Converged(result=result): + return result + case NotConverged(last_result=last): + pytest.fail(f"{failure}; last outcome: {last}") + + +def _chat_poller(transport: Transport, key: str, model: str) -> Poller[StreamingResponse]: + return lambda: transport.send( + "/chat/completions", + headers=transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"say hi {unique_marker()}")], + max_tokens=16, + ), + ) + + +def _create_key(client: ManagementClient, resources: ResourceManager, model: str) -> CreatedKey: + marker: Final = unique_marker() + team_id: Final = client.create_team(TeamNewBody(team_alias=f"e2e-key-lifecycle-team-{marker}")) + resources.defer(lambda: client.delete_team(team_id)) + written: Final = KeyGenerateBody( + key_alias=f"e2e-key-lifecycle-{marker}", + models=[model], + max_budget=MAX_BUDGET, + tpm_limit=TPM_LIMIT, + rpm_limit=RPM_LIMIT, + budget_duration=BUDGET_DURATION, + metadata=KeyMetadata(tag=marker), + team_id=team_id, + ) + response: Final = unwrap(client.generate_key(written)) + resources.defer(lambda: client.proxy.delete_key(response.key)) + return CreatedKey(written=written, response=response) + + +def _key_info_everywhere( + client: ManagementClient, key: str, settled: Callable[[KeyInfo], bool] +) -> Mapping[str, KeyInfo]: + def converged(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, Success) and settled(result.data.info) + + reads: Final = client.proxy.read_back_everywhere( + "/key/info", params=KeyInfoParams(key=key), response_type=KeyInfoResponse, converged=converged + ) + return MappingProxyType({replica: unwrap(read).info for replica, read in reads.items()}) + + +def _is_key_not_found(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, UnknownApiError) and result.status_code == 404 + + +def _assert_reads_back(info: KeyInfo, expected: KeyGenerateBody, replica: str) -> None: + for field, observed, wanted in ( + ("key_alias", info.key_alias, expected.key_alias), + ("models", info.models, expected.models), + ("max_budget", info.max_budget, expected.max_budget), + ("tpm_limit", info.tpm_limit, expected.tpm_limit), + ("rpm_limit", info.rpm_limit, expected.rpm_limit), + ("budget_duration", info.budget_duration, expected.budget_duration), + ("team_id", info.team_id, expected.team_id), + ("metadata", info.metadata, expected.metadata), + ): + assert observed == wanted, f"{replica}: /key/info reports {field}={observed!r}, expected {wanted!r}" + + +def _poll_chat_ok(client: ManagementClient, key: str, model: str) -> None: + _ = _await( + client, + _chat_poller(client.proxy.transport, key, model), + lambda outcome: outcome.ok, + f"chat on {model} never succeeded for the key before the deadline", + ) + + +def _warm_every_replica(client: ManagementClient, key: str, model: str) -> None: + """Serve one call from every replica, so each has the key in its auth cache. Without + this the revocation check below would only prove a replica rejects a key it never + knew, which is true of any random string.""" + for replica, transport in client.proxy.replicas.items(): + _ = _await( + client, + _chat_poller(transport, key, model), + lambda outcome: outcome.ok, + f"{replica}: chat on {model} never succeeded for the key before the deadline", + ) + + +def _assert_chat_rejected_everywhere(client: ManagementClient, key: str, model: str) -> None: + outcomes: Final = await_converged_everywhere( + {replica: _chat_poller(transport, key, model) for replica, transport in client.proxy.replicas.items()}, + converged=lambda outcome: outcome.status_code == 401, + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + for replica, outcome in outcomes.items(): + assert isinstance(outcome, Converged), ( + f"{replica}: the deleted key was still accepted on chat after " + f"{client.proxy.poll_timeout}s, last status {outcome.last_result.status_code}" + ) + + +class TestKeyLifecycle: + def test_create_echoes_every_field_written( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + response: Final = created.response + for field, observed, wanted in ( + ("key_alias", response.key_alias, created.written.key_alias), + ("models", response.models, created.written.models), + ("max_budget", response.max_budget, created.written.max_budget), + ("tpm_limit", response.tpm_limit, created.written.tpm_limit), + ("rpm_limit", response.rpm_limit, created.written.rpm_limit), + ("budget_duration", response.budget_duration, created.written.budget_duration), + ("team_id", response.team_id, created.written.team_id), + ("metadata", response.metadata, created.written.metadata), + ): + assert observed == wanted, f"/key/generate echoed {field}={observed!r}, sent {wanted!r}" + + def test_read_reflects_the_create_on_every_replica( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + infos: Final = _key_info_everywhere( + client, created.key, lambda info: info.key_alias == created.written.key_alias + ) + for replica, info in infos.items(): + _assert_reads_back(info, created.written, replica) + assert info.budget_reset_at is not None, ( + f"{replica}: /key/info reports no budget_reset_at for budget_duration={BUDGET_DURATION!r}" + ) + + @pytest.mark.covers("mgmt.key.update.preserves_unrelated_fields") + def test_partial_update_changes_only_the_named_field( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + before: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == RPM_LIMIT) + + _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, rpm_limit=UPDATED_RPM_LIMIT))) + + after: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == UPDATED_RPM_LIMIT) + for replica, info in after.items(): + _assert_reads_back(info, created.written.model_copy(update={"rpm_limit": UPDATED_RPM_LIMIT}), replica) + assert info.budget_reset_at == before[replica].budget_reset_at, ( + f"{replica}: budget_reset_at moved from {before[replica].budget_reset_at!r} to " + f"{info.budget_reset_at!r} on a /key/update that did not name budget_duration" + ) + + @pytest.mark.covers("mgmt.key.update.clear_persists") + def test_explicit_null_clears_the_budget_and_its_reset_time( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + _ = _key_info_everywhere(client, created.key, lambda info: info.max_budget == MAX_BUDGET) + + _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, max_budget=CLEAR, budget_duration=CLEAR))) + + cleared: Final = _key_info_everywhere(client, created.key, lambda info: info.max_budget is None) + for replica, info in cleared.items(): + assert info.budget_duration is None, ( + f"{replica}: budget_duration={info.budget_duration!r} survived an explicit null" + ) + assert info.budget_reset_at is None, ( + f"{replica}: clearing budget_duration left budget_reset_at={info.budget_reset_at!r}" + ) + _assert_reads_back( + info, created.written.model_copy(update={"max_budget": None, "budget_duration": None}), replica + ) + + def test_key_serves_its_model_and_is_denied_others( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + _poll_chat_ok(client, created.key, mock_deployment) + + denied: Final = client.chat_status(created.key, DENIED_MODEL, f"say hi {unique_marker()}") + assert denied.status_code == 403, ( + f"chat on {DENIED_MODEL!r} outside the key's model list must be denied 403, got " + f"{denied.status_code}: {denied.body[:300]}" + ) + assert MODEL_ACCESS_DENIED_MARKER in denied.body, ( + f"403 body must be a model-access denial, got: {denied.body[:300]}" + ) + + def test_delete_revokes_info_and_chat_on_every_replica( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + """The teardown's deferred delete fires again on the already-deleted key by + design: the deferred cleanup must survive this test failing before the + in-body delete, and a repeat /key/delete is a cheap no-op the warn-only + teardown absorbs.""" + created: Final = _create_key(client, resources, mock_deployment) + _warm_every_replica(client, created.key, mock_deployment) + + client.delete_key_strict(created.key) + + _ = client.proxy.read_back_everywhere( + "/key/info", + params=KeyInfoParams(key=created.key), + response_type=KeyInfoResponse, + converged=_is_key_not_found, + ) + _assert_chat_rejected_everywhere(client, created.key, mock_deployment) diff --git a/tests/e2e/management/test_mcp_lifecycle_e2e.py b/tests/e2e/management/test_mcp_lifecycle_e2e.py new file mode 100644 index 00000000000..9257d697647 --- /dev/null +++ b/tests/e2e/management/test_mcp_lifecycle_e2e.py @@ -0,0 +1,294 @@ +"""Live e2e: the MCP server and toolset management routes' lifecycle contract. + +Two customer defects sit on these routes, and each step here is the read-back that +would have caught one of them: a dashboard edit that took several saves to stick +because the read landed on a replica the write had not reached, and a toolset whose +tools were stored under one name and read back under another, so it granted +nothing. Every read-back therefore polls every replica that serves the route +(ProxyClient.read_back_everywhere) and asserts the exact values written, and both +update routes are held to the same partial-update contract: a field left out of the +payload keeps its stored value, a field sent as null is cleared. The server URL is +unreachable on purpose; only persistence is under test, never a tool call. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from management_client import ManagementClient +from models import ( + McpInfo, + McpServerCreateBody, + McpServerListResponse, + McpServerRow, + McpServerUpdateBody, + ToolsetCreateBody, + ToolsetListResponse, + ToolsetRow, + ToolsetTool, + ToolsetUpdateBody, +) + +pytestmark = pytest.mark.e2e + +UNREACHABLE_URL: Final = "https://e2e-fake-mcp.test.local/mcp" + + +def _create_server(client: ManagementClient, resources: ResourceManager) -> tuple[McpServerCreateBody, str]: + name: Final = f"e2e_mcp_lifecycle_{unique_marker()}" + body: Final = McpServerCreateBody( + server_name=name, + alias=name, + url=UNREACHABLE_URL, + transport="http", + description="e2e lifecycle server", + mcp_info=McpInfo( + server_name=f"{name} (display)", + description="shown on the MCP page", + logo_url="https://e2e.test.local/logo.png", + ), + ) + server_id: Final = client.create_mcp_server(body).server_id + resources.defer(lambda: client.delete_mcp_server(server_id)) + return body, server_id + + +def _assert_server_matches(row: McpServerRow, written: McpServerCreateBody, *, where: str) -> None: + stored: Final = (row.server_name, row.alias, row.url, row.transport, row.description, row.mcp_info) + expected: Final = ( + written.server_name, + written.alias, + written.url, + written.transport, + written.description, + written.mcp_info, + ) + assert stored == expected, f"{where}: stored {stored}, expected {expected}" + + +def _server_everywhere( + client: ManagementClient, server_id: str, *, settled: Callable[[McpServerRow], bool] +) -> Mapping[str, McpServerRow]: + return client.proxy.read_body_back_everywhere(f"/v1/mcp/server/{server_id}", McpServerRow, settled=settled) + + +def _listed_server_everywhere(client: ManagementClient, server_id: str) -> Mapping[str, McpServerRow]: + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/server", + McpServerListResponse, + settled=lambda rows: any(row.server_id == server_id for row in rows.root), + ) + return {replica: next(row for row in rows.root if row.server_id == server_id) for replica, rows in listings.items()} + + +class TestMcpServerLifecycle: + @pytest.mark.covers("mgmt.mcp_server.new.persists") + def test_create_persists_every_field_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + by_id: Final = _server_everywhere(client, server_id, settled=lambda row: row.server_id == server_id) + for replica, row in by_id.items(): + _assert_server_matches(row, body, where=f"GET /v1/mcp/server/{server_id} on {replica}") + + @pytest.mark.skip( + reason=( + "product gap: GET /v1/mcp/server builds each row from the in-memory registry, whose " + "_build_mcp_server_table sets description from mcp_info['description'], so the list " + "reports the mcp_info description while GET /v1/mcp/server/{server_id} reports the " + "stored description column. A server created with both set to different text reads " + "back with two different descriptions depending on the route" + ) + ) + @pytest.mark.covers("mgmt.mcp_server.list.persists") + def test_created_server_is_listed_with_every_field( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + for replica, row in _listed_server_everywhere(client, server_id).items(): + _assert_server_matches(row, body, where=f"GET /v1/mcp/server on {replica}") + + @pytest.mark.covers("mgmt.mcp_server.update.preserves_unrelated_fields") + def test_updating_only_the_alias_keeps_every_other_field_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + renamed: Final = f"{body.alias}_renamed" + + _ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, alias=renamed)) + + after_one_put: Final = _server_everywhere(client, server_id, settled=lambda row: row.alias == renamed) + for replica, row in after_one_put.items(): + _assert_server_matches( + row, + body.model_copy(update={"alias": renamed}), + where=f"GET /v1/mcp/server/{server_id} on {replica} after one PUT of alias", + ) + + @pytest.mark.covers("mgmt.mcp_server.update.clear_persists") + def test_clearing_the_description_with_null_reads_back_null( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + _ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, description=None)) + + cleared: Final = _server_everywhere(client, server_id, settled=lambda row: row.description is None) + for replica, row in cleared.items(): + _assert_server_matches( + row, + body.model_copy(update={"description": None}), + where=f"GET /v1/mcp/server/{server_id} on {replica} after PUT description=null", + ) + + @pytest.mark.covers("mgmt.mcp_server.delete.persists") + def test_delete_removes_the_server_from_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + + _ = unwrap(client.delete_mcp_server(server_id)) + + gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/server/{server_id}") + assert set(gone.values()) == {404}, f"a deleted server must 404 on every replica; got {dict(gone)}" + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/server", + McpServerListResponse, + settled=lambda rows: all(row.server_id != server_id for row in rows.root), + ) + for replica, rows in listings.items(): + assert all(row.server_id != server_id for row in rows.root), ( + f"GET /v1/mcp/server on {replica} still lists the deleted server {server_id}" + ) + + +def _create_toolset( + client: ManagementClient, resources: ResourceManager, server_id: str +) -> tuple[ToolsetCreateBody, str]: + body: Final = ToolsetCreateBody( + toolset_name=f"e2e_toolset_{unique_marker()}", + description="e2e lifecycle toolset", + tools=[ + ToolsetTool(server_id=server_id, tool_name="search_datadog_logs"), + ToolsetTool(server_id=server_id, tool_name="get_datadog_metric"), + ], + ) + toolset_id: Final = client.proxy.create_toolset(body).toolset_id + resources.defer(lambda: client.proxy.delete_toolset(toolset_id)) + return body, toolset_id + + +def _assert_toolset_matches(row: ToolsetRow, written: ToolsetCreateBody, *, where: str) -> None: + stored: Final = (row.toolset_name, row.description, row.tools) + expected: Final = (written.toolset_name, written.description, written.tools) + assert stored == expected, f"{where}: stored {stored}, expected {expected}" + + +def _toolset_everywhere( + client: ManagementClient, toolset_id: str, *, settled: Callable[[ToolsetRow], bool] +) -> Mapping[str, ToolsetRow]: + return client.proxy.read_body_back_everywhere(f"/v1/mcp/toolset/{toolset_id}", ToolsetRow, settled=settled) + + +class TestMcpToolsetLifecycle: + @pytest.mark.covers("mgmt.mcp_toolset.new.persists") + def test_create_persists_both_tools_under_the_exact_names_written( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + by_id: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.toolset_id == toolset_id) + for replica, row in by_id.items(): + _assert_toolset_matches(row, body, where=f"GET /v1/mcp/toolset/{toolset_id} on {replica}") + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/toolset", + ToolsetListResponse, + settled=lambda rows: any(row.toolset_id == toolset_id for row in rows.root), + ) + for replica, rows in listings.items(): + _assert_toolset_matches( + next(row for row in rows.root if row.toolset_id == toolset_id), + body, + where=f"GET /v1/mcp/toolset on {replica}", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.preserves_unrelated_fields") + def test_updating_only_the_description_keeps_the_tools_and_name( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description="edited")) + + edited: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description == "edited") + for replica, row in edited.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"description": "edited"}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of description", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.persists") + def test_updating_the_tools_to_one_entry_reads_back_exactly_that_entry( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + kept: Final = body.tools[:1] + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, tools=kept)) + + narrowed: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.tools == kept) + for replica, row in narrowed.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"tools": kept}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of one tool", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.clear_persists") + def test_clearing_the_description_with_null_reads_back_null( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description=None)) + + cleared: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description is None) + for replica, row in cleared.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"description": None}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT description=null", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.delete.persists") + def test_delete_removes_the_toolset_from_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + _, toolset_id = _create_toolset(client, resources, server_id) + + _ = unwrap(client.proxy.delete_toolset(toolset_id)) + + gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/toolset/{toolset_id}") + assert set(gone.values()) == {404}, f"a deleted toolset must 404 on every replica; got {dict(gone)}" + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/toolset", + ToolsetListResponse, + settled=lambda rows: all(row.toolset_id != toolset_id for row in rows.root), + ) + for replica, rows in listings.items(): + assert all(row.toolset_id != toolset_id for row in rows.root), ( + f"GET /v1/mcp/toolset on {replica} still lists the deleted toolset {toolset_id}" + ) diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index d1ea53a0b3b..352b4446cfd 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from collections.abc import Sequence from e2e_config import datadog_mcp_url, unique_marker from lifecycle import ResourceManager @@ -35,7 +36,11 @@ def register_datadog_mcp( resources: ResourceManager, *, mcp_access_groups: list[str] | None = None, + allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,), ) -> str: + """Register the core Datadog toolset with its credentials from the env. By default + the server exposes only `search_datadog_logs`; pass `allowed_tools=None` to expose + every tool the core toolset serves.""" assert_dd_mcp_creds() name = f"e2e_dd_mcp_{unique_marker()}" server_id = client.register_server( @@ -47,7 +52,7 @@ def register_datadog_mcp( "DD-API-KEY": _dd_api_key(), "DD-APPLICATION-KEY": _dd_app_key(), }, - allowed_tools=[SEARCH_LOGS_TOOL], + allowed_tools=None if allowed_tools is None else list(allowed_tools), mcp_access_groups=mcp_access_groups, ) resources.defer(lambda: client.delete_server(server_id)) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 73453478e5a..210fc7a1e98 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -16,11 +16,11 @@ import time from collections.abc import Mapping from dataclasses import dataclass -from pydantic import BaseModel, ConfigDict, Field, RootModel +from pydantic import BaseModel, ConfigDict, Field from e2e_config import settle_propagation from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap -from models import KeyGenerateBody, ObjectPermission +from models import KeyGenerateBody, McpServerListResponse, McpServerRow, ObjectPermission from proxy_client import ProxyClient McpToolArg = str | int | float | bool | list[str] | dict[str, str] @@ -46,16 +46,6 @@ class McpServerNewResponse(BaseModel): server_id: str -class McpServerRow(BaseModel): - server_id: str - alias: str | None = None - url: str | None = None - - -class McpServersListResponse(RootModel[list[McpServerRow]]): - pass - - class McpToolMcpInfo(BaseModel): server_id: str | None = None alias: str | None = None @@ -193,7 +183,7 @@ class McpClient: "/v1/mcp/server", headers=self.proxy.transport.master, params=NoBody(), - response_type=McpServersListResponse, + response_type=McpServerListResponse, ) ).root @@ -224,11 +214,16 @@ class McpClient: user_id: str, mcp_servers: list[str] | None, mcp_access_groups: list[str] | None = None, + mcp_toolsets: list[str] | None = None, models: list[str] | None = None, ) -> str: object_permission = ( - ObjectPermission(mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups) - if mcp_servers is not None or mcp_access_groups is not None + ObjectPermission( + mcp_servers=mcp_servers, + mcp_access_groups=mcp_access_groups, + mcp_toolsets=mcp_toolsets, + ) + if mcp_servers is not None or mcp_access_groups is not None or mcp_toolsets is not None else None ) return self.proxy.generate_key( @@ -272,6 +267,20 @@ class McpClient: ) time.sleep(self.proxy.poll_interval) + def await_tools(self, key: str, server_id: str, *, expected: frozenset[str]) -> frozenset[str]: + """Poll tools/list until `server_id`'s tools as `key` sees them are exactly + `expected`, and return the last listing either way, so the caller's equality + assertion names the difference. Fails at poll_timeout only when the read + itself never succeeded.""" + deadline = time.monotonic() + self.proxy.poll_timeout + while True: + result = self.list_tools(key) + if isinstance(result, Success) and result.data.tool_names_for_server(server_id) == expected: + return expected + if time.monotonic() >= deadline: + return unwrap(result).tool_names_for_server(server_id) + time.sleep(self.proxy.poll_interval) + def await_call_tool( self, key: str, diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py new file mode 100644 index 00000000000..6b901145eb1 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -0,0 +1,95 @@ +"""Live e2e: a key granted a toolset lists exactly the toolset's tools. + +An admin registers the real Datadog remote MCP server with its whole core toolset +exposed, discovers two of its tool names through a key granted the server outright, +and curates a toolset naming exactly those two. A second key is granted the server +plus that toolset, and its tools/list must come back as exactly those two names: no +more, so the rest of the server's catalog stays hidden behind the toolset, and no +fewer, so a tool stored under one name and read under another (which granted +nothing) fails here first. Requires DD_API_KEY + DD_APP_KEY (the suite's real MCP +upstream). +""" + +from __future__ import annotations + +from typing import Final + +import pytest +from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import ToolsetCreateBody, ToolsetTool + +pytestmark = pytest.mark.e2e + + +def _key( + client: McpClient, + resources: ResourceManager, + label: str, + *, + server_id: str, + toolset_id: str | None = None, +) -> str: + key: Final = client.generate_key( + user_id=f"e2e-mcp-{label}-{unique_marker()}", + mcp_servers=[server_id], + mcp_toolsets=None if toolset_id is None else [toolset_id], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + +def _wire_prefix(wire_name: str, tool_name: str, catalog: frozenset[str]) -> str: + """The prefix tools/list puts in front of one server's tool names, measured off a + tool whose own name is known rather than guessed from the alias. A toolset grants + by the tool's own name, never the wire name, and the prefix is whatever the proxy + is configured to build (the alias, or a short server id), so measuring it is the + only way to cross between the two.""" + assert wire_name.endswith(tool_name), f"tools/list served {wire_name!r}, expected it to end with {tool_name!r}" + prefix: Final = wire_name[: len(wire_name) - len(tool_name)] + unprefixed: Final = frozenset(name for name in catalog if not name.startswith(prefix)) + assert not unprefixed, ( + f"every tool of one server shares the wire prefix {prefix!r}, so {sorted(unprefixed)} " + f"cannot be reduced to the names a toolset grants by" + ) + return prefix + + +class TestMcpToolsetEnforcement: + @pytest.mark.covers("mcp.list_tools.api_key.toolset_scoped") + def test_key_granted_a_toolset_lists_exactly_its_tools(self, client: McpClient, resources: ResourceManager) -> None: + server_id: Final = register_datadog_mcp(client, resources, allowed_tools=None) + client.await_registered(server_id) + + catalog_key: Final = _key(client, resources, "catalog", server_id=server_id) + known_wire: Final = client.await_tool(catalog_key, server_id, SEARCH_LOGS_TOOL) + catalog: Final = unwrap(client.list_tools(catalog_key)).tool_names_for_server(server_id) + assert len(catalog) > 2, ( + f"the Datadog core toolset must serve more tools than the toolset names, or the " + f"restriction has nothing to hide; got {sorted(catalog)}" + ) + prefix: Final = _wire_prefix(known_wire, SEARCH_LOGS_TOOL, catalog) + chosen_wire: Final = frozenset(sorted(catalog)[:2]) + chosen: Final = frozenset(name.removeprefix(prefix) for name in chosen_wire) + + toolset: Final = client.proxy.create_toolset( + ToolsetCreateBody( + toolset_name=f"e2e_toolset_{unique_marker()}", + description="two Datadog tools", + tools=[ToolsetTool(server_id=server_id, tool_name=name) for name in sorted(chosen)], + ) + ) + resources.defer(lambda: client.proxy.delete_toolset(toolset.toolset_id)) + assert frozenset(tool.tool_name for tool in toolset.tools) == chosen, ( + f"toolset stored {toolset.tools}, expected the two names {sorted(chosen)} verbatim" + ) + + scoped_key: Final = _key(client, resources, "toolset", server_id=server_id, toolset_id=toolset.toolset_id) + listed: Final = client.await_tools(scoped_key, server_id, expected=chosen_wire) + assert listed == chosen_wire, ( + f"a key granted the toolset must list exactly its two tools; " + f"got {sorted(listed)}, expected {sorted(chosen_wire)}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 2c6c0e9bbd4..faf8557498b 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -8,9 +8,10 @@ from __future__ import annotations from collections.abc import Sequence from datetime import datetime -from typing import Literal +from typing import Final, Literal -from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_validator +from e2e_http import PartialBody +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator # ---------- keys ---------- @@ -49,11 +50,13 @@ class KeyMetadata(BaseModel): logging: list[KeyLoggingCallback] | None = None priority: str | None = None batch_enqueued_token_limit: int | None = None + tag: str | None = None class ObjectPermission(BaseModel): mcp_servers: list[str] | None = None mcp_access_groups: list[str] | None = None + mcp_toolsets: list[str] | None = None class KeyGenerateBody(BaseModel): @@ -76,11 +79,20 @@ class KeyGenerateBody(BaseModel): allowed_passthrough_routes: list[str] | None = None metadata: KeyMetadata | None = None object_permission: ObjectPermission | None = None - router_settings: "RouterSettingsOverride | None" = None + router_settings: RouterSettingsOverride | None = None class KeyGenerateResponse(BaseModel): key: str + token: str | None = None + key_alias: str | None = None + models: list[str] = [] + max_budget: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + budget_duration: str | None = None + team_id: str | None = None + metadata: KeyMetadata | None = None class KeyRegenerateBody(BaseModel): @@ -122,6 +134,7 @@ class KeyInfo(BaseModel): blocked: bool | None = None spend: float | None = None max_budget: float | None = None + budget_duration: str | None = None budget_reset_at: str | None = None budget_id: str | None = None litellm_budget_table: LiteLLMBudgetTable | None = None @@ -291,6 +304,7 @@ class RouterSettingsOverride(BaseModel): context_window_fallbacks: list[dict[str, list[str]]] | None = None content_policy_fallbacks: list[dict[str, list[str]]] | None = None num_retries: int | None = None + model_group_retry_policy: dict[str, dict[str, int]] | None = None enable_tag_filtering: bool | None = None @@ -505,6 +519,15 @@ class CountTokensResponse(BaseModel): # ---------- mcp servers ---------- +class McpInfo(BaseModel): + """The `mcp_info` display block stored on an MCP server; only the fields the + lifecycle test writes and reads back.""" + + server_name: str | None = None + description: str | None = None + logo_url: str | None = None + + class McpServerCreateBody(BaseModel): """POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is `oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints @@ -519,6 +542,18 @@ class McpServerCreateBody(BaseModel): oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None authorization_url: str | None = None token_url: str | None = None + server_name: str | None = None + description: str | None = None + mcp_info: McpInfo | None = None + + +class McpServerUpdateBody(PartialBody): + """PUT /v1/mcp/server: a field left unset keeps its stored value, a field set + to None is cleared.""" + + server_id: str + alias: str | None = None + description: str | None = None class McpServerInfo(BaseModel): @@ -532,6 +567,54 @@ class McpServerInfo(BaseModel): allow_all_keys: bool | None = None +class McpServerRow(McpServerInfo): + """A stored MCP server as the create, get, and list routes return it: the + fields the lifecycle test asserts survive the round trip.""" + + server_name: str | None = None + transport: str | None = None + description: str | None = None + mcp_info: McpInfo | None = None + + +class McpServerListResponse(RootModel[list[McpServerRow]]): + """GET /v1/mcp/server answers with a bare array of servers.""" + + +class ToolsetTool(BaseModel): + server_id: str + tool_name: str + + +class ToolsetCreateBody(BaseModel): + toolset_name: str + description: str | None = None + tools: list[ToolsetTool] + + +class ToolsetUpdateBody(PartialBody): + """PUT /v1/mcp/toolset: a field left unset keeps its stored value, a field set + to None is cleared.""" + + toolset_id: str + description: str | None = None + tools: list[ToolsetTool] | None = None + + +class ToolsetRow(BaseModel): + """A stored toolset as POST /v1/mcp/toolset, GET /v1/mcp/toolset/{toolset_id}, + and each row of GET /v1/mcp/toolset return it.""" + + toolset_id: str + toolset_name: str + description: str | None = None + tools: list[ToolsetTool] = Field(default_factory=list) + + +class ToolsetListResponse(RootModel[list[ToolsetRow]]): + """GET /v1/mcp/toolset answers with a bare array of toolsets.""" + + class EmbedBody(BaseModel): model: str input: str @@ -573,6 +656,28 @@ class OcrResponse(BaseModel): # ---------- spend logs ---------- +class GuardrailEntityMatch(BaseModel): + entity_type: str + score: float + start: int + end: int + + +class GuardrailRunRecord(BaseModel): + guardrail_name: str | None = None + guardrail_mode: str | None = None + guardrail_status: str | None = None + guardrail_provider: str | None = None + masked_entity_count: dict[str, int] | None = None + guardrail_response: object | None = None + + +class SpendLogMetadata(BaseModel): + user_api_key_alias: str | None = None + applied_guardrails: list[str] | None = None + guardrail_information: list[GuardrailRunRecord] | None = None + + class SpendLogRow(BaseModel): request_id: str | None = None api_key: str | None = None @@ -589,6 +694,7 @@ class SpendLogRow(BaseModel): completion_tokens: int | None = None total_tokens: int | None = None request_tags: list[str] | None = None + metadata: SpendLogMetadata | None = None class SpendLogs(RootModel[list[SpendLogRow]]): @@ -918,12 +1024,34 @@ class CredentialCreateResponse(BaseModel): # ---------- key / team / user / organization management ---------- +class Cleared(BaseModel): + """An explicit JSON null in a merge-patch body. The transport drops `None` fields + before sending (`exclude_none`), so `None` means "leave the stored value alone"; a + field set to `CLEAR` reaches the wire as `null`, which tells the proxy to clear it.""" + + model_config = ConfigDict(frozen=True) + + @model_serializer + def _as_null(self) -> None: + return None + + +CLEAR: Final = Cleared() + + class KeyUpdateBody(BaseModel): + """POST /key/update is a merge patch: a field left `None` is dropped from the body and + keeps its stored value, `CLEAR` sends an explicit null that clears it (`budget_duration` + clears `budget_reset_at` with it), and `metadata` replaces the stored metadata wholesale.""" + key: str models: list[str] | None = None key_alias: str | None = None tpm_limit: int | None = None rpm_limit: int | None = None + max_budget: float | Cleared | None = None + budget_duration: str | Cleared | None = None + metadata: KeyMetadata | None = None class KeyBlockBody(BaseModel): diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 5292f56b324..506edd22d72 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -10,12 +10,18 @@ from __future__ import annotations import time import warnings -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass +from functools import reduce from datetime import datetime +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel from e2e_http import ( AnthropicHeaders, + AuthHeaders, NoBody, ProbeResult, Result, @@ -71,6 +77,9 @@ from models import ( TeamNewResponse, UserDeleteBody, UserDeleteResponse, + ToolsetCreateBody, + ToolsetRow, + ToolsetUpdateBody, ) from e2e_config import ( CONTROL_PLANE_BASE_URL, @@ -78,18 +87,20 @@ from e2e_config import ( POLL_INTERVAL, POLL_TIMEOUT, PROXY_BASE_URL, + PROXY_REPLICA_URLS, REQUEST_TIMEOUT, SLOW_PROVIDER_TIMEOUT_SECONDS, settle_propagation, ) -from transport import HttpTransport, SplitTransport, Transport +from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path 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 @@ -115,6 +126,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]], @@ -140,9 +161,7 @@ def await_servable( last_result: Result[ModelsListResponse] | None = None while True: t = now() - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) + phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds remaining = phase_deadline - t if remaining <= 0: if ( @@ -155,9 +174,7 @@ def await_servable( poll_timeout = min(request_timeout, remaining) last_result = list_models(poll_timeout) - listed = isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ) + listed = isinstance(last_result, Success) and any(entry.id == model_name for entry in last_result.data.data) t = now() if not listed: first_seen_at = None @@ -170,17 +187,47 @@ def await_servable( elif t - first_seen_at >= db_sync_seconds: return Servable() - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) + phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds wait = min(interval, phase_deadline - now()) if wait > 0: 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, @@ -191,16 +238,193 @@ 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}" ) +type ReplicaRead[T] = Callable[[float], T] + + +@dataclass(frozen=True, slots=True) +class EverywhereConverged[T]: + """Every replica answered with something `settled` accepts, keyed by replica.""" + + answers: Mapping[str, T] + + +@dataclass(frozen=True, slots=True) +class NeverConvergedOn[T]: + """`replica` ran out its budget without an answer `settled` accepts; `last` is + its final answer, so the failure can say what that replica still serves.""" + + replica: str + last: T + + +def _last_answer[T]( + read: ReplicaRead[T], + *, + settled: Callable[[T], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> T: + """Poll `read` until `settled` accepts its answer or `timeout` runs out, and + return the last answer either way. Each read's request timeout is clamped to + the budget left, and the final poll runs even when less than an interval + remains, so a deadline never skips the read that would have settled.""" + deadline: Final = now() + timeout + answer = read(min(request_timeout, timeout)) + while not settled(answer): + remaining = deadline - now() + if remaining <= 0: + return answer + sleep(min(interval, remaining)) + answer = read(min(request_timeout, remaining)) + return answer + + +def await_everywhere[T]( + reads: Mapping[str, ReplicaRead[T]], + *, + settled: Callable[[T], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> EverywhereConverged[T] | NeverConvergedOn[T]: + """`_last_answer` against every replica in turn, each with the full budget, so a + write counts as visible only once the last replica reflects it, and stop at the + first replica that never converges. Clock and sleep are injected.""" + def read_replica( + outcome: EverywhereConverged[T] | NeverConvergedOn[T], + item: tuple[str, ReplicaRead[T]], + ) -> EverywhereConverged[T] | NeverConvergedOn[T]: + if isinstance(outcome, NeverConvergedOn): + return outcome + replica, read = item + answer: Final = _last_answer( + read, + settled=settled, + timeout=timeout, + interval=interval, + request_timeout=request_timeout, + now=now, + sleep=sleep, + ) + if not settled(answer): + return NeverConvergedOn(replica=replica, last=answer) + return EverywhereConverged(answers=MappingProxyType({**outcome.answers, replica: answer})) + + initial: Final[EverywhereConverged[T] | NeverConvergedOn[T]] = EverywhereConverged(answers=MappingProxyType({})) + return reduce(read_replica, reads.items(), initial) + + +def _is_not_found[R: BaseModel](result: Result[R]) -> bool: + return isinstance(result, UnknownApiError) and result.status_code == 404 + + +def _status_of[R: BaseModel](result: Result[R]) -> int: + match result: + case Success(status_code=status_code) | UnknownApiError(status_code=status_code): + return status_code + case _: + return -1 + + +type Poller[T] = Callable[[], T] + + +@dataclass(frozen=True, slots=True) +class Converged[T]: + result: T + + +@dataclass(frozen=True, slots=True) +class NotConverged[T]: + """The deadline passed without a read satisfying the predicate; `last_result` is + the final read, so the caller can tell a stale body from a failed request.""" + + last_result: T + + +type ConvergeOutcome[T] = Converged[T] | NotConverged[T] + + +def await_converged[T]( + poll: Poller[T], + *, + converged: Callable[[T], bool], + timeout: float, + interval: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ConvergeOutcome[T]: + """Poll until a read satisfies `converged` or `timeout` elapses. + + Polls before testing the deadline, so a zero or already-spent budget still gets one + attempt, and sleeps only min(interval, time left), so the attempt that lands exactly + on the deadline is taken rather than skipped. Clock and sleep are injected.""" + deadline: Final = now() + timeout + while True: + result = poll() + if converged(result): + return Converged(result=result) + remaining = deadline - now() + if remaining <= 0: + return NotConverged(last_result=result) + sleep(min(interval, remaining)) + + +def await_converged_everywhere[T]( + pollers: Mapping[str, Poller[T]], + *, + converged: Callable[[T], bool], + timeout: float, + interval: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> Mapping[str, ConvergeOutcome[T]]: + """`await_converged` against every replica in turn, each with the full budget, so a + replica that lags behind the one a write landed on is polled until it catches up + rather than failing on its first stale read.""" + return MappingProxyType( + { + replica: await_converged( + poll, converged=converged, timeout=timeout, interval=interval, now=now, sleep=sleep + ) + for replica, poll in pollers.items() + } + ) + + +def first_lagging_replica[T]( + outcomes: Mapping[str, ConvergeOutcome[T]], +) -> tuple[str, NotConverged[T]] | None: + return next( + ((replica, outcome) for replica, outcome in outcomes.items() if isinstance(outcome, NotConverged)), + None, + ) + + +def converge_timeout_message(*, what: str, replica: str, timeout: float, last_result: object) -> str: + return ( + f"{what} on {replica} never converged within {timeout}s of the write " + f"(control/data-plane propagation issue); last read: {last_result}" + ) + + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport + replicas: Mapping[str, Transport] + control_replicas: Mapping[str, Transport] poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT @@ -248,6 +472,52 @@ class ProxyClient: ) ).info + def read_back_everywhere[R: BaseModel]( + self, + path: str, + *, + params: BaseModel, + response_type: type[R], + converged: Callable[[Result[R]], bool], + ) -> Mapping[str, Result[R]]: + """GET `path` under the master key on every replica in PROXY_REPLICA_URLS (the + data-plane URL alone when the stack exports no per-gateway addresses), polling + each to poll_timeout until its read satisfies `converged`. Returns that read per + replica, or fails naming the first replica that never converged and its last + read. Behind a load balancer the single address proves one replica converged, + not all of them; only per-gateway addresses make this a fleet-wide proof.""" + outcomes: Final = await_converged_everywhere( + { + url: self._body_poller(transport, path, params, response_type) + for url, transport in self.replicas.items() + }, + converged=converged, + timeout=self.poll_timeout, + interval=self.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + lagging: Final = first_lagging_replica(outcomes) + if lagging is not None: + replica, outcome = lagging + raise AssertionError( + converge_timeout_message( + what=f"GET {path}", + replica=replica, + timeout=self.poll_timeout, + last_result=outcome.last_result, + ) + ) + return MappingProxyType( + {replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)} + ) + + @staticmethod + def _body_poller[R: BaseModel]( + transport: Transport, path: str, params: BaseModel, response_type: type[R] + ) -> Poller[Result[R]]: + return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type) + def model_info(self) -> list[ModelInfoEntry]: """Every configured deployment with the price the proxy resolved for it (config override merged over cost-map defaults).""" @@ -278,9 +548,7 @@ class ProxyClient: response_type=FileListResponse, ) - def list_fine_tuning_jobs( - self, key: str, params: FineTuningJobsParams - ) -> Result[FineTuningJobsResponse]: + def list_fine_tuning_jobs(self, key: str, params: FineTuningJobsParams) -> Result[FineTuningJobsResponse]: return self.transport.get( "/v1/fine_tuning/jobs", headers=self.transport.bearer(key), @@ -317,12 +585,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", @@ -332,21 +601,19 @@ class ProxyClient: ) ).model_id written_at = time.monotonic() - self._await_model_servable(body.model_name, listed_for) + try: + self._await_model_servable(body.model_name, listed_for) + except BaseException: + self.delete_model(model_id) + raise settle_propagation(written_at) 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, @@ -358,16 +625,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 @@ -395,6 +673,112 @@ class ProxyClient: if not is_ok(result): warnings.warn(f"delete_model({model_id!r}) failed: {result}", stacklevel=2) + # ---- replica read-back ---------------------------------------------- + + def replicas_for(self, path: str) -> Mapping[str, Transport]: + """The replicas that serve `path`: every data-plane replica for an LLM route, + and for a management route the control-plane replicas, since the data-plane + replicas trim management routes and answer them 404. A monolith serves both + from every replica, so a management read-back polls all of them; a split + deployment exposes one control-plane address (there is one backend process + behind it on the stack these suites run against), so it polls that. A + control plane fronting several backends would need its own replica list to + prove each one converged, the way PROXY_REPLICA_URLS does for the gateways. + Never empty: a read-back against no replica would assert nothing and pass.""" + replicas: Final = self.control_replicas if is_control_plane_path(path) else self.replicas + assert replicas, f"no replica is configured to serve {path}, so a read-back there would prove nothing" + return replicas + + def read_body_back_everywhere[R: BaseModel]( + self, path: str, response_type: type[R], *, settled: Callable[[R], bool] + ) -> Mapping[str, R]: + """GET `path` on every replica that serves it, polling each to poll_timeout + until `settled` accepts its body, and fail naming the first replica that + never converged. Returns each replica's settled body, keyed by replica, so + the caller can assert the rest of it.""" + outcome: Final = await_everywhere( + {url: self._reader(transport, path, response_type) for url, transport in self.replicas_for(path).items()}, + settled=lambda result: isinstance(result, Success) and settled(result.data), + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case EverywhereConverged(answers=answers): + return MappingProxyType({url: unwrap(result) for url, result in answers.items()}) + case NeverConvergedOn(replica=replica, last=last): + raise AssertionError( + f"GET {path} on {replica} never converged within {self.poll_timeout}s of the write; " + f"last read: {last}" + ) + + def gone_everywhere(self, path: str) -> Mapping[str, int]: + """Poll GET `path` on every replica that serves it until each stops serving + it, and fail naming the first replica that still does at poll_timeout. + Returns each replica's final status, so the caller asserts the 404 itself.""" + outcome: Final = await_everywhere( + {url: self._reader(transport, path, NoBody) for url, transport in self.replicas_for(path).items()}, + settled=_is_not_found, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case EverywhereConverged(answers=answers): + return MappingProxyType({url: _status_of(result) for url, result in answers.items()}) + case NeverConvergedOn(replica=replica, last=last): + raise AssertionError( + f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}" + ) + + @staticmethod + def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: + return lambda request_timeout: transport.get( + path, + headers=transport.master, + params=NoBody(), + response_type=response_type, + timeout=request_timeout, + ) + + # ---- mcp toolsets --------------------------------------------------- + + def create_toolset(self, body: ToolsetCreateBody) -> ToolsetRow: + return unwrap( + self.transport.post( + "/v1/mcp/toolset", + headers=self.transport.master, + json=body, + response_type=ToolsetRow, + ) + ) + + def update_toolset(self, body: ToolsetUpdateBody) -> ToolsetRow: + """PUT /v1/mcp/toolset: a partial update where a field left unset keeps its + stored value and None clears it.""" + return unwrap( + self.transport.put( + "/v1/mcp/toolset", + headers=self.transport.master, + json=body, + response_type=ToolsetRow, + ) + ) + + def delete_toolset(self, toolset_id: str) -> Result[NoBody]: + """DELETE /v1/mcp/toolset/{toolset_id}. Returns the outcome so the act phase + can unwrap it while a deferred teardown can ignore an already-deleted row.""" + return self.transport.delete( + f"/v1/mcp/toolset/{toolset_id}", + headers=self.transport.master, + json=NoBody(), + response_type=NoBody, + ) + def create_credential(self, body: CredentialCreateBody) -> None: unwrap( self.transport.post( @@ -589,16 +973,23 @@ 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. Management read-backs poll those same + replicas when the two planes share a base URL (a monolith, where every replica + serves every route) and the control plane alone when they differ (a split + deployment, where the data-plane replicas do not serve management routes). 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 @@ -615,8 +1006,19 @@ 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 + } + ) + control_replicas: Final = ( + replicas if control_plane_base_url == base_url else MappingProxyType({control_plane_base_url: split.control}) + ) return ProxyClient( transport=split, + replicas=replicas, + control_replicas=control_replicas, poll_timeout=POLL_TIMEOUT, poll_interval=POLL_INTERVAL, ) diff --git a/tests/e2e/quota_management/spend_tracking/conftest.py b/tests/e2e/quota_management/spend_tracking/conftest.py index 0597c9af400..9c8ffd18144 100644 --- a/tests/e2e/quota_management/spend_tracking/conftest.py +++ b/tests/e2e/quota_management/spend_tracking/conftest.py @@ -36,6 +36,7 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = ( ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY"), ("openai-text-embedding-3-small", "openai/text-embedding-3-small", "OPENAI_API_KEY"), ("openai-responses-codex", "openai/gpt-5.3-codex", "OPENAI_API_KEY"), + ("openai-gpt-4o-mini", "openai/gpt-4o-mini", "OPENAI_API_KEY"), ) diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 056799b8499..9ac97f57f47 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -15,9 +15,12 @@ import time from collections.abc import Callable from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from typing import Final from e2e_config import unique_marker from e2e_http import ( + FileUploadForm, + Headers, NoBody, ProbeResult, Result, @@ -35,6 +38,8 @@ from models import ( DateRangeParams, EmbedBody, EmbedResponse, + KeyGenerateBody, + KeyGenerateResponse, OpenAPISchema, SpendCalculateBody, SpendCalculateResponse, @@ -43,13 +48,27 @@ from models import ( SpendLogsPageParams, SpendTagsResponse, TagSpend, + UserDeleteBody, + UserDeleteResponse, + UserNewBody, + UserNewResponse, + UserRole, ) -from proxy_client import ProxyClient +from proxy_client import Converged, ProxyClient, await_converged +from pydantic import BaseModel, Field __all__ = [ + "BatchCreateBody", + "CallbackLogMetadata", + "CallbackLogPayload", + "BatchObject", + "DailyActivityKeyBreakdown", + "FileObject", "ProbeResult", + "ResponseIdentity", "SpendClient", "SpendLogRow", + "StreamingResponse", "build_client", "is_ok", "unique_marker", @@ -57,6 +76,139 @@ __all__ = [ ] +class GeminiApiKeyHeaders(Headers): + x_goog_api_key: str = Field(serialization_alias="x-goog-api-key") + content_type: str = Field(default="application/json", serialization_alias="Content-Type") + + +class GeminiPart(BaseModel): + text: str + + +class GeminiContent(BaseModel): + parts: list[GeminiPart] + + +class GeminiGenerationConfig(BaseModel): + maxOutputTokens: int + + +class GeminiGenerateBody(BaseModel): + contents: list[GeminiContent] + generationConfig: GeminiGenerationConfig + + +class ResponsesBody(BaseModel): + model: str + input: str + cache: dict[str, bool] | None = {"no-cache": True} + + +class QueuedChatBody(ChatBody): + priority: int = 0 + + +class ResponseIdentity(BaseModel): + id: str | None = None + + +class HealthParams(BaseModel): + model: str + + +class ModelQuery(BaseModel): + model: str + + +class FileObject(BaseModel): + id: str + + +class BatchCreateBody(BaseModel): + input_file_id: str + endpoint: str = "/v1/chat/completions" + completion_window: str = "24h" + model: str + metadata: dict[str, str] + + +class BatchObject(BaseModel): + id: str + status: str + + +class ProviderQuery(BaseModel): + provider: str + + +class CallbackLogMetadata(BaseModel): + user_api_key_hash: str + user_api_key_alias: str + user_api_key_user_id: str + + +class CallbackLogPayload(BaseModel): + id: str + litellm_call_id: str + model: str + call_type: str = "acompletion" + start_time: float = Field(serialization_alias="startTime") + end_time: float = Field(serialization_alias="endTime") + response_cost: float + prompt_tokens: int + completion_tokens: int + total_tokens: int + metadata: CallbackLogMetadata + + +class CallbackLogRecord(BaseModel): + status: str = "success" + standard_logging_payload: CallbackLogPayload + + +class CallbackLogsRequest(BaseModel): + records: list[CallbackLogRecord] + + +class CallbackLogsResponse(BaseModel): + processed: int + failed: int + + +class DailyActivityParams(BaseModel): + start_date: str + end_date: str + api_key: str + + +class DailyActivityKeyMetadata(BaseModel): + key_alias: str | None = None + team_id: str | None = None + user_email: str | None = None + + +class DailyActivityKeyMetrics(BaseModel): + api_requests: int = 0 + + +class DailyActivityKeyBreakdown(BaseModel): + metrics: DailyActivityKeyMetrics + metadata: DailyActivityKeyMetadata + + +class DailyActivityBreakdown(BaseModel): + api_keys: dict[str, DailyActivityKeyBreakdown] = {} + + +class DailyActivityRow(BaseModel): + date: str + breakdown: DailyActivityBreakdown + + +class DailyActivityResponse(BaseModel): + results: list[DailyActivityRow] = [] + + def _chat_body( model: str, content: str, @@ -207,6 +359,166 @@ class SpendClient: def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult: return self.proxy.transport.probe(path, params=params) + def create_user(self, *, email: str, role: UserRole, user_id: str) -> str: + return unwrap( + self.proxy.transport.post( + "/user/new", + headers=self.proxy.transport.master, + json=UserNewBody(user_email=email, user_role=role, user_id=user_id), + response_type=UserNewResponse, + ) + ).user_id + + def delete_user(self, user_id: str) -> None: + _ = unwrap( + self.proxy.transport.post( + "/user/delete", + headers=self.proxy.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=UserDeleteResponse, + ) + ) + + def generate_key_record(self, body: KeyGenerateBody) -> KeyGenerateResponse: + return unwrap( + self.proxy.transport.post( + "/key/generate", + headers=self.proxy.transport.master, + json=body, + response_type=KeyGenerateResponse, + ) + ) + + def send_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + "/chat/completions", + headers=self.proxy.transport.bearer(key), + json=_chat_body(model, content, max_tokens=max_tokens), + ) + + def send_queued_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + "/queue/chat/completions", + headers=self.proxy.transport.bearer(key), + json=QueuedChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + ), + ) + + def send_messages(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + "/v1/messages", + headers=self.proxy.transport.bearer(key), + json=AnthropicMessagesBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + ), + ) + + def send_responses(self, key: str, model: str, content: str) -> StreamingResponse: + return self.proxy.transport.send( + "/v1/responses", + headers=self.proxy.transport.bearer(key), + json=ResponsesBody(model=model, input=content), + ) + + def send_embed(self, key: str, model: str, content: str) -> StreamingResponse: + return self.proxy.transport.send( + "/embeddings", + headers=self.proxy.transport.bearer(key), + json=EmbedBody(model=model, input=content), + ) + + def send_gemini_generate(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + f"/gemini/v1beta/models/{model}:generateContent", + headers=GeminiApiKeyHeaders(x_goog_api_key=key), + json=GeminiGenerateBody( + contents=[GeminiContent(parts=[GeminiPart(text=content)])], + generationConfig=GeminiGenerationConfig(maxOutputTokens=max_tokens), + ), + ) + + def upload_batch_file(self, key: str, model: str, content: bytes) -> FileObject: + return unwrap( + self.proxy.transport.upload( + "/v1/files", + headers=self.proxy.transport.bearer(key), + form=FileUploadForm(purpose="batch"), + filename="key_attribution.jsonl", + content=content, + params=ModelQuery(model=model), + response_type=FileObject, + ) + ) + + def create_batch(self, key: str, body: BatchCreateBody) -> BatchObject: + return unwrap( + self.proxy.transport.post( + "/v1/batches", + headers=self.proxy.transport.bearer(key), + json=body, + response_type=BatchObject, + ) + ) + + def retrieve_batch(self, key: str, batch_id: str, *, provider: str) -> BatchObject: + return unwrap( + self.proxy.transport.get( + f"/v1/batches/{batch_id}", + headers=self.proxy.transport.bearer(key), + params=ProviderQuery(provider=provider), + response_type=BatchObject, + ) + ) + + def replay_callback_log(self, key: str, payload: CallbackLogPayload) -> CallbackLogsResponse: + return unwrap( + self.proxy.transport.post( + "/v1/rust_control_plane/logs", + headers=self.proxy.transport.bearer(key), + json=CallbackLogsRequest(records=[CallbackLogRecord(standard_logging_payload=payload)]), + response_type=CallbackLogsResponse, + ) + ) + + def health(self, model: str) -> ProbeResult: + return self.proxy.transport.probe("/health", params=HealthParams(model=model)) + + def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None: + response: Final = unwrap( + self.proxy.transport.get( + "/user/daily/activity", + headers=self.proxy.transport.master, + params=DailyActivityParams( + start_date=start.strftime("%Y-%m-%d"), + end_date=end.strftime("%Y-%m-%d"), + api_key=token, + ), + response_type=DailyActivityResponse, + ) + ) + return next( + (row.breakdown.api_keys[token] for row in response.results if token in row.breakdown.api_keys), + None, + ) + + def poll_daily_activity_for_key( + self, token: str, *, start: datetime, end: datetime, min_requests: int + ) -> DailyActivityKeyBreakdown | None: + outcome: Final = await_converged( + lambda: self.daily_activity_for_key(token, start=start, end=end), + converged=lambda found: found is not None and found.metrics.api_requests >= min_requests, + timeout=self.proxy.poll_timeout, + interval=self.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + def openapi(self) -> OpenAPISchema: return unwrap( self.proxy.transport.get( diff --git a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py new file mode 100644 index 00000000000..4a2c23927c6 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -0,0 +1,405 @@ +"""Every spend row a live proxy writes joins its virtual key (MAT-180). + +One virtual key with an alias, owned by a user with an email, drives every spend +write path a key can reach: /chat/completions, /queue/chat/completions, +/v1/messages, /v1/responses, /embeddings, the Gemini native passthrough, a batch +input file upload, a batch create, and a replayed callback log (POST +/v1/rust_control_plane/logs, the writer an external gateway feeds). Each row those calls write must carry +`api_key` equal to the key's LiteLLM_VerificationToken.token (the sha256 hash +/key/generate returns as `token`), which is the join /spend/logs?api_key= and +/user/daily/activity rely on to report key_alias and user_email. A row keyed by a +re-hashed token (v1.99.0's regression, #39568 and #39572) shows up as a +key-hash-* row with no alias and no email in the customer's usage exports. + +The health-check service account writes rows too; those must stay keyed by the +literal service-account name, never by a hash of it. A batch's cost row is +written by the retrieve that first sees the batch in a terminal state, so the +batch the run creates is one OpenAI fails at validation within seconds (its one +line targets /v1/embeddings under a /v1/chat/completions batch), and the test +retrieves it by its raw provider id with the same key until it is failed. A raw +id is never owned by the CheckBatchCost poller, so that retrieve prices the batch +inline against the retrieving key and its {provider_batch_id}_batch_cost row +must join the key's token with its alias. A completed batch with a positive +cost is out of a single run's reach (OpenAI's completion window is 24h, and a +stack booted fresh per run lists no earlier run's batches), so the poller's own +row is not asserted here. + +/spend/logs carries no email field, so the email assertion lives on +/user/daily/activity alone; /spend/logs is held to the alias in metadata. +""" + +import base64 +import time +from collections.abc import Iterator +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from models import KeyGenerateBody +from proxy_client import Converged, await_converged +from pydantic import BaseModel +from spend_e2e_client import ( + BatchCreateBody, + BatchObject, + CallbackLogMetadata, + CallbackLogPayload, + DailyActivityKeyBreakdown, + ResponseIdentity, + SpendClient, + SpendLogRow, + StreamingResponse, + unique_marker, +) + +pytestmark = pytest.mark.e2e + +CHAT_MODEL: Final = "gemini-2.5-flash" +MESSAGES_MODEL: Final = "claude-haiku-4-5" +RESPONSES_MODEL: Final = "openai-responses-codex" +EMBED_MODEL: Final = "openai-text-embedding-3-small" +BATCH_MODEL: Final = "openai-gpt-4o-mini" +BATCH_BACKEND_MODEL: Final = "gpt-4o-mini" +BATCH_PROVIDER: Final = "openai" +HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check" +BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"}) +FAILED_BATCH_POLL_SECONDS: Final = 120.0 +FAILED_BATCH_POLL_INTERVAL_SECONDS: Final = 5.0 +MAX_TOKENS: Final = 8 +REPLAY_RESPONSE_COST: Final = 0.0001 +REPLAY_PROMPT_TOKENS: Final = 5 +REPLAY_COMPLETION_TOKENS: Final = 1 +WRITE_PATHS: Final = ( + "chat_completions", + "queue_chat_completions", + "messages", + "responses", + "embeddings", + "gemini_passthrough", + "batch_file_upload", + "batch_create", + "callback_replay", +) + + +class EmbeddingLineBody(BaseModel): + model: str + input: str + + +class EmbeddingLine(BaseModel): + custom_id: str + method: str = "POST" + url: str = "/v1/embeddings" + body: EmbeddingLineBody + + +@dataclass(frozen=True, slots=True) +class AttributedKey: + key: str + token: str + alias: str + email: str + user_id: str + + +@dataclass(frozen=True, slots=True) +class WritePath: + name: str + request_id: str + + +@dataclass(frozen=True, slots=True) +class DrivenKey: + identity: AttributedKey + paths: tuple[WritePath, ...] + started_at: datetime + + +def _body_id(name: str, sent: StreamingResponse) -> WritePath: + assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}" + response_id: Final = ResponseIdentity.model_validate_json(sent.body).id + assert response_id, f"{name} answered without a response id: {sent.body[:300]}" + return WritePath(name=name, request_id=response_id) + + +def _call_id(name: str, sent: StreamingResponse) -> WritePath: + assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}" + assert sent.call_id, f"{name} answered without an x-litellm-call-id header" + return WritePath(name=name, request_id=sent.call_id) + + +def _endpoint_mismatched_jsonl(marker: str) -> bytes: + line: Final = EmbeddingLine(custom_id=marker, body=EmbeddingLineBody(model=BATCH_BACKEND_MODEL, input=marker)) + return f"{line.model_dump_json()}\n".encode() + + +def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> tuple[WritePath, WritePath]: + uploaded: Final = client.upload_batch_file(identity.key, BATCH_MODEL, _endpoint_mismatched_jsonl(marker)) + created: Final = client.create_batch( + identity.key, + BatchCreateBody( + input_file_id=uploaded.id, + model=BATCH_MODEL, + metadata={"run": marker}, + ), + ) + return ( + WritePath(name="batch_file_upload", request_id=uploaded.id), + WritePath(name="batch_create", request_id=created.id), + ) + + +def _drive_callback_replay(client: SpendClient, identity: AttributedKey, marker: str) -> WritePath: + request_id: Final = f"callback-replay-{marker}" + finished_at: Final = time.time() + replayed: Final = client.replay_callback_log( + identity.key, + CallbackLogPayload( + id=request_id, + litellm_call_id=request_id, + model=CHAT_MODEL, + start_time=finished_at - 1, + end_time=finished_at, + response_cost=REPLAY_RESPONSE_COST, + prompt_tokens=REPLAY_PROMPT_TOKENS, + completion_tokens=REPLAY_COMPLETION_TOKENS, + total_tokens=REPLAY_PROMPT_TOKENS + REPLAY_COMPLETION_TOKENS, + metadata=CallbackLogMetadata( + user_api_key_hash=identity.token, + user_api_key_alias=identity.alias, + user_api_key_user_id=identity.user_id, + ), + ), + ) + assert replayed.processed == 1 and replayed.failed == 0, f"callback replay rejected the payload: {replayed}" + return WritePath(name="callback_replay", request_id=request_id) + + +def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tuple[WritePath, ...]: + marker: Final = unique_marker() + prompt: Final = f"Reply with the word ok. {marker}" + key: Final = identity.key + return ( + _body_id("chat_completions", client.send_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("queue_chat_completions", client.send_queued_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("messages", client.send_messages(key, MESSAGES_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("responses", client.send_responses(key, RESPONSES_MODEL, prompt)), + _call_id("embeddings", client.send_embed(key, EMBED_MODEL, prompt)), + _call_id("gemini_passthrough", client.send_gemini_generate(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + *_drive_batch(client, identity, marker), + _drive_callback_replay(client, identity, marker), + ) + + +def _provider_batch_id(unified_batch_id: str) -> str: + encoded: Final = unified_batch_id.removeprefix("batch_") + decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode() + return decoded.removeprefix("litellm:").split(";", 1)[0] + + +def _driven_batch_id(driven: DrivenKey) -> str: + return next(path.request_id for path in driven.paths if path.name == "batch_create") + + +def _await_terminal_batch(client: SpendClient, key: str, provider_batch_id: str) -> BatchObject: + outcome: Final = await_converged( + lambda: client.retrieve_batch(key, provider_batch_id, provider=BATCH_PROVIDER), + converged=lambda batch: batch.status in BATCH_TERMINAL_STATUSES, + timeout=FAILED_BATCH_POLL_SECONDS, + interval=FAILED_BATCH_POLL_INTERVAL_SECONDS, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + + +def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: + return [ + row + for row in client.proxy.spend_logs_window( + start=started_at - timedelta(minutes=1), end=datetime.now(timezone.utc) + timedelta(minutes=1) + ) + if HEALTH_SERVICE_ACCOUNT in (row.request_tags or []) + ] + + +def _health_rows_since(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: + outcome: Final = await_converged( + lambda: _health_rows_between(client, started_at), + converged=lambda rows: bool(rows), + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + + +class TestKeyAttribution: + @pytest.fixture(scope="class") + def driven(self, client: SpendClient) -> Iterator[DrivenKey]: + marker: Final = unique_marker() + user_id: Final = client.create_user( + email=f"key-attribution-{marker}@example.com", + role="proxy_admin", + user_id=f"key-attribution-{marker}", + ) + record: Final = client.generate_key_record( + KeyGenerateBody(models=[], user_id=user_id, key_alias=f"key-attribution-{marker}") + ) + assert record.token, "/key/generate answered without the key's token hash" + assert record.key_alias, "/key/generate dropped the key alias" + identity: Final = AttributedKey( + key=record.key, + token=record.token, + alias=record.key_alias, + email=f"key-attribution-{marker}@example.com", + user_id=user_id, + ) + started_at: Final = datetime.now(timezone.utc) + try: + yield DrivenKey( + identity=identity, + paths=_drive_every_write_path(client, identity), + started_at=started_at, + ) + finally: + client.proxy.delete_key(identity.key) + client.delete_user(identity.user_id) + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.joins_key", + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], + ) + def test_every_write_path_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: + assert tuple(path.name for path in driven.paths) == WRITE_PATHS + found: Final = tuple((path, client.proxy.poll_logs_for_request_id(path.request_id)) for path in driven.paths) + unwritten: Final = [path.name for path, rows in found if not rows] + assert not unwritten, f"write paths that produced no spend row within the poll window: {unwritten}" + unjoined: Final = [ + (path.name, row.call_type, row.api_key) + for path, rows in found + for row in rows + if row.api_key != driven.identity.token + ] + assert not unjoined, ( + "spend rows whose api_key does not join LiteLLM_VerificationToken.token " + f"{driven.identity.token}: {unjoined}" + ) + unaliased: Final = [ + (path.name, row.call_type, row.metadata.user_api_key_alias if row.metadata else None) + for path, rows in found + for row in rows + if row.metadata is None or row.metadata.user_api_key_alias != driven.identity.alias + ] + assert not unaliased, f"spend rows written without key alias {driven.identity.alias!r}: {unaliased}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.reports_alias_and_email", + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], + ) + def test_spend_logs_by_key_return_every_row_with_the_alias(self, client: SpendClient, driven: DrivenKey) -> None: + expected_ids: Final = frozenset(path.request_id for path in driven.paths) + rows: Final = client.poll_logs_for_key( + driven.identity.key, + min_rows=len(driven.paths), + predicate=lambda found: expected_ids <= frozenset(row.request_id or "" for row in found), + ) + missing: Final = expected_ids - frozenset(row.request_id or "" for row in rows) + assert not missing, ( + f"/spend/logs?api_key= does not return {len(missing)} of {len(expected_ids)} rows for the key: " + f"{sorted(path.name for path in driven.paths if path.request_id in missing)}" + ) + aliases: Final = frozenset(row.metadata.user_api_key_alias if row.metadata else None for row in rows) + assert aliases == {driven.identity.alias}, f"/spend/logs rows carry aliases {sorted(map(str, aliases))}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.reports_alias_and_email", + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], + ) + def test_user_daily_activity_reports_alias_and_email(self, client: SpendClient, driven: DrivenKey) -> None: + breakdown: Final[DailyActivityKeyBreakdown | None] = client.poll_daily_activity_for_key( + driven.identity.token, + start=driven.started_at - timedelta(days=1), + end=datetime.now(timezone.utc) + timedelta(days=1), + min_requests=len(driven.paths), + ) + assert breakdown is not None, ( + f"/user/daily/activity?api_key={driven.identity.token} has no api_keys breakdown: " + "the key's rows did not aggregate under its token" + ) + assert breakdown.metrics.api_requests >= len(driven.paths), ( + f"/user/daily/activity counts {breakdown.metrics.api_requests} requests for the key, " + f"expected at least {len(driven.paths)}" + ) + assert breakdown.metadata.key_alias == driven.identity.alias, f"key_alias={breakdown.metadata.key_alias!r}" + assert breakdown.metadata.user_email == driven.identity.email, f"user_email={breakdown.metadata.user_email!r}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.health_rows_keep_service_account", + exercised_on=["chat_completions"], + ) + def test_health_check_rows_keep_the_service_account_key(self, client: SpendClient) -> None: + started_at: Final = datetime.now(timezone.utc) + probe: Final = client.health(CHAT_MODEL) + assert probe.healthy, f"/health?model={CHAT_MODEL} answered {probe.status_code}: {probe.body[:300]}" + rows: Final = _health_rows_since(client, started_at) + assert rows, f"/health?model={CHAT_MODEL} wrote no {HEALTH_SERVICE_ACCOUNT}-tagged spend row" + rehashed: Final = [(row.request_id, row.api_key) for row in rows if row.api_key != HEALTH_SERVICE_ACCOUNT] + assert not rehashed, f"health-check rows keyed by something other than {HEALTH_SERVICE_ACCOUNT!r}: {rehashed}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key", + exercised_on=["batches"], + ) + def test_terminal_batch_cost_row_joins_the_retrieving_key(self, client: SpendClient, driven: DrivenKey) -> None: + provider_batch_id: Final = _provider_batch_id(_driven_batch_id(driven)) + fetched: Final = _await_terminal_batch(client, driven.identity.key, provider_batch_id) + assert fetched.status == "failed", ( + f"endpoint-mismatched batch {provider_batch_id} is {fetched.status!r} after " + f"{FAILED_BATCH_POLL_SECONDS:.0f}s, so its terminal cost row cannot be asserted" + ) + cost_request_id: Final = f"{provider_batch_id}_batch_cost" + rows: Final = client.proxy.poll_logs_for_request_id(cost_request_id) + assert rows, f"retrieving failed batch {provider_batch_id} wrote no cost row under {cost_request_id}" + call_types: Final = tuple(sorted({row.call_type or "" for row in rows})) + assert call_types == ("aretrieve_batch",), f"cost rows under {cost_request_id} carry call types {call_types}" + unjoined: Final = [ + (row.call_type, row.api_key, row.metadata.user_api_key_alias if row.metadata else None) + for row in rows + if row.api_key != driven.identity.token + or row.metadata is None + or row.metadata.user_api_key_alias != driven.identity.alias + ] + assert not unjoined, ( + f"batch cost rows that do not join the retrieving key's token {driven.identity.token} " + f"with alias {driven.identity.alias!r}: {unjoined}" + ) diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 5822058003c..1efcb1a045b 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -73,10 +73,25 @@ def create_always_timing_out_deployment(proxy: ProxyClient, name: str) -> str: ) +def create_always_picked_small_context_deployment(proxy: ProxyClient, name: str) -> str: + """The always-picked half of a retry pair on the smallest-context model OpenAI + still serves: it holds all of the model group's shuffle weight, so an oversized + prompt opens on it and earns a real context-window refusal, which never benches + a deployment, so only the retry itself can steer the request off it.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY, weight=1), + model_info=ModelInfoBody(), + ) + ) + + def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str: """The other half of a retry pair: healthy, but weight 0, so the weighted shuffle - never opens on it. It is reachable only once its sibling is benched and the - weighted pick falls through to a uniform one over what is left.""" + never opens on it. It is reachable only once its sibling is out of the running, + benched by a cooldown or skipped by the retry, and the weighted pick falls through + to a uniform one over what is left.""" return proxy.register_model( ModelNewBody( model_name=name, diff --git a/tests/e2e/router/test_auto_router_regressions_e2e.py b/tests/e2e/router/test_auto_router_regressions_e2e.py index 188db2a8eb5..374badcf5fc 100644 --- a/tests/e2e/router/test_auto_router_regressions_e2e.py +++ b/tests/e2e/router/test_auto_router_regressions_e2e.py @@ -41,6 +41,7 @@ which stores either the registered alias or the provider-prefixed form. import json import os from collections.abc import Iterator +from contextlib import ExitStack from dataclasses import dataclass from typing import Final @@ -120,19 +121,10 @@ class ResponsesApiResponse(BaseModel): @dataclass(frozen=True, slots=True) -class TagSplitDeployments: - """Scenario A mirrors the customer-shaped config from GitHub issue #36619: - plain deployment registered first, tier deployment and marker both tagged. - Scenario B flips both axes for GitHub issue #36621: marker registered first - and its tier deployment left untagged, so routing depends neither on - registration order nor on tier deployments carrying tags.""" - - tag_a: str - shared_a: str - tier_a: str - tag_b: str - shared_b: str - tier_b: str +class TagSplitDeployment: + tag: str + shared: str + tier: str @dataclass(frozen=True, slots=True) @@ -173,9 +165,7 @@ def _uniform_tier_config(tier_model: str) -> dict[str, object]: } -def _key_for( - proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False -) -> str: +def _key_for(proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False) -> str: key: Final = proxy.generate_key( KeyGenerateBody( models=models, @@ -211,46 +201,61 @@ def _assert_served_only_by(rows: list[SpendLogRow], allowed: frozenset[str], con ) -@pytest.fixture(scope="module") -def split(proxy: ProxyClient) -> Iterator[TagSplitDeployments]: +@pytest.fixture(scope="class") +def router_stack() -> Iterator[ExitStack]: + with ExitStack() as stack: + yield stack + + +def _register_models( + proxy: ProxyClient, stack: ExitStack, registrations: tuple[tuple[str, LiteLLMParamsBody], ...] +) -> None: + for name, params in registrations: + stack.callback(proxy.delete_model, proxy.create_model(name, params)) + + +def _tag_split(proxy: ProxyClient, stack: ExitStack, *, marker_first: bool) -> TagSplitDeployment: marker: Final = unique_marker() - deployments: Final = TagSplitDeployments( - tag_a=f"e2e-split-a-{marker}", - shared_a=f"e2e-autoroute-a-{marker}", - tier_a=f"e2e-tier-a-{marker}", - tag_b=f"e2e-split-b-{marker}", - shared_b=f"e2e-autoroute-b-{marker}", - tier_b=f"e2e-tier-b-{marker}", + named: Final = TagSplitDeployment( + tag=f"e2e-split-{marker}", + shared=f"e2e-autoroute-{marker}", + tier=f"e2e-tier-{marker}", ) anthropic_key: Final = _provider_key("ANTHROPIC_API_KEY") - marker_params_a: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_a), - tags=[deployments.tag_a], + marker_registration: Final = ( + named.shared, + LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config=_uniform_tier_config(named.tier), + tags=[named.tag], + ), ) - marker_params_b: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_b), - tags=[deployments.tag_b], + tier_registration: Final = ( + named.tier, + LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=None if marker_first else [named.tag]), ) - registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( - (deployments.shared_a, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), - (deployments.tier_a, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=[deployments.tag_a])), - (deployments.shared_a, marker_params_a), - (deployments.shared_b, marker_params_b), - (deployments.tier_b, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key)), - (deployments.shared_b, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), + plain_registration: Final = (named.shared, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)) + registrations: Final = ( + (marker_registration, tier_registration, plain_registration) + if marker_first + else (plain_registration, tier_registration, marker_registration) ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield deployments - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, stack, registrations) + return named -@pytest.fixture(scope="module") -def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: +@pytest.fixture(scope="class") +def plain_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=False) + + +@pytest.fixture(scope="class") +def marker_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=True) + + +@pytest.fixture(scope="class") +def zero_priced_alias(proxy: ProxyClient, router_stack: ExitStack) -> ZeroPricedAlias: marker: Final = unique_marker() named: Final = ZeroPricedAlias(alias=f"e2e-priced-alias-{marker}", tier=f"e2e-priced-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -263,16 +268,12 @@ def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: +@pytest.fixture(scope="class") +def heuristic_split(proxy: ProxyClient, router_stack: ExitStack) -> HeuristicSplit: marker: Final = unique_marker() named: Final = HeuristicSplit( alias=f"e2e-heuristic-router-{marker}", @@ -289,16 +290,12 @@ def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: (named.strong, LiteLLMParamsBody(model=STRONG_MODEL, api_key=_provider_key("OPENAI_API_KEY"))), (named.alias, LiteLLMParamsBody(model="auto_router/complexity_router", complexity_router_config=config)), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: +@pytest.fixture(scope="class") +def semantic_auto_router(proxy: ProxyClient, router_stack: ExitStack) -> SemanticAutoRouter: marker: Final = unique_marker() named: Final = SemanticAutoRouter( marker=f"e2e-semantic-router-{marker}", @@ -321,16 +318,12 @@ def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: (named.fallback, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.marker, marker_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: +@pytest.fixture(scope="class") +def credentialed_alias(proxy: ProxyClient, router_stack: ExitStack) -> CredentialedAlias: marker: Final = unique_marker() named: Final = CredentialedAlias(alias=f"e2e-cred-alias-{marker}", tier=f"e2e-cred-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -342,104 +335,110 @@ def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named class TestTagSplitRouting: @pytest.mark.covers("reliability.routing.tagged_marker.request_tag_selects_marker") def test_body_tagged_chat_routes_through_the_marker_to_its_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36619: with tag filtering on, a chat request whose body metadata tags match the tagged marker under a shared model name is answered by the marker's tier deployment, not by the plain deployment that was registered under the name first.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a, tags=[split.tag_a]))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + chat: Final = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared, tags=[plain_first_split.tag]))) assert chat.choices, "tagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged chat on the shared name") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_chat_is_always_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620: untagged chat requests to the shared name succeed on every call and are all served by the plain deployment; the tagged marker never captures them, so no intermittent auto-router errors and no tier hijacking.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) for _ in range(5): - chat = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a))) + chat = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared))) assert chat.choices, "untagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=5) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged chat on the shared name") + _assert_served_only_by(rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_messages_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620 on the /v1/messages surface: an untagged Anthropic-native request to the shared name is served by the plain deployment, not captured by the tagged marker.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - answer: Final = unwrap(proxy.messages(key, _hello_messages_body(split.shared_a))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + answer: Final = unwrap(proxy.messages(key, _hello_messages_body(plain_first_split.shared))) assert answer.content or answer.choices, "untagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/messages on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/messages on the shared name" + ) class TestUntaggedTierDeployments: @pytest.mark.covers("reliability.routing.tagged_marker.header_tag_selects_marker") def test_header_tagged_messages_routes_through_the_marker_to_an_untagged_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36621: a /v1/messages request tagged only via the x-litellm-tags header selects the tagged marker, and the rewrite still lands on the tier deployment even though that deployment carries no tags, because the marker consumed the routing tags.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_b) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=marker_first_split.tag) answer: Final = unwrap( proxy.transport.post( "/v1/messages", headers=headers, - json=_hello_messages_body(split.shared_b), + json=_hello_messages_body(marker_first_split.shared), response_type=AnthropicMessagesResponse, ) ) assert answer.content or answer.choices, "header-tagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "header-tagged /v1/messages on the shared name") + _assert_served_only_by( + rows, CHEAP_SERVED | {marker_first_split.tier}, "header-tagged /v1/messages on the shared name" + ) @pytest.mark.covers("reliability.routing.tagged_marker.untagged_tier_deployments_still_served") def test_body_tagged_chat_reaches_the_untagged_tier_after_marker_rewrite( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins the tag-consumption half of GitHub issue #36621: after the tagged marker rewrites the request to its tier model, the consumed routing tags no longer constrain deployment selection, so the untagged tier deployment serves the request instead of a strict-tag denial.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_b, tags=[split.tag_b]))) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + chat: Final = unwrap( + proxy.chat(key, _hello_chat_body(marker_first_split.shared, tags=[marker_first_split.tag])) + ) assert chat.choices, "body-tagged chat through the marker-first shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "body-tagged chat with untagged tier") + _assert_served_only_by(rows, CHEAP_SERVED | {marker_first_split.tier}, "body-tagged chat with untagged tier") @pytest.mark.covers("reliability.routing.tagged_marker.tag_semantics_stay_strict") def test_tagged_call_straight_at_an_untagged_deployment_stays_denied( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """The tag-consumption fix must not loosen strict tag semantics: a tagged request aimed directly at an untagged deployment (no marker involved) is still rejected with the 401 tags-configuration error.""" - key: Final = _key_for(proxy, resources, [split.tier_b], tag_filtering=True) - result: Final = proxy.chat(key, _hello_chat_body(split.tier_b, tags=[split.tag_b])) + key: Final = _key_for(proxy, resources, [marker_first_split.tier], tag_filtering=True) + result: Final = proxy.chat(key, _hello_chat_body(marker_first_split.tier, tags=[marker_first_split.tag])) assert isinstance(result, UnauthorizedError), ( f"expected the tagged direct call to an untagged deployment to be denied with 401, got {result}" ) @@ -451,37 +450,39 @@ class TestUntaggedTierDeployments: class TestResponsesApiTagRouting: @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_header_tagged_responses_with_string_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the /v1/responses surface of the tag split (GitHub issues #36620/#36621): a /v1/responses request with string input, tagged via the x-litellm-tags header, succeeds and routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_a) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=plain_first_split.tag) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post("/v1/responses", headers=headers, json=body, response_type=ResponsesApiResponse) ) assert answer.id, "header-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "header-tagged /v1/responses string input") + _assert_served_only_by( + rows, CHEAP_SERVED | {plain_first_split.tier}, "header-tagged /v1/responses string input" + ) @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_body_tagged_responses_with_list_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the body-tag and list-input combination of the same split: /v1/responses with litellm_metadata.tags and structured input items routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, + model=plain_first_split.shared, input=[ResponsesInputItem(role="user", content=f"say hello {unique_marker()}")], max_output_tokens=64, - litellm_metadata=ResponsesTagMetadata(tags=[split.tag_a]), + litellm_metadata=ResponsesTagMetadata(tags=[plain_first_split.tag]), ) answer: Final = unwrap( proxy.transport.post( @@ -493,18 +494,18 @@ class TestResponsesApiTagRouting: ) assert answer.id, "body-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged /v1/responses list input") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged /v1/responses list input") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_responses_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the untagged half of the /v1/responses tag split: an untagged request to the shared name is served by the plain deployment, matching the chat and messages surfaces.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post( @@ -516,7 +517,9 @@ class TestResponsesApiTagRouting: ) assert answer.id, "untagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/responses on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/responses on the shared name" + ) class TestStrategyAliasPricing: @@ -551,9 +554,7 @@ class TestComplexityHeuristicScope: while the accompanying ~2KB agent system prompt is packed with enough reasoning and complexity keywords that scoring the combined text lands in REASONING; only ask-only scoring keeps this on the cheap tier.""" - key: Final = _key_for( - proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong] - ) + key: Final = _key_for(proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong]) body: Final = ChatBody( model=heuristic_split.alias, messages=[ diff --git a/tests/e2e/router/test_reliability_retries_e2e.py b/tests/e2e/router/test_reliability_retries_e2e.py index 5441412935c..da45cb46a46 100644 --- a/tests/e2e/router/test_reliability_retries_e2e.py +++ b/tests/e2e/router/test_reliability_retries_e2e.py @@ -1,13 +1,17 @@ """Live e2e: a request that fails on its first deployment is retried inside its own model group and still comes back a completion. -The model group is a pair: an always-timing-out deployment that holds all of the -group's shuffle weight, and a healthy backup at weight 0. The weighted pick always -opens on the timing-out one, its first Timeout benches it (an -`allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`), and the retry falls -through to the only deployment left. So the customer sees a completion and the -proxy reports that it took a retry to get there, with no random first pick in the -middle of it. +Each model group is a pair: a deployment that always refuses and holds all of the +group's shuffle weight, plus a healthy backup at weight 0. The weighted pick always +opens on the refusing one, so the customer sees a completion only if the retry +lands on the backup, and the proxy reports that it took a retry to get there, with +no random first pick in the middle of it. + +The timeout pair relies on cooldown: the first Timeout benches the timing-out +deployment (an `allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`) and the +retry falls through to the only deployment left. The context-window pair cannot: +a 400 never benches a deployment, so the retry policy's `BadRequestErrorRetries` +has to steer the retry off the deployment that just refused the prompt. """ from __future__ import annotations @@ -16,20 +20,48 @@ import pytest from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker +from e2e_http import StreamingResponse from lifecycle import ResourceManager from models import RouterSettingsOverride from reliability_support import ( chat_override, completion_tokens_of, content_of, + create_always_picked_small_context_deployment, create_always_timing_out_deployment, create_zero_weight_backup_deployment, finish_reason_of, + oversized_prompt, ) pytestmark = pytest.mark.e2e +def assert_retry_landed_on_backup(resp: StreamingResponse) -> None: + assert resp.status_code == 200, ( + f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}" + ) + + attempted = resp.headers.get("x-litellm-attempted-retries") + assert attempted is not None, "response is missing the x-litellm-attempted-retries header" + assert int(attempted) >= 1, ( + f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never " + "opened on the refusing deployment, so this proves nothing about retries" + ) + + content = content_of(resp) + finish_reason = finish_reason_of(resp) + completion_tokens = completion_tokens_of(resp) or 0 + assert isinstance(content, str), ( + f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})" + ) + assert content or (finish_reason == "length" and completion_tokens > 0), ( + f"the retry returned empty content with finish_reason={finish_reason!r}, " + f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget " + f"was spent on non-visible reasoning (body={resp.body[:300]})" + ) + + class TestReliabilityRetries: @pytest.mark.covers("reliability.retry.timeout.succeeds_within_retries") def test_timeout_on_first_deployment_succeeds_on_retry( @@ -49,25 +81,27 @@ class TestReliabilityRetries: override=RouterSettingsOverride(num_retries=2), ) - assert resp.status_code == 200, ( - f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}" + assert_retry_landed_on_backup(resp) + + @pytest.mark.covers("reliability.retry.context_window.succeeds_within_retries") + def test_context_window_refusal_on_first_deployment_succeeds_on_retry( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + group = f"reliability-retry-{unique_marker()}" + small_context = create_always_picked_small_context_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(small_context)) + backup = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup)) + + resp = chat_override( + client.proxy, + scoped_key, + group, + oversized_prompt(unique_marker()), + override=RouterSettingsOverride( + num_retries=2, + model_group_retry_policy={group: {"BadRequestErrorRetries": 2}}, + ), ) - attempted = resp.headers.get("x-litellm-attempted-retries") - assert attempted is not None, "response is missing the x-litellm-attempted-retries header" - assert int(attempted) >= 1, ( - f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never " - "opened on the timing-out deployment, so this proves nothing about retries" - ) - - content = content_of(resp) - finish_reason = finish_reason_of(resp) - completion_tokens = completion_tokens_of(resp) or 0 - assert isinstance(content, str), ( - f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})" - ) - assert content or (finish_reason == "length" and completion_tokens > 0), ( - f"the retry returned empty content with finish_reason={finish_reason!r}, " - f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget " - f"was spent on non-visible reasoning (body={resp.body[:300]})" - ) + assert_retry_landed_on_backup(resp) diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py index 66841725d1d..81cd6c8d3d1 100644 --- a/tests/e2e/test_e2e_http.py +++ b/tests/e2e/test_e2e_http.py @@ -13,13 +13,24 @@ monkeypatches anything. from __future__ import annotations from collections.abc import Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from types import MappingProxyType from typing import Final import pytest - -from e2e_http import RETRY_ATTEMPTS, TRANSIENT_STATUSES, request_with_retry, streaming_outcome +from e2e_http import ( + RETRY_ATTEMPTS, + TRANSIENT_STATUSES, + NoBody, + PartialBody, + Success, + ValidationError, + classify, + request_with_retry, + streaming_outcome, + wire_body, +) +from pydantic import BaseModel, TypeAdapter @dataclass @@ -33,10 +44,10 @@ class FakeResponse: @dataclass class SleepRecorder: - delays: list[float] = field(default_factory=list) + delays: tuple[float, ...] = () def __call__(self, seconds: float) -> None: - self.delays.append(seconds) + self.delays += (seconds,) def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse]: @@ -55,7 +66,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[0] - assert sleep.delays == [] + assert sleep.delays == () assert responses[0].close_calls == 0 def test_429_is_never_retried(self) -> None: @@ -63,7 +74,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[0] - assert sleep.delays == [] + assert sleep.delays == () assert responses[0].close_calls == 0 def test_overloaded_529_retries_with_backoff_then_returns_the_success(self) -> None: @@ -71,7 +82,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[1] - assert sleep.delays == [0.5] + assert sleep.delays == (0.5,) assert responses[0].close_calls == 1 assert responses[1].close_calls == 0 @@ -80,7 +91,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[RETRY_ATTEMPTS - 1] - assert sleep.delays == [0.5, 1.0] + assert sleep.delays == (0.5, 1.0) assert [r.close_calls for r in responses] == [1, 1, 0, 0] @@ -134,3 +145,65 @@ class TestStreamEventArrivals: assert result.stream_events == [] assert result.stream_event_arrivals == [] assert result.body == "bad request" + + +class _ServerUpdate(PartialBody): + server_id: str + alias: str | None = None + description: str | None = None + + +class _ServerCreate(BaseModel): + alias: str + description: str | None = None + + +class TestWireBody: + """A partial-update body must put exactly the caller's choice on the wire: an + omitted field stays off it so the route keeps the stored value, and an explicit + None goes out as JSON null so the route clears it. Plain bodies keep dropping + None, which is what every create route expects.""" + + def test_partial_body_omits_unset_fields_and_sends_explicit_none_as_null(self) -> None: + assert wire_body(_ServerUpdate(server_id="s1", description=None)) == {"server_id": "s1", "description": None} + assert wire_body(_ServerUpdate(server_id="s1", alias="renamed")) == {"server_id": "s1", "alias": "renamed"} + + def test_plain_body_drops_none_fields(self) -> None: + assert wire_body(_ServerCreate(alias="a", description=None)) == {"alias": "a"} + + +_JSON: Final[TypeAdapter[object]] = TypeAdapter(object) + + +@dataclass +class FakeJsonResponse: + """The `classify` view of a response: a status, the raw body bytes, and the + parse that would raise on an empty one.""" + + status_code: int + content: bytes + + @property + def ok(self) -> bool: + return self.status_code < 400 + + @property + def text(self) -> str: + return self.content.decode() + + def json(self) -> object: + return _JSON.validate_json(self.content) + + +class TestClassifyEmptyBody: + """A delete that answers 202 with no body is a success, not a parse failure: + the MCP server and toolset delete routes both answer that way, and reading it + as a failure would hide a delete that did not happen behind one that did.""" + + def test_empty_2xx_body_is_a_success(self) -> None: + result: Final = classify(FakeJsonResponse(status_code=202, content=b""), NoBody) + assert isinstance(result, Success) and result.status_code == 202 + + def test_body_that_is_not_json_is_still_a_validation_failure(self) -> None: + result: Final = classify(FakeJsonResponse(status_code=200, content=b""), NoBody) + assert isinstance(result, ValidationError) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py new file mode 100644 index 00000000000..3b84a47e3cc --- /dev/null +++ b/tests/e2e/test_proxy_client.py @@ -0,0 +1,276 @@ +"""Harness coverage for the barriers that gate 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, and that a management write only counts as read back once every +replica's read satisfies the caller's predicate, which is what keeps a two-gateway +stack from handing a test a model or a key that one gateway has not caught up on +yet. The fakes are plain pollers standing in for each replica's transport plus 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 types import MappingProxyType +from typing import Final, cast + +import pytest +from e2e_config import parse_replica_urls +from e2e_http import Result, Success +from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from proxy_client import ( + ConvergeOutcome, + Converged, + EverywhereConverged, + ModelsPoller, + NeverConvergedOn, + NotConverged, + NotServableOn, + Poller, + ProxyClient, + ReplicaRead, + Servable, + await_converged_everywhere, + await_everywhere, + await_servable_everywhere, + build_proxy_client, + converge_timeout_message, + first_lagging_replica, +) +from transport import Transport + +MODEL: Final = "gpt-under-test" +_NO_TRANSPORTS: Final = cast(Transport, None) +TIMEOUT: Final = 10.0 +INTERVAL: Final = 2.0 +RPM_BEFORE_UPDATE: Final = 100 +RPM_AFTER_UPDATE: Final = 200 + + +@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() + + +def _key_info(rpm_limit: int) -> Success[KeyInfoResponse]: + return Success(status_code=200, data=KeyInfoResponse(info=KeyInfo(rpm_limit=rpm_limit))) + + +def _reads(results: Iterable[Result[KeyInfoResponse]]) -> Poller[Result[KeyInfoResponse]]: + it: Final = iter(results) + return lambda: next(it) + + +def _updated(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, Success) and result.data.info.rpm_limit == RPM_AFTER_UPDATE + + +def _converge( + pollers: Mapping[str, Poller[Result[KeyInfoResponse]]], clock: FakeClock +) -> Mapping[str, ConvergeOutcome[Result[KeyInfoResponse]]]: + return await_converged_everywhere( + pollers, + converged=_updated, + timeout=TIMEOUT, + interval=INTERVAL, + now=clock.now, + sleep=clock.sleep, + ) + + +class TestAwaitConvergedEverywhere: + def test_waits_for_the_replica_that_lags_behind_the_write(self) -> None: + clock: Final = FakeClock() + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))), + "gateway-2": _reads( + chain(repeat(_key_info(RPM_BEFORE_UPDATE), 2), repeat(_key_info(RPM_AFTER_UPDATE))) + ), + } + ) + outcomes: Final = _converge(pollers, clock) + assert outcomes == { + "gateway-1": Converged(result=_key_info(RPM_AFTER_UPDATE)), + "gateway-2": Converged(result=_key_info(RPM_AFTER_UPDATE)), + } + assert first_lagging_replica(outcomes) is None + assert clock.elapsed == 2 * INTERVAL + + def test_names_the_replica_that_never_converges_with_its_last_read(self) -> None: + clock: Final = FakeClock() + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))), + "gateway-2": _reads(repeat(_key_info(RPM_BEFORE_UPDATE))), + } + ) + outcomes: Final = _converge(pollers, clock) + assert first_lagging_replica(outcomes) == ( + "gateway-2", + NotConverged(last_result=_key_info(RPM_BEFORE_UPDATE)), + ) + assert clock.elapsed == TIMEOUT + message: Final = converge_timeout_message( + what="GET /key/info", + replica="gateway-2", + timeout=TIMEOUT, + last_result=_key_info(RPM_BEFORE_UPDATE), + ) + assert "gateway-2" in message and "/key/info" in message and str(RPM_BEFORE_UPDATE) in message + + def test_each_replica_gets_its_own_full_budget(self) -> None: + """A replica that converges late must not eat into the next replica's budget: both + need most of the timeout here, so one shared deadline would starve the second.""" + clock: Final = FakeClock() + slow: Final = chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE))) + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(slow), + "gateway-2": _reads( + chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE))) + ), + } + ) + outcomes: Final = _converge(pollers, clock) + assert first_lagging_replica(outcomes) is None + assert clock.elapsed == 2 * 3 * INTERVAL + + +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",) + + +def _answers(answers: Iterable[str]) -> ReplicaRead[str]: + it: Final = iter(answers) + return lambda _timeout: next(it) + + +def _await_everywhere(reads: Mapping[str, ReplicaRead[str]]) -> EverywhereConverged[str] | NeverConvergedOn[str]: + clock: Final = FakeClock() + return await_everywhere( + reads, + settled=lambda answer: answer == "renamed", + timeout=TIMEOUT, + interval=INTERVAL, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + +class TestAwaitEverywhere: + def test_waits_for_the_lagging_replica_and_returns_every_settled_answer(self) -> None: + reads: Final = { + "gateway-1": _answers(repeat("renamed")), + "gateway-2": _answers(chain(repeat("stale", 2), repeat("renamed"))), + } + outcome: Final = _await_everywhere(reads) + assert isinstance(outcome, EverywhereConverged) + assert dict(outcome.answers) == {"gateway-1": "renamed", "gateway-2": "renamed"} + + def test_names_the_replica_that_never_converges_with_what_it_last_served(self) -> None: + reads: Final = { + "gateway-1": _answers(repeat("renamed")), + "gateway-2": _answers(repeat("stale")), + } + assert _await_everywhere(reads) == NeverConvergedOn(replica="gateway-2", last="stale") + + def test_polls_until_the_deadline_before_giving_up(self) -> None: + lagging: Final = chain(repeat("stale", int(TIMEOUT / INTERVAL)), repeat("renamed")) + outcome: Final = _await_everywhere({"gateway-1": _answers(lagging)}) + assert isinstance(outcome, EverywhereConverged), outcome + + +class TestReplicasFor: + def test_split_deployment_reads_management_routes_back_from_the_control_plane(self) -> None: + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://backend", + replica_urls=("http://gateway-1", "http://gateway-2"), + ) + assert set(client.replicas_for("/key/info")) == {"http://backend"} + assert set(client.replicas_for("/v1/models")) == {"http://gateway-1", "http://gateway-2"} + + def test_monolith_reads_management_routes_back_from_every_replica(self) -> None: + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://lb", + replica_urls=("http://pod-1", "http://pod-2"), + ) + assert set(client.replicas_for("/key/info")) == {"http://pod-1", "http://pod-2"} + + def test_mcp_admin_routes_read_back_from_every_data_plane_replica(self) -> None: + """/v1/mcp/* is a lazily mounted feature, so a data-plane replica serves it + too and answers from its own in-memory registry. Routing it to the control + plane would leave every replica but that one unproven, and would move the + tools/list barrier in mcp_client off the plane that serves tools/list.""" + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://backend", + replica_urls=("http://gateway-1", "http://gateway-2"), + ) + assert set(client.replicas_for("/v1/mcp/server/abc")) == {"http://gateway-1", "http://gateway-2"} + assert set(client.replicas_for("/v1/mcp/toolset/abc")) == {"http://gateway-1", "http://gateway-2"} + + def test_a_route_no_replica_serves_is_refused_rather_than_read_back_vacuously(self) -> None: + """A read-back over zero replicas would satisfy every predicate and assert + nothing, so asking for one fails instead of passing silently.""" + client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) + with pytest.raises(AssertionError, match="no replica is configured"): + _ = client.replicas_for("/v1/models") diff --git a/tests/e2e/ui/constants.ts b/tests/e2e/ui/constants.ts index 3a62252915c..158dea53b71 100644 --- a/tests/e2e/ui/constants.ts +++ b/tests/e2e/ui/constants.ts @@ -15,6 +15,10 @@ export const UI_BASE_URL = ( // writable path (the image runner already exports TMPDIR) to relocate them. export const ARTIFACT_DIR = process.env.E2E_UI_ARTIFACT_DIR || "."; +export const MOCK_PRESIDIO_URL = (process.env.E2E_MOCK_PRESIDIO_URL || "http://127.0.0.1:8091").replace(/\/+$/, ""); + +export const PROPAGATION_TIMEOUT_MS = 90_000; + const storagePath = (name: string): string => path.join(ARTIFACT_DIR, name); // Storage state paths for each role diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts index 58939ca2b9a..bce09b49e10 100644 --- a/tests/e2e/ui/fixtures/migratedPages.ts +++ b/tests/e2e/ui/fixtures/migratedPages.ts @@ -1,50 +1,142 @@ -/** - * Source of truth for the App Router migration E2E suites. - * - * Add an entry (legacy sidebar page id -> route segment) once a page's migration - * has MERGED to the branch under test. Consumers pick it up automatically: - * - migration smoke (tests/migration/migratedPages.spec.ts), via MIGRATED_E2E_SEGMENTS: - * default mount: npm run e2e:migration - * server-root-path mount: SERVER_ROOT_PATH=/ npm run e2e:migration:root - * - navigation specs that assert per-page URLs (tests/navigation/sidebar.spec.ts) - * - * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts. - */ -export const MIGRATED_E2E_PAGES: Record = { - "api-keys": "api-keys", - models: "models-and-endpoints", - api_ref: "api-reference", - "llm-playground": "playground", - projects: "projects", - "access-groups": "access-groups", - budgets: "budgets", - workflows: "workflows", - "guardrails-monitor": "guardrails-monitor", - "mcp-servers": "mcp-servers", - "search-tools": "search-tools", - "tag-management": "tag-management", - "vector-stores": "vector-stores", - memory: "memory", - policies: "policies", - guardrails: "guardrails", - prompts: "prompts", - "tool-policies": "tool-policies", - skills: "skills", - caching: "caching", - "cost-tracking": "cost-tracking", - "transform-request": "transform-request", - "ui-theme": "ui-theme", - logs: "logs", - "admin-panel": "admin-panel", - "logging-and-alerts": "logging-and-alerts", - "model-hub-table": "model-hub-table", - new_usage: "usage", - usage: "old-usage", - agents: "agents", - "router-settings": "router-settings", - users: "users", - teams: "teams", - organizations: "organizations", -}; +export type MigratedPage = Readonly<{ + segment: string; + linkName: string | RegExp; + group?: string; + content: Readonly<{ role: "heading" | "tab" | "button"; name: string }> | Readonly<{ text: string }>; + unlicensedText?: string; +}>; -export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))]; +export const MIGRATED_E2E_PAGES: Readonly> = { + "api-keys": { segment: "api-keys", linkName: "Virtual Keys", content: { role: "heading", name: "Virtual Keys" } }, + models: { + segment: "models-and-endpoints", + linkName: "Models + Endpoints", + content: { role: "heading", name: "Model Management" }, + }, + api_ref: { + segment: "api-reference", + linkName: "API Reference", + content: { role: "heading", name: "OpenAI Compatible Proxy: API Reference" }, + }, + "llm-playground": { segment: "playground", linkName: "Playground", content: { role: "tab", name: "Chat" } }, + projects: { + segment: "projects", + linkName: /^Projects(?: Beta)?$/, + content: { role: "heading", name: "Projects" }, + }, + "access-groups": { + segment: "access-groups", + linkName: "Access Groups", + content: { role: "heading", name: "Access Groups" }, + }, + budgets: { segment: "budgets", linkName: "Budgets", content: { role: "heading", name: "Budgets" } }, + workflows: { + segment: "workflows", + linkName: "Workflow Runs", + group: "Agentic", + content: { text: "Workflow Runs" }, + }, + "guardrails-monitor": { + segment: "guardrails-monitor", + linkName: "Guardrails Monitor", + content: { role: "heading", name: "Guardrails Monitor" }, + }, + "mcp-servers": { + segment: "mcp-servers", + linkName: "MCP Servers", + content: { role: "heading", name: "MCP Servers" }, + }, + "search-tools": { + segment: "search-tools", + linkName: "Search Tools", + group: "Tools", + content: { role: "heading", name: "Search Tools" }, + }, + "tag-management": { + segment: "tag-management", + linkName: "Tag Management", + group: "Experimental", + content: { role: "heading", name: "Tag Management" }, + }, + "vector-stores": { + segment: "vector-stores", + linkName: "Vector Stores", + group: "Tools", + content: { role: "heading", name: "Vector Store Management" }, + }, + memory: { segment: "memory", linkName: "Memory", group: "Agentic", content: { role: "heading", name: "Memory" } }, + policies: { segment: "policies", linkName: "Policies", content: { role: "tab", name: "Policy Simulator" } }, + guardrails: { segment: "guardrails", linkName: "Guardrails", content: { role: "tab", name: "Guardrails" } }, + prompts: { + segment: "prompts", + linkName: "Prompts", + group: "Experimental", + content: { role: "button", name: "Add New Prompt" }, + }, + "tool-policies": { + segment: "tool-policies", + linkName: "Tool Policies", + group: "Tools", + content: { role: "heading", name: "Tool Policies" }, + }, + skills: { segment: "skills", linkName: "Skills", content: { role: "heading", name: "Skills" } }, + caching: { segment: "caching", linkName: "Response Cache", content: { role: "tab", name: "Cache Settings" } }, + "cost-tracking": { + segment: "cost-tracking", + linkName: "Cost Tracking", + group: "Settings", + content: { text: "Cost Tracking Settings" }, + }, + "transform-request": { + segment: "transform-request", + linkName: "API Playground", + group: "Experimental", + content: { role: "heading", name: "Playground" }, + }, + "ui-theme": { + segment: "ui-theme", + linkName: "UI Theme", + group: "Settings", + content: { role: "heading", name: "UI Theme Customization" }, + }, + logs: { segment: "logs", linkName: "Logs", content: { role: "heading", name: "Request Logs" } }, + "admin-panel": { + segment: "admin-panel", + linkName: "Admin Settings", + group: "Settings", + content: { role: "heading", name: "Admin Access" }, + }, + "logging-and-alerts": { + segment: "logging-and-alerts", + linkName: "Logging & Alerts", + group: "Settings", + content: { role: "tab", name: "Logging Callbacks" }, + }, + "model-hub-table": { + segment: "model-hub-table", + linkName: "AI Hub", + content: { role: "heading", name: "AI Hub" }, + }, + new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage View" } }, + usage: { + segment: "old-usage", + linkName: "Old Usage", + group: "Experimental", + content: { role: "tab", name: "All Up" }, + }, + agents: { segment: "agents", linkName: "Agents", group: "Agentic", content: { role: "heading", name: "Agents" } }, + "router-settings": { + segment: "router-settings", + linkName: "Router Settings", + group: "Settings", + content: { role: "heading", name: "Routing Settings" }, + }, + users: { segment: "users", linkName: "Internal Users", content: { role: "tab", name: "Users" } }, + teams: { segment: "teams", linkName: "Teams", content: { role: "heading", name: "Teams" } }, + organizations: { + segment: "organizations", + linkName: "Organizations", + content: { text: "Click on an organization ID to view its details." }, + unlicensedText: "This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key here.", + }, +}; diff --git a/tests/e2e/ui/fixtures/mock_presidio_server/server.py b/tests/e2e/ui/fixtures/mock_presidio_server/server.py new file mode 100644 index 00000000000..0fc97c4a065 --- /dev/null +++ b/tests/e2e/ui/fixtures/mock_presidio_server/server.py @@ -0,0 +1,107 @@ +""" +Mock Presidio analyzer + anonymizer for UI e2e tests. +Serves POST /analyze and POST /anonymize over a fixed set of regex recognizers. +""" + +import os +import re + +import uvicorn +from fastapi import FastAPI, Request +from fastapi.middleware.cors import CORSMiddleware + + +def _luhn_ok(candidate): + digits = [int(c) for c in candidate if c.isdigit()] + doubled = [d * 2 - 9 if d * 2 > 9 else d * 2 for d in digits[-2::-2]] + return len(digits) >= 13 and (sum(digits[::-1][::2]) + sum(doubled)) % 10 == 0 + + +RECOGNIZERS = ( + ("EMAIL_ADDRESS", re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b"), 1.0, None), + ("CREDIT_CARD", re.compile(r"\b(?:\d[ -]?){13,19}\b"), 1.0, _luhn_ok), + ("US_SSN", re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), 0.85, None), + ("PHONE_NUMBER", re.compile(r"\b(?:\+?\d{1,2}[ .-]?)?(?:\(\d{3}\)|\d{3})[ .-]?\d{3}[ .-]?\d{4}\b"), 0.75, None), + ("IP_ADDRESS", re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), 0.6, None), + ("URL", re.compile(r"\bhttps?://[^\s]+"), 0.5, None), +) + +app = FastAPI(title="Mock Presidio Server") +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_methods=["*"], + allow_headers=["*"], +) + + +def _detect(text, wanted): + found = [ + {"entity_type": name, "start": m.start(), "end": m.end(), "score": score} + for name, pattern, score, validate in RECOGNIZERS + if wanted is None or name in wanted + for m in pattern.finditer(text) + if validate is None or validate(m.group()) + ] + ranked = sorted(found, key=lambda r: (-r["score"], r["start"])) + kept = [] + for candidate in ranked: + overlaps = any(candidate["start"] < k["end"] and k["start"] < candidate["end"] for k in kept) + if not overlaps: + kept.append(candidate) + return sorted(kept, key=lambda r: r["start"]) + + +@app.get("/health") +async def health(): + return {"status": "ok"} + + +@app.post("/analyze") +async def analyze(request: Request): + body = await request.json() + text = body.get("text", "") or "" + entities = body.get("entities") + wanted = set(entities) if entities else None + threshold = body.get("score_threshold") or 0.0 + return [ + {**result, "analysis_explanation": None, "recognition_metadata": {"recognizer_name": "MockRecognizer"}} + for result in _detect(text, wanted) + if result["score"] >= threshold + ] + + +@app.post("/anonymize") +async def anonymize(request: Request): + body = await request.json() + text = body.get("text", "") or "" + results = sorted(body.get("analyzer_results") or [], key=lambda r: r["start"]) + + pieces = [] + items = [] + cursor = 0 + for result in results: + start, end = result["start"], result["end"] + if start < cursor: + continue + placeholder = f"<{result['entity_type']}>" + pieces.append(text[cursor:start]) + masked_start = sum(len(p) for p in pieces) + pieces.append(placeholder) + items.append( + { + "operator": "replace", + "entity_type": result["entity_type"], + "start": masked_start, + "end": masked_start + len(placeholder), + "text": placeholder, + } + ) + cursor = end + pieces.append(text[cursor:]) + + return {"text": "".join(pieces), "items": list(reversed(items))} + + +if __name__ == "__main__": + uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("MOCK_PRESIDIO_PORT", "8091"))) diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index a58ece16f9c..e0e7b4da396 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -1,5 +1,29 @@ import { Page } from "../fixtures/pages"; import { Page as PlaywrightPage, expect } from "@playwright/test"; +import { UI_BASE_URL } from "../constants"; + +export const sidebarLink = (page: PlaywrightPage, name: string | RegExp) => + page.getByRole("complementary").getByRole("link", { name, exact: true }); + +export async function clickSidebarLink(page: PlaywrightPage, name: string | RegExp, groupName?: string): Promise { + const link = sidebarLink(page, name); + if (groupName && !(await link.isVisible())) { + const group = page.getByRole("complementary").getByRole("button", { name: groupName, exact: true }); + await expect(group).toBeVisible(); + if ((await group.getAttribute("aria-expanded")) === "false") { + await group.click(); + } + } + await link.click(); +} + +export async function expectUiRoute(page: PlaywrightPage, segment: string): Promise { + const root = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, ""); + const expected = new URL(`${root}/ui/${segment}`, UI_BASE_URL); + await expect(page, `navigate to ${expected.pathname}`).toHaveURL( + (url) => url.origin === expected.origin && url.pathname.replace(/\/+$/, "") === expected.pathname, + ); +} /** * Navigates to a specific page using the page query parameter. @@ -49,3 +73,13 @@ export async function clickTeamId(page: PlaywrightPage, teamId: string): Promise await cell.click(); await expect(page.getByText("Back to Teams")).toBeVisible({ timeout: 10_000 }); } + +export async function openKeyDetail(page: PlaywrightPage, alias: string): Promise { + await page.getByPlaceholder("Search by key alias or ID").fill(alias); + const row = page.getByRole("row").filter({ hasText: alias }); + await expect(row, `key row "${alias}" never appeared on the Virtual Keys page`).toBeVisible({ timeout: 15_000 }); + await row.getByRole("button", { name: alias }).click(); + await expect(page.getByText("Back to Keys"), `key detail for "${alias}" never opened`).toBeVisible({ + timeout: 15_000, + }); +} diff --git a/tests/e2e/ui/helpers/traffic.ts b/tests/e2e/ui/helpers/traffic.ts index 7f8417cdffb..cb68747b364 100644 --- a/tests/e2e/ui/helpers/traffic.ts +++ b/tests/e2e/ui/helpers/traffic.ts @@ -1,4 +1,4 @@ -import { APIRequestContext, expect } from "@playwright/test"; +import { APIRequestContext, APIResponse, expect } from "@playwright/test"; /** Model names served by fixtures/config.yml, both backed by the mock LLM server. */ export const CHAT_MODEL_A = "fake-openai-gpt-4"; @@ -15,6 +15,9 @@ export const masterKey = (): string => process.env.LITELLM_MASTER_KEY || "sk-123 export const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? ""; +/** Date.now() alone collides: `--repeat-each` starts its copies inside the same millisecond. */ +export const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + interface ChatOptions { model: string; prompt: string; @@ -25,9 +28,8 @@ interface ChatOptions { traceId?: string; } -/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */ -export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { - const res = await request.post(`${rootPath()}/v1/chat/completions`, { +const postChatCompletion = (request: APIRequestContext, opts: ChatOptions): Promise => + request.post(`${rootPath()}/v1/chat/completions`, { headers: { Authorization: `Bearer ${opts.apiKey ?? masterKey()}`, "Content-Type": "application/json", @@ -39,12 +41,26 @@ export async function sendChatCompletion(request: APIRequestContext, opts: ChatO ...(opts.traceId ? { litellm_trace_id: opts.traceId } : {}), }, }); + +/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */ +export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { + const res = await postChatCompletion(request, opts); expect(res.ok(), `chat completion for ${opts.model} failed (${res.status()}): ${await res.text()}`).toBe(true); const body = await res.json(); expect(body.choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); return body.id as string; } +export interface ChatAttempt { + status: number; + body: string; +} + +export async function attemptChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { + const res = await postChatCompletion(request, opts); + return { status: res.status(), body: await res.text() }; +} + /** `key` is the sk- value to authenticate with; `token` is its hash, which spend aggregates are keyed by. */ export async function createVirtualKey( request: APIRequestContext, @@ -66,6 +82,33 @@ export async function createVirtualKey( }; } +export interface KeyInfo { + key_alias: string | null; + max_budget: number | null; + budget_duration: string | null; + budget_reset_at: string | null; + blocked: boolean | null; + models: string[]; + team_id: string | null; +} + +export async function readKeyInfo(request: APIRequestContext, token: string): Promise { + const res = await request.get(`${rootPath()}/key/info?key=${encodeURIComponent(token)}`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(res.ok(), `GET /key/info for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true); + const body = await res.json(); + return body.info as KeyInfo; +} + +export async function deleteVirtualKey(request: APIRequestContext, token: string): Promise { + const res = await request.post(`${rootPath()}/key/delete`, { + headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }, + data: { keys: [token] }, + }); + expect(res.ok(), `key delete for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true); +} + /** Spend logs are flushed on a timer, so an assertion straight after a completion races the writer. */ export async function waitForSpendLog( request: APIRequestContext, diff --git a/tests/e2e/ui/run_e2e.sh b/tests/e2e/ui/run_e2e.sh index 67e3225f668..beb1bc8bf3b 100755 --- a/tests/e2e/ui/run_e2e.sh +++ b/tests/e2e/ui/run_e2e.sh @@ -14,7 +14,7 @@ set -euo pipefail # # Ports default to 4000 / 5432 / 8090 and can be moved when another checkout # already holds them: -# PROXY_PORT=4100 POSTGRES_PORT=5532 MOCK_LLM_PORT=8190 ./run_e2e.sh +# PROXY_PORT=4100 POSTGRES_PORT=5532 MOCK_LLM_PORT=8190 MOCK_PRESIDIO_PORT=8191 ./run_e2e.sh # # In CI (CI=true), expects: # - PostgreSQL already running on 127.0.0.1:5432 @@ -29,6 +29,7 @@ DASHBOARD_DIR="$REPO_ROOT/ui/litellm-dashboard" IS_CI="${CI:-false}" CONTAINER_NAME="litellm-e2e-postgres-$$" MOCK_PID="" +MOCK_PRESIDIO_PID="" PROXY_PID="" PROXY_LOG="" @@ -40,7 +41,8 @@ PROXY_LOG="" PROXY_PORT="${PROXY_PORT:-4000}" POSTGRES_PORT="${POSTGRES_PORT:-5432}" MOCK_LLM_PORT="${MOCK_LLM_PORT:-8090}" -export MOCK_LLM_PORT +MOCK_PRESIDIO_PORT="${MOCK_PRESIDIO_PORT:-8091}" +export MOCK_LLM_PORT MOCK_PRESIDIO_PORT # --- Ensure common tool paths are available (local dev only) --- if [ "$IS_CI" = "false" ]; then @@ -82,6 +84,7 @@ fi cleanup() { echo "Cleaning up..." [ -n "$MOCK_PID" ] && kill "$MOCK_PID" 2>/dev/null || true + [ -n "$MOCK_PRESIDIO_PID" ] && kill "$MOCK_PRESIDIO_PID" 2>/dev/null || true [ -n "$PROXY_PID" ] && kill "$PROXY_PID" 2>/dev/null || true [ -n "$PROXY_LOG" ] && rm -f "$PROXY_LOG" || true if [ "$IS_CI" = "false" ]; then @@ -110,9 +113,9 @@ if [ "$IS_CI" = "false" ]; then # to someone else's :5432 (a psql session, a running app, a Prisma engine # talking to a remote database) aborts the run with "port 5432 is in use" # while nothing is actually bound locally. - for port in "$PROXY_PORT" "$POSTGRES_PORT" "$MOCK_LLM_PORT"; do + for port in "$PROXY_PORT" "$POSTGRES_PORT" "$MOCK_LLM_PORT" "$MOCK_PRESIDIO_PORT"; do if lsof -nP -iTCP:"$port" -sTCP:LISTEN >/dev/null 2>&1; then - echo "Error: port $port is in use (override with PROXY_PORT / POSTGRES_PORT / MOCK_LLM_PORT)" + echo "Error: port $port is in use (override with PROXY_PORT / POSTGRES_PORT / MOCK_LLM_PORT / MOCK_PRESIDIO_PORT)" exit 1 fi done @@ -143,6 +146,7 @@ fi # --- Credentials --- export LITELLM_MASTER_KEY="sk-1234" export MOCK_LLM_URL="http://127.0.0.1:${MOCK_LLM_PORT}/v1" +export E2E_MOCK_PRESIDIO_URL="http://127.0.0.1:${MOCK_PRESIDIO_PORT}" export DISABLE_SCHEMA_UPDATE="true" # The suite resolves its target from E2E_UI_BASE_URL (constants.ts), which # otherwise defaults to :4000 -- so without this a relocated stack would be @@ -151,6 +155,7 @@ export DISABLE_SCHEMA_UPDATE="true" export E2E_UI_BASE_URL="${E2E_UI_BASE_URL:-http://127.0.0.1:${PROXY_PORT}}" # Ensure the proxy serves UI at /ui (not behind a subpath) export SERVER_ROOT_PATH="" +export PROXY_BASE_URL="$E2E_UI_BASE_URL" # Boot with an external logout URL so proxyLogoutUrl.spec.ts can assert the # redirect. This same value is exported to the Playwright process below (the # spec's skip guard reads it). Safe for the rest of the suite — nothing else @@ -203,6 +208,20 @@ for i in $(seq 1 15); do sleep 1 done +echo "=== Starting mock Presidio server ===" +uv run --no-sync python "$SCRIPT_DIR/fixtures/mock_presidio_server/server.py" & +MOCK_PRESIDIO_PID=$! + +PRESIDIO_READY=0 +for i in $(seq 1 15); do + if curl -sf http://127.0.0.1:${MOCK_PRESIDIO_PORT}/health >/dev/null 2>&1; then PRESIDIO_READY=1; break; fi + sleep 1 +done +if [ "$PRESIDIO_READY" -ne 1 ]; then + echo "Mock Presidio server never answered /health on port ${MOCK_PRESIDIO_PORT}" >&2 + exit 1 +fi + # --- LiteLLM proxy --- echo "=== Starting LiteLLM proxy ===" cd "$REPO_ROOT" diff --git a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts index 1e43c7a2b22..0267eeae909 100644 --- a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts +++ b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts @@ -7,6 +7,7 @@ import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traff interface StoredGuardrail { guardrail_id: string; guardrail_name: string | null; + litellm_params?: { pii_entities_config?: Record | null } | null; } async function listGuardrails(page: PlaywrightPage): Promise { @@ -249,6 +250,17 @@ test.describe("Guardrails", () => { const row = page.getByRole("row").filter({ hasText: guardrailName }); await expect(row).toHaveCount(1, { timeout: 15_000 }); + const stored = await findGuardrail(page, guardrailName); + const piiConfig = stored?.litellm_params?.pii_entities_config ?? {}; + expect( + Object.keys(piiConfig).length, + "Select All & Mask persisted no PII entities, so the guardrail would mask nothing", + ).toBeGreaterThan(0); + expect( + Object.entries(piiConfig).filter(([, action]) => action !== "MASK"), + "Select All & Mask must persist every entity with the MASK action", + ).toEqual([]); + await navigateToPage(page, Page.Teams); await dismissFeedbackPopup(page); await clickTeamId(page, E2E_TEAM_NO_ADMIN_ID); diff --git a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts new file mode 100644 index 00000000000..b4dc7c4e8da --- /dev/null +++ b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts @@ -0,0 +1,207 @@ +import { test as base, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, MOCK_PRESIDIO_URL, PROPAGATION_TIMEOUT_MS } from "../../constants"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { Page } from "../../fixtures/pages"; +import { CHAT_MODEL_A, masterKey, rootPath, uniqueSuffix, waitForSpendLog } from "../../helpers/traffic"; +import { openPlayground, selectModel, sendButton, onlyVisible } from "../../helpers/playground"; + +const RAW_EMAIL = "jane.doe@example.com"; +const RAW_PHONE = "555-867-5309"; + +const visibleTestId = (page: PlaywrightPage, id: string) => page.getByTestId(id).filter({ visible: true }); + +const requestLogsRows = (page: PlaywrightPage) => + page.locator("table").filter({ visible: true }).first().locator("tbody tr"); + +async function createPresidioGuardrail(page: PlaywrightPage, guardrailName: string): Promise { + await navigateToPage(page, Page.Guardrails); + await dismissFeedbackPopup(page); + + await page.getByRole("button", { name: /Add New Guardrail/i }).click(); + await page.getByRole("menuitem", { name: "Add Provider Guardrail" }).click(); + + const dialog = page.getByRole("dialog", { name: "Create guardrail" }); + await expect(dialog).toBeVisible({ timeout: 10_000 }); + + await dialog.getByLabel("Guardrail Name").fill(guardrailName); + + const providerSelect = dialog.getByRole("combobox", { name: "Guardrail Provider" }); + await providerSelect.click(); + await providerSelect.fill("Presidio"); + await page.getByRole("option", { name: "Presidio PII" }).click(); + + await dialog.getByLabel("Mode", { exact: true }).click(); + await page.keyboard.type("pre_call"); + await expect(page.getByRole("option", { name: "pre_call" })).toBeAttached({ timeout: 5_000 }); + await page.keyboard.press("Enter"); + await expect(dialog.getByText("pre_call", { exact: true })).toBeVisible({ timeout: 5_000 }); + await dialog.getByText("Create guardrail", { exact: true }).click(); + + await dialog.getByLabel("presidio_analyzer_api_base").fill(MOCK_PRESIDIO_URL); + await dialog.getByLabel("presidio_anonymizer_api_base").fill(MOCK_PRESIDIO_URL); + + await dialog.getByRole("button", { name: "Next" }).click(); + await expect(dialog.getByText("Configure PII Protection")).toBeVisible({ timeout: 10_000 }); + await dialog.getByRole("button", { name: "Select All & Mask" }).click(); + + await dialog.getByRole("button", { name: "Create Guardrail" }).click(); + await expect(page.getByText("Guardrail created successfully").first()).toBeVisible({ timeout: 15_000 }); + await expect(page.getByRole("row").filter({ hasText: guardrailName })).toHaveCount(1, { timeout: 15_000 }); +} + +async function deleteGuardrail(page: PlaywrightPage, guardrailName: string): Promise { + await navigateToPage(page, Page.Guardrails); + await dismissFeedbackPopup(page); + const row = page.getByRole("row").filter({ hasText: guardrailName }); + await expect(row).toHaveCount(1, { timeout: 15_000 }); + await row.getByRole("button", { name: "Open guardrail actions" }).click(); + await page.getByRole("menuitem", { name: "Delete" }).click(); + const deleteModal = page.getByRole("dialog", { name: "Delete Guardrail" }); + await expect(deleteModal).toBeVisible({ timeout: 5_000 }); + await deleteModal.getByRole("button", { name: "Delete", exact: true }).click(); + await expect(page.getByText(`Guardrail "${guardrailName}" deleted successfully`)).toBeVisible({ timeout: 10_000 }); +} + +const test = base.extend<{ guardrailName: string }>({ + guardrailName: async ({ page }, use) => { + const name = `e2e-presidio-story-${uniqueSuffix()}`; + await createPresidioGuardrail(page, name); + try { + await use(name); + } finally { + await deleteGuardrail(page, name); + } + }, +}); + +test.describe("Presidio PII guardrail, end to end from the dashboard", () => { + test.describe.configure({ timeout: 5 * 60_000 }); + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("masks PII sent from the Playground and shows the run in Logs", async ({ page, request, guardrailName }) => { + const marker = `case-ref-${Math.random().toString(36).slice(2, 10)}`; + const prompt = `${marker}. Email me at ${RAW_EMAIL} or call ${RAW_PHONE}.`; + + await expect + .poll( + async () => { + const completion = await request.post(`${rootPath()}/v1/chat/completions`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model: CHAT_MODEL_A, + messages: [ + { role: "user", content: `readiness-${uniqueSuffix()}. Email ${RAW_EMAIL}; phone ${RAW_PHONE}.` }, + ], + guardrails: [guardrailName], + }, + }); + expect(completion.ok(), `guardrail readiness request failed: ${await completion.text()}`).toBe(true); + const { id }: { id: string } = await completion.json(); + expect(id).toBeTruthy(); + await waitForSpendLog(request, id); + const stored = await request.get(`${rootPath()}/spend/logs`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + params: { request_id: id }, + }); + expect(stored.ok(), `guardrail readiness log read failed: ${stored.status()}`).toBe(true); + const body = await stored.text(); + return { + rawEmail: body.includes(RAW_EMAIL), + rawPhone: body.includes(RAW_PHONE), + maskedEmail: body.includes(""), + maskedPhone: body.includes(""), + }; + }, + { + message: `${guardrailName} never masked email and phone data in a completed request`, + timeout: PROPAGATION_TIMEOUT_MS, + intervals: [2_000], + }, + ) + .toEqual({ rawEmail: false, rawPhone: false, maskedEmail: true, maskedPhone: true }); + + await openPlayground(page); + await selectModel(page, CHAT_MODEL_A); + + const guardrailSelect = onlyVisible(page.getByPlaceholder("Select guardrails")); + await expect(guardrailSelect).toBeVisible({ timeout: 20_000 }); + await guardrailSelect.click(); + await guardrailSelect.fill(guardrailName); + await onlyVisible(page.getByRole("option", { name: guardrailName, exact: true })).click({ timeout: 20_000 }); + await page.keyboard.press("Escape"); + + const input = onlyVisible(page.getByPlaceholder("Type your message", { exact: false })); + await expect(input).toBeVisible({ timeout: 15_000 }); + + await input.fill(prompt); + const responsePromise = page.waitForResponse( + (response) => + response.request().method() === "POST" && + new URL(response.url()).pathname.endsWith("/chat/completions") && + (response.request().postData()?.includes(marker) ?? false), + ); + await sendButton(page).click(); + const response = await responsePromise; + expect(response.ok(), `Playground completion failed: ${response.status()}`).toBe(true); + const responseBody = await response.text(); + const chunks: { id?: string; error?: unknown }[] = response.headers()["content-type"]?.includes("text/event-stream") + ? responseBody + .split(/\r?\n/) + .filter((line) => line.startsWith("data: ") && line.trim() !== "data: [DONE]") + .map((line) => JSON.parse(line.slice(6))) + : [JSON.parse(responseBody)]; + expect( + chunks.some((chunk) => chunk.error), + "the Playground stream returned an error", + ).toBe(false); + const requestIds = [...new Set(chunks.map((chunk) => chunk.id).filter((id): id is string => !!id))]; + expect(requestIds, "the Playground response identifies exactly one completion").toHaveLength(1); + const [requestId] = requestIds; + await waitForSpendLog(request, requestId); + + const stored = await request.get(`${rootPath()}/spend/logs?request_id=${requestId}`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(stored.ok(), `spend log read failed: ${stored.status()}`).toBe(true); + const storedBody = JSON.stringify(await stored.json()); + expect(storedBody, "the raw email reached the spend log, so the prompt was stored unmasked").not.toContain( + RAW_EMAIL, + ); + expect(storedBody, "the raw phone number reached the spend log, so the prompt was stored unmasked").not.toContain( + RAW_PHONE, + ); + expect(storedBody, "the stored prompt carries no EMAIL_ADDRESS placeholder, so nothing was masked").toContain( + "", + ); + expect(storedBody, "the stored prompt carries no PHONE_NUMBER placeholder, so nothing was masked").toContain( + "", + ); + + await navigateToPage(page, Page.Logs); + await dismissFeedbackPopup(page); + const search = visibleTestId(page, "datatable-search"); + await expect(search).toBeVisible({ timeout: 20_000 }); + await search.fill(requestId); + const row = requestLogsRows(page).filter({ hasText: requestId }); + await expect(row, `no logs row for request ${requestId}`).toHaveCount(1, { timeout: 30_000 }); + await row.click(); + + const drawer = page.getByRole("dialog").first(); + await expect(onlyVisible(drawer.getByText("Guardrails & Policy Compliance"))).toBeVisible({ timeout: 20_000 }); + await expect(onlyVisible(drawer.getByText(`Pre-call guardrail: ${guardrailName}`))).toBeVisible({ + timeout: 20_000, + }); + const maskedPrompt = drawer.getByText(`${marker}. Email me at or call .`); + await expect(onlyVisible(maskedPrompt)).toBeVisible({ timeout: 20_000 }); + + await onlyVisible(drawer.getByText("2 matched")).click(); + await expect(onlyVisible(drawer.getByText("Detected Entities (2)"))).toBeVisible({ timeout: 10_000 }); + await expect(onlyVisible(drawer.getByText("EMAIL_ADDRESS", { exact: true }))).toBeVisible({ timeout: 10_000 }); + await expect(onlyVisible(drawer.getByText("PHONE_NUMBER", { exact: true }))).toBeVisible({ timeout: 10_000 }); + await expect(onlyVisible(drawer.getByText("Score: 1.00", { exact: true }))).toBeVisible({ timeout: 10_000 }); + await expect(onlyVisible(drawer.getByText("Score: 0.75", { exact: true }))).toBeVisible({ timeout: 10_000 }); + + await expect(drawer.getByText(RAW_EMAIL)).toHaveCount(0); + await expect(drawer.getByText(RAW_PHONE)).toHaveCount(0); + }); +}); diff --git a/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts new file mode 100644 index 00000000000..f923841257a --- /dev/null +++ b/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts @@ -0,0 +1,208 @@ +import { test, expect, type APIRequestContext } from "@playwright/test"; +import { Page } from "../../fixtures/pages"; +import { + dismissFeedbackPopup, + navigateToPage, + openKeyDetail, +} from "../../helpers/navigation"; +import { + CHAT_MODEL_A, + CHAT_MODEL_B, + MOCK_RESPONSE_TEXT, + attemptChatCompletion, + createVirtualKey, + deleteVirtualKey, + masterKey, + readKeyInfo, + rootPath, + uniqueSuffix, +} from "../../helpers/traffic"; + +const MEMBER_PASSWORD = "E2e-Team-Member-Pass-1!"; + +interface CreatedTeam { + readonly team_id: string; +} + +function assertCreatedTeam(body: unknown): asserts body is CreatedTeam { + expect(body, "/team/new returned no team_id").toMatchObject({ + team_id: expect.any(String), + }); +} + +async function postAsMaster( + request: APIRequestContext, + path: string, + data: Record, +): Promise { + const res = await request.post(`${rootPath()}${path}`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data, + }); + expect( + res.ok(), + `POST ${path} failed (${res.status()}): ${await res.text()}`, + ).toBe(true); + return res.json(); +} + +test.describe("Internal User - own team key model scope", () => { + test.use({ storageState: { cookies: [], origins: [] } }); + + test("a team member narrows their own key's models and the proxy enforces it", async ({ + page, + request, + }) => { + const suffix = uniqueSuffix(); + const email = `team-member-${suffix}@test.local`; + const userId = `e2e-key-scope-user-${suffix}`; + const alias = `e2e-key-scope-${suffix}`; + + const team = await postAsMaster(request, "/team/new", { + team_alias: `E2E Key Scope ${suffix}`, + models: [CHAT_MODEL_A, CHAT_MODEL_B], + team_member_permissions: ["/key/generate", "/key/update", "/key/info"], + }); + assertCreatedTeam(team); + const teamId = team.team_id; + + try { + await postAsMaster(request, "/user/new", { + user_id: userId, + user_email: email, + user_role: "internal_user", + auto_create_key: false, + }); + await postAsMaster(request, "/user/update", { + user_id: userId, + password: MEMBER_PASSWORD, + }); + await postAsMaster(request, "/team/member_add", { + team_id: teamId, + member: { role: "user", user_id: userId }, + }); + + const created = await createVirtualKey(request, { + key_alias: alias, + team_id: teamId, + user_id: userId, + models: [], + }); + + try { + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill(email); + await page + .getByPlaceholder("Enter your password") + .fill(MEMBER_PASSWORD); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect( + page.locator("a", { hasText: "Virtual Keys" }), + `${email} never reached the dashboard`, + ).toBeVisible({ timeout: 30_000 }); + await dismissFeedbackPopup(page); + + await navigateToPage(page, Page.ApiKeys); + await openKeyDetail(page, alias); + + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await expect( + page.getByRole("option", { name: CHAT_MODEL_A, exact: true }), + `the Models dropdown does not offer ${CHAT_MODEL_A} to a team member`, + ).toBeVisible({ timeout: 15_000 }); + await expect( + page.getByRole("option", { name: CHAT_MODEL_B, exact: true }), + `the Models dropdown does not offer ${CHAT_MODEL_B} to a team member`, + ).toBeVisible(); + + await page + .getByRole("option", { name: CHAT_MODEL_A, exact: true }) + .click(); + await page.keyboard.press("Escape"); + + const updated = page.waitForResponse( + (res) => + res.url().includes("/key/update") && + res.request().method() === "POST", + ); + await page.getByRole("button", { name: "Save Changes" }).click(); + const updateStatus = (await updated).status(); + expect( + updateStatus, + "a team member's own-key edit was refused", + ).toBeGreaterThanOrEqual(200); + expect( + updateStatus, + "a team member's own-key edit was refused", + ).toBeLessThan(300); + await expect( + page.getByText("Key updated successfully").first(), + ).toBeVisible({ timeout: 15_000 }); + + await expect + .poll( + async () => (await readKeyInfo(request, created.token)).models, + { + message: `the narrowed model scope never reached /key/info for ${alias}`, + timeout: 20_000, + }, + ) + .toEqual([CHAT_MODEL_A]); + + await expect + .poll( + async () => + await attemptChatCompletion(request, { + model: CHAT_MODEL_B, + prompt: `out of scope ${suffix}`, + apiKey: created.key, + }), + { + message: `${CHAT_MODEL_B} was still served after the key was narrowed to ${CHAT_MODEL_A}`, + timeout: 30_000, + }, + ) + .toMatchObject({ + status: 403, + body: expect.stringContaining(CHAT_MODEL_B), + }); + + const inScope = await attemptChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `in scope ${suffix}`, + apiKey: created.key, + }); + expect( + inScope, + `${CHAT_MODEL_A} is no longer served by the narrowed key`, + ).toMatchObject({ + status: 200, + body: expect.stringContaining(MOCK_RESPONSE_TEXT), + }); + } finally { + await deleteVirtualKey(request, created.token); + } + } finally { + await request.post(`${rootPath()}/user/delete`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { user_ids: [userId] }, + }); + await request.post(`${rootPath()}/team/delete`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { team_ids: [teamId] }, + }); + } + }); +}); diff --git a/tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts b/tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts new file mode 100644 index 00000000000..5e2c80b5845 --- /dev/null +++ b/tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts @@ -0,0 +1,196 @@ +import { + test as base, + expect, + type Locator, + type Page as PlaywrightPage, +} from "@playwright/test"; +import { + E2E_TEAM_CRUD_ALIAS, + E2E_TEAM_ORG_ALIAS, + INTERNAL_USER_STORAGE_PATH, +} from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, CHAT_MODEL_B, masterKey } from "../../helpers/traffic"; + +const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; +const CURRENT_TEAM_VIEW = "Current Team Models"; +const ALL_MODELS_VIEW = "All Available Models"; +const PERSONAL_TEAM = "Personal"; + +const teamSelector = (page: PlaywrightPage): Locator => + page.getByRole("combobox", { name: "Current team", exact: true }); +const viewSelector = (page: PlaywrightPage): Locator => + page.getByRole("combobox", { name: "View", exact: true }); + +async function chooseOption( + page: PlaywrightPage, + selector: Locator, + optionName: string, +): Promise { + await selector.click(); + const option = page.getByRole("option", { name: optionName, exact: true }); + await expect(option, `option ${optionName} is offered`).toBeVisible({ + timeout: 10_000, + }); + await option.click(); + await expect( + selector, + `${optionName} is the selection the control now reports`, + ).toContainText(optionName, { + timeout: 10_000, + }); +} + +async function deleteDeployment( + page: PlaywrightPage, + id: string, +): Promise { + const post = () => + page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { id }, + }); + const deleted = await post().catch(() => post()); + expect( + deleted.ok(), + `cleanup: /model/delete ${id} returned ${deleted.status()}`, + ).toBe(true); +} + +function modelRow(page: PlaywrightPage, modelName: string): Locator { + return page.getByRole("row").filter({ hasText: modelName }); +} + +async function isRegistered( + page: PlaywrightPage, + modelName: string, +): Promise { + const body = await readBack<{ data: { model_name?: string }[] }>( + page, + "/v2/model/info", + ); + return body.data.some((row) => row.model_name === modelName); +} + +const uniqueSuffix = (): string => + `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +const test = base.extend<{ ungrantedModelName: string }>({ + ungrantedModelName: async ({ page }, use) => { + const ungrantedModelName = `e2e-ungranted-${uniqueSuffix()}`; + const created = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: ungrantedModelName, + litellm_params: { + model: `openai/${ungrantedModelName}`, + api_base: MOCK_LLM_BASE, + api_key: "fake-key", + }, + model_info: {}, + }, + }); + expect( + created.ok(), + `/model/new failed: ${created.status()} ${await created.text()}`, + ).toBe(true); + const ungrantedModelId = (await created.json()).model_info?.id; + expect(ungrantedModelId, "model id from /model/new").toBeTruthy(); + + try { + await expect + .poll(async () => await isRegistered(page, ungrantedModelName), { + message: `deployment ${ungrantedModelName} never appeared in /v2/model/info after create`, + timeout: 60_000, + }) + .toBe(true); + await use(ungrantedModelName); + } finally { + await deleteDeployment(page, ungrantedModelId); + } + }, +}); + +test.describe("Models and Endpoints for an internal user", () => { + test.use({ storageState: INTERNAL_USER_STORAGE_PATH }); + + test("shows an internal user exactly the models of the team they select", async ({ + page, + ungrantedModelName, + }) => { + await navigateToPage(page, Page.Models); + + await expect( + page.getByRole("tab", { name: "Your Models" }), + "an internal user lands on their own models tab, not an admin-only view", + ).toBeVisible({ timeout: 15_000 }); + await expect( + viewSelector(page), + "the models table opens scoped to the selected team", + ).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 }); + await expect( + modelRow(page, ungrantedModelName), + `the personal view lists ${ungrantedModelName}, so it is on the proxy and reachable from this page`, + ).toHaveCount(1, { timeout: 30_000 }); + + await chooseOption(page, teamSelector(page), E2E_TEAM_CRUD_ALIAS); + await expect( + modelRow(page, CHAT_MODEL_A), + `${E2E_TEAM_CRUD_ALIAS} lists ${CHAT_MODEL_A}`, + ).toHaveCount(1, { + timeout: 15_000, + }); + await expect( + modelRow(page, CHAT_MODEL_B), + `${E2E_TEAM_CRUD_ALIAS} lists ${CHAT_MODEL_B}`, + ).toHaveCount(1, { + timeout: 15_000, + }); + await expect( + modelRow(page, ungrantedModelName), + `${ungrantedModelName} is on the proxy but not granted to ${E2E_TEAM_CRUD_ALIAS}, so it must not be listed`, + ).toHaveCount(0); + + await chooseOption(page, teamSelector(page), E2E_TEAM_ORG_ALIAS); + await expect( + modelRow(page, CHAT_MODEL_A), + `${E2E_TEAM_ORG_ALIAS} lists ${CHAT_MODEL_A}`, + ).toHaveCount(1, { + timeout: 15_000, + }); + await expect( + page.getByTestId("pagination-range"), + `${E2E_TEAM_ORG_ALIAS} lists the one model it grants and nothing else`, + ).toHaveText("Showing 1-1 of 1", { timeout: 15_000 }); + await expect( + modelRow(page, CHAT_MODEL_B), + `${CHAT_MODEL_B} belongs to another team and must not leak into ${E2E_TEAM_ORG_ALIAS}`, + ).toHaveCount(0); + await expect( + modelRow(page, ungrantedModelName), + `${ungrantedModelName} is granted to no team and must not leak into ${E2E_TEAM_ORG_ALIAS}`, + ).toHaveCount(0); + + await chooseOption(page, viewSelector(page), ALL_MODELS_VIEW); + await expect( + modelRow(page, CHAT_MODEL_A), + `switching to ${ALL_MODELS_VIEW} leaves the table populated rather than blanking it`, + ).toHaveCount(1, { timeout: 15_000 }); + + await page.reload(); + await expect( + teamSelector(page), + "the team selection is not persisted across a reload, so the table returns to the personal view", + ).toContainText(PERSONAL_TEAM, { timeout: 15_000 }); + await expect( + viewSelector(page), + "the view selection is not persisted across a reload either", + ).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 }); + await expect( + modelRow(page, ungrantedModelName), + "the personal view still renders models after a reload rather than coming back empty", + ).toHaveCount(1, { timeout: 30_000 }); + }); +}); diff --git a/tests/e2e/ui/tests/migration/README.md b/tests/e2e/ui/tests/migration/README.md index d6b33598ec4..59933502463 100644 --- a/tests/e2e/ui/tests/migration/README.md +++ b/tests/e2e/ui/tests/migration/README.md @@ -1,17 +1,25 @@ # App Router migration smoke A growing E2E smoke for pages migrated from the legacy `?page=` switch to App -Router path routes. For each migrated page it clicks the page's sidebar link, checks -the URL is the path route and the page renders, reloads it, then clicks off to a -legacy page and back to confirm navigation still works. It runs in two situations: -the default mount and a non-root `SERVER_ROOT_PATH` mount. +Router path routes. For each page it clicks the sidebar link by its accessible +name, verifies the destination's content, reloads it, then visits Virtual Keys +and returns. It runs at the default mount and a non-root `SERVER_ROOT_PATH` mount + +Link selection does not depend on `href` formatting. URL assertions compare the +origin and pathname, allowing a trailing slash, query string, and fragment while +rejecting another route or mount. Reloads must return a successful document, +and each journey must finish without uncaught browser errors ## Adding a page -When a page's migration merges, add its route segment to -`tests/e2e/ui/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` -in `ui/litellm-dashboard/src/utils/migratedPages.ts`). Both suites pick it up -automatically. +Add an entry to `tests/e2e/ui/fixtures/migratedPages.ts`, keyed by the legacy page +ID. Specify its route segment, accessible link name, sidebar group if collapsed, +and distinctive visible content such as a heading or tab. Keep expectations +independent of the application's route table so an incorrect destination fails +the test. Both navigation suites use this fixture + +For a licensed-only page, `unlicensedText` describes the expected upgrade notice. +The authenticated session's license claim determines which content must render ## Running @@ -31,4 +39,9 @@ SERVER_ROOT_PATH=/litellm npm run e2e:migration:root ``` `globalSetup` logs in once per role; the admin storage state is reused for these -tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login`. +tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login` + +`tests/navigation/sidebar.spec.ts` also checks the navigation helpers against +equivalent link formats on the live dashboard and a deep link containing a query +string and fragment. The link-format cases change only the rendered `href` +attribute to exercise the locator contract; destination pages and APIs remain live diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts index 547330190bd..f2dbd2dbc77 100644 --- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts +++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts @@ -1,105 +1,73 @@ import { test, expect, type Page } from "@playwright/test"; -import { MIGRATED_E2E_SEGMENTS } from "../../fixtures/migratedPages"; +import { MIGRATED_E2E_PAGES, type MigratedPage } from "../../fixtures/migratedPages"; import { ADMIN_STORAGE_PATH } from "../../constants"; -import { dismissFeedbackPopup } from "../../helpers/navigation"; +import { clickSidebarLink, dismissFeedbackPopup, expectUiRoute, sidebarLink } from "../../helpers/navigation"; +import { proxyIsPremium } from "../../helpers/premium"; -/** - * App Router migration smoke as a user journey: start where the proxy lands you, - * click a migrated page in the sidebar, confirm it routed and rendered, reload it - * (the check a wrong server_root_path breaks), bounce to a legacy page and back, - * and, once two pages are migrated, navigate directly between two migrated pages. - * - * Driven by MIGRATED_E2E_SEGMENTS, so it grows as pages are migrated. Set - * SERVER_ROOT_PATH (e.g. "/litellm") to exercise the non-root mount; leave it - * unset for the default mount. Boot the proxy with the matching value first. - */ -const ROOT = process.env.SERVER_ROOT_PATH ?? ""; +const ROOT = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, ""); +const apiKeys = MIGRATED_E2E_PAGES["api-keys"]; -const esc = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); -const pathRe = (segment: string) => new RegExp(`${esc(ROOT)}/ui/${esc(segment)}/?($|\\?)`); -// Scope nav lookups to the sidebar (a `complementary` landmark). The top bar -// now renders a breadcrumb whose current-page item is also a "Virtual Keys" -// link, so an unscoped locator would match two elements. -const sidebar = (page: Page) => page.getByRole("complementary"); -const virtualKeysLink = (page: Page) => sidebar(page).getByRole("link", { name: "Virtual Keys", exact: true }); - -/** The dashboard shell is present (sidebar rendered); page didn't 404 / crash. */ -async function expectRendered(page: Page) { - await expect(virtualKeysLink(page)).toBeVisible({ timeout: 20_000 }); +async function expectContent(page: Page, destination: MigratedPage): Promise { + await expect(sidebarLink(page, apiKeys.linkName)).toBeVisible({ timeout: 20_000 }); + const main = page.getByRole("main"); + if (destination.unlicensedText && !proxyIsPremium()) { + await expect(main.getByText(destination.unlicensedText, { exact: true })).toBeVisible(); + return; + } + const content = destination.content; + const landmark = + "role" in content + ? main.getByRole(content.role, { name: content.name, exact: true }) + : main.getByText(content.text, { exact: true }); + await expect(landmark).toBeVisible(); } -/** - * Click a migrated page's sidebar link. Migrated items render as ; - * nested ones live under collapsible groups whose children only render while the - * group is open, so expand collapsed groups until the link is clickable. - */ -async function clickSidebar(page: Page, segment: string) { - const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first(); - const collapsedGroups = sidebar(page).getByRole("button", { expanded: false }); - for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { - const stillCollapsed = await collapsedGroups.count(); - if (stillCollapsed === 0) break; - await collapsedGroups.first().click(); - await expect(collapsedGroups).toHaveCount(stillCollapsed - 1); - } - await link.click(); +async function navigateToDestination(page: Page, destination: MigratedPage): Promise { + await clickSidebarLink(page, destination.linkName, destination.group); + await expectUiRoute(page, destination.segment); + await dismissFeedbackPopup(page); + await expectContent(page, destination); } test.use({ storageState: ADMIN_STORAGE_PATH }); test.describe("App Router migrated pages", () => { - for (const segment of MIGRATED_E2E_SEGMENTS) { - test(`${segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => { + for (const destination of Object.values(MIGRATED_E2E_PAGES)) { + test(`${destination.segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => { const pageErrors: string[] = []; - page.on("pageerror", (e) => pageErrors.push(String(e))); + page.on("pageerror", (error) => pageErrors.push(String(error))); - // 1. Start where the proxy lands us. - await page.goto(`${ROOT}/ui/`); + const landing = await page.goto(`${ROOT}/ui/`); + expect(landing?.ok(), "dashboard document loads successfully").toBe(true); await dismissFeedbackPopup(page); - await expectRendered(page); + await expectContent(page, apiKeys); - // 2. Click the migrated page in the sidebar -> path route + rendered. - await clickSidebar(page, segment); - await expect(page).toHaveURL(pathRe(segment)); - await expectRendered(page); - // 3. Reload the path route directly; a wrong server_root_path 404s here. - await page.reload(); + await navigateToDestination(page, destination); + + const reloaded = await page.reload(); + expect(reloaded?.ok(), `${destination.segment} document loads on reload`).toBe(true); + await expectUiRoute(page, destination.segment); await dismissFeedbackPopup(page); - await expect(page).toHaveURL(pathRe(segment)); - await expectRendered(page); - // 4. Click the Virtual Keys sidebar link to the api-keys landing (now a path route), then back. - await virtualKeysLink(page).click(); - await expect(page).toHaveURL(pathRe("api-keys")); - await dismissFeedbackPopup(page); - await expectRendered(page); - // 5. Click back to the migrated page. - await clickSidebar(page, segment); - await expect(page).toHaveURL(pathRe(segment)); - await expectRendered(page); - expect(pageErrors, `page errors during ${segment} journey`).toEqual([]); + await expectContent(page, destination); + + await navigateToDestination(page, apiKeys); + await navigateToDestination(page, destination); + expect(pageErrors, `page errors during ${destination.segment} journey`).toEqual([]); }); } test("navigates directly between two migrated pages", async ({ page }) => { - test.skip(MIGRATED_E2E_SEGMENTS.length < 2, "needs >= 2 migrated pages"); - const [first, second] = MIGRATED_E2E_SEGMENTS; const pageErrors: string[] = []; - page.on("pageerror", (e) => pageErrors.push(String(e))); + page.on("pageerror", (error) => pageErrors.push(String(error))); - await page.goto(`${ROOT}/ui/`); + const landing = await page.goto(`${ROOT}/ui/`); + expect(landing?.ok(), "dashboard document loads successfully").toBe(true); await dismissFeedbackPopup(page); + await expectContent(page, apiKeys); - await clickSidebar(page, first); - await expect(page).toHaveURL(pathRe(first)); - await expectRendered(page); - await clickSidebar(page, second); - await expect(page).toHaveURL(pathRe(second)); - await expectRendered(page); - // Back to the first migrated page. - await clickSidebar(page, first); - await expect(page).toHaveURL(pathRe(first)); - await expectRendered(page); - - expect(pageErrors, "page errors during migrated -> migrated nav").toEqual([]); + for (const destination of [apiKeys, MIGRATED_E2E_PAGES.models, apiKeys]) { + await navigateToDestination(page, destination); + } + expect(pageErrors, "page errors during migrated page navigation").toEqual([]); }); }); diff --git a/tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts b/tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts new file mode 100644 index 00000000000..4480515ae59 --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts @@ -0,0 +1,252 @@ +import { + test as base, + expect, + type Page as PlaywrightPage, +} from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { captureRequestBody, readBack } from "../../helpers/roundTrip"; +import { masterKey, sendChatCompletion } from "../../helpers/traffic"; + +const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; +const CUSTOM_PARAM = "extra_headers"; +const CUSTOM_PARAM_VALUE = { "X-E2E-Edit-Probe": "one" }; + +type StoredParams = Record; + +async function readStoredParams( + page: PlaywrightPage, + modelId: string, +): Promise { + const body = await readBack<{ data: { litellm_params: StoredParams }[] }>( + page, + `/model/info?litellm_model_id=${modelId}`, + ); + return body.data[0]?.litellm_params ?? {}; +} + +function paramsEditor(page: PlaywrightPage) { + return page.getByPlaceholder('"rpm": 100'); +} + +async function editParams( + page: PlaywrightPage, + mutate: (params: StoredParams) => StoredParams, +): Promise { + await page.getByRole("button", { name: "Edit Settings" }).click(); + const editor = paramsEditor(page); + await expect( + editor, + "the LiteLLM Params editor is reachable on every visit to the edit form", + ).toBeVisible({ + timeout: 15_000, + }); + const shown = JSON.parse(await editor.inputValue()) as StoredParams; + await editor.fill(JSON.stringify(mutate(shown), null, 2)); +} + +async function deleteDeployment( + page: PlaywrightPage, + id: string, +): Promise { + const post = () => + page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { id }, + }); + const deleted = await post().catch(() => post()); + expect( + deleted.ok(), + `cleanup: /model/delete ${id} returned ${deleted.status()}`, + ).toBe(true); +} + +const uniqueSuffix = (): string => + `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +const test = base.extend<{ + deployment: { readonly modelName: string; readonly createdModelId: string }; +}>({ + deployment: async ({ page, request }, use) => { + const modelName = `e2e-edit-params-${uniqueSuffix()}`; + const created = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: modelName, + litellm_params: { + model: `openai/${modelName}`, + api_base: MOCK_LLM_BASE, + api_key: "fake-key", + }, + model_info: {}, + }, + }); + expect( + created.ok(), + `/model/new failed: ${created.status()} ${await created.text()}`, + ).toBe(true); + const createdModelId = (await created.json()).model_info?.id; + expect(createdModelId, "model id from /model/new").toBeTruthy(); + + try { + await expect + .poll( + async () => { + try { + await sendChatCompletion(request, { + model: modelName, + prompt: `warmup ${modelName}`, + }); + return true; + } catch { + return false; + } + }, + { + message: `deployment ${modelName} never became routable after /model/new`, + timeout: 60_000, + }, + ) + .toBe(true); + await use({ modelName, createdModelId }); + } finally { + await deleteDeployment(page, createdModelId); + } + }, +}); + +test.describe("Edit LiteLLM Params on a deployment", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("params added on a deployment can be re-edited, and the deployment keeps serving", async ({ + page, + request, + deployment: { modelName, createdModelId }, + }) => { + await navigateToPage(page, Page.Models); + const modelIdCell = page.getByTestId(`model-id-${createdModelId}`); + await expect( + modelIdCell, + `the Models table lists ${modelName}`, + ).toBeVisible({ timeout: 15_000 }); + await modelIdCell.click(); + await expect(page.getByText("Back to Models").first()).toBeVisible({ + timeout: 15_000, + }); + + await editParams(page, (params) => ({ + ...params, + temperature: 0.2, + [CUSTOM_PARAM]: CUSTOM_PARAM_VALUE, + })); + const firstSave = await captureRequestBody( + page, + { method: "PATCH", urlIncludes: `/model/${createdModelId}/update` }, + async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }, + ); + expect( + firstSave.litellm_params?.temperature, + "the added temperature goes on the wire", + ).toBe(0.2); + expect( + firstSave.litellm_params?.[CUSTOM_PARAM], + `the added ${CUSTOM_PARAM} goes on the wire`, + ).toEqual(CUSTOM_PARAM_VALUE); + expect( + firstSave.litellm_params?.model, + "a params edit does not rewrite the upstream model", + ).toBe(`openai/${modelName}`); + expect( + firstSave.litellm_params?.api_base, + "a params edit does not rewrite the api base", + ).toBe(MOCK_LLM_BASE); + expect( + firstSave.litellm_params, + "the credential is never re-sent, so a masked placeholder cannot overwrite the stored key", + ).not.toHaveProperty("api_key"); + + await expect + .poll( + async () => (await readStoredParams(page, createdModelId)).temperature, + { + message: "the added temperature never reached the stored deployment", + timeout: 20_000, + }, + ) + .toBe(0.2); + const afterFirstSave = await readStoredParams(page, createdModelId); + expect( + afterFirstSave[CUSTOM_PARAM], + `the added ${CUSTOM_PARAM} reached the stored deployment`, + ).toEqual(CUSTOM_PARAM_VALUE); + expect( + afterFirstSave.model, + "the stored upstream model survived the edit", + ).toBe(`openai/${modelName}`); + expect( + afterFirstSave.api_base, + "the stored api base survived the edit", + ).toBe(MOCK_LLM_BASE); + + await editParams(page, (params) => ({ + ...Object.fromEntries( + Object.entries(params).filter(([key]) => key !== CUSTOM_PARAM), + ), + temperature: 0.7, + })); + const secondSave = await captureRequestBody( + page, + { method: "PATCH", urlIncludes: `/model/${createdModelId}/update` }, + async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }, + ); + expect( + secondSave.litellm_params?.temperature, + "a param set by an earlier save can be edited again", + ).toBe(0.7); + expect( + secondSave.litellm_params, + `dropping ${CUSTOM_PARAM} from the editor drops it from the request the UI sends`, + ).not.toHaveProperty(CUSTOM_PARAM); + expect( + secondSave.litellm_params?.model, + "a second params edit still leaves the upstream model alone", + ).toBe(`openai/${modelName}`); + expect( + secondSave.litellm_params?.api_base, + "a second params edit still leaves the api base alone", + ).toBe(MOCK_LLM_BASE); + expect( + secondSave.litellm_params, + "the credential is still never re-sent", + ).not.toHaveProperty("api_key"); + + await expect + .poll( + async () => (await readStoredParams(page, createdModelId)).temperature, + { + message: + "the re-edited temperature never reached the stored deployment", + timeout: 20_000, + }, + ) + .toBe(0.7); + + await page.reload(); + await expect( + page + .getByRole("tabpanel", { name: "Overview" }) + .getByText('"temperature": 0.7'), + "reopening the deployment renders the re-edited value, not the one from the first save", + ).toBeVisible({ timeout: 20_000 }); + + await sendChatCompletion(request, { + model: modelName, + prompt: `still serving ${modelName}`, + }); + }); +}); diff --git a/tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts b/tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts new file mode 100644 index 00000000000..7bd1caa6756 --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts @@ -0,0 +1,263 @@ +import { + test as base, + expect, + type Locator, + type Page as PlaywrightPage, +} from "@playwright/test"; +import { ADMIN_STORAGE_PATH, PROPAGATION_TIMEOUT_MS } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { masterKey } from "../../helpers/traffic"; + +const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; +const UNREACHABLE_BASE = "http://127.0.0.1:9/v1"; + +async function isRegistered( + page: PlaywrightPage, + modelName: string, +): Promise { + const body = await readBack<{ data: { model_name?: string }[] }>( + page, + "/v2/model/info", + ); + return body.data.some((row) => row.model_name === modelName); +} + +function healthRow(page: PlaywrightPage, modelName: string): Locator { + return page.getByRole("row").filter({ hasText: modelName }); +} + +function pageOf(label: string): { current: number; total: number } { + const [current, total] = label + .replace("Page ", "") + .split(" of ") + .map((part) => Number(part.trim())); + return { current, total }; +} + +async function locateHealthRow( + page: PlaywrightPage, + modelName: string, +): Promise { + const pageLabel = page.getByTestId("pagination-page"); + await expect( + pageLabel, + "the health table reports which page it is showing", + ).toBeVisible({ timeout: 20_000 }); + + const deadline = Date.now() + 60_000; + while (Date.now() < deadline) { + const row = healthRow(page, modelName); + const onThisPage = await row + .first() + .waitFor({ state: "visible", timeout: 3_000 }) + .then(() => true) + .catch(() => false); + if (onThisPage) return row; + + const { current, total } = pageOf(await pageLabel.innerText()); + const goTo = current < total ? current + 1 : 1; + if (total === 1) continue; + await page + .getByRole("button", { + name: current < total ? "Go to next page" : "Go to first page", + }) + .click(); + await expect(pageLabel).toContainText(`Page ${goTo} of`, { + timeout: 15_000, + }); + } + return healthRow(page, modelName); +} + +async function openHealthTab(page: PlaywrightPage): Promise { + await page.getByRole("tab", { name: "Health Status" }).click(); + await expect( + page.getByRole("heading", { name: "Model Health Status" }), + ).toBeVisible({ timeout: 15_000 }); +} + +async function expectStatus( + page: PlaywrightPage, + modelName: string, + status: string, +): Promise { + const row = await locateHealthRow(page, modelName); + await expect(row, `${modelName} has one row in the health table`).toHaveCount( + 1, + { timeout: 20_000 }, + ); + await expect( + row.getByText(status, { exact: true }), + `the Health Status cell for ${modelName} reads ${status}`, + ).toHaveCount(1, { timeout: 60_000 }); +} + +async function deleteDeployment( + page: PlaywrightPage, + id: string, +): Promise { + const post = () => + page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { id }, + }); + const deleted = await post().catch(() => post()); + expect( + deleted.ok(), + `cleanup: /model/delete ${id} returned ${deleted.status()}`, + ).toBe(true); +} + +const uniqueSuffix = (): string => + `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +async function withDeployment( + page: PlaywrightPage, + prefix: string, + apiBase: string, + use: (name: string) => Promise, +): Promise { + const name = `${prefix}-${uniqueSuffix()}`; + const created = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: name, + litellm_params: { + model: `openai/${name}`, + api_base: apiBase, + api_key: "fake-key", + }, + model_info: {}, + }, + }); + expect( + created.ok(), + `/model/new for ${name} failed: ${created.status()} ${await created.text()}`, + ).toBe(true); + const id = (await created.json()).model_info?.id; + expect(id, `model id from /model/new for ${name}`).toBeTruthy(); + try { + await expect + .poll(() => isRegistered(page, name), { + message: `deployment ${name} never appeared in /v2/model/info after create`, + timeout: 60_000, + }) + .toBe(true); + await expect + .poll( + async () => { + const response = await page.request.get("/v1/models", { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(response.ok(), `/v1/models failed: ${response.status()}`).toBe(true); + const body: { data: { id: string }[] } = await response.json(); + return body.data.some((model) => model.id === name); + }, + { + message: `deployment ${name} never appeared on the serving path`, + timeout: PROPAGATION_TIMEOUT_MS, + intervals: [2_000], + }, + ) + .toBe(true); + await use(name); + } finally { + await deleteDeployment(page, id); + } +} + +const test = base.extend<{ reachableName: string; unreachableName: string }>({ + reachableName: async ({ page }, use) => { + await withDeployment(page, "e2e-health-up", MOCK_LLM_BASE, use); + }, + unreachableName: async ({ page }, use) => { + await withDeployment(page, "e2e-health-down", UNREACHABLE_BASE, use); + }, +}); + +test.describe("Model health status", () => { + test.describe.configure({ timeout: 8 * 60_000 }); + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Run Health Check reports a reachable deployment healthy and an unreachable one unhealthy", async ({ + page, + reachableName, + unreachableName, + }) => { + await navigateToPage(page, Page.Models); + await openHealthTab(page); + + for (const name of [reachableName, unreachableName]) { + const row = await locateHealthRow(page, name); + await expect(row, `${name} has one row in the health table`).toHaveCount( + 1, + { timeout: 20_000 }, + ); + await row + .getByRole("button", { name: "Run Health Check", exact: true }) + .click(); + } + + await expectStatus(page, reachableName, "healthy"); + await expect( + healthRow(page, reachableName).getByText("unhealthy", { exact: true }), + "a reachable deployment is never reported unhealthy", + ).toHaveCount(0); + await expectStatus(page, unreachableName, "unhealthy"); + + const successDetail = ( + await locateHealthRow(page, reachableName) + ).getByRole("button", { + name: "View response details", + }); + await expect( + successDetail, + `${reachableName} offers its health check response for inspection`, + ).toBeVisible({ timeout: 60_000 }); + await successDetail.click(); + const successDialog = page.getByRole("dialog"); + await expect( + successDialog.getByRole("heading", { + name: `Health Check Response - ${reachableName}`, + }), + "the healthy deployment's detail opens its own response dialog", + ).toBeVisible({ timeout: 10_000 }); + await successDialog.getByRole("button", { name: "Close" }).last().click(); + await expect(successDialog).toBeHidden({ timeout: 10_000 }); + + const errorDetail = ( + await locateHealthRow(page, unreachableName) + ).getByRole("button", { + name: "View full error details", + }); + await expect( + errorDetail, + `${unreachableName} offers its health check error for inspection`, + ).toBeVisible({ timeout: 60_000 }); + await errorDetail.click(); + const errorDialog = page.getByRole("dialog"); + await expect( + errorDialog.getByRole("heading", { + name: `Health Check Error - ${unreachableName}`, + }), + "the unreachable deployment's detail opens its own error dialog", + ).toBeVisible({ timeout: 10_000 }); + await expect( + errorDialog, + "the error dialog carries the upstream connection failure, not a generic message", + ).toContainText(/connection error/i, { timeout: 10_000 }); + await expect( + errorDialog, + "the error dialog names the endpoint that could not be reached", + ).toContainText(UNREACHABLE_BASE); + await errorDialog.getByRole("button", { name: "Close" }).last().click(); + await expect(errorDialog).toBeHidden({ timeout: 10_000 }); + + await page.reload(); + await openHealthTab(page); + await expectStatus(page, reachableName, "healthy"); + await expectStatus(page, unreachableName, "unhealthy"); + }); +}); diff --git a/tests/e2e/ui/tests/navigation/sidebar.spec.ts b/tests/e2e/ui/tests/navigation/sidebar.spec.ts index b220dc09ae2..eaa4985c659 100644 --- a/tests/e2e/ui/tests/navigation/sidebar.spec.ts +++ b/tests/e2e/ui/tests/navigation/sidebar.spec.ts @@ -3,7 +3,13 @@ import { Role } from "../../fixtures/roles"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { Page } from "../../fixtures/pages"; import { menuLabelToPage } from "../../fixtures/menuMappings"; -import { navigateToPage } from "../../helpers/navigation"; +import { + clickSidebarLink, + dismissFeedbackPopup, + expectUiRoute, + navigateToPage, + sidebarLink, +} from "../../helpers/navigation"; import { MIGRATED_E2E_PAGES } from "../../fixtures/migratedPages"; import type { Page as PlaywrightPage } from "@playwright/test"; @@ -11,7 +17,7 @@ const sidebarButtons = { [Role.ProxyAdmin]: [ "Virtual Keys", "Playground", - "Models", + "Models + Endpoints", "Usage", "Teams", "Internal Users", @@ -22,9 +28,9 @@ const sidebarButtons = { /** Migrated pages live at a path route; legacy pages keep the ?page= query param. */ async function expectPageUrl(page: PlaywrightPage, pageKey: string): Promise { - const migratedSegment = MIGRATED_E2E_PAGES[pageKey]; - if (migratedSegment) { - await expect(page).toHaveURL(new RegExp(`/ui/${migratedSegment}/?($|\\?)`)); + const migratedPage = MIGRATED_E2E_PAGES[pageKey]; + if (migratedPage) { + await expectUiRoute(page, migratedPage.segment); } else { await expect(page).toHaveURL(new RegExp(`[?&]page=${pageKey}(&|$)`)); } @@ -51,12 +57,7 @@ for (const { role, storage } of roles) { throw new Error(`No page mapping found for menu label: ${buttonLabel}`); } - // Sidebar items are links inside the `complementary` landmark; scoping - // there avoids the top-bar breadcrumb, which also links the page name. - const tab = page.getByRole("complementary").getByRole("link", { name: buttonLabel }); - await expect(tab).toBeVisible(); - - await tab.click(); + await clickSidebarLink(page, buttonLabel); await expectPageUrl(page, expectedPage); } @@ -81,5 +82,41 @@ for (const { role, storage } of roles) { await navigateToPage(page, Page.LlmPlayground); await expectPageUrl(page, Page.LlmPlayground); }); + + for (const format of ["without trailing slash", "absolute with query and fragment", "relative"] as const) { + test(`sidebar locator tolerates hrefs ${format}`, async ({ page }) => { + await page.goto("/ui/"); + await dismissFeedbackPopup(page); + const link = sidebarLink(page, "Models + Endpoints"); + await expect(link).toBeVisible(); + const destination = new URL("/ui/models-and-endpoints/", page.url()); + const href = + format === "without trailing slash" + ? destination.pathname.replace(/\/$/, "") + : format === "relative" + ? "./models-and-endpoints/" + : `${destination.href}?source=navigation-smoke#overview`; + + await link.evaluate((element, value) => element.setAttribute("href", value), href); + await expect(link).toHaveAttribute("href", href); + await clickSidebarLink(page, "Models + Endpoints"); + + await expectUiRoute(page, "models-and-endpoints"); + await expect( + page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }), + ).toBeVisible(); + }); + } + + test("route assertion tolerates a query string and fragment on a deep link", async ({ page }) => { + const response = await page.goto("/ui/models-and-endpoints/?source=navigation-smoke#overview"); + expect(response?.ok()).toBe(true); + await expectUiRoute(page, "models-and-endpoints"); + await expect( + page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }), + ).toBeVisible(); + expect(new URL(page.url()).search).toBe("?source=navigation-smoke"); + expect(new URL(page.url()).hash).toBe("#overview"); + }); }); } diff --git a/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts new file mode 100644 index 00000000000..788f4ca1284 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts @@ -0,0 +1,113 @@ +import { test as base, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, PROPAGATION_TIMEOUT_MS } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation"; +import { + CHAT_MODEL_A, + MOCK_RESPONSE_TEXT, + attemptChatCompletion, + createVirtualKey, + deleteVirtualKey, + readKeyInfo, + sendChatCompletion, + uniqueSuffix, +} from "../../helpers/traffic"; + +interface ScopedKey { + alias: string; + token: string; + apiKey: string; +} + +const test = base.extend<{ scopedKey: ScopedKey }>({ + scopedKey: async ({ page }, use) => { + const alias = `e2e-block-key-${uniqueSuffix()}`; + const created = await createVirtualKey(page.request, { + key_alias: alias, + models: [CHAT_MODEL_A], + }); + await use({ alias, token: created.token, apiKey: created.key }); + await deleteVirtualKey(page.request, created.token); + }, +}); + +test.describe("Proxy Admin - Key blocking", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("blocking a key stops it serving and unblocking restores it", async ({ page, scopedKey }) => { + test.setTimeout(5 * 60_000); + const { alias, token, apiKey } = scopedKey; + + await sendChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: `pre-block ${alias}`, + apiKey, + }); + + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await openKeyDetail(page, alias); + + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Block Key" }).click(); + const blockDialog = page.getByRole("dialog", { name: "Block Key" }); + await expect(blockDialog, "the Block Key confirmation never opened").toBeVisible({ timeout: 10_000 }); + await blockDialog.getByRole("button", { name: "Block", exact: true }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).blocked, { + message: "the key never came back blocked from /key/info", + timeout: 20_000, + }) + .toBe(true); + + await expect + .poll( + async () => + await attemptChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: "blocked", + apiKey, + }), + { + message: "a blocked key was still served by /v1/chat/completions", + timeout: PROPAGATION_TIMEOUT_MS, + }, + ) + .toMatchObject({ status: 401, body: expect.stringContaining("blocked") }); + + await page.reload(); + await expect( + page.getByText("Blocked", { exact: true }), + "the reloaded key detail does not show the key as blocked", + ).toBeVisible({ timeout: 15_000 }); + + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Unblock Key" }).click(); + const unblockDialog = page.getByRole("dialog", { name: "Unblock Key" }); + await expect(unblockDialog, "the Unblock Key confirmation never opened").toBeVisible({ timeout: 10_000 }); + await unblockDialog.getByRole("button", { name: "Unblock", exact: true }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).blocked, { + message: "the key never came back unblocked from /key/info", + timeout: 20_000, + }) + .toBe(false); + + await expect + .poll( + async () => + await attemptChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: "unblocked", + apiKey, + }), + { + message: "an unblocked key is still refused by /v1/chat/completions", + timeout: PROPAGATION_TIMEOUT_MS, + }, + ) + .toMatchObject({ status: 200, body: expect.stringContaining(MOCK_RESPONSE_TEXT) }); + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts new file mode 100644 index 00000000000..4e4d0a395c3 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts @@ -0,0 +1,101 @@ +import { test as base, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation"; +import { captureRequestBody } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, createVirtualKey, deleteVirtualKey, readKeyInfo, uniqueSuffix } from "../../helpers/traffic"; + +interface ScopedKey { + alias: string; + token: string; +} + +const test = base.extend<{ scopedKey: ScopedKey }>({ + scopedKey: async ({ page }, use) => { + const alias = `e2e-budget-window-${uniqueSuffix()}`; + const created = await createVirtualKey(page.request, { + key_alias: alias, + team_id: E2E_TEAM_CRUD_ID, + models: [CHAT_MODEL_A], + }); + await use({ alias, token: created.token }); + await deleteVirtualKey(page.request, created.token); + }, +}); + +test.describe("Proxy Admin - Key budget window", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("a monthly spend cap survives a reload, and clearing the window keeps the cap", async ({ page, scopedKey }) => { + const { alias, token } = scopedKey; + + const before = await readKeyInfo(page.request, token); + expect(before.max_budget, "a freshly generated key starts with no budget").toBeNull(); + + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await openKeyDetail(page, alias); + + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + await page.getByLabel("Reset Budget", { exact: true }).click(); + await page.getByRole("option", { name: "monthly", exact: true }).click(); + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).max_budget, { + message: "the $12.50 cap never reached /key/info", + timeout: 20_000, + }) + .toBe(12.5); + await expect + .poll(async () => (await readKeyInfo(page.request, token)).budget_duration, { + message: "the monthly reset window never reached /key/info", + timeout: 20_000, + }) + .toBe("30d"); + + const capped = await readKeyInfo(page.request, token); + const resetAt = new Date(capped.budget_reset_at ?? ""); + expect(Number.isNaN(resetAt.getTime()), "a monthly window left the key with no budget_reset_at").toBe(false); + expect(resetAt.getTime(), "budget_reset_at was set in the past").toBeGreaterThan(Date.now()); + expect(resetAt.getUTCDate(), "a monthly window resets on the 1st, a daily one would not").toBe(1); + + await page.reload(); + await expect( + page.getByRole("paragraph").filter({ hasText: "of $12.50" }), + "the reloaded key detail does not render the $12.50 cap", + ).toBeVisible({ timeout: 15_000 }); + + await page.getByRole("tab", { name: "Settings" }).click(); + await expect( + page.getByTestId("budget-reset-value"), + "the reloaded key detail does not name the 30d reset window", + ).toHaveText(/Every 30d/, { timeout: 15_000 }); + + await page.getByRole("button", { name: "Edit Settings" }).click(); + await page.getByLabel("Reset Budget", { exact: true }).click(); + await page.getByRole("option", { name: "Never resets", exact: true }).click(); + + const cleared = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/update" }, async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }); + expect(cleared).toHaveProperty("budget_duration"); + expect(cleared.budget_duration, "clearing the window must send budget_duration: null explicitly").toBeNull(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).budget_duration, { + message: "the reset window was never cleared on /key/info", + timeout: 20_000, + }) + .toBeNull(); + + const after = await readKeyInfo(page.request, token); + expect(after.budget_reset_at, "clearing the reset window left a stale next-reset timestamp").toBeNull(); + expect(after.max_budget, "clearing the reset window also wiped the spend cap").toBe(12.5); + expect(after.models, "editing the budget left the key's models untouched").toEqual(before.models); + expect(after.team_id, "editing the budget left the key's team untouched").toEqual(before.team_id); + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts b/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts new file mode 100644 index 00000000000..f7930378b1e --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts @@ -0,0 +1,145 @@ +import { test, expect, type APIRequestContext, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { proxyIsPremium } from "../../helpers/premium"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +const auth = () => ({ Authorization: `Bearer ${masterKey()}` }); + +async function guardrailId(request: APIRequestContext, name: string): Promise { + const res = await request.get("/v2/guardrails/list", { headers: auth() }); + expect(res.ok(), `GET /v2/guardrails/list (${res.status()})`).toBe(true); + const rows = (await res.json()).guardrails as { guardrail_id: string; guardrail_name: string | null }[]; + return rows.find((row) => row.guardrail_name === name)?.guardrail_id; +} + +async function teamGuardrails(page: PlaywrightPage, teamId: string): Promise { + const body = await readBack<{ team_info: { metadata: { guardrails?: string[] } | null } }>( + page, + `/team/info?team_id=${encodeURIComponent(teamId)}`, + ); + return body.team_info.metadata?.guardrails ?? []; +} + +async function keywordPromptStatus(request: APIRequestContext, apiKey: string, keyword: string): Promise { + const res = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" }, + data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `please tell me about ${keyword}` }] }, + }); + return res.status(); +} + +test.describe("Proxy Admin - Team guardrail removal", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Clearing a team's only guardrail on the Settings tab lets blocked traffic through again", async ({ + page, + request, + }) => { + test.skip(!proxyIsPremium(), "proxy under test is unlicensed, so team guardrails are premium-gated"); + + const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`; + const guardrailName = `e2e-team-guardrail-${stamp}`; + const bannedKeyword = `e2eteamban${stamp}`; + const teamAlias = `e2e-guardrail-team-${stamp}`; + + let teamId = ""; + let teamKey = ""; + try { + const guardrailRes = await request.post("/guardrails", { + headers: auth(), + data: { + guardrail: { + guardrail_name: guardrailName, + litellm_params: { + guardrail: "litellm_content_filter", + mode: "pre_call", + default_on: false, + blocked_words: [{ keyword: bannedKeyword, action: "BLOCK" }], + }, + }, + }, + }); + expect( + guardrailRes.ok(), + `POST /guardrails failed (${guardrailRes.status()}): ${await guardrailRes.text()}`, + ).toBe(true); + + const teamRes = await request.post("/team/new", { + headers: auth(), + data: { team_alias: teamAlias, models: [CHAT_MODEL_A], metadata: { guardrails: [guardrailName] } }, + }); + expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + teamId = (await teamRes.json()).team_id as string; + + const keyRes = await request.post("/key/generate", { headers: auth(), data: { team_id: teamId } }); + expect(keyRes.ok(), `POST /key/generate failed (${keyRes.status()}): ${await keyRes.text()}`).toBe(true); + teamKey = (await keyRes.json()).key as string; + + await expect + .poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), { + message: "the team's guardrail never started refusing the banned keyword", + timeout: 60_000, + }) + .toBe(400); + + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + await clickTeamId(page, teamId); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + const chip = page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName }); + await expect(chip).toBeVisible({ timeout: 10_000 }); + await chip.locator('[data-slot="combobox-chip-remove"]').click(); + await expect(chip).toHaveCount(0, { timeout: 10_000 }); + + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll(async () => teamGuardrails(page, teamId), { + message: "the team still carries a guardrail in /team/info after the save", + timeout: 20_000, + }) + .toEqual([]); + + await page.reload(); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + await expect(page.getByRole("combobox", { name: "Select guardrails" })).toBeVisible({ timeout: 15_000 }); + await expect( + page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName }), + "the removed guardrail is gone from the Settings tab after a reload", + ).toHaveCount(0); + + await expect + .poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), { + message: "the team key is still refused for a keyword whose guardrail was removed", + timeout: 60_000, + }) + .toBe(200); + + const served = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${teamKey}`, "Content-Type": "application/json" }, + data: { + model: CHAT_MODEL_A, + messages: [{ role: "user", content: `please tell me about ${bannedKeyword}` }], + }, + }); + expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); + } finally { + if (teamKey) { + await request.post("/key/delete", { headers: auth(), data: { keys: [teamKey] } }); + } + if (teamId) { + await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } }); + } + const id = await guardrailId(request, guardrailName); + if (id) { + await request.delete(`/guardrails/${id}`, { headers: auth() }); + } + } + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts b/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts new file mode 100644 index 00000000000..c7325e78f75 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts @@ -0,0 +1,129 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, masterKey } from "../../helpers/traffic"; + +interface TeamInfoResponse { + team_info: { + models: string[]; + members_with_roles: { user_id?: string; role?: string }[]; + }; + team_memberships: { + user_id: string; + litellm_budget_table: { max_budget: number | null } | null; + }[]; +} + +const auth = () => ({ Authorization: `Bearer ${masterKey()}` }); + +async function teamInfo(page: PlaywrightPage, teamId: string): Promise { + return readBack(page, `/team/info?team_id=${encodeURIComponent(teamId)}`); +} + +function roleOf(info: TeamInfoResponse, userId: string): string | undefined { + return info.team_info.members_with_roles.find((member) => member.user_id === userId)?.role; +} + +function budgetOf(info: TeamInfoResponse, userId: string): number | null | undefined { + return info.team_memberships.find((membership) => membership.user_id === userId)?.litellm_budget_table?.max_budget; +} + +function otherMembers(info: TeamInfoResponse, userId: string): string[] { + return info.team_info.members_with_roles + .filter((member) => member.user_id !== userId) + .map((member) => `${member.user_id}:${member.role}`) + .sort(); +} + +test.describe("Proxy Admin - Team member edit", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + const createdTeams: string[] = []; + const createdUsers: string[] = []; + + test.afterEach(async ({ request }) => { + for (const teamId of createdTeams.splice(0)) { + await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } }); + } + for (const userId of createdUsers.splice(0)) { + await request.post("/user/delete", { headers: auth(), data: { user_ids: [userId] } }); + } + }); + + test("Editing a member's role and per-member budget persists and survives a reload", async ({ page, request }) => { + const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`; + const memberId = `e2e-member-edit-${stamp}`; + const teamAlias = `e2e-member-edit-team-${stamp}`; + + createdUsers.push(memberId); + const created = await request.post("/user/new", { + headers: auth(), + data: { user_id: memberId, user_role: "internal_user", auto_create_key: false }, + }); + expect(created.ok(), `POST /user/new failed (${created.status()}): ${await created.text()}`).toBe(true); + + const teamRes = await request.post("/team/new", { + headers: auth(), + data: { + team_alias: teamAlias, + models: [CHAT_MODEL_A], + members_with_roles: [{ user_id: memberId, role: "admin" }], + }, + }); + expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + const teamId = (await teamRes.json()).team_id as string; + createdTeams.push(teamId); + + const before = await teamInfo(page, teamId); + expect(roleOf(before, memberId), "the member starts out as a team admin").toBe("admin"); + expect(budgetOf(before, memberId) ?? null, "the member starts out with no per-member budget").toBeNull(); + expect( + otherMembers(before, memberId).length, + "the team has another member for the edit to leave alone", + ).toBeGreaterThan(0); + + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + await clickTeamId(page, teamId); + await page.getByRole("tab", { name: "Members" }).click(); + + const memberRow = page.locator("tr", { hasText: memberId }).first(); + await expect(memberRow).toBeVisible({ timeout: 10_000 }); + await memberRow.getByTestId("edit-member").click(); + + const modal = page.getByRole("dialog", { name: "Edit Member" }); + await expect(modal).toBeVisible({ timeout: 10_000 }); + + await modal.getByLabel(/^Role/).click(); + await page.getByRole("option", { name: "User", exact: true }).click(); + await modal.getByLabel(/Team Member Budget \(USD\)/).fill("5"); + await modal.getByRole("button", { name: "Save Changes" }).click(); + + await expect(page.getByText("Team member updated successfully").first()).toBeVisible({ timeout: 10_000 }); + + await expect + .poll( + async () => { + const info = await teamInfo(page, teamId); + return [roleOf(info, memberId), budgetOf(info, memberId)]; + }, + { message: "the member's role and budget never landed in /team/info", timeout: 20_000 }, + ) + .toEqual(["user", 5]); + + await page.reload(); + await page.getByRole("tab", { name: "Members" }).click(); + const reloadedRow = page.locator("tr", { hasText: memberId }).first(); + await expect(reloadedRow).toBeVisible({ timeout: 15_000 }); + await expect(reloadedRow.getByText("user", { exact: true }), "role shown after a reload").toBeVisible(); + await expect(reloadedRow.getByText("$5.00"), "per-member budget shown after a reload").toBeVisible(); + + const after = await teamInfo(page, teamId); + expect(after.team_info.models, "model access untouched by a member edit").toEqual(before.team_info.models); + expect(otherMembers(after, memberId), "the rest of the roster untouched by a member edit").toEqual( + otherMembers(before, memberId), + ); + }); +}); diff --git a/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts new file mode 100644 index 00000000000..75fb3be9b64 --- /dev/null +++ b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts @@ -0,0 +1,177 @@ +import { test, expect, type Browser, type BrowserContext, type Page as PlaywrightPage } from "@playwright/test"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +const PASSWORD = "E2e-Member-Perms-Pass-1!"; + +const auth = () => ({ Authorization: `Bearer ${masterKey()}` }); + +async function sessionKey(page: PlaywrightPage): Promise { + const cookie = (await page.context().cookies()).find((candidate) => candidate.name === "token"); + expect(cookie?.value, "logged-in session carries a token cookie").toBeTruthy(); + const payload = JSON.parse(Buffer.from(cookie!.value.split(".")[1], "base64url").toString("utf-8")) as { + key?: string; + }; + expect(payload.key, "session JWT carries the virtual key the dashboard calls with").toMatch(/^sk-/); + return payload.key!; +} + +async function signIn(browser: Browser, email: string): Promise { + const context = await browser.newContext({ storageState: { cookies: [], origins: [] } }); + const page = await context.newPage(); + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill(email); + await page.getByPlaceholder("Enter your password").fill(PASSWORD); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 }); + await dismissFeedbackPopup(page); + return context; +} + +test.describe("Team Admin - Member permissions", () => { + test.use({ storageState: { cookies: [], origins: [] } }); + + test("Granting /key/generate lets a plain member create a team key that serves traffic", async ({ + browser, + request, + }) => { + const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`; + const adminId = `e2e-perm-admin-${stamp}`; + const memberId = `e2e-perm-member-${stamp}`; + const adminEmail = `${adminId}@test.local`; + const memberEmail = `${memberId}@test.local`; + const teamAlias = `e2e-perm-team-${stamp}`; + + const createUser = async (userId: string, email: string): Promise => { + const created = await request.post("/user/new", { + headers: auth(), + data: { user_id: userId, user_email: email, user_role: "internal_user", auto_create_key: false }, + }); + expect(created.ok(), `POST /user/new for ${userId} (${created.status()}): ${await created.text()}`).toBe(true); + const password = await request.post("/user/update", { + headers: auth(), + data: { user_id: userId, password: PASSWORD }, + }); + expect(password.ok(), `POST /user/update for ${userId} (${password.status()})`).toBe(true); + }; + + let teamId = ""; + const createdKeys: string[] = []; + const contexts: BrowserContext[] = []; + try { + await createUser(adminId, adminEmail); + await createUser(memberId, memberEmail); + + const teamRes = await request.post("/team/new", { + headers: auth(), + data: { + team_alias: teamAlias, + models: [CHAT_MODEL_A], + members_with_roles: [ + { user_id: adminId, role: "admin" }, + { user_id: memberId, role: "user" }, + ], + }, + }); + expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + teamId = (await teamRes.json()).team_id as string; + + const memberContext = await signIn(browser, memberEmail); + contexts.push(memberContext); + const memberPage = memberContext.pages()[0]; + const memberSessionKey = await sessionKey(memberPage); + + const refused = await memberPage.request.post("/key/generate", { + headers: { Authorization: `Bearer ${memberSessionKey}`, "Content-Type": "application/json" }, + data: { team_id: teamId, key_alias: `e2e-perm-denied-${stamp}` }, + }); + expect(refused.status(), "a plain member cannot mint a team key before the grant").toBe(401); + expect(await refused.text()).toContain("/key/generate"); + + const adminContext = await signIn(browser, adminEmail); + contexts.push(adminContext); + const adminPage = adminContext.pages()[0]; + await navigateToPage(adminPage, Page.Teams); + await clickTeamId(adminPage, teamId); + await adminPage.getByRole("tab", { name: "Member Permissions" }).click(); + + for (const route of ["/key/generate", "/key/update"]) { + await adminPage.getByRole("row").filter({ hasText: route }).getByRole("checkbox").check(); + } + await adminPage.getByRole("button", { name: "Save Changes" }).click(); + await expect(adminPage.getByText("Permissions updated successfully").first()).toBeVisible({ timeout: 10_000 }); + + await expect + .poll( + async () => { + const res = await request.get(`/team/permissions_list?team_id=${encodeURIComponent(teamId)}`, { + headers: auth(), + }); + if (!res.ok()) return []; + return ((await res.json()).team_member_permissions ?? []) as string[]; + }, + { message: "the granted permissions never landed in /team/permissions_list", timeout: 20_000 }, + ) + .toEqual(expect.arrayContaining(["/key/generate", "/key/update"])); + + const keyAlias = `e2e-perm-key-${stamp}`; + await navigateToPage(memberPage, Page.ApiKeys); + await memberPage.getByRole("button", { name: /Create New Key/i }).click(); + await expect(memberPage.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + await memberPage.getByLabel(/Key Name/).fill(keyAlias); + + const teamSelect = memberPage.getByTestId("team-dropdown").getByRole("combobox"); + await teamSelect.click(); + await memberPage.keyboard.type(teamAlias); + await memberPage.getByRole("option", { name: teamAlias }).first().click(); + + await memberPage.getByRole("combobox", { name: "Select models" }).click(); + await memberPage.getByRole("option", { name: "All Team Models", exact: true }).click(); + await memberPage.keyboard.press("Escape"); + + await memberPage.getByRole("button", { name: "Create Key", exact: true }).click(); + const saveDialog = memberPage.getByRole("dialog", { name: "Save your Key" }); + await expect(saveDialog).toBeVisible({ timeout: 15_000 }); + const apiKey = (await saveDialog.locator("pre").innerText()).trim(); + expect(apiKey).toMatch(/^sk-/); + createdKeys.push(apiKey); + await memberPage.keyboard.press("Escape"); + + await expect + .poll( + async () => { + const res = await request.get( + `/key/list?team_id=${encodeURIComponent(teamId)}&return_full_object=true&size=100`, + { headers: auth() }, + ); + if (!res.ok()) return null; + const row = ((await res.json()).keys as Record[]).find( + (candidate) => candidate.key_alias === keyAlias, + ); + return row ? [row.user_id, row.team_id] : null; + }, + { message: `key ${keyAlias} never appeared on the team with the member as its owner`, timeout: 20_000 }, + ) + .toEqual([memberId, teamId]); + + const served = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" }, + data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `member key ping ${stamp}` }] }, + }); + expect(served.status(), "the delegated key is a real key the gateway serves").toBe(200); + expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); + } finally { + for (const context of contexts) { + await context.close(); + } + for (const key of createdKeys) { + await request.post("/key/delete", { headers: auth(), data: { keys: [key] } }); + } + if (teamId) { + await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } }); + } + await request.post("/user/delete", { headers: auth(), data: { user_ids: [adminId, memberId] } }); + } + }); +}); diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 3fab20a28ad..e5826b18668 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -2,6 +2,7 @@ import glob import os import re import sys +from pathlib import Path import pytest @@ -870,3 +871,69 @@ class TestMigrateDeployAttemptAccounting: harness.run() assert len(harness.deploy_calls) == 1 assert harness.resolved == [] + + +class TestJWTKeyMappingCascade: + """Regression tests for issue #33702. + + A virtual key referenced by a LiteLLM_JWTKeyMapping row could not be deleted + because LiteLLM_JWTKeyMapping_token_fkey was created ON DELETE RESTRICT, so + deleting the key (Admin UI, /key/delete, team delete, ...) raised a foreign + key violation. The mapping must be removed automatically when its key is + deleted, which the FK now enforces via ON DELETE CASCADE. + """ + + _FK_NAME = "LiteLLM_JWTKeyMapping_token_fkey" + + def _effective_on_delete(self): + """Replay every migration in order and return the last ON DELETE action + declared for the JWT key mapping FK.""" + action = None + for _migration_name, sql in _get_all_migrations(): + for match in re.finditer( + rf'ADD\s+CONSTRAINT\s+"{re.escape(self._FK_NAME)}".*?' + r"ON\s+DELETE\s+(CASCADE|RESTRICT|SET\s+NULL|NO\s+ACTION|SET\s+DEFAULT)", + sql, + re.IGNORECASE | re.DOTALL, + ): + action = re.sub(r"\s+", " ", match.group(1).upper()) + return action + + def test_fk_effective_on_delete_is_cascade(self): + """The final FK definition across all migrations must cascade deletes.""" + assert self._effective_on_delete() == "CASCADE", ( + f"{self._FK_NAME} must end up ON DELETE CASCADE so deleting a " + "virtual key removes its JWT key mapping (issue #33702)" + ) + + def test_schema_declares_cascade_on_relation(self): + """schema.prisma must declare onDelete: Cascade on the mapping relation + so the generated client and DB agree.""" + schema_paths = glob.glob( + os.path.abspath( + os.path.join( + os.path.dirname(__file__), "../../**/schema.prisma" + ) + ), + recursive=True, + ) + declaring = tuple( + (path, schema) + for path, schema in ((p, Path(p).read_text()) for p in schema_paths) + if "model LiteLLM_JWTKeyMapping" in schema + ) + assert declaring, "No schema.prisma declaring LiteLLM_JWTKeyMapping found" + for path, schema in declaring: + match = re.search( + r"litellm_verification_token\s+LiteLLM_VerificationToken\s+@relation\(([^)]*)\)", + schema, + ) + assert match is not None, ( + f"{path} declares LiteLLM_JWTKeyMapping but its verification token " + "relation could not be parsed, so this test cannot vouch for it " + "(issue #33702)" + ) + assert "onDelete: Cascade" in match.group(1), ( + f"{path} must declare onDelete: Cascade on the JWT key mapping " + "relation (issue #33702)" + ) diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 9f6e1f4c3f7..4103536950d 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -401,7 +401,7 @@ async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance() enduser_writes = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"]["user_id"]["in"] == [f"user{i}" for i in range(1, 7)] + assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} assert enduser_writes[0]["data"] == {"spend": 0} budget_writes = [c for c in batch_calls if c["table"] == "budget"] @@ -602,7 +602,7 @@ async def test_reset_budget_continues_other_categories_on_failure(): assert len([c for c in batch_calls if c["table"] == "team_membership"]) == 1 enduser_writes = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1"]}} + assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} assert enduser_writes[0]["data"] == {"spend": 0} # Check the new batch write path: 2 keys + 1 user (user1 failed) + 2 teams. @@ -1031,7 +1031,7 @@ async def test_service_logger_endusers_success(): enduser_writes = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1", "user2"]}} + assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once() ( diff --git a/tests/llm_translation/test_azure_o_series.py b/tests/llm_translation/test_azure_o_series.py index 1a2d672af71..ce7e614cbe2 100644 --- a/tests/llm_translation/test_azure_o_series.py +++ b/tests/llm_translation/test_azure_o_series.py @@ -159,15 +159,23 @@ def test_azure_o_series_routing(): def test_openai_o_series_max_retries_0(mock_get_openai_client): import litellm + mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} + mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( + ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + ) litellm.set_verbose = True response = litellm.completion( model="azure/o1-preview", messages=[{"role": "user", "content": "hi"}], max_retries=0, + api_key="fake-key", + api_base="https://fake-azure.openai.azure.com", + api_version="2024-10-21", ) mock_get_openai_client.assert_called_once() assert mock_get_openai_client.call_args.kwargs["max_retries"] == 0 + assert response.choices[0].message.content == "Hello" @pytest.mark.asyncio diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index 0fa72b45ed8..e6528e77749 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -335,6 +335,10 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): ] with patch.object(client.chat.completions.with_raw_response, "create") as mock_post: + mock_post.return_value.headers = {} + mock_post.return_value.parse.return_value = litellm.ModelResponse( + choices=[{"message": {"role": "assistant", "content": InvestigationOutput().model_dump_json()}}] + ) response = litellm.completion( model="azure/gpt-4.1-mini", messages=[ @@ -362,6 +366,7 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): assert "response_format" in mock_post.call_args.kwargs else: assert "response_format" not in mock_post.call_args.kwargs + assert response.choices[0].message.content == InvestigationOutput().model_dump_json() def test_map_openai_params(): diff --git a/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py b/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py index dad2fdbf065..6c059423f74 100644 --- a/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py +++ b/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py @@ -194,6 +194,7 @@ class DummyCredentials: ("aws_web_identity_token", "dummy_web_identity_token"), ("aws_sts_endpoint", "dummy_sts_endpoint"), ("aws_external_id", "dummy_external_id"), + ("aws_session_tags", [{"Key": "team", "Value": "genai"}]), ], ) def test_dynamic_aws_params_propagation(model, param_name, param_value): diff --git a/tests/llm_translation/test_bedrock_embedding.py b/tests/llm_translation/test_bedrock_embedding.py index 56baed141da..1fc05b43b23 100644 --- a/tests/llm_translation/test_bedrock_embedding.py +++ b/tests/llm_translation/test_bedrock_embedding.py @@ -1,14 +1,12 @@ import json import os -from datetime import datetime -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import Mock, patch import pytest import base64 -import httpx import litellm -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import HTTPHandler titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} @@ -394,8 +392,6 @@ def test_bedrock_embedding_uses_correct_region_when_specified(): os.environ["AWS_REGION_NAME"] = original_region_name else: os.environ.pop("AWS_REGION_NAME", None) - - def test_bedrock_embedding_region_bug_reproduction(): """ Reproduces the bug where aws_region_name is ignored when passed explicitly. @@ -458,13 +454,3 @@ def test_bedrock_embedding_region_bug_reproduction(): os.environ["AWS_REGION_NAME"] = original_region_name else: os.environ.pop("AWS_REGION_NAME", None) - - -def test_bedrock_titan_g1_text_02_model_info(): - """Test that amazon.titan-embed-g1-text-02 has correct pricing metadata""" - model_info = litellm.get_model_info("amazon.titan-embed-g1-text-02") - assert model_info is not None, "Model info should not be None" - assert model_info["litellm_provider"] == "bedrock" - assert model_info["mode"] == "embedding" - assert model_info["input_cost_per_token"] == 1e-07 - assert model_info["max_input_tokens"] == 8192 diff --git a/tests/llm_translation/test_bedrock_embedding_pricing.py b/tests/llm_translation/test_bedrock_embedding_pricing.py deleted file mode 100644 index 099d73fed87..00000000000 --- a/tests/llm_translation/test_bedrock_embedding_pricing.py +++ /dev/null @@ -1,34 +0,0 @@ -""" -Tests for AWS Bedrock embedding model pricing in the model cost map. - -Regression test for the Amazon Titan Text Embeddings V2 commercial price, -which was previously set 10x too high (2e-07 instead of 2e-08). -AWS lists Titan Text Embeddings V2 at $0.02 per 1M input tokens -(= $0.00002 per 1K tokens = 2e-08 per token). -""" - -import importlib - - -class TestBedrockEmbeddingPricing: - """Test suite for Bedrock embedding model pricing in the cost map.""" - - def test_titan_embed_v2_commercial_input_cost(self, monkeypatch): - """Titan Text Embeddings V2 should be priced at $0.02 / 1M tokens (2e-08).""" - # Scope the local-cost-map flag to this test only, so it does not leak - # into sibling tests. monkeypatch restores the environment on teardown. - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - - import litellm.litellm_core_utils.get_model_cost_map - import litellm - - # Reload so the cost map is re-read from the local file with the flag set. - importlib.reload(litellm.litellm_core_utils.get_model_cost_map) - importlib.reload(litellm) - - model = litellm.model_cost["amazon.titan-embed-text-v2:0"] - - assert model["input_cost_per_token"] == 2e-08 - assert model["output_cost_per_token"] == 0.0 - assert model["litellm_provider"] == "bedrock" - assert model["mode"] == "embedding" diff --git a/tests/llm_translation/test_bedrock_govcloud.py b/tests/llm_translation/test_bedrock_govcloud.py index e69a95c714d..3ac1fa7cf2e 100644 --- a/tests/llm_translation/test_bedrock_govcloud.py +++ b/tests/llm_translation/test_bedrock_govcloud.py @@ -40,38 +40,6 @@ class TestBedrockGovCloudSupport: assert "us-gov-east-1" in all_regions assert "us-gov-west-1" in all_regions - def test_govcloud_models_in_model_cost(self): - """Test that GovCloud models are present in model cost configuration""" - from litellm import model_cost - - # Test Claude models in GovCloud - assert ( - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0" - in model_cost - ) - assert ( - "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0" - in model_cost - ) - assert ( - "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost - ) - assert ( - "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost - ) - assert "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0" in model_cost - assert "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0" in model_cost - - # Test Llama models in GovCloud - assert "bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0" in model_cost - assert "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0" in model_cost - assert "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0" in model_cost - assert "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0" in model_cost - - # Test Titan models in GovCloud - assert "bedrock/us-gov-east-1/amazon.titan-text-lite-v1" in model_cost - assert "bedrock/us-gov-west-1/amazon.titan-text-lite-v1" in model_cost - def test_govcloud_model_routing(self): """Test that GovCloud models are routed correctly""" # Test Claude model routing @@ -148,135 +116,6 @@ class TestBedrockGovCloudSupport: assert not any("us-gov-east-1" in model for model in litellm.bedrock_models) assert not any("us-gov-west-1" in model for model in litellm.bedrock_models) - def test_govcloud_model_cost_properties(self): - """Test that GovCloud models have proper cost configuration""" - from litellm import model_cost - - # Check a specific GovCloud model has all required properties - govcloud_model = model_cost[ - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0" - ] - - assert "max_tokens" in govcloud_model - assert "max_input_tokens" in govcloud_model - assert "max_output_tokens" in govcloud_model - assert "input_cost_per_token" in govcloud_model - assert "output_cost_per_token" in govcloud_model - assert govcloud_model["litellm_provider"] == "bedrock" - assert govcloud_model["mode"] == "chat" - - def test_govcloud_model_pricing_verification(self): - """Test that GovCloud models have correct pricing that differs from base models""" - from litellm import model_cost - - # Claude Haiku 4.5 commercial list pricing is under the us.* inference profile id - base_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0" - gov_east_model = ( - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0" - ) - gov_west_model = ( - "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0" - ) - - # Verify base model pricing (us.* inference profile: $1.10/$5.50 per MTok) - base_pricing = model_cost[base_model] - assert base_pricing["input_cost_per_token"] == 1.1e-06 - assert base_pricing["output_cost_per_token"] == 5.5e-06 - - # Verify GovCloud models have different (higher) pricing - gov_east_pricing = model_cost[gov_east_model] - gov_west_pricing = model_cost[gov_west_model] - - # GovCloud models should have ~20% higher pricing than base models - assert gov_east_pricing["input_cost_per_token"] == 1.2e-06 - assert gov_east_pricing["output_cost_per_token"] == 6e-06 - assert gov_west_pricing["input_cost_per_token"] == 1.2e-06 - assert gov_west_pricing["output_cost_per_token"] == 6e-06 - - # Verify the pricing difference is approximately 20% - assert ( - abs( - gov_east_pricing["input_cost_per_token"] - / base_pricing["input_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - assert ( - abs( - gov_east_pricing["output_cost_per_token"] - / base_pricing["output_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - assert ( - abs( - gov_west_pricing["input_cost_per_token"] - / base_pricing["input_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - assert ( - abs( - gov_west_pricing["output_cost_per_token"] - / base_pricing["output_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - - # Test Claude 3 Haiku pricing - base_haiku_model = "anthropic.claude-3-haiku-20240307-v1:0" - gov_east_haiku_model = ( - "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" - ) - gov_west_haiku_model = ( - "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" - ) - - # Verify base Haiku model pricing - base_haiku_pricing = model_cost[base_haiku_model] - assert base_haiku_pricing["input_cost_per_token"] == 2.5e-07 # 0.00000025 - assert base_haiku_pricing["output_cost_per_token"] == 1.25e-06 # 0.00000125 - - # Verify GovCloud Haiku models have different (higher) pricing - gov_east_haiku_pricing = model_cost[gov_east_haiku_model] - gov_west_haiku_pricing = model_cost[gov_west_haiku_model] - - # GovCloud Haiku models should have 20% higher pricing than base models - assert ( - gov_east_haiku_pricing["input_cost_per_token"] == 3e-07 - ) # 0.0000003 (20% higher) - assert ( - gov_east_haiku_pricing["output_cost_per_token"] == 1.5e-06 - ) # 0.0000015 (20% higher) - assert ( - gov_west_haiku_pricing["input_cost_per_token"] == 3e-07 - ) # 0.0000003 (20% higher) - assert ( - gov_west_haiku_pricing["output_cost_per_token"] == 1.5e-06 - ) # 0.0000015 (20% higher) - - # Verify the pricing difference is exactly 20% - assert ( - gov_east_haiku_pricing["input_cost_per_token"] - == base_haiku_pricing["input_cost_per_token"] * 1.2 - ) - assert ( - gov_east_haiku_pricing["output_cost_per_token"] - == base_haiku_pricing["output_cost_per_token"] * 1.2 - ) - assert ( - gov_west_haiku_pricing["input_cost_per_token"] - == base_haiku_pricing["input_cost_per_token"] * 1.2 - ) - assert ( - gov_west_haiku_pricing["output_cost_per_token"] - == base_haiku_pricing["output_cost_per_token"] * 1.2 - ) - @patch("litellm.completion") def test_govcloud_completion_cost_calculation(self, mock_completion): """Test that completion requests use correct pricing for GovCloud models""" diff --git a/tests/llm_translation/test_crusoe.py b/tests/llm_translation/test_crusoe.py index 56aa4e4cd42..576428684fc 100644 --- a/tests/llm_translation/test_crusoe.py +++ b/tests/llm_translation/test_crusoe.py @@ -4,7 +4,6 @@ Tests for Crusoe provider integration import os from unittest import mock -import litellm CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1" @@ -71,38 +70,3 @@ def test_get_llm_provider_crusoe(): ) assert model == "meta-llama/Llama-3.3-70B-Instruct" assert provider == "crusoe" - - -def test_crusoe_models_configuration(): - """Test that Crusoe models are configured correctly""" - from litellm import get_model_info - - original_model_cost = litellm.model_cost - original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") - try: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - crusoe_models = [ - "crusoe/meta-llama/Llama-3.3-70B-Instruct", - "crusoe/deepseek-ai/DeepSeek-R1-0528", - "crusoe/deepseek-ai/DeepSeek-V3-0324", - "crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507", - "crusoe/moonshotai/Kimi-K2-Thinking", - "crusoe/openai/gpt-oss-120b", - "crusoe/google/gemma-3-12b-it", - ] - - for model in crusoe_models: - model_info = get_model_info(model) - assert model_info is not None, f"Model info not found for {model}" - assert model_info.get("litellm_provider") == "crusoe", ( - f"{model} should have crusoe as provider" - ) - assert model_info.get("mode") == "chat", f"{model} should be in chat mode" - finally: - litellm.model_cost = original_model_cost - if original_env is None: - os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) - else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env diff --git a/tests/llm_translation/test_hyperbolic.py b/tests/llm_translation/test_hyperbolic.py index 78817fbd902..b7206e40a4e 100644 --- a/tests/llm_translation/test_hyperbolic.py +++ b/tests/llm_translation/test_hyperbolic.py @@ -1,8 +1,3 @@ -import os -from datetime import datetime -from unittest.mock import MagicMock - -import pytest import litellm @@ -69,34 +64,6 @@ def test_hyperbolic_in_provider_lists(): assert "https://api.hyperbolic.xyz/v1" in openai_compatible_endpoints -def test_hyperbolic_models_configuration(): - """Test that Hyperbolic models are properly configured""" - import json - - # Load model configuration directly from the JSON file - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path, "r") as f: - model_data = json.load(f) - - # Test a few key models - test_models = [ - "hyperbolic/deepseek-ai/DeepSeek-V3", - "hyperbolic/Qwen/Qwen2.5-Coder-32B-Instruct", - "hyperbolic/deepseek-ai/DeepSeek-R1", - ] - - for model in test_models: - assert model in model_data - model_info = model_data[model] - assert model_info["litellm_provider"] == "hyperbolic" - assert model_info["mode"] == "chat" - assert "max_tokens" in model_info - assert "input_cost_per_token" in model_info - assert "output_cost_per_token" in model_info - - def test_hyperbolic_supported_params(): """Test that supported OpenAI parameters are correctly configured""" from litellm.llms.hyperbolic.chat.transformation import HyperbolicChatConfig diff --git a/tests/llm_translation/test_lambda_ai.py b/tests/llm_translation/test_lambda_ai.py index 7ae18828d3f..edba459b352 100644 --- a/tests/llm_translation/test_lambda_ai.py +++ b/tests/llm_translation/test_lambda_ai.py @@ -8,7 +8,6 @@ from unittest import mock import pytest import litellm -from litellm import completion from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig @@ -103,48 +102,6 @@ async def test_lambda_ai_completion_call(): raise -def test_lambda_ai_models_configuration(): - """Test that Lambda AI models are configured correctly""" - from litellm import get_model_info - - # Reload model cost map to pick up local changes - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - # Clear and repopulate lambda_ai_models list after reloading model_cost - litellm.lambda_ai_models = set() - litellm.add_known_models() - - # Some Lambda AI models to test - lambda_ai_models = [ - "lambda_ai/deepseek-llama3.3-70b", - "lambda_ai/hermes3-8b", - "lambda_ai/llama3.1-8b-instruct", - "lambda_ai/llama3.2-11b-vision-instruct", - "lambda_ai/qwen25-coder-32b-instruct", - ] - - for model in lambda_ai_models: - model_info = get_model_info(model) - assert model_info is not None, f"Model info not found for {model}" - assert ( - model_info.get("litellm_provider") == "lambda_ai" - ), f"{model} should have lambda_ai as provider" - assert model_info.get("mode") == "chat", f"{model} should be in chat mode" - assert ( - model_info.get("supports_function_calling") is True - ), f"{model} should support function calling" - assert ( - model_info.get("supports_system_messages") is True - ), f"{model} should support system messages" - - # Check vision support for vision models - if "vision" in model: - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - - def test_lambda_ai_model_list_populated(): """Test that lambda_ai_models list is populated correctly""" # Ensure we're using local model cost map and repopulate models diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index b6e30ddc711..31c554985a7 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -1623,15 +1623,11 @@ class TestMissingChoicesGuard: assert "no 'choices'" in exc_info.value.message - def test_convert_to_model_response_object_empty_choices_raises_api_error(self): - """Empty choices list raises APIError, same as missing/null choices. + def test_convert_to_model_response_object_empty_choices_returns_empty_list(self): + """An empty choices list is a real provider answer, so it converts to choices=[] instead of raising. - Provider-specific repair (e.g. github_copilot synthesizing choices for - Anthropic-native responses) happens before this guard, in the provider - config; the core utility keeps treating empty choices as an error. + See: https://github.com/BerriAI/litellm/issues/40276 """ - from litellm.exceptions import APIError - response_object = { "id": "msg_123", "model": "some-model", @@ -1639,16 +1635,17 @@ class TestMissingChoicesGuard: "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, } - with pytest.raises(APIError) as exc_info: - convert_to_model_response_object( - response_object=response_object, - model_response_object=ModelResponse(), - ) + result = convert_to_model_response_object( + response_object=response_object, + model_response_object=ModelResponse(), + ) - assert "no 'choices'" in exc_info.value.message + assert isinstance(result, ModelResponse) + assert result.choices == [] + assert result.usage.prompt_tokens == 10 def test_convert_to_model_response_object_null_choices_raises_api_error(self): - """choices=None raises APIError.""" + """choices=None raises APIError that names the type instead of claiming the key is missing.""" from litellm.exceptions import APIError response_object = { @@ -1664,7 +1661,7 @@ class TestMissingChoicesGuard: model_response_object=ModelResponse(), ) - assert "no 'choices'" in exc_info.value.message + assert "'choices' that is not a list (NoneType)" in exc_info.value.message def test_convert_to_streaming_response_no_choices_raises_api_error(self): """Missing choices in streaming cache-hit path raises APIError.""" diff --git a/tests/llm_translation/test_morph.py b/tests/llm_translation/test_morph.py index b91d1810d38..752fb3b9083 100644 --- a/tests/llm_translation/test_morph.py +++ b/tests/llm_translation/test_morph.py @@ -68,24 +68,6 @@ def test_morph_in_provider_lists(): ) -def test_morph_model_info(): - """Test that morph models have correct configuration.""" - import litellm - - model_info = litellm.get_model_info("morph/morph-v3-large") - - assert model_info["litellm_provider"] == "morph" - assert model_info["mode"] == "chat" - assert model_info["max_tokens"] == 16000 - assert model_info["max_input_tokens"] == 16000 - assert model_info["max_output_tokens"] == 16000 - assert model_info["input_cost_per_token"] == 9e-07 # $0.9/1M tokens - assert model_info["output_cost_per_token"] == 1.9e-06 # $1.9/1M tokens - assert model_info["supports_function_calling"] is False - assert model_info["supports_vision"] is False - assert model_info["supports_system_messages"] is True - - def test_morph_supported_params(): """Test that MorphChatConfig returns correct supported parameters.""" config = MorphChatConfig() diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index 2b9abdec5d0..af4ba85d58e 100644 --- a/tests/llm_translation/test_openai.py +++ b/tests/llm_translation/test_openai.py @@ -292,15 +292,21 @@ class TestOpenAIChatCompletion(BaseLLMChatTest): def test_openai_max_retries_0(mock_get_openai_client): import litellm + mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} + mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( + ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + ) litellm.set_verbose = True response = litellm.completion( model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], max_retries=0, + api_key="fake-key", ) mock_get_openai_client.assert_called_once() assert mock_get_openai_client.call_args.kwargs["max_retries"] == 0 + assert response.choices[0].message.content == "Hello" @patch("litellm.main.openai_chat_completions._get_openai_client") diff --git a/tests/llm_translation/test_openai_o1.py b/tests/llm_translation/test_openai_o1.py index e188a3af647..fd25e04d67d 100644 --- a/tests/llm_translation/test_openai_o1.py +++ b/tests/llm_translation/test_openai_o1.py @@ -1,15 +1,11 @@ -import json import os -from datetime import datetime -from unittest.mock import AsyncMock, patch, MagicMock +from unittest.mock import patch - -import httpx import pytest import litellm -from litellm import Choices, Message, ModelResponse +from litellm import ModelResponse from base_llm_unit_tests import BaseLLMChatTest, BaseOSeriesModelsTest @@ -74,7 +70,6 @@ async def test_o1_handle_tool_calling_optional_params( - max_tokens is translated to 'max_completion_tokens' - role 'system' is translated to 'user' """ - from openai import AsyncOpenAI from litellm.utils import ProviderConfigManager from litellm.types.utils import LlmProviders @@ -186,15 +181,6 @@ class TestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest): pass -def test_o1_supports_vision(): - """Test that o1 supports vision""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - for k, v in litellm.model_cost.items(): - if k.startswith("o1") and v.get("litellm_provider") == "openai": - assert v.get("supports_vision") is True, f"{k} does not support vision" - - def test_o3_reasoning_effort(): resp = litellm.completion( model="o3-mini", diff --git a/tests/llm_translation/test_v0.py b/tests/llm_translation/test_v0.py index 95708dd855a..e96022e1e22 100644 --- a/tests/llm_translation/test_v0.py +++ b/tests/llm_translation/test_v0.py @@ -8,7 +8,6 @@ from unittest import mock import pytest import litellm -from litellm import completion from litellm.llms.v0.chat.transformation import V0ChatConfig @@ -111,33 +110,3 @@ def test_v0_supported_params(): ] assert set(supported_params) == set(expected_params) - - -def test_v0_models_configuration(): - """Test that v0 models are configured correctly""" - from litellm import get_model_info - - # Reload model cost map to pick up local changes - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - # All v0 models - v0_models = ["v0/v0-1.0-md", "v0/v0-1.5-md", "v0/v0-1.5-lg"] - - for model in v0_models: - model_info = get_model_info(model) - assert model_info is not None, f"Model info not found for {model}" - # All v0 models support vision (multimodal) - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - assert ( - model_info.get("litellm_provider") == "v0" - ), f"{model} should have v0 as provider" - assert model_info.get("mode") == "chat", f"{model} should be in chat mode" - assert ( - model_info.get("supports_function_calling") is True - ), f"{model} should support function calling" - assert ( - model_info.get("supports_system_messages") is True - ), f"{model} should support system messages" diff --git a/tests/llm_translation/test_voyage_ai.py b/tests/llm_translation/test_voyage_ai.py index 208e01110da..800751be115 100644 --- a/tests/llm_translation/test_voyage_ai.py +++ b/tests/llm_translation/test_voyage_ai.py @@ -139,6 +139,7 @@ class TestVoyageContextualEmbeddings: # Test contextual model detection assert config.is_contextualized_embeddings("voyage-context-3") is True + assert config.is_contextualized_embeddings("voyage-context-4") is True assert config.is_contextualized_embeddings("voyage-context-2") is True assert config.is_contextualized_embeddings("context-model") is True diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index f8f23ea015a..43ed57f63af 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -3999,10 +3999,14 @@ def test_completion_novita_ai(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client.chat.completions.with_raw_response, "create" ) as mock_call: + mock_call.return_value.headers = {} + mock_call.return_value.parse.return_value = litellm.ModelResponse( + choices=[{"message": {"role": "assistant", "content": "Hello"}}] + ) try: - completion( + response = completion( model="novita/meta-llama/llama-3.3-70b-instruct", messages=messages, client=openai_client, @@ -4010,6 +4014,7 @@ def test_completion_novita_ai(): ) mock_call.assert_called_once() + assert response.choices[0].message.content == "Hello" # Verify model is passed correctly assert ( diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index f0f24a6e6b2..834570091bd 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -1076,7 +1076,7 @@ def test_standard_logging_payload(model, turn_off_message_logging): ) ) - keys_list = list(StandardLoggingPayload.__annotations__.keys()) + keys_list = list(StandardLoggingPayload.__required_keys__) for k in keys_list: assert ( @@ -1190,7 +1190,7 @@ def test_standard_logging_payload_audio(turn_off_message_logging, stream): ) ) - keys_list = list(StandardLoggingPayload.__annotations__.keys()) + keys_list = list(StandardLoggingPayload.__required_keys__) for k in keys_list: assert ( diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index 2de83778f1c..38ccfd91f95 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -1,8 +1,6 @@ # What is this? ## Unit testing for the 'get_model_info()' function import os -import traceback -import json from typing import List, Dict, Any @@ -11,7 +9,7 @@ import pytest import litellm from litellm import get_model_info -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch def test_get_model_info_simple_model_name(): @@ -49,34 +47,6 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch): assert model_info["input_cost_per_token"] == 0.0 -def test_get_model_info_shows_correct_supports_vision(): - info = litellm.get_model_info("gemini/gemini-2.0-flash") - print("info", info) - assert info["supports_vision"] is True - - -def test_get_model_info_shows_assistant_prefill(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - info = litellm.get_model_info("deepseek/deepseek-chat") - print("info", info) - assert info.get("supports_assistant_prefill") is True - - -def test_get_model_info_shows_supports_prompt_caching(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - info = litellm.get_model_info("deepseek/deepseek-chat") - print("info", info) - assert info.get("supports_prompt_caching") is True - - -def test_get_model_info_finetuned_models(): - info = litellm.get_model_info("ft:gpt-3.5-turbo:my-org:custom_suffix:id") - print("info", info) - assert info["input_cost_per_token"] == 0.000003 - - def test_get_model_info_gemini_pro(): info = litellm.get_model_info("gemini-2.0-flash") print("info", info) @@ -219,7 +189,7 @@ def test_model_info_bedrock_converse_enforcement(monkeypatch): def test_get_model_info_custom_provider(): # Custom provider example copied from https://docs.litellm.ai/docs/providers/custom_llm_server: import litellm - from litellm import CustomLLM, completion, get_llm_provider + from litellm import CustomLLM, completion class MyCustomLLM(CustomLLM): def completion(self, *args, **kwargs) -> litellm.ModelResponse: diff --git a/tests/local_testing/test_least_busy_routing.py b/tests/local_testing/test_least_busy_routing.py index 18ab8bf779d..fb83d4e601f 100644 --- a/tests/local_testing/test_least_busy_routing.py +++ b/tests/local_testing/test_least_busy_routing.py @@ -33,8 +33,8 @@ def test_model_added(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"gpt-3.5-turbo_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = "gpt-3.5-turbo_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 def test_get_available_deployments(): @@ -52,8 +52,8 @@ def test_get_available_deployments(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"{model_group}_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = f"{model_group}_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 # test_get_available_deployments() @@ -104,15 +104,20 @@ async def test_router_get_available_deployments(async_test): router.leastbusy_logger.test_flag = True model_group = "azure-model" - request_count_dict = {1: 10, 2: 54, 3: 100} - cache_key = f"{model_group}_request_count" + request_count_dict = {"1": 10, "2": 54, "3": 100} + cache_keys = { + deployment_id: f"{model_group}_request_count:{deployment_id}" + for deployment_id in request_count_dict + } if async_test is True: - await router.cache.async_set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + await router.cache.async_set_cache(key=cache_keys[deployment_id], value=count) deployment = await router.async_get_available_deployment( model=model_group, messages=None, request_kwargs={} ) else: - router.cache.set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + router.cache.set_cache(key=cache_keys[deployment_id], value=count) deployment = router.get_available_deployment(model=model_group, messages=None) print(f"deployment: {deployment}") assert deployment["model_info"]["id"] == "1" @@ -124,15 +129,18 @@ async def test_router_get_available_deployments(async_test): messages=[{"role": "user", "content": "Hey, how's it going?"}], ) - return_dict = router.cache.get_cache(key=cache_key) - # wait 2 seconds time.sleep(2) + return_dict = { + deployment_id: router.cache.get_cache(key=cache_key) + for deployment_id, cache_key in cache_keys.items() + } + assert router.leastbusy_logger.logged_success == 1 - assert return_dict[1] == 10 - assert return_dict[2] == 54 - assert return_dict[3] == 100 + assert return_dict["1"] == 10 + assert return_dict["2"] == 54 + assert return_dict["3"] == 100 ## Test with Real calls ## @@ -192,9 +200,11 @@ async def test_router_atext_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.atext_completion(model=model, prompt=prompt, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" @@ -259,8 +269,10 @@ async def test_router_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.acompletion(model=model, messages=messages, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" diff --git a/tests/local_testing/test_lowest_latency_routing.py b/tests/local_testing/test_lowest_latency_routing.py index 598b1dbcaf9..aba4500199a 100644 --- a/tests/local_testing/test_lowest_latency_routing.py +++ b/tests/local_testing/test_lowest_latency_routing.py @@ -1077,73 +1077,6 @@ async def test_latency_list_trimming_discards_oldest_entry_async(): ), f"Oldest latency {oldest_latency} should have been discarded" -def test_ttft_list_trimming_discards_oldest_entry(): - """ - The time_to_first_token list trims the oldest entry when full, matching - the behavior of the latency list. - """ - max_size = 3 - test_cache = DualCache() - lowest_latency_logger = LowestLatencyLoggingHandler( - router_cache=test_cache, routing_args={"max_latency_list_size": max_size} - ) - - model_group = "gpt-3.5-turbo" - deployment_id = "test-deployment" - - ttft_values = [] - for i in range(max_size + 1): - start_time = time.time() - expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4 - completion_start_time = start_time + expected_ttft - end_time = start_time + float(i + 1) - ttft_values.append(expected_ttft) - - kwargs = { - "litellm_params": { - "metadata": { - "model_group": model_group, - "deployment": "azure/gpt-4.1-mini", - }, - "model_info": {"id": deployment_id}, - }, - "stream": True, - "completion_start_time": completion_start_time, - } - # TTFT is only recorded when response_obj is a ModelResponse. - response_obj = litellm.ModelResponse( - usage=litellm.Usage(completion_tokens=1, total_tokens=1) - ) - - lowest_latency_logger.log_success_event( - response_obj=response_obj, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) - - latency_key = f"{model_group}_map" - cached_data = test_cache.get_cache(key=latency_key) - ttft_list = cached_data[deployment_id].get("time_to_first_token", []) - - assert ( - len(ttft_list) == max_size - ), f"Expected {max_size} entries, got {len(ttft_list)}" - - newest_ttft = ttft_values[-1] - oldest_ttft = ttft_values[0] - tolerance = 0.05 - - assert ( - abs(ttft_list[-1] - newest_ttft) < tolerance - ), f"Newest TTFT {newest_ttft} should be at end of list" - - for ttft in ttft_list: - assert ( - abs(ttft - oldest_ttft) > tolerance - ), f"Oldest TTFT {oldest_ttft} should have been discarded" - - @pytest.mark.asyncio async def test_timeout_penalty_discards_oldest_entry(): """ @@ -1269,72 +1202,3 @@ def test_list_order_preserved_after_multiple_trims(): assert ( abs(latency_list[i] - expected) < tolerance ), f"At index {i}, expected ~{expected}, got {latency_list[i]}" - - -@pytest.mark.asyncio -async def test_ttft_list_trimming_discards_oldest_entry_async(): - """ - Async counterpart: the time_to_first_token list trims the oldest entry - when full. Exercises the async_log_success_event TTFT path, which only - runs when response_obj is a ModelResponse and the call is marked as - streaming with a completion_start_time. - """ - max_size = 3 - test_cache = DualCache() - lowest_latency_logger = LowestLatencyLoggingHandler( - router_cache=test_cache, routing_args={"max_latency_list_size": max_size} - ) - - model_group = "gpt-3.5-turbo" - deployment_id = "test-deployment" - - ttft_values = [] - for i in range(max_size + 1): - start_time = time.time() - expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4 - completion_start_time = start_time + expected_ttft - end_time = start_time + float(i + 1) - ttft_values.append(expected_ttft) - - kwargs = { - "litellm_params": { - "metadata": { - "model_group": model_group, - "deployment": "azure/gpt-4.1-mini", - }, - "model_info": {"id": deployment_id}, - }, - "stream": True, - "completion_start_time": completion_start_time, - } - response_obj = litellm.ModelResponse( - usage=litellm.Usage(completion_tokens=1, total_tokens=1) - ) - - await lowest_latency_logger.async_log_success_event( - response_obj=response_obj, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) - - latency_key = f"{model_group}_map" - cached_data = await test_cache.async_get_cache(key=latency_key) - ttft_list = cached_data[deployment_id].get("time_to_first_token", []) - - assert ( - len(ttft_list) == max_size - ), f"Expected {max_size} entries, got {len(ttft_list)}" - - newest_ttft = ttft_values[-1] - oldest_ttft = ttft_values[0] - tolerance = 0.05 - - assert ( - abs(ttft_list[-1] - newest_ttft) < tolerance - ), f"Newest TTFT {newest_ttft} should be at end of list" - - for ttft in ttft_list: - assert ( - abs(ttft - oldest_ttft) > tolerance - ), f"Oldest TTFT {oldest_ttft} should have been discarded" diff --git a/tests/local_testing/test_redis_increment_with_floor.py b/tests/local_testing/test_redis_increment_with_floor.py new file mode 100644 index 00000000000..e358d5f31e0 --- /dev/null +++ b/tests/local_testing/test_redis_increment_with_floor.py @@ -0,0 +1,80 @@ +"""Least-busy routing keeps its in-flight counters in Redis, and the clamp at zero plus the +create-once TTL both live inside a Lua script. Nothing but a real Redis runs that script, so +these are the only tests that fail when the script itself is wrong.""" + +import os +import uuid +from typing import Final + +import pytest +from dotenv import load_dotenv + +load_dotenv() + +from litellm.caching.redis_cache import RedisCache + +TTL: Final = 600 + + +@pytest.fixture +def counter(): + cache: Final = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT")) + key: Final = f"lit7039-{uuid.uuid4()}" + yield cache, key, cache.check_and_fix_namespace(key=key) + cache.delete_cache(key) + + +def test_a_counter_adds_every_increment_and_reads_back_what_it_holds(counter): + cache, key, _ = counter + + assert cache.increment_with_floor(key, 3, TTL) == 3 + assert cache.increment_with_floor(key, 2, TTL) == 5 + assert cache.batch_get_counts([key]) == (5,) + + +def test_a_decrement_past_zero_leaves_the_counter_at_zero(counter): + """A worker whose counter expired mid-request decrements a key that is no longer there. + Without the clamp that deployment reads negative, and least-busy pins every later request + on it until the count climbs back to zero.""" + cache, key, _ = counter + + assert cache.increment_with_floor(key, 1, TTL) == 1 + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.batch_get_counts([key]) == (0,) + + +def test_traffic_never_pushes_a_counters_expiry_back_out(counter): + """The TTL is what releases a count whose worker died mid-request. Rewriting it on every + touch would keep that stuck count alive for as long as the group takes traffic.""" + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + assert cache.redis_client.ttl(namespaced_key) > TTL - 60 + + cache.redis_client.expire(namespaced_key, 30) + cache.increment_with_floor(key, 1, TTL) + + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +def test_clamping_to_zero_keeps_the_expiry_it_already_had(counter): + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + cache.redis_client.expire(namespaced_key, 30) + + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +@pytest.mark.asyncio +async def test_the_async_counter_behaves_the_same_way(counter): + cache, key, namespaced_key = counter + + assert await cache.async_increment_with_floor(key, 2, TTL) == 2 + assert await cache.async_batch_get_counts([key]) == (2,) + + cache.redis_client.expire(namespaced_key, 30) + + assert await cache.async_increment_with_floor(key, -9, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 diff --git a/tests/logging_callback_tests/test_datadog.py b/tests/logging_callback_tests/test_datadog.py index 83a652e8884..7ac9ac0b5ad 100644 --- a/tests/logging_callback_tests/test_datadog.py +++ b/tests/logging_callback_tests/test_datadog.py @@ -270,7 +270,7 @@ async def test_datadog_logging_http_request(): message = json.loads(body[0]["message"]) print("logged message", json.dumps(message, indent=4)) - expected_message_fields = StandardLoggingPayload.__annotations__.keys() + expected_message_fields = StandardLoggingPayload.__required_keys__ for field in expected_message_fields: assert field in message, f"Field '{field}' is missing from the message" diff --git a/tests/mcp_tests/mcp_server.py b/tests/mcp_tests/mcp_server.py index bc6accbb721..eba7cae1bca 100644 --- a/tests/mcp_tests/mcp_server.py +++ b/tests/mcp_tests/mcp_server.py @@ -1,10 +1,12 @@ # math_server.py import argparse import os +from typing import Final -from mcp.server.fastmcp import FastMCP +from mcp.server.fastmcp import Context, FastMCP mcp = FastMCP("Math") +ADD_OFFSET: Final = int(os.getenv("MCP_ADD_OFFSET", "0")) def _parse_args() -> argparse.Namespace: @@ -31,7 +33,7 @@ def _parse_args() -> argparse.Namespace: @mcp.tool() def add(a: int, b: int) -> int: """Add two numbers""" - return a + b + return a + b + ADD_OFFSET @mcp.tool() @@ -40,6 +42,15 @@ def multiply(a: int, b: int) -> int: return a * b +@mcp.tool() +def request_headers(ctx: Context) -> dict[str, str]: + request: Final = ctx.request_context.request + return { + "authorization": request.headers.get("authorization", "") if request is not None else "", + "x-request-tag": request.headers.get("x-request-tag", "") if request is not None else "", + } + + def main() -> None: args = _parse_args() transport = (args.transport or "stdio").lower() diff --git a/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml b/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml index ad68a03781d..19fad3d1393 100644 --- a/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml +++ b/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml @@ -23,3 +23,6 @@ mcp_servers: transport: http url: http://127.0.0.1:0/mcp allow_all_keys: true + math_restricted: + transport: http + url: http://127.0.0.1:0/mcp diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index 2dd57e13d3b..e1099fe0a62 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -1,26 +1,39 @@ import asyncio +import json import os +import queue import socket import subprocess import sys +import tempfile import threading import time import typing +from contextlib import asynccontextmanager, contextmanager +from dataclasses import dataclass +from datetime import datetime from pathlib import Path +import httpx import pytest import uvicorn import yaml from mcp import ClientSession from mcp.client.streamable_http import streamablehttp_client +from mcp.types import CallToolResult +from starlette.requests import Request +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._experimental.mcp_server.tool_search import handle_mcp_proxy_tool +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, ProxyException, UserAPIKeyAuth from litellm.proxy.proxy_server import ( app as proxy_app, +) +from litellm.proxy.proxy_server import ( cleanup_router_config_variables, initialize, ) - CONFIG_TEMPLATE_PATH = Path("tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml") MCP_SERVER_SCRIPT = Path("tests/mcp_tests/mcp_server.py") PROJECT_ROOT = Path(__file__).resolve().parents[2] @@ -45,28 +58,49 @@ def _clear_proxy_database_env() -> typing.Iterator[None]: mp.undo() -def _initialize_proxy(config_path: str) -> None: +async def _initialize_proxy(config_path: str) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + cleanup_router_config_variables() - asyncio.run(initialize(config=config_path, debug=True)) + await initialize(config=config_path, debug=True) + for server_id, upstream in tuple(global_mcp_server_manager.registry.items()): + if upstream.server_name != "math_restricted": + continue + global_mcp_server_manager.registry[server_id] = upstream.model_copy( + update={"tool_name_to_display_name": {"add": "Add Numbers"}} + ) + + +@dataclass(frozen=True) +class ProxyRig: + url: str + config_path: str + loop: asyncio.AbstractEventLoop def _start_proxy_server( config_path: str, -) -> tuple[str, uvicorn.Server, threading.Thread, socket.socket]: - _initialize_proxy(config_path) - +) -> tuple[ProxyRig, uvicorn.Server, threading.Thread, socket.socket]: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.bind(("127.0.0.1", 0)) host, port = sock.getsockname() - config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning") + config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning", lifespan="off") server = uvicorn.Server(config) + loop = asyncio.new_event_loop() + + async def _serve() -> None: + from litellm.proxy._experimental.mcp_server import server as mcp_server + + await _initialize_proxy(config_path) + async with proxy_app.router.lifespan_context(proxy_app), mcp_server.lifespan(proxy_app): + await server.serve(sockets=[sock]) + def _run() -> None: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - loop.run_until_complete(server.serve(sockets=[sock])) + with asyncio.Runner(loop_factory=lambda: loop) as runner: + runner.run(_serve()) thread = threading.Thread(target=_run, daemon=True) thread.start() @@ -79,79 +113,93 @@ def _start_proxy_server( raise TimeoutError("Proxy server did not start in time") time.sleep(0.05) - return f"http://{host}:{port}", server, thread, sock + return ProxyRig(f"http://{host}:{port}", config_path, loop), server, thread, sock -@pytest.fixture(scope="session") -def math_streamable_http_server() -> str: +@contextmanager +def _math_http_server(offset: int) -> typing.Iterator[str]: host = "127.0.0.1" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind((host, 0)) _, port = sock.getsockname() - cmd = [ - sys.executable, - str(MCP_SERVER_SCRIPT), - "--transport", - "http", - "--host", - host, - "--port", - str(port), - ] - - env = os.environ.copy() - server_process = subprocess.Popen( - cmd, - cwd=str(PROJECT_ROOT), - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - - start_time = time.time() - while True: - if server_process.poll() is not None: - stdout, stderr = server_process.communicate() - raise RuntimeError( - f"Streamable HTTP MCP server exited early.\nSTDOUT: {stdout.decode()}\nSTDERR: {stderr.decode()}" - ) + with tempfile.TemporaryFile() as server_log: + process = subprocess.Popen( + [sys.executable, str(MCP_SERVER_SCRIPT), "--transport", "http", "--host", host, "--port", str(port)], + cwd=str(PROJECT_ROOT), + stdout=server_log, + stderr=subprocess.STDOUT, + env={**os.environ, "MCP_ADD_OFFSET": str(offset)}, + ) try: - with socket.create_connection((host, port), timeout=0.1): - break - except OSError: - if time.time() - start_time > PROXY_START_TIMEOUT: - server_process.terminate() - raise TimeoutError("Streamable HTTP MCP server did not start in time") - time.sleep(0.05) - - yield f"http://{host}:{port}" - - server_process.terminate() - try: - server_process.wait(timeout=5) - except subprocess.TimeoutExpired: - server_process.kill() + start_time = time.monotonic() + while True: + if process.poll() is not None: + server_log.seek(0) + raise RuntimeError(f"MCP upstream exited early: {server_log.read().decode()}") + try: + with socket.create_connection((host, port), timeout=0.1): + break + except OSError: + if time.monotonic() - start_time > PROXY_START_TIMEOUT: + raise TimeoutError("Streamable HTTP MCP server did not start in time") + time.sleep(0.05) + yield f"http://{host}:{port}" + finally: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) @pytest.fixture(scope="session") -def proxy_server_url( - tmp_path_factory: pytest.TempPathFactory, math_streamable_http_server: str +def math_streamable_http_server() -> typing.Iterator[str]: + with _math_http_server(100) as url: + yield url + + +@pytest.fixture(scope="session") +def math_restricted_server() -> typing.Iterator[str]: + with _math_http_server(200) as url: + yield url + + +@pytest.fixture(scope="session") +def _proxy_server( + tmp_path_factory: pytest.TempPathFactory, + math_streamable_http_server: str, + math_restricted_server: str, ): config_dir = tmp_path_factory.mktemp("mcp_e2e") config_path = config_dir / "config.yaml" config = yaml.safe_load(CONFIG_TEMPLATE_PATH.read_text()) - config["mcp_servers"]["math_streamable_http"][ - "url" - ] = f"{math_streamable_http_server}/mcp" + config["mcp_servers"]["math_stdio"]["command"] = sys.executable + config["mcp_servers"]["math_streamable_http"]["url"] = f"{math_streamable_http_server}/mcp" + config["mcp_servers"]["math_restricted"]["url"] = f"{math_restricted_server}/mcp" + config["general_settings"]["custom_auth"] = f"{__name__}.authorize_proxy_key" + config["litellm_settings"]["callbacks"] = [f"{__name__}.proxy_call_recorder"] + config["mcp_servers"]["math_restricted"]["mcp_info"] = {"mcp_server_cost_info": {"default_cost_per_query": 0.25}} config_path.write_text(yaml.safe_dump(config)) - server_url, server, thread, sock = _start_proxy_server(str(config_path)) + rig, server, thread, sock = _start_proxy_server(str(config_path)) - yield server_url + try: + yield rig + finally: + server.should_exit = True + thread.join(timeout=10) + sock.close() + assert not thread.is_alive(), "Proxy did not shut down" - server.should_exit = True - thread.join(timeout=10) - sock.close() + +@pytest.fixture +def proxy_server_url(_proxy_server: ProxyRig, setup_and_teardown: None) -> str: + asyncio.run_coroutine_threadsafe(_initialize_proxy(_proxy_server.config_path), _proxy_server.loop).result( + timeout=30 + ) + return _proxy_server.url class TestProxyMcpSimpleConnections: @@ -177,9 +225,7 @@ class TestProxyMcpSimpleConnections: assert text == "7" @pytest.mark.asyncio - async def test_proxy_mcp_streamable_http_roundtrip( - self, proxy_server_url: str - ) -> None: + async def test_proxy_mcp_streamable_http_roundtrip(self, proxy_server_url: str) -> None: async with asyncio.timeout(20): async with streamablehttp_client( url=f"{proxy_server_url}/mcp", @@ -197,12 +243,10 @@ class TestProxyMcpSimpleConnections: assert result.content first_content = result.content[0] text = getattr(first_content, "text", None) - assert text == "11" + assert text == "111" @pytest.mark.asyncio - async def test_proxy_mcp_lists_all_servers_without_header( - self, proxy_server_url: str - ) -> None: + async def test_proxy_mcp_lists_all_servers_without_header(self, proxy_server_url: str) -> None: async with asyncio.timeout(20): async with streamablehttp_client( url=f"{proxy_server_url}/mcp", @@ -220,22 +264,16 @@ class TestProxyMcpSimpleConnections: } assert expected_tool_names <= tool_names - async def _call_and_get_text( - tool_name: str, *, a: int, b: int - ) -> str | None: - result = await session.call_tool( - tool_name, arguments={"a": a, "b": b} - ) + async def _call_and_get_text(tool_name: str, *, a: int, b: int) -> str | None: + result = await session.call_tool(tool_name, arguments={"a": a, "b": b}) assert result.content first_content = result.content[0] return getattr(first_content, "text", None) stdio_result = await _call_and_get_text("math_stdio-add", a=2, b=3) - streamable_result = await _call_and_get_text( - "math_streamable_http-add", a=4, b=5 - ) + streamable_result = await _call_and_get_text("math_streamable_http-add", a=4, b=5) assert stdio_result == "5" - assert streamable_result == "9" + assert streamable_result == "109" class TestProxyMcpStatelessBehavior: @@ -254,9 +292,7 @@ class TestProxyMcpStatelessBehavior: """ @pytest.mark.asyncio - async def test_independent_clients_no_shared_session( - self, proxy_server_url: str - ) -> None: + async def test_independent_clients_no_shared_session(self, proxy_server_url: str) -> None: """Two independent clients connect and operate without sharing session state.""" async with asyncio.timeout(30): # --- Client A: connect, initialize, call tool --- @@ -269,9 +305,7 @@ class TestProxyMcpStatelessBehavior: ) as (read_a, write_a, _get_sid_a): async with ClientSession(read_a, write_a) as session_a: await session_a.initialize() - result_a = await session_a.call_tool( - "add", arguments={"a": 10, "b": 20} - ) + result_a = await session_a.call_tool("add", arguments={"a": 10, "b": 20}) assert result_a.content text_a = getattr(result_a.content[0], "text", None) assert text_a == "30" @@ -293,9 +327,394 @@ class TestProxyMcpStatelessBehavior: await session_b.initialize() tools = await session_b.list_tools() assert any(t.name.endswith("add") for t in tools.tools) - result_b = await session_b.call_tool( - "add", arguments={"a": 100, "b": 200} - ) + result_b = await session_b.call_tool("add", arguments={"a": 100, "b": 200}) assert result_b.content text_b = getattr(result_b.content[0], "text", None) assert text_b == "300" + + +PROXY_MODE_TOOLS = frozenset({"search_tools", "get_tool_schema", "call_tool"}) + + +def _payload(result: typing.Any) -> typing.Any: + assert result.content, f"empty tool result: {result}" + return json.loads(result.content[0].text) + + +def _proxy_session(proxy_server_url: str, **extra_headers: str): + return streamablehttp_client( + url=f"{proxy_server_url}/mcp/proxy", + headers={"Authorization": PROXY_AUTHORIZATION_HEADER, **extra_headers}, + ) + + +class TestProxyMcpSchemaDiscoveryMode: + """Drive /mcp/proxy over the real streamable-HTTP transport with the MCP SDK client: + the fixed three-tool surface, opaque-id discovery, schema-validated execution against + two upstreams that expose the same tool name, and the operations the surface refuses.""" + + @pytest.mark.asyncio + async def test_initialize_and_list_expose_only_discovery_tools(self, proxy_server_url: str) -> None: + async with asyncio.timeout(20): + async with _proxy_session(proxy_server_url) as (read, write, _sid): + async with ClientSession(read, write) as session: + init = await session.initialize() + assert init.capabilities.tools is not None + assert init.capabilities.prompts is None + assert init.capabilities.resources is None + + listed = await session.list_tools() + assert {tool.name for tool in listed.tools} == PROXY_MODE_TOOLS + + @pytest.mark.asyncio + async def test_search_schema_and_call_round_trip_keeps_server_identity(self, proxy_server_url: str) -> None: + async with asyncio.timeout(30): + async with _proxy_session(proxy_server_url) as (read, write, _sid): + async with ClientSession(read, write) as session: + await session.initialize() + + hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"})) + by_name = {hit["name"]: hit for hit in hits} + assert {"math_stdio-add", "math_streamable_http-add"} <= set(by_name) + assert all("inputSchema" not in hit for hit in hits) + assert by_name["math_stdio-add"]["tool_id"] != by_name["math_streamable_http-add"]["tool_id"] + + schema = _payload( + await session.call_tool( + "get_tool_schema", arguments={"tool_id": by_name["math_stdio-add"]["tool_id"]} + ) + ) + assert schema["name"] == "math_stdio-add" + assert set(schema["inputSchema"]["required"]) == {"a", "b"} + assert schema["outputSchema"]["properties"]["result"]["type"] == "integer" + + stdio = await session.call_tool( + "call_tool", + arguments={"tool_id": by_name["math_stdio-add"]["tool_id"], "arguments": {"a": 3, "b": 4}}, + ) + http = await session.call_tool( + "call_tool", + arguments={ + "tool_id": by_name["math_streamable_http-add"]["tool_id"], + "arguments": {"a": 5, "b": 6}, + }, + ) + assert stdio.isError is False and stdio.content[0].text == "7" + assert http.isError is False and http.content[0].text == "111" + + @pytest.mark.asyncio + async def test_server_scope_header_narrows_discovery(self, proxy_server_url: str) -> None: + async with asyncio.timeout(20): + async with _proxy_session(proxy_server_url, **{"x-mcp-servers": "math_streamable_http"}) as ( + read, + write, + _sid, + ): + async with ClientSession(read, write) as session: + await session.initialize() + hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"})) + assert {hit["name"] for hit in hits} == {"math_streamable_http-add"} + + @pytest.mark.asyncio + async def test_rejections_never_reach_upstream(self, proxy_server_url: str) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import METHOD_NOT_FOUND + + async with asyncio.timeout(30): + async with _proxy_session(proxy_server_url) as (read, write, _sid): + async with ClientSession(read, write) as session: + await session.initialize() + hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"})) + tool_id = next(hit["tool_id"] for hit in hits if hit["name"] == "math_stdio-add") + + bad_args = await session.call_tool( + "call_tool", arguments={"tool_id": tool_id, "arguments": {"a": "three", "b": 4}} + ) + assert bad_args.isError is True and "Invalid arguments" in bad_args.content[0].text + + stale = await session.call_tool("get_tool_schema", arguments={"tool_id": "0" * 32}) + assert stale.isError is True and "unauthorized tool_id" in stale.content[0].text + + for not_an_object in ("wrong", False): + refused_args = await session.call_tool( + "call_tool", arguments={"tool_id": tool_id, "arguments": not_an_object} + ) + assert refused_args.isError is True and "object" in refused_args.content[0].text + + direct = await session.call_tool("math_stdio-add", arguments={"a": 1, "b": 2}) + assert direct.isError is True and "unavailable on /mcp/proxy" in direct.content[0].text + + for operation in (session.list_prompts, session.list_resources): + with pytest.raises(McpError) as refused: + await operation() + assert refused.value.error.code == METHOD_NOT_FOUND + + +async def authorize_proxy_key(request: Request, api_key: str) -> UserAPIKeyAuth: + permissions = { + "sk-1234": LiteLLM_ObjectPermissionTable(object_permission_id="open", mcp_servers=["math_stdio"]), + "sk-restricted": LiteLLM_ObjectPermissionTable( + object_permission_id="restricted", mcp_servers=["math_restricted"] + ), + "sk-none": LiteLLM_ObjectPermissionTable(object_permission_id="none", mcp_servers=["no-mcp-servers"]), + "sk-add-only": LiteLLM_ObjectPermissionTable( + object_permission_id="add-only", mcp_servers=["math_stdio"], mcp_tool_permissions={"math_stdio": ["add"]} + ), + } + permission = permissions.get(api_key) + if permission is None: + raise ProxyException(message="Unknown test key", type="authentication_error", param=None, code=401) + return UserAPIKeyAuth(api_key=api_key, user_id=api_key, object_permission=permission) + + +class ProxyCallRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: queue.Queue[str] = queue.Queue() + self.failures: queue.Queue[str] = queue.Queue() + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + payload = kwargs.get("standard_logging_object") + if isinstance(payload, dict) and payload.get("call_type") == "call_mcp_tool": + self.events.put(json.dumps(payload, default=str)) + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + payload = kwargs.get("standard_logging_object") + if isinstance(payload, dict) and payload.get("call_type") == "call_mcp_tool": + self.failures.put(json.dumps(payload, default=str)) + + +proxy_call_recorder = ProxyCallRecorder() + + +@asynccontextmanager +async def _scoped_session(url: str, key: str = "sk-1234", **headers: str) -> typing.AsyncIterator[ClientSession]: + async with asyncio.timeout(30): + async with _proxy_session(url, Authorization=f"Bearer {key}", **headers) as (read, write, _sid): + async with ClientSession(read, write) as session: + await session.initialize() + yield session + + +async def _search(session: ClientSession, query: str) -> dict[str, str]: + result = await session.call_tool("search_tools", arguments={"query": query}) + assert result.isError is False, result + return {hit["name"]: hit["tool_id"] for hit in _payload(result)} + + +async def _call(session: ClientSession, tool_id: str, a: int = 3, b: int = 4) -> CallToolResult: + return await session.call_tool("call_tool", arguments={"tool_id": tool_id, "arguments": {"a": a, "b": b}}) + + +async def _raw_rpc( + proxy_server_url: str, key: str | None, method: str, params: dict[str, object], **headers: str +) -> httpx.Response: + async with httpx.AsyncClient() as client: + return await client.post( + f"{proxy_server_url}/mcp/proxy", + headers={ + "Accept": "application/json, text/event-stream", + **({"Authorization": f"Bearer {key}"} if key else {}), + **headers, + }, + json={"jsonrpc": "2.0", "id": 1, "method": method, "params": params}, + ) + + +async def _raw_initialize(proxy_server_url: str, key: str | None) -> httpx.Response: + return await _raw_rpc( + proxy_server_url, + key, + "initialize", + {"protocolVersion": "2025-03-26", "capabilities": {}, "clientInfo": {"name": "auth-test", "version": "1"}}, + ) + + +def _rpc_result(response: httpx.Response) -> dict[str, typing.Any]: + if response.headers["content-type"].startswith("text/event-stream"): + data_line = next(line for line in response.text.splitlines() if line.startswith("data:")) + return json.loads(data_line.removeprefix("data:"))["result"] + return response.json()["result"] + + +def _assert_unauthorized(result: CallToolResult) -> None: + assert result.isError is True + assert result.content[0].text == "Unknown or unauthorized tool_id" + + +class TestProxyMcpAuthorizationScope: + @pytest.mark.asyncio + async def test_server_grant_bounds_search_and_blocks_foreign_ids(self, proxy_server_url: str) -> None: + async with _scoped_session(proxy_server_url, "sk-restricted") as granted: + restricted_id = (await _search(granted, "add"))["math_restricted-add"] + assert (await _call(granted, restricted_id)).content[0].text == "207" + async with _scoped_session(proxy_server_url) as ungranted: + assert set(await _search(ungranted, "add")) == {"math_stdio-add", "math_streamable_http-add"} + _assert_unauthorized(await ungranted.call_tool("get_tool_schema", {"tool_id": restricted_id})) + _assert_unauthorized(await _call(ungranted, restricted_id)) + + @pytest.mark.asyncio + async def test_no_mcp_servers_sentinel_rejects_initialize_and_hides_every_tool(self, proxy_server_url: str) -> None: + async with _scoped_session(proxy_server_url) as granted: + tool_id = (await _search(granted, "add"))["math_stdio-add"] + response = await _raw_initialize(proxy_server_url, "sk-none") + assert response.status_code == 403, response.text + assert "no MCP servers granted" in response.json()["detail"]["error"] + + async def raw_call(name: str, arguments: dict[str, object]) -> dict[str, typing.Any]: + call = await _raw_rpc(proxy_server_url, "sk-none", "tools/call", {"name": name, "arguments": arguments}) + assert call.status_code == 200, call.text + return _rpc_result(call) + + listed = await _raw_rpc(proxy_server_url, "sk-none", "tools/list", {}) + assert listed.status_code == 200, listed.text + assert {tool["name"] for tool in _rpc_result(listed)["tools"]} == {"search_tools", "get_tool_schema", "call_tool"} + search = await raw_call("search_tools", {"query": "add"}) + assert search["isError"] is False, search + assert json.loads(search["content"][0]["text"]) == [] + for name, arguments in ( + ("get_tool_schema", {"tool_id": tool_id}), + ("call_tool", {"tool_id": tool_id, "arguments": {"a": 3, "b": 4}}), + ): + denied = await raw_call(name, arguments) + assert denied["isError"] is True, denied + assert denied["content"][0]["text"] == "Unknown or unauthorized tool_id" + + @pytest.mark.asyncio + async def test_tool_grant_hides_ungranted_tools_and_blocks_their_ids(self, proxy_server_url: str) -> None: + async with _scoped_session(proxy_server_url) as granted: + multiply_id = (await _search(granted, "multiply"))["math_stdio-multiply"] + async with _scoped_session(proxy_server_url, "sk-add-only", **{"x-mcp-servers": "math_stdio"}) as session: + ids = await _search(session, "add multiply request_headers") + assert set(ids) == {"math_stdio-add"} + assert (await _call(session, ids["math_stdio-add"])).content[0].text == "7" + _assert_unauthorized(await session.call_tool("get_tool_schema", {"tool_id": multiply_id})) + _assert_unauthorized(await _call(session, multiply_id)) + + @pytest.mark.asyncio + async def test_same_named_tools_keep_distinct_ids_and_reach_their_own_upstream(self, proxy_server_url: str) -> None: + async with _scoped_session(proxy_server_url, "sk-restricted") as session: + ids = await _search(session, "add") + assert set(ids) == {"math_stdio-add", "math_streamable_http-add", "math_restricted-add"} + assert len(set(ids.values())) == 3 + assert all(len(tool_id) == 32 for tool_id in ids.values()) + for name, expected in ( + ("math_stdio-add", "7"), + ("math_streamable_http-add", "107"), + ("math_restricted-add", "207"), + ): + schema = _payload(await session.call_tool("get_tool_schema", {"tool_id": ids[name]})) + assert schema["name"] == name + assert schema["tool_id"] == ids[name] + result = await _call(session, ids[name]) + assert result.isError is False + assert result.content[0].text == expected + + @pytest.mark.asyncio + async def test_server_scope_header_narrows_grants_and_blocks_out_of_scope_ids(self, proxy_server_url: str) -> None: + async with _scoped_session(proxy_server_url, "sk-restricted") as unscoped: + other_id = (await _search(unscoped, "add"))["math_stdio-add"] + async with _scoped_session( + proxy_server_url, "sk-restricted", **{"x-mcp-servers": "math_restricted"} + ) as session: + ids = await _search(session, "add") + assert set(ids) == {"math_restricted-add"} + _assert_unauthorized(await session.call_tool("get_tool_schema", {"tool_id": other_id})) + _assert_unauthorized(await _call(session, other_id)) + assert (await _call(session, ids["math_restricted-add"])).content[0].text == "207" + + @pytest.mark.asyncio + @pytest.mark.parametrize("key", [None, "sk-invalid"]) + async def test_missing_or_invalid_key_cannot_initialize(self, proxy_server_url: str, key: str | None) -> None: + response = await _raw_initialize(proxy_server_url, key) + assert response.status_code == 401, response.text + + @pytest.mark.asyncio + async def test_server_headers_are_forwarded_only_to_the_named_upstream(self, proxy_server_url: str) -> None: + for tag in ("first-request", "second-request"): + async with _scoped_session( + proxy_server_url, + "sk-restricted", + **{ + "x-mcp-math_restricted-authorization": f"Bearer {tag}", + "x-mcp-math_restricted-x-request-tag": tag, + }, + ) as session: + ids = await _search(session, "request_headers") + for name, expected in ( + ("math_restricted", {"authorization": f"Bearer {tag}", "x-request-tag": tag}), + ("math_streamable_http", {"authorization": "", "x-request-tag": ""}), + ): + result = await session.call_tool( + "call_tool", {"tool_id": ids[f"{name}-request_headers"], "arguments": {}} + ) + assert result.isError is False + assert _payload(result) == expected + + @pytest.mark.asyncio + async def test_proxy_call_emits_spend_log(self, proxy_server_url: str) -> None: + async with _scoped_session(proxy_server_url, "sk-restricted") as session: + tool_id = (await _search(session, "add"))["math_restricted-add"] + result = await _call(session, tool_id, 123, 456) + assert result.isError is False and result.content[0].text == "779" + async with asyncio.timeout(10): + while True: + payload = json.loads(await asyncio.to_thread(proxy_call_recorder.events.get, True, 5)) + if payload.get("metadata", {}).get("mcp_tool_call_metadata", {}).get("arguments") == { + "a": 123, + "b": 456, + }: + break + assert payload["call_type"] == "call_mcp_tool" + assert payload["response_cost"] == 0.25 + assert payload["status"] == "success" + assert payload["metadata"]["mcp_tool_call_metadata"]["mcp_server_name"] == "math_restricted" + assert payload["metadata"]["mcp_tool_call_metadata"]["name"] == "add" + assert payload["metadata"]["mcp_tool_call_metadata"]["namespaced_tool_name"] == "math_restricted/add" + + @pytest.mark.asyncio + async def test_proxy_scope_exception_returns_iserror_and_emits_failure_log(self, proxy_server_url: str) -> None: + response = await _raw_rpc( + proxy_server_url, + "sk-none", + "tools/call", + {"name": "call_tool", "arguments": {"tool_id": "denied-scope", "arguments": {}}}, + **{"x-mcp-servers": "math_restricted", "x-litellm-call-id": "proxy-scope-denial"}, + ) + assert response.status_code == 200, response.text + result = _rpc_result(response) + assert result["isError"] is True + assert result["content"][0]["text"] == ( + "Error: The key is not allowed to access the requested MCP servers: math_restricted" + ) + async with asyncio.timeout(10): + while True: + payload = json.loads(await asyncio.to_thread(proxy_call_recorder.failures.get, True, 5)) + if payload["id"] == "proxy-scope-denial": + break + assert payload["call_type"] == "call_mcp_tool" + assert payload["status"] == "failure" + assert payload["response_cost"] == 0 + assert "math_restricted" in payload["error_str"] + + @pytest.mark.parametrize("arguments", ["wrong", False, None, [], 0]) + def test_handler_rejects_non_object_arguments( + self, proxy_server_url: str, _proxy_server: ProxyRig, arguments: object + ) -> None: + async def check() -> None: + auth = UserAPIKeyAuth( + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="validation", mcp_servers=["math_stdio"] + ) + ) + hits = _payload(await handle_mcp_proxy_tool("search_tools", {"query": "add"}, auth)) + tool_id = next(hit["tool_id"] for hit in hits if hit["name"] == "math_stdio-add") + result = await handle_mcp_proxy_tool("call_tool", {"tool_id": tool_id, "arguments": arguments}, auth) + assert result.isError is True + assert result.content[0].text == "arguments must be an object" + + asyncio.run_coroutine_threadsafe(check(), _proxy_server.loop).result(timeout=30) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index 7e338dafb86..4a392042d63 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -127,14 +127,14 @@ def test_bad_request_bad_param_error(): ) -def test_anthropic_with_responses_api(): - client = get_test_client() - response = client.responses.create( - model="anthropic/claude-sonnet-4-5-20250929", +def test_anthropic_with_responses_api() -> None: + client: Final = get_test_client() + response: Final = client.responses.create( + model="anthropic/claude-sonnet-5", input="just respond with the word 'ping'", - previous_response_id="hi", ) - print("anthropic response=", response) + assert response.status == "completed" + assert response.output_text.strip() def test_cancel_response(): diff --git a/tests/pass_through_unit_tests/test_claude_code_marketplace.py b/tests/pass_through_unit_tests/test_claude_code_marketplace.py index 2ca81f1d5d3..bedb8830559 100644 --- a/tests/pass_through_unit_tests/test_claude_code_marketplace.py +++ b/tests/pass_through_unit_tests/test_claude_code_marketplace.py @@ -216,7 +216,7 @@ async def test_get_marketplace(mock_prisma_client): ) # Now get the marketplace - response = await get_marketplace() + response = await get_marketplace(request=MagicMock()) # Response is a JSONResponse, get the body body = json.loads(response.body.decode()) diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index 7bc61c40ea0..e8caa241a53 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -40,20 +40,38 @@ async def _turn( tokens: int = 100, spend: float = 0.01, saved: float = 0.02, + classifier_cost: float = 0.0, tier: "str | None" = None, + baseline: "str | None" = None, ) -> None: touched: Final = 1 if (hit or ttl is not None or not covered) else 0 await db.execute_raw( UPSERT_AUTOROUTER_SESSION_SQL, - key, session_id, router, router_type, model, at.isoformat(), tokens, spend, saved, covered, hit, ttl, touched, + key, + session_id, + router, + router_type, + model, + at.isoformat(), + tokens, + spend, + saved, + classifier_cost, + covered, + hit, + ttl, + touched, tier, + baseline, ) async def _row(db, key: str, session_id: str = "s1", router: str = "auto-1") -> dict: rows = await db.query_raw( 'SELECT * FROM "LiteLLM_AutoRouterSession" WHERE api_key = $1 AND session_id = $2 AND router_name = $3', - key, session_id, router, + key, + session_id, + router, ) assert len(rows) == 1 return rows[0] @@ -143,9 +161,9 @@ async def test_out_of_order_turns_do_not_rewind_the_session(db): async def test_concurrent_writers_compose_without_losing_turns(db): key = f"k-{uuid.uuid4()}" - await _turn(db, key, "A", T0) + await _turn(db, key, "A", T0, classifier_cost=0.001) await asyncio.gather( - *(_turn(db, key, "A", T0 + timedelta(seconds=1 + offset), hit=1) for offset in range(30)) + *(_turn(db, key, "A", T0 + timedelta(seconds=1 + offset), hit=1, classifier_cost=0.002) for offset in range(30)) ) row = await _row(db, key) assert row["turns"] == 31 @@ -154,6 +172,51 @@ async def test_concurrent_writers_compose_without_losing_turns(db): == row["turns"] ) assert row["spend"] == pytest.approx(0.31) + assert row["saved_spend"] == pytest.approx(0.62) + assert row["classifier_cost"] == pytest.approx(0.061) + assert row["classifier_cost_recorded_turns"] == 31 + + +async def _legacy_turn(db, key: str, at: datetime, session_id: str = "s1", router: str = "auto-1") -> None: + await db.execute_raw( + """INSERT INTO "LiteLLM_AutoRouterSession" AS t ( + api_key, session_id, router_name, router_type, first_turn_at, last_turn_at, last_model, turns, spend, saved_spend + ) VALUES ($1, $2, $3, 'complexity', $4::timestamp, $4::timestamp, 'A', 1, 0.01, 0.02) + ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET + turns = t.turns + 1, spend = t.spend + EXCLUDED.spend, saved_spend = t.saved_spend + EXCLUDED.saved_spend, + last_turn_at = EXCLUDED.last_turn_at""", + key, + session_id, + router, + at.isoformat(), + ) + + +@pytest.mark.parametrize("writers", [(False,), (True,), (False, True), (True, False)]) +async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers: tuple[bool, ...]): + key: Final = f"k-{uuid.uuid4()}" + for offset, records_cost in enumerate(writers): + at: Final = T0 + timedelta(seconds=offset) + if records_cost: + await _turn(db, key, "A", at, classifier_cost=0.004) + else: + await _legacy_turn(db, key, at) + + row: Final = await _row(db, key) + assert row["turns"] == len(writers) + assert row["spend"] == pytest.approx(0.01 * len(writers)) + assert row["saved_spend"] == pytest.approx(0.02 * len(writers)) + assert row["classifier_cost"] == pytest.approx(0.004 * sum(writers)) + assert row["classifier_cost_recorded_turns"] == sum(writers) + groups: Final = await db.query_raw( + AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key + ) + assert len(groups) == 1 + assert groups[0]["classifier_cost"] == row["classifier_cost"] + assert groups[0]["classifier_cost_recorded_turns"] == sum(writers) + assert groups[0]["turns"] == len(writers) + assert groups[0]["spend"] == row["spend"] + assert groups[0]["saved_spend"] == row["saved_spend"] async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): @@ -161,14 +224,25 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): router = f"r-{uuid.uuid4()}" in_window = f"s-{uuid.uuid4()}" out_of_window = f"s-{uuid.uuid4()}" - await _turn(db, key, "A", T0, session_id=in_window, router=router, saved=0.5, spend=0.25) - await _turn(db, key, "B", T0 + timedelta(seconds=60), session_id=in_window, router=router, saved=0.5, spend=0.25) - await _turn(db, key, "A", T0 - timedelta(days=40), session_id=out_of_window, router=router) + await _turn( + db, + key, + "B", + T0 + timedelta(seconds=60), + session_id=in_window, + router=router, + saved=0.5, + spend=0.25, + classifier_cost=0.01, + ) + await _turn(db, key, "A", T0, session_id=in_window, router=router, saved=0.5, spend=0.25, classifier_cost=0.02) + await _turn(db, key, "A", T0 - timedelta(days=40), session_id=out_of_window, router=router, classifier_cost=9.0) rows = await db.query_raw( AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) matching = [row for row in rows if row["router_name"] == router] assert len(matching) == 1 @@ -178,19 +252,54 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): assert grouped["turns"] == 2 assert grouped["spend"] == pytest.approx(0.5) assert grouped["saved_spend"] == pytest.approx(1.0) + assert grouped["classifier_cost"] == pytest.approx(0.03) + assert grouped["classifier_cost_recorded_turns"] == 2 + assert grouped["unordered_turns"] == 1 assert grouped["session_seconds"] == pytest.approx(60.0) +async def test_the_benchmarks_aggregate_can_filter_to_one_key(db): + router = f"r-{uuid.uuid4()}" + first_key = f"k-{uuid.uuid4()}" + second_key = f"k-{uuid.uuid4()}" + await _turn(db, first_key, "A", T0, router=router, saved=0.5, classifier_cost=0.01) + await _turn(db, second_key, "A", T0, router=router, saved=9.0, classifier_cost=0.09) + + rows = await db.query_raw( + AUTOROUTER_BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + first_key, + ) + matching = [row for row in rows if row["router_name"] == router] + assert len(matching) == 1 + assert matching[0]["sessions"] == 1 + assert matching[0]["saved_spend"] == pytest.approx(0.5) + assert matching[0]["classifier_cost"] == pytest.approx(0.01) + assert matching[0]["classifier_cost_recorded_turns"] == 1 + + unknown_key_rows = await db.query_raw( + AUTOROUTER_BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + f"k-{uuid.uuid4()}", + ) + assert [row for row in unknown_key_rows if row["router_name"] == router] == [] + + async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db): key = f"k-{uuid.uuid4()}" router = f"r-{uuid.uuid4()}" await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity") - await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality") + await _turn( + db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality" + ) rows = await db.query_raw( AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) matching = sorted( (row for row in rows if row["router_name"] == router), @@ -231,6 +340,27 @@ async def test_a_mid_session_router_type_change_keeps_foreign_tier_names_out_of_ assert row["turns"] == 3 +async def test_baseline_models_count_the_turns_priced_against_each_baseline(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, baseline="opus") + await _turn(db, key, "B", T0 + timedelta(seconds=10), baseline="opus") + await _turn(db, key, "A", T0 + timedelta(seconds=20), baseline="sonnet") + + assert (await _row(db, key))["baseline_models"] == {"opus": 2, "sonnet": 1} + + +async def test_a_turn_priced_against_no_baseline_leaves_the_map_alone(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, baseline=None) + assert (await _row(db, key))["baseline_models"] == {} + + await _turn(db, key, "A", T0 + timedelta(seconds=10), baseline="opus") + await _turn(db, key, "A", T0 + timedelta(seconds=20), baseline=None) + row = await _row(db, key) + assert row["baseline_models"] == {"opus": 1} + assert row["turns"] == 3 + + async def test_an_out_of_order_turn_still_counts_toward_its_tier(db): key = f"k-{uuid.uuid4()}" await _turn(db, key, "A", T0 + timedelta(seconds=60), tier="simple") @@ -253,6 +383,7 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db): AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) grouped = next(row for row in rows if row["router_name"] == router) assert grouped["tier_turns"] == {"simple": 2, "complex": 1} @@ -280,6 +411,7 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router} assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}} @@ -294,6 +426,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db): AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) grouped = next(row for row in rows if row["router_name"] == router) assert grouped["tier_turns"] == {} diff --git a/tests/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py new file mode 100644 index 00000000000..741fa7386df --- /dev/null +++ b/tests/proxy_migration_tests/test_invalid_index_repair.py @@ -0,0 +1,264 @@ +import os +import threading +import uuid +from collections.abc import Iterator, Mapping +from types import MappingProxyType +from typing import Final + +import pytest +from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager + +psycopg = pytest.importorskip("psycopg") + +pytestmark = pytest.mark.timeout(120) + +requires_db: Final = pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) + +HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable" +HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx" +HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC' +LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable" +LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx" +PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable" +PARTITIONED_INDEX: Final = "LiteLLM_PartitionedTable_id_idx" + + +def _base_url() -> str: + return os.environ["DATABASE_URL"].split("?")[0] + + +def _index_validity(schema: str) -> Mapping[str, bool]: + with psycopg.connect(_base_url(), autocommit=True) as conn: + rows = conn.execute( + "SELECT c.relname, i.indisvalid FROM pg_index i " + "JOIN pg_class c ON c.oid = i.indexrelid " + "JOIN pg_namespace n ON n.oid = c.relnamespace " + "WHERE n.nspname = %s", + (schema,), + ).fetchall() + return MappingProxyType(dict(rows)) + + +def _interrupt_concurrent_build(schema: str, table: str, statement: str) -> None: + """Abort a CONCURRENTLY build while it waits on an older snapshot, the same + spot the deadlock loser dies at, so it leaves its index INVALID.""" + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"') + with psycopg.connect(_base_url(), autocommit=True) as builder: + builder.execute("SET statement_timeout = '1s'") + with pytest.raises(psycopg.errors.QueryCanceled): + builder.execute(statement) + + +def _leave_invalid_index(schema: str, table: str, index: str, columns: str) -> None: + _interrupt_concurrent_build( + schema, table, f'CREATE INDEX CONCURRENTLY "{index}" ON "{schema}"."{table}" ({columns})' + ) + + +def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None: + _interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"') + + +@pytest.fixture +def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + schema: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE SCHEMA "{schema}"') + conn.execute( + f'CREATE TABLE "{schema}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)' + ) + conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)') + + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") + yield schema + + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP SCHEMA "{schema}" CASCADE') + + +@pytest.fixture +def fresh_database(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + """A brand-new database, what a first deploy sees. A scratch schema would + not do: the migrations guard on pg_constraint by name across every schema, + so a LiteLLM schema already pushed into public makes them skip and then + fail, which is exactly what CI's database looks like.""" + admin_url: Final = _base_url() + name: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" + with psycopg.connect(admin_url, autocommit=True) as conn: + conn.execute(f'CREATE DATABASE "{name}"') + + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", f"{admin_url.rsplit('/', 1)[0]}/{name}") + yield "public" + + with psycopg.connect(admin_url, autocommit=True) as conn: + conn.execute(f'DROP DATABASE "{name}" WITH (FORCE)') + + +@requires_db +def test_repair_rebuilds_invalid_litellm_indexes_and_leaves_lookalike_tables_alone(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + _leave_invalid_index(scratch_schema, LOOKALIKE_TABLE, LOOKALIKE_INDEX, "id") + assert _index_validity(scratch_schema) == {HEALTH_INDEX: False, LOOKALIKE_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, LOOKALIKE_INDEX: False} + + +@requires_db +def test_repair_drops_leftovers_of_interrupted_rebuilds(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + _leave_invalid_reindex_leftover(scratch_schema, HEALTH_TABLE, HEALTH_INDEX) + _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccold", '"model_id"') + _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccnew1", '"model_id"') + before: Final = _index_validity(scratch_schema) + assert len(before) == 4 + assert set(before.values()) == {False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_is_a_no_op_when_every_index_is_valid(scratch_schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE INDEX "{HEALTH_INDEX}" ON "{scratch_schema}"."{HEALTH_TABLE}" ({HEALTH_INDEX_COLUMNS})') + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_leaves_partitioned_parent_indexes_alone(scratch_schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}" (id INT) PARTITION BY RANGE (id)') + conn.execute( + f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}_p0" ' + f'PARTITION OF "{scratch_schema}"."{PARTITIONED_TABLE}" FOR VALUES FROM (0) TO (10)' + ) + conn.execute(f'CREATE INDEX "{PARTITIONED_INDEX}" ON ONLY "{scratch_schema}"."{PARTITIONED_TABLE}" (id)') + assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False} + + +@requires_db +def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + with psycopg.connect(_base_url(), autocommit=True) as other_replica: + other_replica.execute("SELECT pg_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)) + assert ProxyExtrasDBManager.repair_invalid_indexes() is False + assert _index_validity(scratch_schema) == {HEALTH_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{scratch_schema}"."{HEALTH_TABLE}"') + assert ProxyExtrasDBManager.repair_invalid_indexes(lock_timeout="1s") is False + blocked: Final = _index_validity(scratch_schema) + assert blocked[HEALTH_INDEX] is False + assert [name for name in blocked if name.endswith("_ccnew")] + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +def _hold_snapshot(schema: str, table: str, pinned: threading.Event, seconds: float) -> None: + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"') + pinned.set() + pin.execute("SELECT pg_sleep(%s)", (seconds,)) + + +@requires_db +def test_repair_outlives_a_statement_timeout_passed_through_database_url_options( + scratch_schema: str, monkeypatch: pytest.MonkeyPatch +) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={scratch_schema}&options=-c%20statement_timeout%3D2000") + pinned: Final = threading.Event() + holder: Final = threading.Thread(target=_hold_snapshot, args=(scratch_schema, HEALTH_TABLE, pinned, 5.0)) + holder.start() + pinned.wait() + try: + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + finally: + holder.join() + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) -> None: + table: Final = f"LiteLLM_ScratchTable_{uuid.uuid4().hex[:8]}" + index: Final = f"{table}_id_idx" + monkeypatch.setenv("DATABASE_URL", _base_url()) + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE TABLE public."{table}" (id TEXT)') + try: + _leave_invalid_index("public", table, index, "id") + assert _index_validity("public")[index] is False + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity("public")[index] is True + finally: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP TABLE public."{table}"') + + +def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") + + assert ProxyExtrasDBManager.repair_invalid_indexes() is False + + +@requires_db +def test_repair_connects_over_direct_url_but_looks_in_the_schema_database_url_names(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + with pytest.MonkeyPatch.context() as env: + env.setenv("DIRECT_URL", f"{_base_url()}?schema=public") + env.setenv("DATABASE_URL", f"postgresql://u:p@127.0.0.1:9/x?schema={scratch_schema}") + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +def _invalidate_deployed_index(schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP INDEX "{schema}"."{HEALTH_INDEX}"') + _leave_invalid_index(schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + +@requires_db +@pytest.mark.timeout(300) +@pytest.mark.parametrize("use_v2_resolver", [True, False]) +def test_setup_database_repairs_the_index_after_a_recovered_deploy(fresh_database: str, use_v2_resolver: bool) -> None: + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + _invalidate_deployed_index(fresh_database) + assert _index_validity(fresh_database)[HEALTH_INDEX] is False + + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + + assert _index_validity(fresh_database)[HEALTH_INDEX] is True diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 9a6ab08e9b6..6417c7c8aa6 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -583,6 +583,90 @@ class TestCheckBatchCost: assert passed_model_info["input_cost_per_token_batches"] == 2e-06 assert passed_model_info["output_cost_per_token_batches"] == 4e-06 + @pytest.mark.asyncio + async def test_poller_masks_api_base_credentials_before_logging( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """Request rows mask `key=` query credentials out of api_base before it is + logged, but the poller skips that pre-call step, so an unmasked deployment + api_base would land verbatim on the batch cost row: regression test for the + poller masking the same way. + """ + import base64 + from unittest.mock import patch + + import httpx + import respx + + from litellm.litellm_core_utils.litellm_logging import Logging + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + mock_job = MagicMock() + mock_job.id = "job-masked-api-base-1" + mock_job.unified_object_id = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() + mock_job.created_by = "user-1" + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "file-output-123" + mock_response.error_file_id = None + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "openai" + mock_deployment.litellm_params.model = "gpt-5.4-mini" + mock_deployment.litellm_params.api_base = "https://gateway.example.com/v1?key=AIzaSyVERYSECRET7890" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + output_line = json.dumps( + { + "custom_id": "req-1", + "response": { + "status_code": 200, + "body": { + "id": "chatcmpl-1", + "object": "chat.completion", + "model": "gpt-5.4-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, + "error": None, + } + ) + + with ( + respx.mock(assert_all_called=True) as provider, + patch.object( # test-quality-ok: the poller builds Logging inline, the only seam to the row it logs + Logging, "async_success_handler", autospec=True + ) as success_handler, + ): + provider.get("https://api.openai.com/v1/files/file-output-123/content").mock( + return_value=httpx.Response(200, content=f"{output_line}\n".encode()) + ) + await check_batch_cost_instance.check_batch_cost() + + cost_row_calls = [call for call in success_handler.await_args_list if "batch_cost" in call.kwargs] + assert len(cost_row_calls) == 1 + logged_api_base = cost_row_calls[0].args[0].litellm_params["api_base"] + assert logged_api_base == "https://gateway.example.com/v1?key=*****7890" + assert "VERYSECRET" not in logged_api_base + @pytest.mark.asyncio async def test_primary_path_completion_update_includes_batch_processed( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router @@ -1760,6 +1844,35 @@ class TestUnmanagedVertexRouting: ) router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") + def test_flag_on_routes_fine_tuned_endpoint_to_vertex_deployment(self): + """A fine-tuned Gemini batch stores `endpoints/` in the gs:// path; the bare model + (the endpoint id) must round-trip to the deployment configured as + `vertex_ai/gemini/` (LIT-6899).""" + endpoint_id = "7768560373388541952" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = None + router.get_model_list.return_value = [ + { + "model_name": "gemini-2.5-flash-dts-usc1", + "litellm_params": { + "model": f"vertex_ai/gemini/{endpoint_id}", + "custom_llm_provider": "vertex_ai", + }, + "model_info": {"id": "deploy-ft"}, + }, + ] + instance = self._instance(track_unmanaged=True, router=router) + job = self._job( + file_object=_unmanaged_vertex_file_object( + input_file_id=f"gs://bucket/litellm-vertex-files/endpoints/{endpoint_id}/abc.jsonl" + ) + ) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-ft", "8823717160934178816") + def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): """Flag on, but the only deployment for the model group is a non-vertex_ai provider: must not be selected, even though the model group name matches.""" @@ -2554,6 +2667,85 @@ class TestBatchCostAttribution: assert metadata["user_api_key_alias"] == "prod-key" + @pytest.mark.asyncio + async def test_org_id_snapshotted_on_the_row_wins(self): + """The org_id column captures the creating key's organization at submission time, + like team_id, so a key later moved to another org still bills the original one.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", organization_id="org-moved-to"), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + + metadata = await instance._build_creator_attribution_metadata( + self._job(org_id="org-at-creation"), "batch-1" + ) + + assert metadata["user_api_key_org_id"] == "org-at-creation" + + @pytest.mark.asyncio + async def test_org_id_comes_from_the_creating_key(self): + """The spend update writer increments organization spend from user_api_key_org_id. + A legacy row without the org_id column falls back to the creating key's org.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", organization_id="org-42"), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_org_id"] == "org-42" + + @pytest.mark.asyncio + async def test_org_id_falls_back_to_the_team_organization(self): + """A key with no org of its own still books batch spend against its team's + organization, matching how the request path resolves org attribution.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", organization_id=None), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_org_id"] == "org-team" + + @pytest.mark.asyncio + async def test_key_lookup_failure_still_bills_the_team_org(self): + """A key-table error while resolving a legacy row's org must not drop the team's + organization: the two lookups fail independently, so org spend still lands.""" + from types import SimpleNamespace + + instance = self._instance( + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=Exception("db down") + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_org_id"] == "org-team" + + @pytest.mark.asyncio + async def test_no_org_leaves_the_key_unset(self): + """Without any org the key is absent entirely, so the spend writer's org update + stays skipped instead of matching an empty-string organization.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", organization_id=None), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert "user_api_key_org_id" not in metadata + @pytest.mark.asyncio async def test_metadata_provenance_keeps_spend_log_api_key_joinable(self): """ diff --git a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py index 89899d3e762..c4b1f4f3afd 100644 --- a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py +++ b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py @@ -23,9 +23,9 @@ from litellm.proxy.proxy_server import token_counter def _fake_hf_tokenizer(num_tokens: int) -> MagicMock: encoding = MagicMock() - encoding.ids = list(range(num_tokens)) + encoding.__len__.return_value = num_tokens tokenizer = MagicMock() - tokenizer.encode.return_value = encoding + tokenizer.encode_batch_fast.return_value = [encoding] return tokenizer @@ -68,13 +68,11 @@ async def test_custom_tokenizer_from_model_info_is_used(monkeypatch): ) ) - mock_tokenizer_cls.from_pretrained.assert_called_once_with( - "my-org/custom-tokenizer", revision="v2", auth_token=None - ) + mock_tokenizer_cls.from_pretrained.assert_called_once_with("my-org/custom-tokenizer", revision="v2", token=None) assert response.tokenizer_type == "huggingface_tokenizer" assert response.request_model == "my-embedding-model" assert response.model_used == "self-hosted-embedder" - assert response.total_tokens > 0 + assert response.total_tokens >= 7 @pytest.mark.asyncio diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 54cce9cdd78..d06eb0426c9 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2920,7 +2920,7 @@ async def test_get_config_callbacks_with_all_types(client_no_auth): assert result["status"] == "success" assert "callbacks" in result - callbacks = result["callbacks"] + callbacks = [cb for cb in result["callbacks"] if not cb.get("read_only", False)] # Verify we have all 5 callbacks (2 success + 1 failure + 2 success_and_failure) assert len(callbacks) == 5 diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 772d3622745..467c1332325 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -14,13 +14,14 @@ These tests ensure the polling handler correctly manages response state following the OpenAI Response API format. """ +import asyncio import json from datetime import datetime, timezone from typing import Any, Dict, Optional from unittest.mock import AsyncMock, Mock, patch import pytest - +from fastapi import Request from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler @@ -1414,7 +1415,7 @@ def _make_background_streaming_kwargs( polling_id=polling_id, data={"model": "gpt-4o", "stream": False, "background": True}, polling_handler=polling_handler, - request=Mock(), + request=Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}), fastapi_response=Mock(), user_api_key_dict=Mock(), general_settings={}, @@ -1663,6 +1664,63 @@ class TestBackgroundStreamingTerminalEvents: final_call = handler.update_state.call_args_list[-1] assert final_call.kwargs["status"] == "completed" + @pytest.mark.asyncio + async def test_polling_client_disconnect_does_not_cancel_upstream_call(self): + """The polling client hangs up right after getting its polling id. The detached task + must still stream the upstream response through the client-disconnect guards.""" + from litellm.proxy.common_request_processing import create_response + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + async def client_already_left(): + return {"type": "http.disconnect"} + + async def slow_upstream_stream(): + await asyncio.sleep(0.05) + for event in ( + {"type": "response.in_progress"}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "status": "completed", + "usage": {"input_tokens": 13, "output_tokens": 10}, + "model": "gpt-4o", + "output": [{"id": "item_1", "type": "message"}], + }, + }, + ): + yield f"data: {json.dumps(event)}\n\n" + + async def upstream_call_behind_disconnect_guard(**kwargs): + return await create_response( + slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"] + ) + + handler = AsyncMock(spec=ResponsePollingHandler) + kwargs = _make_background_streaming_kwargs("poll_7", handler) + kwargs["request"] = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"x-litellm-call-id", b"call-123")], + "query_string": b"", + }, + client_already_left, + ) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = upstream_call_behind_disconnect_guard + await background_streaming_task(**kwargs) + + final_call = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "completed" + assert final_call.kwargs["usage"] == {"input_tokens": 13, "output_tokens": 10} + class TestEdgeCases: """Test edge cases and error scenarios""" diff --git a/tests/router_unit_tests/test_router_cooldown_per_deployment.py b/tests/router_unit_tests/test_router_cooldown_per_deployment.py index b8ae8a8c013..964782c348b 100644 --- a/tests/router_unit_tests/test_router_cooldown_per_deployment.py +++ b/tests/router_unit_tests/test_router_cooldown_per_deployment.py @@ -204,8 +204,8 @@ class TestExceptionTypeCountersTrackedIndependently: cache_key_suffix="RateLimitError", ) - rl_counter = router.failed_calls.get_cache(key="primary:RateLimitError") or 0 - generic_counter = router.failed_calls.get_cache(key="primary:generic") or 0 + rl_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0 + generic_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0 assert rl_counter == 3, "RateLimitError counter should be 3" assert generic_counter == 0, "generic counter must be untouched by RateLimitError increments" @@ -218,8 +218,8 @@ class TestExceptionTypeCountersTrackedIndependently: cache_key_suffix="generic", ) - generic_counter_after = router.failed_calls.get_cache(key="primary:generic") or 0 - rl_counter_after = router.failed_calls.get_cache(key="primary:RateLimitError") or 0 + generic_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0 + rl_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0 assert generic_counter_after == 1, "generic counter should now be 1" assert rl_counter_after == 3, "RateLimitError counter must remain unchanged after InternalServerError" @@ -246,12 +246,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired cooldown entry must not appear in active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" + assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" def test_active_entry_is_returned(self): """ @@ -267,7 +267,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -290,14 +290,14 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - (60.0 - remaining), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + cc.in_memory_cache.set_cache(key, value, ttl=600) - before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + before_expiry = cc.in_memory_cache.ttl_dict.get(key) assert before_expiry is not None cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry is not None corrected_remaining = after_expiry - time.time() assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" @@ -318,12 +318,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired entry must not appear in async active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None class TestFallbackDeploymentCooldown: diff --git a/tests/rust-python-harness/shared/parity/compare.py b/tests/rust-python-harness/shared/parity/compare.py index adf85e5c8d7..88239ed042a 100644 --- a/tests/rust-python-harness/shared/parity/compare.py +++ b/tests/rust-python-harness/shared/parity/compare.py @@ -78,3 +78,4 @@ def assert_parity(baseline: Execution, candidate: Execution, baseline_user_agent validate_harness(baseline, candidate, baseline_user_agent) assert_request_parity(baseline.requests, candidate.requests) assert_value_parity(baseline.report, candidate.report) + assert_value_parity(baseline.callbacks, candidate.callbacks, path="$.callbacks") diff --git a/tests/rust-python-harness/shared/parity/models.py b/tests/rust-python-harness/shared/parity/models.py index 898b58d23ee..5612e218824 100644 --- a/tests/rust-python-harness/shared/parity/models.py +++ b/tests/rust-python-harness/shared/parity/models.py @@ -37,6 +37,25 @@ class SDKError(BaseModel): llm_provider: str | None +class CallbackObservation(BaseModel): + model_config = ConfigDict(frozen=True) + + hook: Literal[ + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", + ] + phase: Literal["success", "failure"] + model: str | None + call_type: str | None + litellm_call_id: str | None + metadata: JsonValue + kwargs: JsonValue + payload: JsonValue + error: SDKError | None + + class SDKJsonChunk(BaseModel): model_config = ConfigDict(frozen=True) @@ -119,6 +138,7 @@ class Execution(BaseModel): requests: tuple[CapturedRequest, ...] report: SDKReport + callbacks: tuple[CallbackObservation, ...] | None = None class SDKCommand(BaseModel): @@ -133,6 +153,7 @@ class WorkerSuccess(BaseModel): status: Literal["ok"] = "ok" report: SDKReport + callbacks: tuple[CallbackObservation, ...] | None = None class WorkerFailure(BaseModel): diff --git a/tests/rust-python-harness/shared/parity/runner.py b/tests/rust-python-harness/shared/parity/runner.py index 43a583382cb..feae0201aee 100644 --- a/tests/rust-python-harness/shared/parity/runner.py +++ b/tests/rust-python-harness/shared/parity/runner.py @@ -39,9 +39,7 @@ class SubprocessRunner: return ( sys.executable, "-m", - ".".join( - self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts - ), + ".".join(self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts), "--parity-worker", provider_url, ) @@ -113,7 +111,11 @@ class SubprocessWorker: ) assert isinstance(result, WorkerSuccess) try: - return Execution(requests=self.provider.take_requests(len(responses)), report=result.report) + return Execution( + requests=self.provider.take_requests(len(responses)), + report=result.report, + callbacks=result.callbacks, + ) except AssertionError: self.provider.reset() raise diff --git a/tests/rust-python-harness/shared/parity/test_parity.py b/tests/rust-python-harness/shared/parity/test_parity.py index 83daccdf8ba..7c8355fbfd8 100644 --- a/tests/rust-python-harness/shared/parity/test_parity.py +++ b/tests/rust-python-harness/shared/parity/test_parity.py @@ -7,7 +7,7 @@ import pytest from pydantic import BaseModel, ConfigDict, JsonValue, PrivateAttr from .compare import assert_model_parity, assert_parity -from .models import CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report +from .models import CallbackObservation, CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report SENTINEL: Final = "python-parity-fallback" @@ -69,6 +69,56 @@ def test_parity_rejects_response_difference() -> None: assert_parity(python, rust, SENTINEL) +def test_parity_distinguishes_unobserved_callbacks_from_zero_events() -> None: + python: Final = _execution(user_agent=SENTINEL) + rust: Final = _execution(user_agent="litellm-rust").model_copy(update={"callbacks": ()}) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + +def test_parity_rejects_callback_payload_difference() -> None: + observation: Final = CallbackObservation( + hook="log_success_event", + phase="success", + model="test-model", + call_type="ocr", + litellm_call_id="test-call", + metadata={"profile": "success"}, + kwargs={"model": "test-model"}, + payload={"model": "test-model", "pages": []}, + error=None, + ) + python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)}) + rust: Final = _execution(user_agent="litellm-rust").model_copy( + update={"callbacks": (observation.model_copy(update={"payload": {"model": "changed", "pages": []}}),)} + ) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + +def test_parity_rejects_callback_kwargs_difference() -> None: + observation: Final = CallbackObservation( + hook="log_success_event", + phase="success", + model="test-model", + call_type="ocr", + litellm_call_id="test-call", + metadata={"profile": "success"}, + kwargs={"model": "test-model"}, + payload={"model": "test-model", "pages": []}, + error=None, + ) + python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)}) + rust: Final = _execution(user_agent="litellm-rust").model_copy( + update={"callbacks": (observation.model_copy(update={"kwargs": {"model": "changed"}}),)} + ) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + def test_parity_rejects_error_difference() -> None: python: Final = Execution( requests=(), diff --git a/tests/rust-python-harness/shared/tracing/native.py b/tests/rust-python-harness/shared/tracing/native.py index 4f988f65294..688995cbc4b 100644 --- a/tests/rust-python-harness/shared/tracing/native.py +++ b/tests/rust-python-harness/shared/tracing/native.py @@ -19,7 +19,8 @@ class _TraceEventPayload(BaseModel): class TraceResponsePayload(BaseModel): model_config = ConfigDict(strict=True, extra="forbid") - response: object + response: object = None + error: str | None = None trace: tuple[_TraceEventPayload, ...] | list[_TraceEventPayload] diff --git a/tests/rust-python-harness/strategies/e2e_parity/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/__init__.py index f668e178eef..95346bbd496 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/__init__.py +++ b/tests/rust-python-harness/strategies/e2e_parity/__init__.py @@ -18,8 +18,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity", note=( - "Recorded sync/async SDK parity; invalid-model provider errors differ, " - "and Reducto lacks a Rust contract." + "Recorded sync/async SDK parity with focused success/error callback profiles; " + "Reducto lacks a Rust contract, and known provider parity gaps remain." ), ), surface="sdk", diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py index e72980752f2..5c4abc78081 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py @@ -1,23 +1,31 @@ from __future__ import annotations import asyncio +import datetime +import queue import sys import tempfile +import time import traceback -from collections.abc import Callable, Coroutine, Generator +from collections.abc import Callable, Coroutine, Generator, Mapping from contextlib import contextmanager from enum import Enum from functools import partial from pathlib import Path from typing import Annotated, Final, Literal, cast +from urllib.parse import urlsplit, urlunsplit from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse from .....shared.parity.compare import assert_parity from .....shared.parity.fixtures.store import fixture_id, recorded_fixtures from .....shared.parity.models import ( + JSON_VALUE_ADAPTER, + CallbackObservation, + Execution, SDKCommand, SDKError, SDKReport, @@ -39,6 +47,9 @@ from .fixtures.config import configured_fixture_directory from .fixtures.models import OcrParityCase, OcrSdkInput API_KEY: Final = "test-key" +CALLBACK_DELAY_SECONDS: Final = 0.05 +CALLBACK_DRAIN_TIMEOUT_SECONDS: Final = 10.0 +CALLBACK_TERMINALS: Final[tuple[Literal["success", "failure"], ...]] = ("success", "failure") PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback" PYTHON_VARIANT: Final = ExecutionVariant(name="Python", environment=(("LITELLM_RUST", "0"),)) RUST_VARIANT: Final = ExecutionVariant(name="Rust", environment=(("LITELLM_RUST", "1"),)) @@ -75,10 +86,148 @@ class InvalidOcrWorkerCase(BaseModel): case: InvalidOcrCase -OcrWorkerCase = Annotated[RecordedOcrWorkerCase | InvalidOcrWorkerCase, Field(discriminator="kind")] +class CallbackOcrWorkerCase(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["callback"] = "callback" + case: OcrParityCase + terminal: Literal["success", "failure"] + + +OcrWorkerCase = Annotated[ + RecordedOcrWorkerCase | InvalidOcrWorkerCase | CallbackOcrWorkerCase, + Field(discriminator="kind"), +] OCR_WORKER_CASE_ADAPTER: Final[TypeAdapter[OcrWorkerCase]] = TypeAdapter(OcrWorkerCase) +class RecordingCallback(CustomLogger): + def __init__(self) -> None: + self.message_logging: Final = True + self.turn_off_message_logging: Final = False + self._observations: Final[queue.SimpleQueue[CallbackObservation]] = queue.SimpleQueue() + + def _normalized_kwargs(self, value: object, key: str | None = None) -> JsonValue: + if value is None or isinstance(value, (bool, int, float)): + return value + if isinstance(value, str): + if key != "api_base": + return value + parsed: Final = urlsplit(value) + return urlunsplit(("", "", parsed.path, parsed.query, parsed.fragment)) + if isinstance(value, datetime.datetime): + return "datetime" + if isinstance(value, Exception): + return sdk_error_report(value).model_dump(mode="json") + if isinstance(value, BaseModel): + return self._normalized_kwargs(value.model_dump(mode="json"), key) + if isinstance(value, Mapping): + if any(not isinstance(map_key, str) for map_key in value): + raise TypeError("callback kwarg mappings must use string keys") + return { + map_key: self._normalized_kwargs(map_value, map_key) + for map_key, map_value in value.items() + } + if isinstance(value, (list, tuple)): + return [self._normalized_kwargs(item) for item in value] + raise TypeError(f"unsupported callback kwarg type: {type(value)}") + + def _record( + self, + hook: Literal[ + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", + ], + phase: Literal["success", "failure"], + kwargs: dict[str, object], + response_obj: object, + ) -> None: + raw_litellm_params: Final = kwargs.get("litellm_params") + litellm_params: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], raw_litellm_params) if isinstance(raw_litellm_params, Mapping) else {} + ) + raw_metadata: Final = litellm_params.get("metadata") + metadata_mapping: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], raw_metadata) if isinstance(raw_metadata, Mapping) else {} + ) + metadata: Final = JSON_VALUE_ADAPTER.validate_python( + {key: metadata_mapping[key] for key in ("callback_profile", "sdk_route") if key in metadata_mapping} + ) + raw_error: Final = kwargs.get("exception") + error: Final = sdk_error_report(raw_error) if isinstance(raw_error, Exception) else None + normalized_kwargs: Final = self._normalized_kwargs(kwargs) + payload_source: Final = ( + response_obj.model_dump(mode="json") if isinstance(response_obj, BaseModel) else response_obj + ) + payload: Final = JSON_VALUE_ADAPTER.validate_python(payload_source) + raw_model: Final = kwargs.get("model") + raw_call_type: Final = kwargs.get("call_type") + raw_call_id: Final = kwargs.get("litellm_call_id") + self._observations.put( + CallbackObservation( + hook=hook, + phase=phase, + model=raw_model if isinstance(raw_model, str) else None, + call_type=str(raw_call_type) if raw_call_type is not None else None, + litellm_call_id=raw_call_id if isinstance(raw_call_id, str) else None, + metadata=metadata, + kwargs=normalized_kwargs, + payload=payload, + error=error, + ) + ) + + def log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + time.sleep(CALLBACK_DELAY_SECONDS) + self._record("log_success_event", "success", kwargs, response_obj) + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + await asyncio.sleep(CALLBACK_DELAY_SECONDS) + self._record("async_log_success_event", "success", kwargs, response_obj) + + def log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + time.sleep(CALLBACK_DELAY_SECONDS) + self._record("log_failure_event", "failure", kwargs, response_obj) + + async def async_log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + await asyncio.sleep(CALLBACK_DELAY_SECONDS) + self._record("async_log_failure_event", "failure", kwargs, response_obj) + + def observations(self) -> tuple[CallbackObservation, ...]: + observations: Final = tuple(self._observations.get_nowait() for _ in range(self._observations.qsize())) + return tuple(sorted(observations, key=lambda observation: observation.hook)) + + INVALID_OCR_CASES: Final = ( InvalidOcrCase( name="unsupported_provider", @@ -229,6 +378,106 @@ def _execute_invalid_sdk_case( return _execute_sdk_call(call_kwargs, route, event_loop) +def _callback_call_id(route: SDKRoute, terminal: Literal["success", "failure"]) -> str: + return f"ocr-callback-{route.value}-{terminal}" + + +def _callback_metadata(route: SDKRoute, terminal: Literal["success", "failure"]) -> dict[str, str]: + return {"callback_profile": terminal, "sdk_route": route.value} + + +def _drain_callback_delivery(route: SDKRoute, event_loop: asyncio.AbstractEventLoop) -> None: + if route is SDKRoute.AOCR: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + async def drain_async_callbacks() -> None: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=CALLBACK_DRAIN_TIMEOUT_SECONDS) + await GLOBAL_LOGGING_WORKER.stop() + + event_loop.run_until_complete(drain_async_callbacks()) + + from litellm.litellm_core_utils.thread_pool_executor import executor + + executor.shutdown(wait=True, cancel_futures=False) + + +def _execute_callback_sdk_case( + case: OcrParityCase, + route: SDKRoute, + terminal: Literal["success", "failure"], + mock_url: str, + event_loop: asyncio.AbstractEventLoop, +) -> WorkerSuccess: + callback: Final = RecordingCallback() + call_kwargs: Final = { + **_call_kwargs(case.litellm_input, mock_url, route), + "callbacks": [callback], + "litellm_call_id": _callback_call_id(route, terminal), + "litellm_trace_id": _callback_call_id(route, terminal), + "metadata": _callback_metadata(route, terminal), + } + report: Final = _execute_sdk_call(call_kwargs, route, event_loop) + _drain_callback_delivery(route, event_loop) + return WorkerSuccess(report=report, callbacks=callback.observations()) + + +def _assert_callback_lifecycle( + execution: Execution, + route: SDKRoute, + terminal: Literal["success", "failure"], +) -> None: + observations: Final = execution.callbacks + assert observations is not None, f"{route.value} {terminal} callbacks were not observed" + expected_hooks: Final = ( + (f"log_{terminal}_event",) + if route is SDKRoute.OCR + else ("async_log_success_event",) + if terminal == "success" + else ("async_log_failure_event", "log_failure_event") + ) + actual_hooks: Final = tuple(observation.hook for observation in observations) + assert actual_hooks == expected_hooks, ( + f"{route.value} {terminal} expected callback hooks {expected_hooks}, received {actual_hooks}" + ) + expected_call_id: Final = _callback_call_id(route, terminal) + expected_metadata: Final = _callback_metadata(route, terminal) + for observation in observations: + assert observation.phase == terminal + assert observation.call_type == route.value + assert observation.litellm_call_id == expected_call_id + assert observation.metadata == expected_metadata + assert observation.model + if terminal == "success": + assert isinstance(execution.report, SDKSuccess) + assert observation.payload == execution.report.response + assert observation.error is None + else: + assert isinstance(execution.report, SDKError) + assert observation.payload is None + assert observation.error is not None + assert observation.error.exception_type + assert observation.error.message + assert observation.error.status_code is not None + assert observation.error.status_code >= 400 + + +def _check_callback_ocr_sdk_parity( + case: OcrParityCase, + route: SDKRoute, + terminal: Literal["success", "failure"], + case_file: Path, + runner: SubprocessRunner, +) -> None: + with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: + python_worker, rust_worker = workers + python: Final = python_worker.execute(case_file, route.value, case.provider_responses) + rust: Final = rust_worker.execute(case_file, route.value, case.provider_responses) + + _assert_callback_lifecycle(python, route, terminal) + _assert_callback_lifecycle(rust, route, terminal) + assert_parity(python, rust, PYTHON_HTTP_SENTINEL) + + def _check_recorded_ocr_sdk_parity( ocr_fixture: OcrParityCase, route: SDKRoute, @@ -276,6 +525,25 @@ def _write_worker_case(directory: Path, index: int, case: OcrWorkerCase) -> Path return case_file +def _callback_fixture( + fixtures: tuple[OcrParityCase, ...], + terminal: Literal["success", "failure"], +) -> OcrParityCase: + matching: Final = tuple( + fixture + for fixture in fixtures + if fixture.litellm_input.contract == "mistral" + and ( + all(response.status_code < 400 for response in fixture.provider_responses) + if terminal == "success" + else any(response.status_code >= 400 for response in fixture.provider_responses) + ) + ) + if not matching: + raise AssertionError(f"no recorded Mistral OCR {terminal} fixture is available for callback parity") + return min(matching, key=lambda fixture: fixture_id(fixture.litellm_input, fixture.litellm_input.model)) + + @contextmanager def parity_checks() -> Generator[tuple[E2ECheck, ...]]: fixtures: Final = tuple( @@ -298,6 +566,17 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]: _write_worker_case(directory, len(recorded_files) + index, InvalidOcrWorkerCase(case=case)) for index, case in enumerate(INVALID_OCR_CASES) ) + callback_cases: Final[tuple[tuple[Literal["success", "failure"], OcrParityCase], ...]] = tuple( + (terminal, _callback_fixture(fixtures, terminal)) for terminal in CALLBACK_TERMINALS + ) + callback_files: Final = tuple( + _write_worker_case( + directory, + len(recorded_files) + len(invalid_files) + index, + CallbackOcrWorkerCase(case=case, terminal=terminal), + ) + for index, (terminal, case) in enumerate(callback_cases) + ) with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: recorded: Final = tuple( E2ECheck( @@ -315,7 +594,15 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]: for case, case_file in zip(INVALID_OCR_CASES, invalid_files, strict=True) for route in SDKRoute ) - yield (*recorded, *invalid) + callbacks: Final = tuple( + E2ECheck( + f"callback:{route.value}:{terminal}", + partial(_check_callback_ocr_sdk_parity, case, route, terminal, case_file, runner), + ) + for (terminal, case), case_file in zip(callback_cases, callback_files, strict=True) + for route in SDKRoute + ) + yield (*recorded, *invalid, *callbacks) def _execute_worker_command( @@ -333,6 +620,8 @@ def _execute_worker_command( return WorkerSuccess(report=_execute_sdk_case(recorded.litellm_input, route, mock_url, event_loop)) case InvalidOcrWorkerCase(case=invalid): return WorkerSuccess(report=_execute_invalid_sdk_case(invalid, route, mock_url, event_loop)) + case CallbackOcrWorkerCase(case=callback_case, terminal=terminal): + return _execute_callback_sdk_case(callback_case, route, terminal, mock_url, event_loop) except Exception: return WorkerFailure(error=traceback.format_exc()) diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index da560f99730..04659b25382 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -16,6 +16,7 @@ TraceFailureSource = Literal["python", "rust", "harness"] class RouteFixture: kwargs: dict[str, object] provider_responses: tuple[RecordedHttpResponse, ...] + expected_failure: bool = False @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index f8d7c55d4e2..eb6c9233565 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -2,15 +2,16 @@ from __future__ import annotations import asyncio from collections.abc import Awaitable +from dataclasses import dataclass from pathlib import Path from typing import Final, Protocol, cast from ....shared.parity.replay import replay_server from ....shared.reporting.models import Surface -from ....shared.tracing.native import native_trace_events +from ....shared.tracing.native import TraceResponsePayload, native_trace_events from ....shared.tracing.profiler import FunctionTraceEvent, profile_python from ....shared.tracing.steps import Engine, pipeline_projection -from ..models import RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario +from ..models import RouteFixture, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario from ..reporting import TraceComparisonArtifact @@ -18,9 +19,22 @@ class SdkCall(Protocol): def __call__(self, **kwargs: object) -> object: ... +@dataclass(frozen=True, slots=True) +class _CollectedTrace: + events: tuple[FunctionTraceEvent, ...] + error: str | None = None + + def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object: async def invoke_async() -> object: - return await cast(Awaitable[object], function(**kwargs)) + try: + return await cast(Awaitable[object], function(**kwargs)) + finally: + await asyncio.sleep(0) + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + await GLOBAL_LOGGING_WORKER.stop() if asynchronous: return asyncio.run(invoke_async()) @@ -48,16 +62,30 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)])) +def _python_invocation_error(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> str | None: + try: + _invoke(function, kwargs, asynchronous=asynchronous) + except Exception as error: + return f"{type(error).__name__}: {error}" + return None + + def _collect( - function: SdkCall, kwargs: dict[str, object], engine: Engine, *, asynchronous: bool -) -> tuple[FunctionTraceEvent, ...]: + function: SdkCall, + fixture: RouteFixture, + engine: Engine, + *, + asynchronous: bool, +) -> _CollectedTrace: + kwargs: Final = fixture.kwargs if engine == "rust": - return native_trace_events(_invoke(function, kwargs, asynchronous=asynchronous)) + payload: Final = TraceResponsePayload.model_validate(_invoke(function, kwargs, asynchronous=asynchronous)) + return _CollectedTrace(native_trace_events(payload), payload.error) import litellm with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: - _invoke(function, kwargs, asynchronous=asynchronous) - return tuple(profiler.events) + error: Final = _python_invocation_error(function, kwargs, asynchronous=asynchronous) + return _CollectedTrace(tuple(profiler.events), error) def collect_trace( @@ -68,22 +96,30 @@ def collect_trace( return function try: with replay_server() as provider: - fixture: Final = spec.fixture(engine, provider.url) - for response in fixture.provider_responses: + base_fixture: Final = spec.fixture(engine, provider.url) + for response in base_fixture.provider_responses: provider.enqueue_response(response) - kwargs: Final = { - **fixture.kwargs, - "api_key": "test-key", - "api_base": provider.url, - **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), - } - events: Final = _collect(function, kwargs, engine, asynchronous=asynchronous) + fixture: Final = RouteFixture( + kwargs={ + **base_fixture.kwargs, + "api_key": "test-key", + "api_base": provider.url, + **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), + }, + provider_responses=base_fixture.provider_responses, + expected_failure=base_fixture.expected_failure, + ) + collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) except Exception as error: return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}") - if not events: + if fixture.expected_failure and collected.error is None: + return TraceExecutionFailure(engine, "call succeeded but the scenario expects failure") + if not fixture.expected_failure and collected.error is not None: + return TraceExecutionFailure(engine, collected.error) + if not collected.events: return TraceExecutionFailure(engine, "trace is empty") - return events + return collected.events def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFailure) -> str | None: diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index fe214f45339..2a4a1b3a152 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -19,6 +19,20 @@ COMMON_MAPPINGS: Final = ( mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), ) +SUCCESS_CALLBACK_SYNC_MAPPING: Final = mapping( + rust_span="success_callback", + python_frame=r"BoundedLoggingThreadPoolExecutor\.submit$", +) +SUCCESS_CALLBACK_ASYNC_MAPPING: Final = mapping( + rust_span="success_callback", + python_frame=r"Logging\.async_success_handler$", +) +FAILURE_CALLBACK_MAPPING: Final = mapping( + rust_span="failure_callback", + python_frame=r"Logging\.(?:async_)?failure_handler$", +) +IGNORED_SUCCESS_CALLBACK_MAPPING: Final = mapping(rust_span="success_callback") + SYNC_MAPPINGS: Final = ( *COMMON_MAPPINGS, mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"), @@ -39,6 +53,21 @@ ASYNC_MAPPINGS: Final = ( ), ) +CALLBACK_SUCCESS_SYNC_MAPPINGS: Final = (*SYNC_MAPPINGS, SUCCESS_CALLBACK_SYNC_MAPPING) +CALLBACK_SUCCESS_ASYNC_MAPPINGS: Final = (*ASYNC_MAPPINGS, SUCCESS_CALLBACK_ASYNC_MAPPING) +CALLBACK_FAILURE_SYNC_MAPPINGS: Final = ( + *COMMON_MAPPINGS, + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + FAILURE_CALLBACK_MAPPING, +) +CALLBACK_FAILURE_ASYNC_MAPPINGS: Final = ( + *COMMON_MAPPINGS, + mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"), + FAILURE_CALLBACK_MAPPING, +) + + AZURE_COMMON_MAPPINGS: Final = ( *COMMON_MAPPINGS[:7], mapping( @@ -99,6 +128,34 @@ def _mistral_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture(engine, "mistral/mistral-ocr-latest") +def _callback_fixture(engine: Engine, *, failure: bool) -> RouteFixture: + fixture: Final = _fixture(engine, "mistral/mistral-ocr-latest") + provider_responses: Final = ( + ( + RecordedHttpResponse.from_bytes( + 400, + (HttpHeader(name="content-type", value="application/json"),), + b'{"message":"trace callback provider failure"}', + ), + ) + if failure + else fixture.provider_responses + ) + return RouteFixture( + kwargs=fixture.kwargs, + provider_responses=provider_responses, + expected_failure=failure, + ) + + +def _mistral_callback_success_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _callback_fixture(engine, failure=False) + + +def _mistral_callback_failure_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _callback_fixture(engine, failure=True) + + def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture( engine, @@ -304,36 +361,50 @@ TRACE_SUITE: Final = TraceSuite( name="mistral", fixture=_mistral_fixture, mappings=COMMON_MAPPINGS, - sync_mappings=SYNC_MAPPINGS, - async_mappings=ASYNC_MAPPINGS, + sync_mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + ), + TraceScenario( + name="mistral-callback-success", + fixture=_mistral_callback_success_fixture, + mappings=COMMON_MAPPINGS, + sync_mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS, + async_mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS, + ), + TraceScenario( + name="mistral-callback-failure", + fixture=_mistral_callback_failure_fixture, + mappings=(*COMMON_MAPPINGS, FAILURE_CALLBACK_MAPPING), + sync_mappings=CALLBACK_FAILURE_SYNC_MAPPINGS, + async_mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS, ), TraceScenario( name="azure-ai", fixture=_azure_fixture, mappings=AZURE_COMMON_MAPPINGS, - sync_mappings=AZURE_SYNC_MAPPINGS, - async_mappings=AZURE_ASYNC_MAPPINGS, + sync_mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), TraceScenario( name="azure-document-intelligence", fixture=_azure_document_intelligence_fixture, mappings=DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS, - sync_mappings=DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, - async_mappings=DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, + sync_mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), TraceScenario( name="vertex-ai", fixture=_vertex_fixture, mappings=VERTEX_COMMON_MAPPINGS, - sync_mappings=VERTEX_SYNC_MAPPINGS, - async_mappings=VERTEX_ASYNC_MAPPINGS, + sync_mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), TraceScenario( name="vertex-deepseek", fixture=_vertex_deepseek_fixture, mappings=DEEPSEEK_COMMON_MAPPINGS, - sync_mappings=DEEPSEEK_SYNC_MAPPINGS, - async_mappings=DEEPSEEK_ASYNC_MAPPINGS, + sync_mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), ), ) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py index 3e6c4060134..0e771f0dc17 100644 --- a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py @@ -230,6 +230,10 @@ _HOST_ONLY_BRIDGE_EXCLUSIONS: Final = tuple( "test_ocr_exception_type_uses_resolved_provider_context", "Python wraps bridge exceptions into public errors.", ), + ( + "test_rust_upstream_error_uses_ocr_provider_error_mapping", + "Python maps native upstream errors through the selected OCR provider config.", + ), ("test_aocr_routes_to_async_rust_when_enabled", "Python selects and invokes the async native bridge."), ("test_aocr_exception_type_uses_resolved_provider_context", "Python wraps async bridge exceptions."), ("test_ocr_forwards_timeout_to_rust", "Python converts and forwards explicit timeouts."), diff --git a/tests/store_model_in_db_tests/test_openai_error_handling.py b/tests/store_model_in_db_tests/test_openai_error_handling.py index d3f38f93bf3..5de7c427d19 100644 --- a/tests/store_model_in_db_tests/test_openai_error_handling.py +++ b/tests/store_model_in_db_tests/test_openai_error_handling.py @@ -193,7 +193,7 @@ async def test_chat_completion_bad_model_with_spend_logs(): # Verify the structure of the log entry assert log_entry["request_id"] == litellm_call_id - assert log_entry["model"] == "non-existent-model" + assert log_entry["model"] == "unknown-model" assert log_entry["model_group"] in ("", "non-existent-model") assert log_entry["spend"] == 0.0 assert log_entry["total_tokens"] == 0 diff --git a/tests/test_gateway/test_launch.py b/tests/test_gateway/test_launch.py new file mode 100644 index 00000000000..a783ce6ac7e --- /dev/null +++ b/tests/test_gateway/test_launch.py @@ -0,0 +1,202 @@ +import os +import socket +import sys +import textwrap +import urllib.parse +from collections.abc import Iterator +from pathlib import Path +from typing import Final, cast +from unittest.mock import MagicMock, patch + +import pytest +from uvicorn.importer import import_from_string +from uvicorn.main import main as uvicorn_main + +import gateway.main +from gateway.launch import GATEWAY_APP, main, pool_database_url, uvicorn_argv +from litellm.proxy.db.db_url_settings import DatabaseURLSettings +from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR, PgBouncerError, PgBouncerSettings + +DB_ENV: Final = { + "DATABASE_HOST": "db.internal", + "DATABASE_PORT": "5432", + "DATABASE_USER": "litellm_pool", + "DATABASE_NAME": "litellm", + "DATABASE_PASSWORD": "p@ss", +} + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return cast(tuple[str, int], probe.getsockname())[1] + + +def _fake_pooler(tmp_path: Path) -> Path: + script: Final = tmp_path / "fake-pgbouncer" + script.write_text( + textwrap.dedent( + f"""\ + #!{sys.executable} + import configparser, select, socket, sys + if sys.argv[1:] == ["--version"]: + print("PgBouncer 1.25.2") + sys.exit(0) + ini = configparser.ConfigParser() + ini.read(sys.argv[1]) + port = ini.getint("pgbouncer", "listen_port") + tcp = socket.socket() + tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + tcp.bind(("127.0.0.1", port)) + tcp.listen() + unix = socket.socket(socket.AF_UNIX) + unix.bind(ini.get("pgbouncer", "unix_socket_dir") + f"/.s.PGSQL.{{port}}") + unix.listen() + while True: + for ready in select.select([tcp, unix], [], [])[0]: + ready.accept()[0].close() + """ + ) + ) + script.chmod(0o700) + return script + + +def _query(url: str) -> dict[str, str]: + return dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query)) + + +@pytest.fixture +def password_env(monkeypatch: pytest.MonkeyPatch) -> Iterator[dict[str, str]]: + for var in ( + "DATABASE_URL", + "IAM_TOKEN_DB_AUTH", + "AZURE_POSTGRESQL_AUTH", + "DATABASE_HOST_READ_REPLICA", + PGBOUNCER_POOLED_ENV_VAR, + ): + monkeypatch.setenv(var, "") + monkeypatch.delenv(var) + for var, value in DB_ENV.items(): + monkeypatch.setenv(var, value) + yield dict(DB_ENV) + os.environ.pop("DATABASE_URL", None) + + +def _minted_iam_token(token: str): + rds: Final = MagicMock() + rds.generate_db_auth_token.return_value = token + return patch("boto3.client", return_value=rds) + + +def _uvicorn_params(argv: tuple[str, ...]) -> dict[str, object]: + return uvicorn_main.make_context("uvicorn", list(argv)).params + + +class TestUvicornArgv: + def test_keepalive_env_reaches_uvicorn(self): + params: Final = _uvicorn_params(uvicorn_argv(("--workers", "4"), {"KEEPALIVE_TIMEOUT": "75"})) + assert params["app"] == GATEWAY_APP + assert params["workers"] == 4 + assert params["timeout_keep_alive"] == 75 + + def test_unset_env_keeps_the_uvicorn_default(self): + assert _uvicorn_params(uvicorn_argv(("--workers", "4"), {}))["timeout_keep_alive"] == 5 + + def test_an_explicit_flag_wins_over_the_env(self): + argv: Final = uvicorn_argv(("--timeout-keep-alive", "30"), {"KEEPALIVE_TIMEOUT": "75"}) + assert _uvicorn_params(argv)["timeout_keep_alive"] == 30 + + def test_the_app_uvicorn_is_told_to_serve_is_the_trimmed_gateway(self): + assert import_from_string(cast(str, _uvicorn_params(uvicorn_argv((), {}))["app"])) is gateway.main.app + + +class TestPoolDatabaseUrl: + def test_a_disabled_pooler_yields_no_url_to_install(self, password_env: dict[str, str]): + settings: Final = DatabaseURLSettings.from_env() + settings.apply_to_env() + environ: Final = {"DATABASE_URL": "postgresql://litellm_pool:p%40ss@db.internal:5432/litellm"} + assert pool_database_url(settings, PgBouncerSettings(enabled=False), environ) is None + + def test_a_missing_upstream_url_is_reported(self, password_env: dict[str, str]): + environ: Final[dict[str, str]] = {} + outcome: Final = pool_database_url(DatabaseURLSettings.from_env(), PgBouncerSettings(enabled=True), environ) + assert isinstance(outcome, PgBouncerError) + assert "DATABASE_URL" in outcome.reason + + def test_token_auth_hands_the_workers_the_pool_user_not_the_token( + self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ): + monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + port: Final = _free_port() + environ: Final = {"DATABASE_URL": "postgresql://litellm:MINTED_TOKEN@db.internal:5432/litellm"} + with _minted_iam_token("MINTED_TOKEN"): + outcome: Final = pool_database_url( + DatabaseURLSettings.from_env(), + PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path))), + environ, + ) + assert isinstance(outcome, str), outcome + pooled: Final = urllib.parse.urlsplit(outcome) + assert (pooled.username, pooled.hostname, pooled.port) == ("litellm_pgbouncer", "127.0.0.1", port) + assert "MINTED_TOKEN" not in outcome + + +class TestMain: + def test_workers_inherit_the_loopback_url_the_supervisor_installed( + self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ): + port: Final = _free_port() + monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") + monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", str(port)) + monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(tmp_path))) + monkeypatch.setenv("KEEPALIVE_TIMEOUT", "75") + served: Final[list[tuple[str, ...]]] = [] + main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv))) + + pooled: Final = os.environ["DATABASE_URL"] + assert urllib.parse.urlsplit(pooled).hostname == "127.0.0.1" + assert urllib.parse.urlsplit(pooled).port == port + assert urllib.parse.urlsplit(pooled).username == "litellm_pgbouncer" + assert "p%40ss" not in pooled + assert _query(pooled)["pgbouncer"] == "true" + assert _uvicorn_params(served[0])["timeout_keep_alive"] == 75 + + DatabaseURLSettings.from_env().apply_to_env() + assert urllib.parse.urlsplit(os.environ["DATABASE_URL"]).netloc == urllib.parse.urlsplit(pooled).netloc + assert _query(os.environ["DATABASE_URL"])["pgbouncer"] == "true" + + def test_iam_workers_keep_the_loopback_url_instead_of_minting_their_own( + self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ): + port: Final = _free_port() + monkeypatch.delenv("DATABASE_PASSWORD") + monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") + monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", str(port)) + monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(tmp_path))) + served: Final[list[tuple[str, ...]]] = [] + with _minted_iam_token("SUPERVISOR_TOKEN"): + main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv))) + pooled: Final = os.environ["DATABASE_URL"] + assert urllib.parse.urlsplit(pooled).netloc.endswith(f"@127.0.0.1:{port}") + assert "SUPERVISOR_TOKEN" not in pooled + assert os.environ[PGBOUNCER_POOLED_ENV_VAR] == "true" + assert len(served) == 1 + + with _minted_iam_token("WORKER_TOKEN"): + DatabaseURLSettings.from_env().apply_to_env() + assert os.environ["DATABASE_URL"] == pooled + + def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn( + self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ): + monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") + monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(tmp_path / "missing-pgbouncer")) + served: Final[list[tuple[str, ...]]] = [] + with pytest.raises(SystemExit) as stopped: + main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv))) + assert "missing-pgbouncer" in str(stopped.value) + assert served == [] diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index b2cf253d164..a30474245c6 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -126,16 +126,6 @@ def test_load_rust_messages_returns_injected_impl(): assert rust_messages.load_rust_messages() is bridge -def test_bare_rust_still_toggles_ocr(): - from litellm.rust_bridge.ocr import rust_ocr_enabled - - litellm.rust(True) - assert rust_ocr_enabled() is True - - litellm.rust(False) - assert rust_ocr_enabled() is False - - def test_load_rust_amessages_returns_injected_impl(): bridge = RecordingAsyncMessages() litellm.rust(True) @@ -214,7 +204,7 @@ async def test_amessages_wrapper_forwards_args(): def _gate(**overrides): kwargs = { "custom_llm_provider": "azure_ai", - "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), + "litellm_params": GenericLiteLLMParams(api_key="sk-azure"), "has_agentic_hook": False, "model": "claude-sonnet-4-5", "api_key": "sk-azure", @@ -282,18 +272,6 @@ async def test_gate_uses_process_enable_without_request_override(): assert bridge.calls[0]["custom_llm_provider"] == "azure_ai" -@pytest.mark.asyncio -async def test_gate_skips_rust_when_flag_false(): - bridge = ExplodingAsyncMessages() - litellm.rust(True) - rust_messages.set_rust_messages(amessages=bridge) - - response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False)) - - assert response is None - assert bridge.calls == 0 - - @pytest.mark.asyncio async def test_gate_invokes_rust_for_native_anthropic_provider(): bridge = RecordingAsyncMessages() @@ -302,7 +280,7 @@ async def test_gate_invokes_rust_for_native_anthropic_provider(): response = await _gate( custom_llm_provider="anthropic", - litellm_params=GenericLiteLLMParams(api_key="sk-ant", rust=True), + litellm_params=GenericLiteLLMParams(api_key="sk-ant"), api_key="sk-ant", api_base="https://api.anthropic.com", headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"}, diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index c86c7c4df03..768ea332677 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -21,11 +21,12 @@ from types import MappingProxyType import httpx import pytest import respx +from openai.types.batch import BatchRequestCounts import litellm import litellm.batches.batch_utils as bu -from litellm.types.utils import Usage +from litellm.types.utils import LiteLLMBatch, Usage # --------------------------------------------------------------------------- # # Builders for batch OUTPUT file rows. @@ -489,7 +490,9 @@ def test_aggregate_counts_successful_and_failed_requests(monkeypatch): def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0) + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.4, 0.6)) result = bu._aggregate_batch_cost_usage_models( entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai" ) @@ -500,6 +503,7 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): 1, 0, ) + assert (result.prompt_cost, result.completion_cost) == (0.4, 0.6) # =========================================================================== # @@ -507,15 +511,17 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): # =========================================================================== # -def test_cost_from_content_completion_cost_path(monkeypatch): - # model_info is None -> litellm.completion_cost per successful row. +def test_cost_without_model_info_prices_each_row_by_its_response_model(monkeypatch): + # model_info is None -> batch_cost_calculator per successful row, model from the response body. + import litellm.cost_calculator as cc + calls = [] - def _completion_cost(**kw): + def _batch_cost(**kw): calls.append(kw) - return 0.5 + return (0.3, 0.2) - monkeypatch.setattr(litellm, "completion_cost", _completion_cost) + monkeypatch.setattr(cc, "batch_cost_calculator", _batch_cost) rows = [ _success_row(usage=_usage(10, 5)), _failed_row(), # excluded -> not costed @@ -524,8 +530,10 @@ def test_cost_from_content_completion_cost_path(monkeypatch): result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert result.cost == 1.0 # 2 successful * 0.5 + assert result.cost == pytest.approx(1.0) # 2 successful * (0.3 + 0.2) + assert (result.prompt_cost, result.completion_cost) == (pytest.approx(0.6), pytest.approx(0.4)) assert len(calls) == 2 # failed row not costed + assert all(call["model"] == "gpt-4o" and call["model_info"] is None for call in calls) assert result.successful_requests == 2 assert result.failed_requests == 1 @@ -578,7 +586,9 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch): """A one-shot generator: any implementation that iterates the entries twice (e.g. separate cost and usage passes) sees nothing on the second pass and returns wrong totals for at least one of cost/usage/models.""" - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5) + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.25, 0.25)) one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))]) result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai") @@ -753,12 +763,15 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): @pytest.mark.asyncio async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch): + import litellm.cost_calculator as cc + rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5) + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (1.5, 1.0)) result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai") assert result.cost == 2.5 + assert (result.prompt_cost, result.completion_cost) == (1.5, 1.0) assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) assert result.models == ["gpt-4o"] @@ -1107,8 +1120,10 @@ async def test_handle_completed_batch_orchestration(monkeypatch): async def fake_fetch(batch, custom_llm_provider, litellm_params=None): return _vertex_jsonl(rows) + import litellm.cost_calculator as cc + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3) + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3)) result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") @@ -1718,3 +1733,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 diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index c3edb40c819..b87f9489250 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -158,6 +158,14 @@ def test_create__vertex_ai_dispatch(seams): _assert_only(seams.vertex.create_batch, seams, "create_batch") +def test_create__vertex_ai_forwards_custom_endpoint(seams): + """The vertex handler owns the custom_endpoint batch rejection (LIT-6899), so the dispatcher + must forward the flag for the handler to act on.""" + bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True) + + assert seams.vertex.create_batch.call_args.kwargs["custom_endpoint"] is True + + def test_create__provider_config_routes_to_base_http_handler(seams): """model + a provider batches config (bedrock-style) routes to the generic base_llm_http_handler, NOT the per-provider instance.""" diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 8e0bc200012..071b99850f6 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -658,3 +658,38 @@ def test_async_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatc asyncio.run(_short_lived_script()) assert len(writes) == 1 + + +@pytest.mark.asyncio +async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monkeypatch): + """The spend log for a cache hit must reuse the key the lookup already computed instead of hashing again.""" + import litellm + from litellm.caching.caching import Cache + from litellm.types.utils import CallTypes + + async def acompletion(**kwargs): + return None + + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + kwargs = {"model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}], "caching": True} + await litellm.cache.async_add_cache( + litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}]), **kwargs + ) + handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now()) + logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False) + logging_obj.async_success_handler = AsyncMock() + + hit = await handler._async_get_cache( + model="gpt-5.4", + original_function=acompletion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.acompletion.value, + kwargs=kwargs, + args=(), + ) + + assert hit is not None and hit.cached_result is not None + assert handler.preset_cache_key is not None + assert logging_obj.litellm_params["preset_cache_key"] == handler.preset_cache_key + assert hit.cached_result._hidden_params["cache_key"] == handler.preset_cache_key diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index ded3be26630..4c9068722b8 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -1,4 +1,5 @@ import asyncio +import logging import time import uuid from unittest.mock import AsyncMock, MagicMock, patch @@ -7,7 +8,8 @@ import pytest from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync +from litellm.types.caching import RedisPipelineIncrementOperation @pytest.mark.asyncio @@ -576,3 +578,129 @@ async def test_dual_cache_late_attach_redis_wires_writes_and_ttl_async(): assert mock_redis.async_set_cache.call_args[0][:2] == (key_after, val_after) assert in_memory.get_cache(key_after) == val_after + + +class _OpenBreakerRedis: + def __init__(self) -> None: + from litellm.caching.redis_cache import RedisCircuitBreaker + + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + for _ in range(3): + self._circuit_breaker.record_failure() + + @_redis_circuit_breaker_guard + async def async_get_cache(self, key, **kwargs): + raise AssertionError("never reached") + + @_redis_circuit_breaker_guard + async def async_batch_get_cache(self, key_list, **kwargs): + raise AssertionError("never reached") + + @_redis_circuit_breaker_guard + async def async_set_cache(self, key, value, **kwargs): + raise AssertionError("never reached") + + @_redis_circuit_breaker_guard + async def async_set_cache_pipeline(self, cache_list, **kwargs): + raise AssertionError("never reached") + + @_redis_circuit_breaker_guard + async def async_increment_pipeline(self, increment_list, **kwargs): + raise AssertionError("never reached") + + @_redis_circuit_breaker_guard + async def async_increment(self, key, value, **kwargs): + raise AssertionError("never reached") + + @_redis_circuit_breaker_guard_sync + def get_cache(self, key, **kwargs): + raise AssertionError("never reached") + + @_redis_circuit_breaker_guard_sync + def batch_get_cache(self, key_list, **kwargs): + raise AssertionError("never reached") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call", + [ + lambda cache: cache.async_get_cache("k"), + lambda cache: cache.async_batch_get_cache(["k1", "k2"]), + lambda cache: cache.async_set_cache("k", "v"), + lambda cache: cache.async_set_cache_pipeline([("k", "v")]), + lambda cache: cache.async_increment_cache_pipeline( + increment_list=[RedisPipelineIncrementOperation(key="k", increment_value=1.0, ttl=60)] + ), + lambda cache: cache.async_increment_cache("k", 1.0), + ], + ids=["get", "batch_get", "set", "set_pipeline", "increment_pipeline", "increment"], +) +async def test_an_open_circuit_breaker_is_not_an_error_per_request(caplog, call): + cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + caplog.clear() + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + await call(cache) + + assert [record.levelno for record in caplog.records if record.levelno >= logging.WARNING] == [] + assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) + + +@pytest.mark.parametrize( + "call", + [lambda cache: cache.get_cache("k"), lambda cache: cache.batch_get_cache(["k1", "k2"])], + ids=["get", "batch_get"], +) +def test_an_open_circuit_breaker_is_not_an_error_per_sync_request(caplog, call): + cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + caplog.clear() + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + call(cache) + + assert [record.levelno for record in caplog.records if record.levelno >= logging.WARNING] == [] + assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_a_real_redis_failure_still_logs_an_error(caplog): + class _BrokenRedis: + async def async_get_cache(self, key, **kwargs): + raise ConnectionError("redis is down") + + cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_BrokenRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + assert await cache.async_get_cache("k") is None + + errors = [record for record in caplog.records if record.levelno == logging.ERROR] + assert [record.getMessage() for record in errors] == ["LiteLLM Cache: exception in async_get_cache: redis is down"] + assert errors[0].exc_info is not None + + +def _dual_cache_with_open_breaker_and_a_memory_hit() -> DualCache: + in_memory = InMemoryCache() + in_memory.set_cache("k1", "v1") + return DualCache(in_memory_cache=in_memory, redis_cache=_OpenBreakerRedis(), default_redis_batch_cache_expiry=10) # pyright: ignore[reportArgumentType] # duck-typed Redis double + + +def test_open_breaker_keeps_sync_batch_read_memory_hits_and_releases_reservations(): + """A refused Redis batch read must still answer with the in-memory hits and hold no reservation. + + The refusal was logged and turned into a bare None, so a caller lost its in-memory hits + for as long as the breaker stayed open, and the reserved keys stayed throttled until + the batch expiry passed even though nothing was ever read for them. + """ + cache = _dual_cache_with_open_breaker_and_a_memory_hit() + + assert list(cache.batch_get_cache(["k1", "k2"])) == ["v1", None] + assert "k2" not in cache.last_redis_batch_access_time + + +@pytest.mark.asyncio +async def test_open_breaker_keeps_async_batch_read_memory_hits_and_releases_reservations(): + cache = _dual_cache_with_open_breaker_and_a_memory_hit() + + assert list(await cache.async_batch_get_cache(["k1", "k2"])) == ["v1", None] + assert "k2" not in cache.last_redis_batch_access_time diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 2f412e7382b..bcae33b976e 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1,11 +1,12 @@ import asyncio +import time from collections.abc import Iterator from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm._service_logger import ServiceLogging -from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError @pytest.fixture @@ -515,14 +516,90 @@ async def test_circuit_breaker_opens_when_method_swallows_redis_failure(call_met await call_method(cache) -def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_redis_cache): - """An open breaker must preserve the sync batch read's dictionary fallback.""" +def test_circuit_breaker_open_makes_sync_batch_get_cache_fast_fail(sync_batch_redis_cache, caplog): + """Once the breaker is open the sync batch read refuses with the typed error instead of a miss. + + Swallowing the refusal into `{}` made every sync batch read on an open breaker emit an ERROR + log and a service failure event per call, and the DualCache caller could not tell the + refusal from a dead Redis, so it dropped its in-memory hits too. + """ from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} - assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} + caplog.clear() + with caplog.at_level("INFO"): + with pytest.raises(RedisCircuitBreakerOpenError): + sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) + sync_batch_redis_cache.redis_client.mget.assert_called() + assert caplog.records == [] + + +def test_sync_get_cache_failure_feeds_the_breaker_and_logs_a_well_formed_record(sync_batch_redis_cache, caplog): + """The sync get path swallowed its Redis error without recording it, and its log call was malformed. + + `verbose_logger.error("...: ", e)` passes the exception as a format argument to a message + with no placeholder, so the record carried no error text. Nothing fed the breaker either, + so a dead Redis read through this path never opened it. + """ + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + + sync_batch_redis_cache.redis_client.get.side_effect = OSError("redis unavailable") + + with caplog.at_level("ERROR"): + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + assert sync_batch_redis_cache.get_cache("lit7468") is None + + assert all("redis unavailable" in record.getMessage() for record in caplog.records) + assert len(caplog.records) == REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + assert sync_batch_redis_cache._circuit_breaker.is_open() is True + with pytest.raises(RedisCircuitBreakerOpenError): + sync_batch_redis_cache.get_cache("lit7468") + + +def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache): + """A caller that must fall back when Redis is unreachable needs the failure, not zeros. + + The batch read answers a dead Redis with an empty dict, which a counting caller cannot tell + apart from "every counter is unset". Least-busy routing read that as an idle deployment and + kept sending traffic to it instead of falling back to this worker's own in-flight counts. + """ + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + sync_batch_redis_cache.batch_get_counts(["lit7039"]) + + +@pytest.mark.asyncio +async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports_a_miss(redis_no_ping: None): + """Async twin: the async batch read hides the same failure behind an empty dict.""" + failing_client = AsyncMock() + failing_client.mget.side_effect = OSError("redis unavailable") + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + + with patch.object(cache, "init_async_client", return_value=failing_client): + assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + await cache.async_batch_get_counts(["lit7039"]) + + +@pytest.mark.parametrize("stored", [b"3", "3"]) +def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(stored, redis_no_ping: None): + """Counters come back positionally, so an unset key has to stay a hole rather than shift the + rest of the row onto the wrong deployments, and a count has to survive whether the client + hands it back as bytes or as text.""" + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + cache.redis_client.mget.return_value = [stored, None, b"0"] + + assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0) @pytest.fixture @@ -617,7 +694,8 @@ def test_sync_batch_get_cache_survives_a_service_callback_that_raises( with ThreadPoolExecutor(max_workers=1) as pool: assert pool.submit(cache.batch_get_cache, key_list=["lit6729"]).result() == {} - assert cache.batch_get_cache(key_list=["lit6729"]) == {} + with pytest.raises(RedisCircuitBreakerOpenError): + cache.batch_get_cache(key_list=["lit6729"]) def test_call_stack_info_skips_breaker_guard_frames(): @@ -966,6 +1044,161 @@ async def test_breaker_metrics_track_state_and_failure_class(): assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before + 1 assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + breaker._opened_at = time.time() - 9999 + assert breaker.is_open() is False breaker.record_success() assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + 1 + + +def test_sync_guard_counts_a_timeout_as_a_timeout(): + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker_sync + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + + def timing_out_call() -> str: + raise RedisTimeoutError("read timed out") + + for _ in range(6): + with pytest.raises(RedisTimeoutError): + _run_under_circuit_breaker_sync(breaker, "op", timing_out_call) + + assert breaker.is_open() is False + + +def test_success_admitted_before_the_breaker_opened_cannot_close_it(): + """A stale in-flight success must not close a breaker that opened while it ran. + + Calls admitted while the breaker was still closed finish after later failures opened it. + Recording their success unconditionally closed the breaker again, skipping the recovery + timeout and the single half-open probe, so the breaker flapped between open and closed + on every straggler while Redis was still down. + """ + from litellm.caching.redis_cache import RedisCircuitBreaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + for _ in range(3): + breaker.record_failure() + assert breaker._state == breaker.OPEN + + breaker.record_success() + + assert breaker._state == breaker.OPEN + assert breaker.is_open() is True + + +def test_recovery_probe_still_closes_the_breaker(): + from litellm.caching.redis_cache import RedisCircuitBreaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + for _ in range(3): + breaker.record_failure() + breaker._opened_at = time.time() - 9999 + assert breaker.is_open() is False + assert breaker._state == breaker.HALF_OPEN + + breaker.record_success() + + assert breaker._state == breaker.CLOSED + assert breaker.is_open() is False + + +@pytest.mark.asyncio +async def test_stale_success_during_the_recovery_probe_leaves_the_breaker_to_the_probe(): + """A call admitted before the trip that finishes while HALF_OPEN must not close the breaker. + + Only the one call designated as the recovery probe has actually reached Redis after the + outage, so closing on the straggler's success resumed full Redis traffic before the probe + had proven anything. + """ + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + stale_admitted = asyncio.Event() + stale_release = asyncio.Event() + probe_admitted = asyncio.Event() + probe_release = asyncio.Event() + + async def stale_call() -> str: + stale_admitted.set() + await stale_release.wait() + return "stale" + + async def probe_call() -> str: + probe_admitted.set() + await probe_release.wait() + return "probe" + + stale = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", stale_call)) + await stale_admitted.wait() + for _ in range(3): + breaker.record_failure() + assert breaker._state == breaker.OPEN + breaker._opened_at = time.time() - 9999 + probe = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", probe_call)) + await probe_admitted.wait() + assert breaker._state == breaker.HALF_OPEN + + stale_release.set() + assert await stale == "stale" + + assert breaker._state == breaker.HALF_OPEN, "the straggler must not close the breaker for the probe" + assert breaker.is_open() is True + + probe_release.set() + assert await probe == "probe" + + assert breaker._state == breaker.CLOSED + assert breaker.is_open() is False + + +@pytest.mark.asyncio +async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new_probe(): + """A probe still in flight when a late failure reopens the breaker must not close it for the next probe. + + Once the breaker has reopened, only the probe admitted after that outage has reached + Redis, so the older probe's success no longer says anything about whether Redis recovered. + """ + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + old_probe_admitted = asyncio.Event() + old_probe_release = asyncio.Event() + new_probe_admitted = asyncio.Event() + new_probe_release = asyncio.Event() + + async def old_probe_call() -> str: + old_probe_admitted.set() + await old_probe_release.wait() + return "old probe" + + async def new_probe_call() -> str: + new_probe_admitted.set() + await new_probe_release.wait() + return "new probe" + + for _ in range(3): + breaker.record_failure() + breaker._opened_at = time.time() - 9999 + old_probe = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", old_probe_call)) + await old_probe_admitted.wait() + assert breaker._state == breaker.HALF_OPEN + + breaker.record_failure() + assert breaker._state == breaker.OPEN + breaker._opened_at = time.time() - 9999 + new_probe = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", new_probe_call)) + await new_probe_admitted.wait() + assert breaker._state == breaker.HALF_OPEN + + old_probe_release.set() + assert await old_probe == "old probe" + + assert breaker._state == breaker.HALF_OPEN, "the overtaken probe must not close the breaker for the new probe" + assert breaker.is_open() is True + + new_probe_release.set() + assert await new_probe == "new probe" + assert breaker._state == breaker.CLOSED diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 6590718878d..e4ada0a9b31 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3952,3 +3952,338 @@ def test_convert_chat_completion_messages_to_responses_api_tool_result_with_tool function_call_output = next(item for item in response if item.get("type") == "function_call_output") assert function_call_output["output"] == [{"type": "input_text", "text": "1 tool found"}] + + +def _litellm_encoded_response_id(upstream_id: str) -> str: + from litellm.responses.utils import ResponsesAPIRequestUtils + + return ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="azure", model_id="deployment-1", response_id=upstream_id + ) + + +def test_transform_response_keeps_upstream_id_and_provider_extras(): + from unittest.mock import Mock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse, Usage + + content_filters = [ + {"blocked": False, "source_type": "prompt", "content_filter_results": {"hate": {"filtered": False}}} + ] + raw_response = ResponsesAPIResponse.model_validate( + { + "id": _litellm_encoded_response_id("resp_azure_123"), + "created_at": 1734366691, + "object": "response", + "model": "gpt-5.6", + "status": "completed", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_weather", + "arguments": '{"city": "Seattle"}', + "status": "completed", + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "service_tier": "default", + "content_filters": content_filters, + "max_tool_calls": None, + "background": False, + "top_logprobs": 0, + "store": True, + } + ) + model_response = ModelResponse( + id="chatcmpl-local", + created=1734366691, + model=None, + object="chat.completion", + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ) + + result = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-5.6", + raw_response=raw_response, + model_response=model_response, + logging_obj=Mock(), + request_data={"model": "gpt-5.6"}, + messages=[{"role": "user", "content": "What is the weather in Seattle?"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + dumped = result.model_dump() + + assert dumped["id"] == "resp_azure_123" + assert dumped["object"] == "chat.completion" + assert dumped["service_tier"] == "default" + assert dumped["content_filters"] == content_filters + assert "max_tool_calls" not in dumped, "a null provider field must not appear as a null top-level key" + assert "output" not in dumped and "status" not in dumped, ( + "Responses schema fields must not leak into the chat response" + ) + assert not {"background", "top_logprobs", "store"} & dumped.keys(), ( + "Responses API bookkeeping must not ride along as chat metadata" + ) + assert dumped["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "lookup_weather" + + +def test_bridged_response_is_priced_by_the_reported_service_tier(): + from unittest.mock import Mock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse + + raw_response = ResponsesAPIResponse.model_validate( + { + "id": "resp_flex", + "created_at": 1734366691, + "object": "response", + "model": "gpt-5.4", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + "service_tier": "flex", + } + ) + + result = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + pricing = litellm.model_cost["gpt-5.4"] + flex_cost = 1000 * pricing["input_cost_per_token_flex"] + 100 * pricing["output_cost_per_token_flex"] + standard_cost = 1000 * pricing["input_cost_per_token"] + 100 * pricing["output_cost_per_token"] + + cost = litellm.completion_cost(completion_response=result, custom_llm_provider="openai") + + assert cost == pytest.approx(flex_cost) + assert cost < standard_cost + + +def test_streaming_chunks_carry_the_upstream_response_id(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + encoded_id = _litellm_encoded_response_id("resp_azure_stream") + events = [ + {"type": "response.created", "response": {"id": encoded_id, "output": []}}, + {"type": "response.output_text.delta", "delta": "Hel"}, + {"type": "response.completed", "response": {"id": encoded_id, "output": [{"type": "message"}]}}, + ] + + ids = [iterator.chunk_parser(event).id for event in events] + + assert ids == ["resp_azure_stream"] * len(events), f"streamed chunks did not carry the upstream id: {ids}" + + +def test_streaming_final_chunk_carries_provider_metadata(): + from unittest.mock import MagicMock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + content_filters = [{"blocked": False, "source_type": "completion", "content_filter_results": {}}] + events = [ + {"type": "response.created", "response": {"id": "resp_azure_stream", "output": []}}, + {"type": "response.output_text.delta", "delta": "Hello"}, + { + "type": "response.completed", + "response": { + "id": "resp_azure_stream", + "output": [{"type": "message"}], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + "service_tier": "default", + "content_filters": content_filters, + "background": False, + }, + }, + ] + stream = CustomStreamWrapper( + completion_stream=iter([iterator.chunk_parser(event) for event in events]), + model="gpt-5.6", + custom_llm_provider="azure", + logging_obj=MagicMock(), + ) + + chunks = [chunk.model_dump() for chunk in stream] + + assert chunks[-1]["choices"][0]["finish_reason"] == "stop" + assert chunks[-1]["service_tier"] == "default" + assert chunks[-1]["content_filters"] == content_filters + assert "background" not in chunks[-1] + assert all("service_tier" not in chunk for chunk in chunks[:-1]) + + +def _system_input_item(text: str) -> dict[str, object]: + return {"type": "message", "role": "system", "content": [{"type": "input_text", "text": text}]} + + +def test_mid_conversation_system_string_stays_in_input_after_a_user_turn(): + handler: Final = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Read the file."}, + {"role": "system", "content": "14982391 tokens left"}, + {"role": "user", "content": "Now summarize it."}, + ] + ) + + assert instructions == "You are a helpful assistant." + assert input_items == [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Read the file."}]}, + _system_input_item("14982391 tokens left"), + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Now summarize it."}]}, + ] + + +def test_leading_system_strings_still_join_instructions_without_a_following_turn(): + handler: Final = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "system", "content": "Be brief."}, + {"role": "system", "content": "Answer in French."}, + ] + ) + + assert instructions == "Be brief. Answer in French." + assert input_items == [] + + +def test_mid_conversation_system_reminder_as_string_and_as_text_block_produce_identical_input_items(): + handler: Final = LiteLLMResponsesTransformationHandler() + reminder: Final = "14982391 tokens left" + + as_string, string_instructions = handler.convert_chat_completion_messages_to_responses_api( + [{"role": "user", "content": "Read the file."}, {"role": "system", "content": reminder}] + ) + as_block, block_instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "user", "content": "Read the file."}, + { + "role": "system", + "content": [{"type": "text", "text": reminder, "cache_control": {"type": "ephemeral"}}], + }, + ] + ) + + assert string_instructions is None + assert block_instructions is None + assert json.dumps(as_string) == json.dumps(as_block) + assert as_string[1] == _system_input_item(reminder) + + +def test_claude_code_shaped_history_keeps_a_byte_stable_input_prefix_across_requests(): + handler: Final = LiteLLMResponsesTransformationHandler() + top_level_system: Final = [{"type": "text", "text": "You are Claude Code.", "cache_control": {"type": "ephemeral"}}] + first_reminder: Final = "27k chars of deferred tools" + second_reminder: Final = "14982391 tokens left" + first_request_messages: Final = [ + {"role": "system", "content": top_level_system}, + {"role": "user", "content": "Read inventory.py."}, + { + "role": "system", + "content": [{"type": "text", "text": first_reminder, "cache_control": {"type": "ephemeral"}}], + }, + ] + second_request_messages: Final = [ + {"role": "system", "content": top_level_system}, + {"role": "user", "content": "Read inventory.py."}, + {"role": "system", "content": first_reminder}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "Read", "arguments": '{"file_path": "inventory.py"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "ITEMS = []"}, + { + "role": "system", + "content": [{"type": "text", "text": second_reminder, "cache_control": {"type": "ephemeral"}}], + }, + ] + + first_request: Final = handler.transform_request( + model="gpt-5.6-luna", + messages=first_request_messages, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + second_request: Final = handler.transform_request( + model="gpt-5.6-luna", + messages=second_request_messages, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert "instructions" not in first_request + assert "instructions" not in second_request + assert first_request["input"][0] == _system_input_item("You are Claude Code.") + assert json.dumps(second_request["input"][: len(first_request["input"])]) == json.dumps(first_request["input"]) + assert second_request["input"][len(first_request["input"]) :] == [ + {"type": "function_call", "call_id": "call_1", "name": "Read", "arguments": '{"file_path": "inventory.py"}'}, + {"type": "function_call_output", "call_id": "call_1", "output": [{"type": "input_text", "text": "ITEMS = []"}]}, + _system_input_item(second_reminder), + ] + + +def test_system_string_after_a_developer_message_stays_in_input_in_client_order(): + handler: Final = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "developer", "content": "Always answer in French."}, + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Bonjour"}, + ] + ) + + assert instructions is None + assert [item["role"] for item in input_items] == ["developer", "system", "user"] + assert input_items[1] == _system_input_item("Be brief.") diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index 62c95cb100b..a4f32df46ae 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -7,9 +7,12 @@ # 4. Added proper cleanup in fixtures # 5. Added worker-specific isolation for parallel execution +import base64 import importlib import os from pathlib import Path +from types import SimpleNamespace +import httpx import pytest import asyncio @@ -595,3 +598,43 @@ def pytest_sessionfinish(session, exitstatus): _close_handler_if_needed(getattr(litellm, "aclient", None)) _close_handler_if_needed(getattr(litellm, "client", None)) _run_coroutine_if_needed(close_litellm_async_clients()) + + +ONE_PIXEL_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +@pytest.fixture +def async_only_image_fetch(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import factory, image_handling + from litellm.llms.gemini.chat import transformation as gemini_chat_transformation + + fetch = SimpleNamespace( + fetched=[], + base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(), + data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(), + ) + + def forbid_sync_fetch(client, url, **kwargs): + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client, url, **kwargs): + fetch.fetched.append(url) + return httpx.Response( + 200, + content=ONE_PIXEL_PNG, + headers={"content-type": "image/png"}, + request=httpx.Request("GET", url), + ) + + def forbid_sync_convert(url, *args, **kwargs): + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + for module in (image_handling, factory, gemini_chat_transformation): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) + return fetch diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py index f46df5baadf..c1c3569c0e2 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -262,7 +262,7 @@ def test_proxied_traffic_stays_on_native_hooks(): never sees ``data["prompt"]``.""" guardrail = _guardrail() assert guardrail.uses_apply_guardrail_interface() is True - assert guardrail._deployment_pre_call_target() is guardrail + assert guardrail._deployment_hook_target() is guardrail @pytest.mark.asyncio diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index ebd33aa2e53..d3e668b8987 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -375,6 +375,7 @@ def _in_memory_managed_files(): table.upsert = AsyncMock(side_effect=_upsert) prisma = MagicMock() prisma.db.litellm_managedobjecttable = table + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) cache = MagicMock() cache.async_set_cache = AsyncMock() @@ -390,7 +391,7 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create(): """Regression (spend loss): the batch create persists the creating key hash and tags so CheckBatchCost can write an attributed spend row instead of a blank one the DB drops.""" instance, store = _in_memory_managed_files() - creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice") + creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice", org_id="org-acme") await instance.store_unified_object_id( unified_object_id="unified-b", @@ -407,9 +408,70 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create(): assert row["api_key"] == "hash-alice" assert row["created_by"] == "alice" assert row["team_id"] == "team-alpha" + assert row["org_id"] == "org-acme" assert row["request_tags"].data == ["env:prod"] +@pytest.mark.asyncio +async def test_store_unified_object_id_resolves_org_through_the_cached_team(): + """Most keys belong to an org only through their team, so the auth object carries no + org_id. The create reads the team that auth already cached, so org spend is snapshotted + at submission time without a database query in the request path.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + from litellm.proxy.proxy_server import user_api_key_cache + + instance, store = _in_memory_managed_files() + creator = UserAPIKeyAuth(user_id="alice", team_id="team-cached", api_key="hash-alice") + await user_api_key_cache.async_set_cache( + key="team_id:team-cached", + value=LiteLLM_TeamTableCachedObj(team_id="team-cached", organization_id="org-via-team"), + model_type=LiteLLM_TeamTableCachedObj, + ) + try: + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + ) + finally: + user_api_key_cache.delete_cache(key="team_id:team-cached") + + assert store["unified-b"]["org_id"] == "org-via-team" + instance.prisma_client.db.litellm_teamtable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_store_unified_object_id_resolves_org_from_the_db_when_the_team_is_not_cached(): + """A team no request has run under yet is absent from the auth cache; its organization + still comes back from the table so the org is billed rather than dropped.""" + from litellm.models.team import LiteLLM_TeamTable + from litellm.proxy.proxy_server import user_api_key_cache + + instance, store = _in_memory_managed_files() + instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=LiteLLM_TeamTable(team_id="team-uncached", organization_id="org-via-db") + ) + creator = UserAPIKeyAuth(user_id="alice", team_id="team-uncached", api_key="hash-alice") + try: + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + ) + finally: + user_api_key_cache.delete_cache(key="team_id:team-uncached") + + assert store["unified-b"]["org_id"] == "org-via-db" + + @pytest.mark.asyncio async def test_store_unified_object_id_omits_key_and_tags_without_persist_attribution(): """Regression (spend redirect): a caller that is not the batch create (a poll, or the @@ -471,6 +533,7 @@ async def test_store_unified_object_id_attribution_columns_are_write_once(): upsert_data = instance.prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"] assert "api_key" not in upsert_data["update"] assert "request_tags" not in upsert_data["update"] + assert "org_id" not in upsert_data["update"] @pytest.mark.asyncio diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 091b958d7c3..48fceb50403 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -1095,6 +1095,110 @@ async def test_afile_content_passes_trusted_model_credentials_to_router(): assert trusted_credentials["s3_bucket_name"] == "my-bucket" +def _managed_deletion_file_id(provider_file_id): + from litellm.types.utils import SpecialEnums + + value = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "test-file", "batch-model", provider_file_id, "model-123" + ) + return base64.urlsafe_b64encode(value.encode()).decode().rstrip("=") + + +def _managed_files_with_deletion_row(unified_file_id, provider_file_id, file_object): + from litellm.caching import DualCache + from litellm.models.managed_files import LiteLLM_ManagedFileTable + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + row = LiteLLM_ManagedFileTable( + unified_file_id=unified_file_id, + model_mappings={"model-123": provider_file_id}, + flat_model_file_ids=[provider_file_id], + file_object=file_object, + ) + table = MagicMock( + find_first=AsyncMock(return_value=row), + delete=AsyncMock(), + ) + return _PROXY_LiteLLMManagedFiles( + internal_usage_cache=DualCache(), + prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=table)), + ), table + + +@pytest.mark.asyncio +async def test_afile_delete_bedrock_uses_deployment_bucket_and_signed_s3_delete(monkeypatch): + import httpx + import respx + + from litellm import Router + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + router = Router( + model_list=[ + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-bucket", + }, + "model_info": {"id": "model-123"}, + } + ], + num_retries=0, + ) + s3_uri = "s3://my-bucket/litellm-bedrock-files/input.jsonl" + unified_file_id = _managed_deletion_file_id(s3_uri) + managed_files, table = _managed_files_with_deletion_row(unified_file_id, s3_uri, None) + with respx.mock: + route = respx.delete( + "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/input.jsonl" + ).mock(return_value=httpx.Response(204)) + response = await managed_files.afile_delete( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=router, + _litellm_internal_model_credentials={"s3_bucket_name": "request-bucket"}, + ) + + assert len(route.calls) == 1 + assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert response.id == unified_file_id + assert response.deleted is True + table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id}) + + +@pytest.mark.asyncio +async def test_afile_delete_returns_managed_id_for_stored_provider_output(): + from openai.types import FileDeleted + + provider_file_id = "file-error-output" + unified_file_id = _managed_deletion_file_id(provider_file_id) + stored_file = _make_file_object(provider_file_id) + managed_files, table = _managed_files_with_deletion_row(unified_file_id, provider_file_id, stored_file) + router = MagicMock( + get_deployment_credentials_with_provider=MagicMock(return_value=None), + afile_delete=AsyncMock(return_value=FileDeleted(id=provider_file_id, object="file", deleted=True)), + ) + response = await managed_files.afile_delete( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=router, + _litellm_internal_model_credentials={"s3_bucket_name": "request-bucket"}, + ) + + assert response.id == unified_file_id + assert response.object == "file" + assert response.filename == stored_file.filename + assert stored_file.id == provider_file_id + router.afile_delete.assert_awaited_once_with(model="model-123", file_id=provider_file_id) + table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id}) + + @pytest.mark.asyncio async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch): """ diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 4db131da62c..3283af26cf7 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1,25 +1,33 @@ import asyncio import base64 +import json import os import sys +from collections.abc import AsyncIterator from importlib import metadata from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import anyio import httpx import pytest +import respx from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth from mcp import McpError +from mcp.client.streamable_http import streamable_http_client +from pydantic import ValidationError from mcp.shared.message import SessionMessage from mcp.types import ( LATEST_PROTOCOL_VERSION, + CallToolResult, ErrorData, Implementation, InitializeResult, JSONRPCError, JSONRPCMessage, JSONRPCResponse, + LoggingMessageNotificationParams, ServerCapabilities, ) @@ -29,8 +37,9 @@ import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import ( MCP_STREAMABLE_HTTP_REQUIREMENT, MCPClient, - _as_read_timeout, _first_non_cancelled_cause, + _TransportContext, + as_mcp_read_timeout, missing_streamable_http_client_error, strip_auth_scheme, ) @@ -859,25 +868,25 @@ def _raise_mcp_error_while_handling_a_timeout(code: int, message: str) -> McpErr return raised -def test_as_read_timeout_separates_the_sdk_timeout_from_a_relayed_upstream_error(): +def test_as_mcp_read_timeout_separates_the_sdk_timeout_from_a_relayed_upstream_error(): """Neither signal alone is enough. The code alone cannot separate the SDK's own timeout from an upstream JSON-RPC error that happens to use 408, and the context chain alone cannot separate it from any other relayed error that surfaces while a timeout is being handled, so both must hold. """ timeout_code = int(httpx.codes.REQUEST_TIMEOUT) - translated = _as_read_timeout(_raise_mcp_error_while_handling_a_timeout(timeout_code, "Timed out while waiting")) + translated = as_mcp_read_timeout(_raise_mcp_error_while_handling_a_timeout(timeout_code, "Timed out while waiting")) assert isinstance(translated, TimeoutError) assert str(translated) == "Timed out while waiting" relayed_408 = McpError(ErrorData(code=timeout_code, message="upstream said 408")) - assert _as_read_timeout(relayed_408) is None, "an upstream 408 with no elapsed timeout is not our timeout" + assert as_mcp_read_timeout(relayed_408) is None, "an upstream 408 with no elapsed timeout is not our timeout" relayed_other = _raise_mcp_error_while_handling_a_timeout(-32603, "upstream internal error") - assert _as_read_timeout(relayed_other) is None, "a non-timeout code is not our timeout, whatever the chain" + assert as_mcp_read_timeout(relayed_other) is None, "a non-timeout code is not our timeout, whatever the chain" - assert _as_read_timeout(McpError(ErrorData(code=-32603, message="boom"))) is None - assert _as_read_timeout(RuntimeError("not an McpError")) is None + assert as_mcp_read_timeout(McpError(ErrorData(code=-32603, message="boom"))) is None + assert as_mcp_read_timeout(RuntimeError("not an McpError")) is None @pytest.mark.asyncio @@ -1224,14 +1233,14 @@ def test_without_a_configured_slot_the_existing_precedence_is_unchanged(): _REDIRECT_CASES = [ - ("https://upstream.example.com/mcp", "https://upstream.example.com/other"), # same origin + ("https://upstream.example.com/mcp", "https://upstream.example.com/other"), # same origin ("https://upstream.example.com/mcp", "https://upstream.example.com:443/other"), # explicit default port - ("https://upstream.example.com/mcp", "https://attacker.example.com/collect"), # different host - ("https://upstream.example.com/mcp", "http://upstream.example.com/collect"), # scheme downgrade - ("https://upstream.example.com/mcp", "https://upstream.example.com:8443/other"), # different port - ("https://upstream.example.com/mcp", "https://sub.upstream.example.com/x"), # different host - ("http://upstream.example.com/mcp", "https://upstream.example.com/other"), # http -> https upgrade - ("http://upstream.example.com/mcp", "http://upstream.example.com/other"), # same origin, plain http + ("https://upstream.example.com/mcp", "https://attacker.example.com/collect"), # different host + ("https://upstream.example.com/mcp", "http://upstream.example.com/collect"), # scheme downgrade + ("https://upstream.example.com/mcp", "https://upstream.example.com:8443/other"), # different port + ("https://upstream.example.com/mcp", "https://sub.upstream.example.com/x"), # different host + ("http://upstream.example.com/mcp", "https://upstream.example.com/other"), # http -> https upgrade + ("http://upstream.example.com/mcp", "http://upstream.example.com/other"), # same origin, plain http ] @@ -1283,3 +1292,618 @@ def test_a_differently_cased_injected_header_cannot_shadow_the_slot() -> None: headers = client._get_auth_headers() assert [v for k, v in headers.items() if k.lower() == "esb-oauth"] == ["Bearer minted-token"] assert headers["X-Trace"] == "keep" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("content_type", "body", "expected_type"), + [ + ("text/html", b"secret-page", ValueError), + ("application/json", b"secret-invalid-json", ValidationError), + ("application/json", b"", ValidationError), + ("application/json", b'{"secret":"invalid-rpc"}', ValidationError), + ("application/json", b'{"jsonrpc":"2.0","id":0,"result":{"secret":"invalid-schema"}}', ValidationError), + ], +) +async def test_invalid_http_response_surfaces_without_waiting_for_timeout( + content_type: str, body: bytes, expected_type: type[Exception] +) -> None: + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": content_type}, content=body) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(expected_type) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + message: Final = _connection_error_message(caught.value, client.server_url, 30) + assert "unsupported content type" in message or "invalid MCP response" in message + assert "secret" not in message + assert "timed out" not in message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [200, 401, 503]) +async def test_http_response_handler_preserves_success_and_http_errors(status_code: int) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": []} + ) + return httpx.Response(status_code, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + operation: Final = client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ) + if status_code == 200: + result: Final = await asyncio.wait_for(operation, timeout=3) + assert result.tools == [] + else: + with pytest.raises(httpx.HTTPStatusError) as caught: + await asyncio.wait_for(operation, timeout=3) + assert caught.value.response.status_code == status_code + + +@pytest.mark.asyncio +async def test_http_response_handler_preserves_notifications_and_tool_listing() -> None: + notification: Final = { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": {"level": "info", "data": "Listing tools"}, + } + logging_callback: Final = AsyncMock() + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + if payload["method"] == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload["id"], + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"logging": {}, "tools": {}}, + "serverInfo": {"name": "test", "version": "1"}, + }, + }, + ) + response: Final = { + "jsonrpc": "2.0", + "id": payload["id"], + "result": {"tools": [{"name": "search", "inputSchema": {"type": "object"}}]}, + } + return httpx.Response( + 200, + headers={"Content-Type": "text/event-stream"}, + content="".join(f"event: message\ndata: {json.dumps(message)}\n\n" for message in (notification, response)), + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30, logging_callback=logging_callback) + result: Final = await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ), + timeout=3, + ) + + assert [tool.name for tool in result.tools] == ["search"] + logging_callback.assert_awaited_once_with(LoggingMessageNotificationParams(level="info", data="Listing tools")) + + +@pytest.mark.asyncio +async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response() -> None: + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": "secret-invalid-tools"} + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(ValidationError) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + message: Final = _connection_error_message(caught.value, client.server_url, 30) + assert "invalid MCP response" in message + assert "secret" not in message + + +class _DiagnosticSSEStream(httpx.AsyncByteStream): + def __init__(self, messages: asyncio.Queue[bytes | Exception | None]) -> None: + self.messages = messages + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"event: endpoint\ndata: /messages\n\n" + while True: + message: Final = await self.messages.get() + if message is None: + return + if isinstance(message, Exception): + raise message + yield b"event: message\ndata: " + message + b"\n\n" + + +_DIAGNOSTIC_STDIO_SERVER: Final = """ +import json, sys +mode, failure_method = sys.argv[1:] +for line in sys.stdin: + request = json.loads(line) + if "method" not in request or "id" not in request: + continue + if request["method"] == failure_method: + if mode == "bad-json": + print("secret-invalid-json", flush=True) + continue + if mode == "closed": + sys.exit(0) + if mode == "silent": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Waiting"}}), flush=True) + continue + if request["method"] == "initialize": + result = {"protocolVersion": request["params"]["protocolVersion"], "capabilities": {"tools": {}, "logging": {}}, "serverInfo": {"name": "diagnostic", "version": "1"}} + elif request["method"] == "tools/list": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Listing tools"}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "unmatched", "result": {}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}), flush=True) + result = {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + else: + result = {"content": [{"type": "text", "text": "pong"}], "isError": False} + print(json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": result}), flush=True) +""" + + +def _diagnostic_transport(transport: MCPTransport, mode: str, failure_method: str) -> _TransportContext: + from mcp import StdioServerParameters + from mcp.client.sse import sse_client + from mcp.client.stdio import stdio_client + + if transport == MCPTransport.stdio: + return stdio_client( + StdioServerParameters( + command=sys.executable, args=["-u", "-c", _DIAGNOSTIC_STDIO_SERVER, mode, failure_method] + ) + ) + messages: Final[asyncio.Queue[bytes | Exception | None]] = asyncio.Queue() + + async def respond(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, headers={"Content-Type": "text/event-stream"}, stream=_DiagnosticSSEStream(messages) + ) + payload: Final = json.loads(request.content) + if "method" not in payload or "id" not in payload: + return httpx.Response(202) + if payload["method"] == failure_method and mode != "ok": + if mode == "bad-json": + await messages.put(b"secret-invalid-json") + elif mode == "io-error": + await messages.put(httpx.ReadError("secret-read-error")) + elif mode == "closed": + await messages.put(None) + elif mode == "silent": + await messages.put( + b'{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Waiting"}}' + ) + return httpx.Response(202) + if payload["method"] == "tools/list": + for message in ( + { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": {"level": "info", "data": "Listing tools"}, + }, + {"jsonrpc": "2.0", "id": "unmatched", "result": {}}, + {"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}, + ): + await messages.put(json.dumps(message).encode()) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"tools": {}, "logging": {}}, + "serverInfo": {"name": "diagnostic", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + if payload["method"] == "tools/list" + else {"content": [{"type": "text", "text": "pong"}], "isError": False} + ) + await messages.put(json.dumps({"jsonrpc": "2.0", "id": payload["id"], "result": result}).encode()) + return httpx.Response(202) + + def factory( + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, + ) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(respond), headers=headers, timeout=timeout, auth=auth) + + return sse_client("https://example.com/sse", httpx_client_factory=factory) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("failure_method", ["initialize", "tools/list"]) +async def test_transport_parsing_failure_is_preserved(transport: MCPTransport, failure_method: str) -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=transport, timeout=0.2) + with pytest.raises(ValidationError): + await asyncio.wait_for( + client._execute_session_operation( + _diagnostic_transport(transport, "bad-json", failure_method), lambda session: session.list_tools() + ), + timeout=3, + ) + + +@pytest.mark.asyncio +async def test_sse_read_failure_is_preserved() -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=MCPTransport.sse, timeout=0.2) + with pytest.raises(httpx.ReadError, match="secret-read-error"): + await asyncio.wait_for( + client._execute_session_operation( + _diagnostic_transport(MCPTransport.sse, "io-error", "tools/list"), lambda session: session.list_tools() + ), + timeout=3, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("mode", ["ok", "closed", "silent"]) +async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str) -> None: + from mcp import ClientSession + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + logging_callback: Final = AsyncMock() + client: Final = MCPClient( + server_url="https://example.com/sse", transport_type=transport, timeout=0.2, logging_callback=logging_callback + ) + + async def operation(session: ClientSession) -> CallToolResult: + tools: Final = await session.list_tools() + assert [tool.name for tool in tools.tools] == ["ping"] + return await session.call_tool("ping", {}) + + pending: Final = client._execute_session_operation(_diagnostic_transport(transport, mode, "tools/list"), operation) + if mode == "ok": + result: Final = await asyncio.wait_for(pending, timeout=3) + assert result.isError is False + assert result.content[0].text == "pong" + logging_callback.assert_awaited_once_with(LoggingMessageNotificationParams(level="info", data="Listing tools")) + else: + with pytest.raises(McpError) as caught: + await asyncio.wait_for(pending, timeout=3) + if mode == "closed": + assert "connection was closed" in _connection_error_message(caught.value, client.server_url, 0.2) + else: + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +async def test_transport_cancellation_cleans_up_a_pending_request(transport: MCPTransport) -> None: + ready: Final = asyncio.Event() + + async def on_log(message: LoggingMessageNotificationParams) -> None: + if message.data == "Waiting": + ready.set() + + client: Final = MCPClient( + server_url="https://example.com/sse", transport_type=transport, timeout=30, logging_callback=on_log + ) + task: Final = asyncio.create_task( + client._execute_session_operation( + _diagnostic_transport(transport, "silent", "tools/list"), lambda session: session.list_tools() + ) + ) + try: + await asyncio.wait_for(ready.wait(), timeout=3) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) + + +class _InterruptedHTTPBody(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b'{"jsonrpc":' + raise httpx.RemoteProtocolError("secret-incomplete-response") + + +@pytest.mark.asyncio +async def test_interrupted_http_response_preserves_the_transport_failure() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": "application/json"}, stream=_InterruptedHTTPBody()) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(httpx.RemoteProtocolError, match="secret-incomplete-response"): + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + +@pytest.mark.asyncio +async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=b"") + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=0.2) + with pytest.raises(McpError) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +@pytest.mark.parametrize( + "outcome", + ( + "absent", + "other_capability", + "supported", + "method_not_found", + "internal_error", + "unauthorized", + "timeout", + "initialize_not_found", + ), +) +async def test_optional_discovery_capabilities_and_errors( + method: str, outcome: str, caplog: pytest.LogCaptureFixture +) -> None: + import logging + from unittest.mock import Mock + + from mcp.types import JSONRPCRequest + + capability: Final = "prompts" if method == "prompts/list" else "resources" + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + advertised: Final = "resources" if capability == "prompts" else "prompts" + entry: Final = { + "prompts/list": {"name": "example"}, + "resources/list": {"name": "example", "uri": "test://example"}, + "resources/templates/list": {"name": "example", "uriTemplate": "test://{name}"}, + }[method] + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + if outcome == "initialize_not_found": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "error": {"code": -32601, "message": "Initialization rejected"}, + }, + ) + if payload.method == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {} + if outcome == "absent" + else {advertised if outcome == "other_capability" else capability: {}}, + "serverInfo": {"name": "discovery", "version": "1"}, + }, + }, + ) + if outcome == "timeout": + raise httpx.ReadTimeout("Optional list timed out", request=request) + if outcome == "unauthorized": + return httpx.Response(401) + if outcome in ("method_not_found", "internal_error", "absent", "other_capability"): + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "error": { + "code": -32603 if outcome == "internal_error" else -32601, + "message": "Optional list rejected", + }, + }, + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry]}}) + + responder: Final = Mock(side_effect=respond) + caplog.set_level(logging.DEBUG, logger="LiteLLM") + with respx.mock(base_url="https://example.com") as router: + router.route().mock(side_effect=responder) + client: Final = MCPClient(server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + result: Final = await operation() + + requests: Final = tuple( + JSONRPCMessage.model_validate_json(call.args[0].content).root + for call in responder.call_args_list + if call.args[0].method == "POST" + ) + assert sum(isinstance(request, JSONRPCRequest) and request.method == method for request in requests) == ( + 0 if outcome in ("absent", "other_capability", "initialize_not_found") else 1 + ) + assert [item.name for item in result] == (["example"] if outcome == "supported" else []) + failures: Final = tuple( + record for record in caplog.records if record.name == "LiteLLM" and record.levelno >= logging.WARNING + ) + if outcome in ("internal_error", "unauthorized", "timeout", "initialize_not_found"): + assert any(record.levelno == logging.ERROR and "failed" in record.message for record in failures) + else: + assert failures == () + if outcome == "method_not_found": + assert any( + record.levelno == logging.DEBUG and "Optional list rejected" in record.message for record in caplog.records + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("supports_first", (True, False)) +async def test_optional_discovery_uses_each_sessions_capabilities(supports_first: bool) -> None: + from unittest.mock import Mock + from mcp.types import JSONRPCRequest + + capabilities: Final = iter(({"resources": {}}, {}) if supports_first else ({}, {"resources": {}})) + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": next(capabilities), + "serverInfo": {"name": "changing", "version": "1"}, + } + if payload.method == "initialize" + else {"resources": [{"name": "example", "uri": "test://example"}]} + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + + responder: Final = Mock(side_effect=respond) + with respx.mock(base_url="https://example.com") as router: + router.route().mock(side_effect=responder) + client: Final = MCPClient(server_url="https://example.com/mcp") + first: Final = await client.list_resources() + second: Final = await client.list_resources() + + assert [item.name for item in first] == (["example"] if supports_first else []) + assert [item.name for item in second] == ([] if supports_first else ["example"]) + requests: Final = tuple( + JSONRPCMessage.model_validate_json(call.args[0].content).root + for call in responder.call_args_list + if call.args[0].method == "POST" + ) + assert sum(isinstance(request, JSONRPCRequest) and request.method == "resources/list" for request in requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +async def test_optional_discovery_preserves_cancellation(method: str) -> None: + from mcp.types import JSONRPCRequest + + ready: Final = asyncio.Event() + pending: Final = asyncio.Event() + + async def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + if payload.method == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"resources": {}, "prompts": {}}, + "serverInfo": {"name": "pending", "version": "1"}, + }, + }, + ) + ready.set() + await pending.wait() + return httpx.Response(202) + + with respx.mock(base_url="https://example.com") as router: + router.route().mock(side_effect=respond) + client: Final = MCPClient(server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + task: Final = asyncio.create_task(operation()) + try: + await asyncio.wait_for(ready.wait(), timeout=3) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) + + + +def test_client_import_before_proxy_credentials_succeeds_in_fresh_process(): + import subprocess + + result = subprocess.run( + [sys.executable, "-c", "import litellm.experimental_mcp_client.client; from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager; print(MCPServerManager.__name__)"], + capture_output=True, text=True, timeout=60, check=False, + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "MCPServerManager" diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py index 0555447e34f..a2f81091893 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py @@ -541,24 +541,166 @@ def _span_json(logger_under_test: DataDogLLMObsLogger, payload: dict[str, Any]) return json.loads(safe_dumps(span)) +SECRET_TOOL_RESULT: Final = '{"city": "Paris", "temp_c": 18, "account_secret": "SECRET-7545"}' +TOOL_CONVERSATION: Final[list[dict[str, Any]]] = [ + {"role": "user", "content": "secret prompt"}, + {"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, + {"role": "tool", "tool_call_id": "call_abc123", "content": SECRET_TOOL_RESULT}, +] + + +def _redacted_span_as_the_proxy_builds_it(payload: dict[str, Any]) -> dict[str, Any]: + logger_under_test = _redacting_logger(turn_off_message_logging=True) + return _span_json( + logger_under_test, logger_under_test.redact_standard_logging_payload_from_model_call_details(payload) + ) + + def test_redaction_keeps_the_conversation_shape_without_its_content() -> None: - """Roles and message count survive so the trace stays legible; contents and tool payloads do not.""" - result = _span_json( - _redacting_logger(turn_off_message_logging=True), + result = _redacted_span_as_the_proxy_builds_it( + build_payload( + messages=TOOL_CONVERSATION, + response_message={"role": "assistant", "content": "secret response", "tool_calls": [ASSISTANT_TOOL_CALL]}, + ) + ) + + redacted_call = { + "name": "get_weather", + "arguments": "redacted-by-litellm", + "tool_id": "call_abc123", + "type": "function", + } + assert result["meta"]["input"]["messages"] == [ + {"role": "user", "content": "redacted-by-litellm"}, + {"role": "assistant", "content": "redacted-by-litellm", "tool_calls": [redacted_call]}, + { + "role": "tool", + "content": "redacted-by-litellm", + "tool_results": [ + {"name": "get_weather", "result": "redacted-by-litellm", "tool_id": "call_abc123", "type": "function"} + ], + }, + ] + assert result["meta"]["output"]["messages"] == [ + {"role": "assistant", "content": "redacted-by-litellm", "tool_calls": [redacted_call]} + ] + serialized = safe_dumps(result) + assert "SECRET-7545" not in serialized + assert "Paris" not in serialized + assert "secret" not in serialized + + +def test_redaction_counts_tool_result_tokens_before_replacing_them() -> None: + payload = build_payload(messages=TOOL_CONVERSATION) + payload["standard_logging_object"]["model"] = "claude-sonnet-5" + + result = _redacted_span_as_the_proxy_builds_it(payload) + + expected_tokens = litellm.token_counter(model="claude-sonnet-5", text=SECRET_TOOL_RESULT) + assert expected_tokens > 0 + assert result["metrics"]["tool_output_tokens"] == float(expected_tokens) + assert result["metrics"]["input_tokens"] == 4447.0 + + +def test_tool_output_tokens_sum_every_result_in_the_request(logger: DataDogLLMObsLogger) -> None: + payload = build( + logger, + messages=[ + {"role": "tool", "tool_call_id": "call_1", "content": "one two three"}, + {"role": "tool", "tool_call_id": "call_2", "content": "four five six seven"}, + ], + ) + + assert payload["metrics"]["tool_output_tokens"] == float( + litellm.token_counter(text="one two three") + litellm.token_counter(text="four five six seven") + ) + + +def test_a_request_without_tool_results_reports_no_tool_output_tokens(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, messages=[{"role": "user", "content": "hi"}]) + + assert "tool_output_tokens" not in payload["metrics"] + assert "tool_output_tokens" not in _redacted_span_as_the_proxy_builds_it(build_payload())["metrics"] + + +def test_a_tool_that_returned_nothing_still_counts_as_zero_tool_output_tokens(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, messages=[{"role": "tool", "tool_call_id": "call_1", "content": ""}]) + + assert payload["metrics"]["tool_output_tokens"] == 0.0 + + +def test_redaction_keeps_anthropic_tool_blocks_as_structure_only() -> None: + result = _redacted_span_as_the_proxy_builds_it( build_payload( messages=[ - {"role": "user", "content": "secret prompt"}, - {"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, - ], - response_message={"role": "assistant", "content": "secret response"}, - ), + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": SECRET_TOOL_RESULT}], + }, + ] + ) ) assert result["meta"]["input"]["messages"] == [ - {"role": "user", "content": "redacted-by-litellm"}, - {"role": "assistant", "content": "redacted-by-litellm"}, + { + "role": "assistant", + "content": "redacted-by-litellm", + "tool_calls": [ + {"name": "get_weather", "arguments": "redacted-by-litellm", "tool_id": "toolu_1", "type": "tool_use"} + ], + }, + { + "role": "user", + "content": "redacted-by-litellm", + "tool_results": [ + {"name": "get_weather", "result": "redacted-by-litellm", "tool_id": "toolu_1", "type": "function"} + ], + }, ] - assert result["meta"]["output"]["messages"] == [{"role": "assistant", "content": "redacted-by-litellm"}] + assert "Paris" not in safe_dumps(result) + + +def test_redaction_blanks_tool_identifiers_that_are_not_strings() -> None: + result = _redacted_span_as_the_proxy_builds_it( + build_payload( + messages=[ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": {"leak": "SECRET-7545"}, "type": ["SECRET-7545"], "function": {"name": ["SECRET-7545"]}} + ], + } + ] + ) + ) + + assert result["meta"]["input"]["messages"][0]["tool_calls"] == [ + {"name": "", "arguments": "redacted-by-litellm", "tool_id": "", "type": ""} + ] + assert "SECRET-7545" not in safe_dumps(result) + + +def test_the_shared_hook_still_strips_what_redaction_governs_besides_messages() -> None: + payload = build_payload(messages=TOOL_CONVERSATION) + payload["standard_logging_object"]["classifier_input"] = {"system": "SECRET-7545"} + logger_under_test = _redacting_logger(turn_off_message_logging=True) + + with patch.object( # test-quality-ok: the hook reads this module global with no injection seam + litellm, "standard_logging_payload_excluded_fields", ["response"] + ): + redacted = logger_under_test.redact_standard_logging_payload_from_model_call_details(payload) + + assert "classifier_input" not in redacted["standard_logging_object"] + assert "response" not in redacted["standard_logging_object"] + assert redacted["standard_logging_object"]["messages"] == TOOL_CONVERSATION + assert payload["standard_logging_object"]["classifier_input"] == {"system": "SECRET-7545"} def test_redaction_drops_unrecognized_and_malformed_message_roles() -> None: diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py b/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py index 9c75e0b0a47..a9074e89806 100644 --- a/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py +++ b/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py @@ -24,6 +24,8 @@ from litellm.types.integrations.newrelic import ( NEWRELIC_METRIC_PROMPT_TOKENS, NEWRELIC_METRIC_REQUEST_DURATION_MS, NEWRELIC_METRIC_REQUESTS, + NEWRELIC_METRIC_TEAM_MAX_BUDGET, + NEWRELIC_METRIC_TEAM_REMAINING_BUDGET, NEWRELIC_METRIC_TOTAL_TOKENS, NewRelicMetricRecord, ) @@ -40,6 +42,8 @@ def _record( completion_tokens=20, total_tokens=30, duration_ms=100.0, + team_max_budget=None, + team_spend=None, ) -> NewRelicMetricRecord: return NewRelicMetricRecord( team_id=team_id, @@ -53,12 +57,28 @@ def _record( completion_tokens=completion_tokens, total_tokens=total_tokens, duration_ms=duration_ms, + team_max_budget=team_max_budget, + team_spend=team_spend, ) -def _standard_logging_object(team_id="team-a", response_cost=0.25) -> dict: +def _standard_logging_object( + team_id="team-a", response_cost=0.25, team_max_budget: float | None = None, team_spend: float | None = None +) -> dict: + budget_metadata = { + key: value + for key, value in ( + ("user_api_key_team_max_budget", team_max_budget), + ("user_api_key_team_spend", team_spend), + ) + if value is not None + } return { - "metadata": {"user_api_key_team_id": team_id, "user_api_key_team_alias": f"{team_id}-alias"}, + "metadata": { + "user_api_key_team_id": team_id, + "user_api_key_team_alias": f"{team_id}-alias", + **budget_metadata, + }, "model_group": "gpt-4o-group", "model": "gpt-4o", "custom_llm_provider": "openai", @@ -204,6 +224,74 @@ class TestBuildMetricPayload: assert "model_group" not in attributes +class TestTeamBudgetGauges: + def test_latest_record_per_team_drives_one_gauge_pair(self): + records = ( + _record(team_id="team-a", model="gpt-4o", response_cost=0.5, team_max_budget=100.0, team_spend=10.0), + _record(team_id="team-a", model="claude-4", response_cost=2.0, team_max_budget=100.0, team_spend=10.5), + _record(team_id="team-b", response_cost=1.0, team_max_budget=None, team_spend=3.0), + _record(team_id="", team_alias="", response_cost=1.0, team_max_budget=50.0, team_spend=1.0), + ) + payload = build_metric_payload(records, window_start=1_000.0, now=1_005.0) + + max_budget_gauges = _metrics_by_name(payload, NEWRELIC_METRIC_TEAM_MAX_BUDGET) + remaining_gauges = _metrics_by_name(payload, NEWRELIC_METRIC_TEAM_REMAINING_BUDGET) + assert [(m["type"], m["value"], m["attributes"]) for m in max_budget_gauges] == [ + ("gauge", 100.0, {"team_id": "team-a", "team_alias": "team-a-alias"}) + ] + assert [(m["type"], m["attributes"]) for m in remaining_gauges] == [ + ("gauge", {"team_id": "team-a", "team_alias": "team-a-alias"}) + ] + assert remaining_gauges[0]["value"] == pytest.approx(100.0 - 10.5 - 2.0) + assert len(_metrics_by_name(payload, NEWRELIC_METRIC_COST_USD)) == 4 + + def test_missing_team_spend_counts_only_this_request(self): + payload = build_metric_payload( + (_record(response_cost=0.25, team_max_budget=10.0, team_spend=None),), window_start=1_000.0, now=1_005.0 + ) + + assert _metrics_by_name(payload, NEWRELIC_METRIC_TEAM_REMAINING_BUDGET)[0]["value"] == pytest.approx(9.75) + + @pytest.mark.asyncio + async def test_budget_gauges_reach_the_metric_api_from_standard_logging_metadata(self): + logger = _make_logger() + logger.async_client.post = AsyncMock(return_value=_response(202)) + slo = _standard_logging_object(response_cost=0.25, team_max_budget=20.0, team_spend=4.5) + + await logger.async_log_success_event( + kwargs={"standard_logging_object": slo}, response_obj={}, start_time=None, end_time=None + ) + await logger.flush_queue() + + body = json.loads(gzip.decompress(logger.async_client.post.await_args.kwargs["data"]).decode("utf-8")) + by_name = {m["name"]: m for m in body[0]["metrics"]} + assert by_name[NEWRELIC_METRIC_TEAM_MAX_BUDGET] == { + "name": NEWRELIC_METRIC_TEAM_MAX_BUDGET, + "type": "gauge", + "value": 20.0, + "attributes": {"team_id": "team-a", "team_alias": "team-a-alias"}, + } + assert by_name[NEWRELIC_METRIC_TEAM_REMAINING_BUDGET]["type"] == "gauge" + assert by_name[NEWRELIC_METRIC_TEAM_REMAINING_BUDGET]["value"] == pytest.approx(15.25) + assert by_name[NEWRELIC_METRIC_COST_USD]["value"] == 0.25 + + @pytest.mark.asyncio + async def test_no_budget_metadata_sends_no_gauges(self): + logger = _make_logger() + logger.async_client.post = AsyncMock(return_value=_response(202)) + + await logger.async_log_success_event( + kwargs={"standard_logging_object": _standard_logging_object()}, + response_obj={}, + start_time=None, + end_time=None, + ) + await logger.flush_queue() + + body = json.loads(gzip.decompress(logger.async_client.post.await_args.kwargs["data"]).decode("utf-8")) + assert {m["type"] for m in body[0]["metrics"]} == {"count", "summary"} + + class TestQueueAndFlush: @pytest.mark.asyncio async def test_log_event_queues_record_from_standard_logging_object(self): diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index 4aa28b5abfd..ae41c74944d 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -3,14 +3,21 @@ baggage helpers, metrics, the typed coercion helpers, mapper branches, span-name builders, and the registry validator's failure paths. Needs the OTel SDK.""" import json +import threading +from collections.abc import Iterator from dataclasses import replace +from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer import pytest pytest.importorskip("opentelemetry") +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( # noqa: E402 + ExportTraceServiceRequest, +) from opentelemetry.sdk.metrics import MeterProvider # noqa: E402 from opentelemetry.sdk.metrics.export import InMemoryMetricReader # noqa: E402 +from opentelemetry.sdk.trace import TracerProvider # noqa: E402 from opentelemetry.sdk.trace.export import ( # noqa: E402 BatchSpanProcessor, ConsoleSpanExporter, @@ -538,6 +545,175 @@ def test_build_span_exporter_variants(): assert "OTLPSpanExporter" in type(http_exporter).__name__ +def _export_one_trace_to_local_collector(exporter_kind: str) -> tuple[list[dict], tuple[int, int, int]]: + """Run a parent/child trace through the configured exporter against a + throwaway HTTP collector. Returns the requests as the collector saw them + (child first, since it ends first) and (trace_id, parent span_id, child span_id).""" + received: list[dict] = [] + + class Collector(BaseHTTPRequestHandler): + def do_POST(self): + body = self.rfile.read(int(self.headers["Content-Length"])) + received.append({"path": self.path, "headers": dict(self.headers), "body": body}) + self.send_response(200) + self.end_headers() + + def log_message(self, *_args): + pass + + server = HTTPServer(("127.0.0.1", 0), Collector) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + config = OpenTelemetryV2Config( + exporter=exporter_kind, + endpoint=f"http://127.0.0.1:{server.server_port}", + headers="x-collector-token=secret", + ) + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(providers.build_span_exporter(config))) + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("parent", kind=SpanKind.SERVER) as parent: + with tracer.start_as_current_span("child") as child: + ids = ( + parent.get_span_context().trace_id, + parent.get_span_context().span_id, + child.get_span_context().span_id, + ) + provider.shutdown() + finally: + server.shutdown() + server.server_close() + assert len(received) == 2 + return received, ids + + +def _only_span(request: dict) -> dict: + scope_spans = json.loads(request["body"])["resourceSpans"][0]["scopeSpans"][0]["spans"] + assert len(scope_spans) == 1 + return scope_spans[0] + + +def test_http_json_exporter_posts_otlp_json_to_traces_endpoint(): + """``http/json`` must put the OTLP/JSON mapping on the wire (camelCase + fields, integer enums, hex ids) with a JSON content type, so collectors that + cannot decode protobuf can ingest the trace. Headers still travel.""" + (child_request, parent_request), (trace_id, parent_id, child_id) = _export_one_trace_to_local_collector("http/json") + + assert parent_request["path"] == "/v1/traces" + assert parent_request["headers"]["Content-Type"] == "application/json" + assert parent_request["headers"]["x-collector-token"] == "secret" + parent = _only_span(parent_request) + assert parent["name"] == "parent" + assert parent["kind"] == 2 + assert parent["traceId"] == format(trace_id, "032x") + assert parent["spanId"] == format(parent_id, "016x") + assert "parentSpanId" not in parent + child = _only_span(child_request) + assert child["traceId"] == format(trace_id, "032x") + assert child["spanId"] == format(child_id, "016x") + assert child["parentSpanId"] == format(parent_id, "016x") + + +def test_http_protobuf_exporter_still_posts_protobuf(): + (_child_request, parent_request), (trace_id, _parent_id, _child_id) = _export_one_trace_to_local_collector( + "http/protobuf" + ) + + assert parent_request["path"] == "/v1/traces" + assert parent_request["headers"]["Content-Type"] == "application/x-protobuf" + assert format(trace_id, "032x").encode() not in parent_request["body"] + decoded = ExportTraceServiceRequest.FromString(parent_request["body"]) + span = decoded.resource_spans[0].scope_spans[0].spans[0] + assert span.name == "parent" + assert span.trace_id == trace_id.to_bytes(16, "big") + + +@pytest.fixture +def otlp_collector() -> Iterator[tuple[str, list[str]]]: + received_paths: list[str] = [] + + class RecordingHandler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + self.rfile.read(int(self.headers.get("Content-Length", "0"))) + received_paths.append(self.path) + self.send_response(200) + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), RecordingHandler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + yield f"http://127.0.0.1:{server.server_port}", received_paths + finally: + server.shutdown() + server.server_close() + + +def _export_one_span(cfg: OpenTelemetryV2Config) -> None: + provider = providers.build_tracer_provider(cfg) + provider.get_tracer("probe").start_span("probe").end() + assert provider.force_flush() + provider.shutdown() + + +def test_traces_endpoint_env_posts_to_the_configured_url_verbatim(monkeypatch, otlp_collector): + base_url, received_paths = otlp_collector + for var in ("OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_EXPORTER_OTLP_ENDPOINT"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("OTEL_ENDPOINT", f"{base_url}/services/collector") + monkeypatch.setenv("OTEL_TRACES_ENDPOINT", f"{base_url}/services/collector/traces") + + cfg = OpenTelemetryV2Config.from_env() + assert cfg.exporter == "otlp_http" + _export_one_span(cfg) + assert received_paths == ["/services/collector/traces"] + + +def test_traces_endpoint_alias_alone_implies_otlp_http(monkeypatch, otlp_collector): + base_url, received_paths = otlp_collector + for var in ("OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", f"{base_url}/custom/traces") + + cfg = OpenTelemetryV2Config.from_env() + assert cfg.exporter == "otlp_http" + _export_one_span(cfg) + assert received_paths == ["/custom/traces"] + + +def test_traces_endpoint_per_exporter_coexists_with_default_normalization(otlp_collector): + base_url, received_paths = otlp_collector + cfg = OpenTelemetryV2Config( + exporters=[ + {"kind": "otlp_http", "endpoint": base_url}, + { + "kind": "otlp_http", + "endpoint": f"{base_url}/services/collector", + "traces_endpoint": f"{base_url}/services/collector/traces", + }, + ] + ) + _export_one_span(cfg) + assert sorted(received_paths) == ["/services/collector/traces", "/v1/traces"] + + +def test_http_json_exporter_honors_traces_endpoint(otlp_collector): + base_url, received_paths = otlp_collector + cfg = OpenTelemetryV2Config( + exporters=[ + { + "kind": "http/json", + "endpoint": base_url, + "traces_endpoint": f"{base_url}/services/collector/traces", + } + ] + ) + _export_one_span(cfg) + assert received_paths == ["/services/collector/traces"] + + def test_otlp_metric_exporter_uses_cumulative_histogram_temporality(): """Histograms must export as cumulative, not delta. diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py b/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py new file mode 100644 index 00000000000..67695d5aed8 --- /dev/null +++ b/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py @@ -0,0 +1,2744 @@ +"""Key/team OTLP destinations override the operator's exporters for that backend.""" + +import contextvars +import time +from base64 import b64encode +from collections.abc import Mapping +from functools import reduce +from types import MappingProxyType + +import pytest +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import Status, StatusCode + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.otel import logger as otel_logger +from litellm.integrations.otel.logger import ( + OpenTelemetryV2, + build_otel_v2_logger, + fan_out_provider, + publish_global_otel_v2_provider, +) +from litellm.integrations.otel.model.config import ( + ExporterOwner, + ExporterSpec, + OpenTelemetryV2Config, + is_otel_v2_enabled, +) +from litellm.integrations.otel.model.destination import OtelDestination +from litellm.integrations.otel.plumbing import providers as otel_providers +from litellm.integrations.otel.plumbing.context import ( + destination_backends, + request_destinations, + set_request_destinations, +) +from litellm.integrations.otel.plumbing.providers import ( + TenantFanOutSpanProcessor, + _OverriddenBackendFilter, + _sink_key, + build_tracer_provider, + deliverable_destinations, + operator_sink_keys, +) +from litellm.integrations.otel.plumbing.routing import TenantTracerCache, get_tracer +from litellm.integrations.otel.presets.arize import arize_preset +from litellm.integrations.otel.presets.destinations import ( + destination_capable_backends, + destination_for, +) +from litellm.integrations.otel.presets.langfuse import langfuse_preset +from litellm.proxy._types import AddTeamCallback, UserAPIKeyAuth +from litellm.proxy.litellm_pre_call_utils import ( + convert_key_logging_metadata_to_callback, + resolve_tenant_otel_destinations, +) +from litellm.types.utils import StandardCallbackDynamicParams + +LANGFUSE_DEST = OtelDestination( + endpoint="http://tenant.local/api/public/otel", + headers={"Authorization": "Basic dGVuYW50"}, + callback_name="langfuse_otel", +) + + +@pytest.fixture +def allow_test_hosts(monkeypatch): + """A tenant-supplied host must be allowlisted by the operator. Allowlist the ones + these fixtures name so the resolution tests stay about resolution; + ``TestTenantHostSsrfGuard`` covers the guard itself.""" + monkeypatch.setattr( + litellm, "provider_url_destination_allowed_hosts", ["team.local", "key.local", "x"], raising=False + ) + + +@pytest.fixture(autouse=True) +def isolate_published_provider(monkeypatch): + """Publishing records the fan-out carrier in module state; one test's publish must + not become the next test's provider.""" + monkeypatch.setattr(otel_logger, "_published_v2_provider", None) + + +def in_fresh_context(fn, *args): + """Run ``fn`` in its own context so one test's destinations never leak.""" + return contextvars.copy_context().run(fn, *args) + + +def emit(provider: TracerProvider, name: str = "chat gpt-4") -> None: + with get_tracer(provider, "litellm").start_as_current_span(name): + pass + + +def wired_provider(dest_exporter: InMemorySpanExporter, global_exporter: InMemorySpanExporter) -> TracerProvider: + """The operator's provider: one owned exporter plus the tenant fan-out.""" + provider = TracerProvider() + provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(global_exporter), "langfuse_otel")) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) + return provider + + +class TestOverrideSuppression: + def test_operator_exporter_keeps_the_span_when_no_destination_is_resolved(self): + global_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = wired_provider(dest_exporter, global_exporter) + + in_fresh_context(emit, provider) + + assert [s.name for s in global_exporter.get_finished_spans()] == ["chat gpt-4"] + assert dest_exporter.get_finished_spans() == () + + def test_operator_exporter_is_skipped_once_the_backend_is_overridden(self): + global_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = wired_provider(dest_exporter, global_exporter) + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + emit(provider) + + in_fresh_context(run) + + assert global_exporter.get_finished_spans() == () + assert [s.name for s in dest_exporter.get_finished_spans()] == ["chat gpt-4"] + + def test_a_backend_the_request_did_not_override_still_exports(self): + arize_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(arize_exporter), "arize")) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + emit(provider) + + in_fresh_context(run) + + assert [s.name for s in arize_exporter.get_finished_spans()] == ["chat gpt-4"] + + +class TestRoutingMode: + """The operator's choice between replacing its own exporter and exporting alongside it. + + One org-wide backend across every team is a real deployment, and losing it the + moment a team configures its own is what ``additive`` exists to prevent. + """ + + OPERATOR_SINK = ("https://cloud.langfuse.com/api/public/otel/v1/traces", (("authorization", "Basic op"),)) + #: What a tenant destination for that same project looks like before normalizing: + #: no signal path yet, and the header name cased the way the backend writes it. + SAME_ACCOUNT_ENDPOINT = "https://cloud.langfuse.com/api/public/otel" + + @staticmethod + def _additive(monkeypatch): + monkeypatch.setattr(litellm, "otel_tenant_destination_mode", "additive", raising=False) + + @staticmethod + def _tree(provider): + tracer = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("auth /v1/chat/completions"): + pass + with tracer.start_as_current_span("chat gpt-4"): + pass + + def _run(self, provider, destinations=(LANGFUSE_DEST,)): + def run(): + set_request_destinations(destinations) + self._tree(provider) + + in_fresh_context(run) + + def test_global_only_keeps_every_span_and_delivers_to_nobody(self): + """No team destination resolved, so the operator's backbone is untouched.""" + global_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = wired_provider(dest_exporter, global_exporter) + + self._run(provider, destinations=()) + + assert len(global_exporter.get_finished_spans()) == 3 + assert dest_exporter.get_finished_spans() == () + + def test_team_only_gets_the_whole_tree_with_no_operator_exporter(self): + """A deployment with no operator credentials still gives the team its trace.""" + dest_exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) + + self._run(provider) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "auth /v1/chat/completions", + "chat gpt-4", + } + + def test_additive_gives_the_operator_and_the_team_the_same_tree(self, monkeypatch): + self._additive(monkeypatch) + global_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = wired_provider(dest_exporter, global_exporter) + + self._run(provider) + + names = {"POST /v1/chat/completions", "auth /v1/chat/completions", "chat gpt-4"} + assert {s.name for s in global_exporter.get_finished_spans()} == names + assert {s.name for s in dest_exporter.get_finished_spans()} == names + assert len(global_exporter.get_finished_spans()) == 3, "the operator must not get a span twice" + + def test_override_moves_the_tree_off_the_operator(self): + """The default, unchanged: the tenant's traffic reaches the tenant and nowhere else.""" + global_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = wired_provider(dest_exporter, global_exporter) + + self._run(provider) + + assert global_exporter.get_finished_spans() == () + assert len(dest_exporter.get_finished_spans()) == 3 + + def test_a_team_naming_the_operators_own_project_is_written_once(self, monkeypatch): + """Fanning out to two accounts is the point. Writing the same account twice + is a duplicate the operator would see in their own project.""" + self._additive(monkeypatch) + shared = InMemorySpanExporter() + same = OtelDestination( + endpoint=self.SAME_ACCOUNT_ENDPOINT, + headers=MappingProxyType({"Authorization": "Basic op"}), + callback_name="langfuse_otel", + ) + provider = TracerProvider() + provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(shared), "langfuse_otel")) + provider.add_span_processor( + TenantFanOutSpanProcessor( + processor_factory=lambda _d: SimpleSpanProcessor(shared), + operator_sinks=frozenset({self.OPERATOR_SINK}), + ) + ) + + self._run(provider, destinations=(same,)) + + assert len(shared.get_finished_spans()) == 3, "the same account received the trace twice" + + def test_in_override_a_team_naming_the_operators_project_still_gets_the_trace(self): + """Override suppresses the operator's own exporter, so the fan-out is the only + thing left delivering. Skipping it on a matching account leaves the team with + nothing at all.""" + shared = InMemorySpanExporter() + same = OtelDestination( + endpoint=self.SAME_ACCOUNT_ENDPOINT, + headers=MappingProxyType({"Authorization": "Basic op"}), + callback_name="langfuse_otel", + ) + provider = TracerProvider() + provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(shared), "langfuse_otel")) + provider.add_span_processor( + TenantFanOutSpanProcessor( + processor_factory=lambda _d: SimpleSpanProcessor(shared), + operator_sinks=frozenset({self.OPERATOR_SINK}), + ) + ) + + self._run(provider, destinations=(same,)) + + assert len(shared.get_finished_spans()) == 3, "the team's own destination received nothing" + + def test_a_team_naming_a_different_project_still_gets_its_copy(self, monkeypatch): + """The dedup keys on the account, so a second project is still a second copy.""" + self._additive(monkeypatch) + global_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(global_exporter), "langfuse_otel")) + provider.add_span_processor( + TenantFanOutSpanProcessor( + processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter), + operator_sinks=frozenset({self.OPERATOR_SINK}), + ) + ) + + self._run(provider) + + assert len(global_exporter.get_finished_spans()) == 3 + assert len(dest_exporter.get_finished_spans()) == 3 + + @pytest.mark.parametrize("additive", [True, False]) + def test_a_failing_team_destination_leaves_the_operator_alone(self, monkeypatch, additive): + """A tenant collector that raises on every span must not cost the operator + its own telemetry, nor take the request down with it.""" + if additive: + self._additive(monkeypatch) + global_exporter, arize_exporter = InMemorySpanExporter(), InMemorySpanExporter() + + class Exploding(SimpleSpanProcessor): + def on_end(self, span): + raise RuntimeError("tenant collector is down") + + provider = TracerProvider() + provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(global_exporter), "langfuse_otel")) + provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(arize_exporter), "arize")) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: Exploding(InMemorySpanExporter())) + ) + + self._run(provider) + + assert len(arize_exporter.get_finished_spans()) == 3, "an unrelated backend lost spans" + assert len(global_exporter.get_finished_spans()) == (3 if additive else 0) + + def test_the_env_var_turns_additive_on_without_a_config_file(self, monkeypatch): + monkeypatch.setattr(litellm, "otel_tenant_destination_mode", None, raising=False) + monkeypatch.setenv("LITELLM_OTEL_TENANT_DESTINATION_MODE", "Additive") + global_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = wired_provider(dest_exporter, global_exporter) + + self._run(provider) + + assert len(global_exporter.get_finished_spans()) == 3 + assert len(dest_exporter.get_finished_spans()) == 3 + + def test_an_unrecognized_mode_stays_on_override(self, monkeypatch): + monkeypatch.setattr(litellm, "otel_tenant_destination_mode", "both", raising=False) + global_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = wired_provider(dest_exporter, global_exporter) + + self._run(provider) + + assert global_exporter.get_finished_spans() == () + + def test_operator_sink_keys_skips_an_exporter_with_no_endpoint_of_its_own(self): + """Such an exporter resolves its endpoint from the environment at export + time, so it has no identity to compare a destination against.""" + config = OpenTelemetryV2Config( + exporters=( + ExporterSpec(kind="otlp_http", endpoint=self.OPERATOR_SINK[0], headers="authorization=Basic op"), + ExporterSpec(kind="otlp_http", endpoint=None, headers="authorization=Basic other"), + ) + ) + + assert operator_sink_keys(config) == frozenset({self.OPERATOR_SINK}) + + def test_operator_sink_keys_skips_exporters_that_never_reach_the_wire(self): + """A console kind ignores the endpoint and a header-gated spec with no + credentials is dropped when the provider is built, so treating either as an + account the operator writes to would silently withhold a team's own spans + under additive.""" + config = OpenTelemetryV2Config( + exporters=( + ExporterSpec(kind="otlp_http", endpoint=self.OPERATOR_SINK[0], headers="authorization=Basic op"), + ExporterSpec(kind="console", endpoint="http://team.local/v1/traces"), + ExporterSpec(kind="otlp_http", endpoint="http://gated.local/v1/traces", requires_headers=True), + ) + ) + + assert operator_sink_keys(config) == frozenset({self.OPERATOR_SINK}) + + def test_operator_sink_keys_spans_every_config_it_is_handed(self): + first = OpenTelemetryV2Config( + exporters=( + ExporterSpec( + kind="otlp_http", + endpoint=self.OPERATOR_SINK[0], + headers="authorization=Basic op", + ), + ) + ) + second = OpenTelemetryV2Config( + exporters=( + ExporterSpec( + kind="otlp_http", + endpoint="https://otlp.arize.com/v1/traces", + headers="space_id=s,api_key=k", + ), + ) + ) + + assert operator_sink_keys(first, second) == { + self.OPERATOR_SINK, + _sink_key("https://otlp.arize.com/v1/traces", {"space_id": "s", "api_key": "k"}), + } + + def test_a_team_pointing_at_a_credential_less_operator_exporter_still_gets_its_spans(self, monkeypatch): + """Under additive the fan-out skips a destination the operator already writes + to. An exporter the provider never built writes nothing, so skipping it would + cost the team every span.""" + monkeypatch.setenv("LITELLM_OTEL_TENANT_DESTINATION_MODE", "additive") + gated_endpoint = "http://gated.local/v1/traces" + destination = OtelDestination(endpoint=gated_endpoint, callback_name="newrelic") + config = OpenTelemetryV2Config( + exporters=(ExporterSpec(kind="otlp_http", endpoint=gated_endpoint, requires_headers=True),) + ) + dest_exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor( + TenantFanOutSpanProcessor( + processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter), + operator_sinks=operator_sink_keys(config), + ) + ) + + def run(): + set_request_destinations((destination,)) + emit(provider) + + in_fresh_context(run) + + assert [s.name for s in dest_exporter.get_finished_spans()] == ["chat gpt-4"] + + def test_the_operators_own_langfuse_and_a_team_naming_it_are_one_account(self, monkeypatch): + """The two sides are built by different code that writes the endpoint and the + header names differently, so comparing them raw silently never matches.""" + monkeypatch.setenv("LANGFUSE_HOST", "https://lf.internal") + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-op") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-op") + monkeypatch.setattr(litellm, "provider_url_destination_allowed_hosts", ["lf.internal"], raising=False) + operator = operator_sink_keys(langfuse_preset()) + + def sink(public_key, secret_key): + destination = destination_for( + "langfuse_otel", + StandardCallbackDynamicParams( + langfuse_public_key=public_key, + langfuse_secret_key=secret_key, + langfuse_host="https://lf.internal", + ), + ) + assert destination is not None + return _sink_key(destination.endpoint, destination.headers) + + assert sink("pk-op", "sk-op") in operator, "a team naming the operator's own project" + assert sink("pk-team", "sk-team") not in operator, "a different project on the same server" + + def test_two_accounts_holding_the_same_strings_in_different_roles_are_not_one(self): + """The values alone are not the identity. Two accounts can hold the same pair + of strings with the space id and the api key the other way round, and folding + them together would leave the second one's team with no trace at all.""" + endpoint = "https://otlp.arize.com/v1" + + assert _sink_key(endpoint, {"space_id": "a", "api_key": "b"}) != _sink_key( + endpoint, {"space_id": "b", "api_key": "a"} + ) + + def test_the_operators_own_arize_space_and_a_team_naming_it_are_one_account(self, monkeypatch): + """One account answers to two header names here: the operator's exporter sends + ``space_id`` and a team destination sends ``arize-space-id``. Keyed on the names, + additive would write the operator's own space twice for every request.""" + monkeypatch.setenv("ARIZE_SPACE_ID", "space-op") + monkeypatch.setenv("ARIZE_API_KEY", "key-op") + monkeypatch.delenv("ARIZE_SPACE_KEY", raising=False) + operator = operator_sink_keys(arize_preset()) + + def sink(space, api_key): + destination = destination_for( + "arize", + StandardCallbackDynamicParams(arize_space_key=space, arize_api_key=api_key), + ) + assert destination is not None + return _sink_key(destination.endpoint, destination.headers) + + assert sink("space-op", "key-op") in operator, "a team naming the operator's own space" + assert sink("space-team", "key-team") not in operator, "a different Arize space" + + +class TestFanOut: + def test_every_span_of_the_request_reaches_the_destination_in_one_trace(self): + """The whole tree, gen-AI span included, parented as the operator would see it.""" + dest_exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) + tracer = get_tracer(provider, "litellm") + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("auth /v1/chat/completions"): + pass + with tracer.start_as_current_span("chat gpt-4"): + pass + + in_fresh_context(run) + + spans = dest_exporter.get_finished_spans() + by_name = {s.name: s for s in spans} + assert set(by_name) == {"POST /v1/chat/completions", "auth /v1/chat/completions", "chat gpt-4"} + root = by_name["POST /v1/chat/completions"] + assert len({s.context.trace_id for s in spans}) == 1, "the tenant must receive one connected trace" + for child in ("auth /v1/chat/completions", "chat gpt-4"): + assert by_name[child].parent.span_id == root.context.span_id + + def test_a_team_naming_two_backends_gets_the_trace_at_both(self): + """The fan-out rides one provider, so it cannot skip a destination on the + grounds that some other backend owns it: nothing else would deliver it.""" + langfuse, arize = InMemorySpanExporter(), InMemorySpanExporter() + by_endpoint = {"http://a.local": langfuse, "http://b.local": arize} + provider = TracerProvider() + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda d: SimpleSpanProcessor(by_endpoint[d.endpoint])) + ) + + def run(): + set_request_destinations( + ( + OtelDestination(endpoint="http://a.local", callback_name="langfuse_otel"), + OtelDestination(endpoint="http://b.local", callback_name="arize"), + ) + ) + emit(provider) + + in_fresh_context(run) + + assert [s.name for s in langfuse.get_finished_spans()] == ["chat gpt-4"] + assert [s.name for s in arize.get_finished_spans()] == ["chat gpt-4"] + + def test_a_destination_carries_the_tenants_service_name(self): + """An overridden backend skips per-request tracer routing, so the service name + that route used to apply has to travel on the destination instead.""" + dest = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest))) + + def run(): + set_request_destinations( + ( + OtelDestination( + endpoint="http://a.local", + callback_name="langfuse_otel", + resource_attributes={"service.name": "team-checkout"}, + ), + ) + ) + emit(provider) + + in_fresh_context(run) + + assert {s.resource.attributes["service.name"] for s in dest.get_finished_spans()} == {"team-checkout"} + + def test_the_operators_database_endpoint_does_not_ride_along_to_the_tenant(self): + """A database span describes the proxy's own Postgres, so the tenant gets the + span and its timing without the host, the port, the schema or the error text + that names them. The operator's own copy keeps everything.""" + dest_exporter, operator_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(operator_exporter)) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) + tracer = get_tracer(provider, "litellm") + unreachable = "Can't reach database server at db.internal.example:15400" + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + with tracer.start_as_current_span("postgres get_data") as db_span: + db_span.set_attributes( + { + "db.system.name": "postgresql", + "db.system": "postgresql", + "db.operation.name": "get_data", + "server.address": "db.internal.example", + "server.port": 15400, + "db.namespace": "litellm", + "error.type": "PrismaError", + "error.message": unreachable, + "error": unreachable, + "litellm.provider.error.stack_trace": f"Traceback: {unreachable}", + } + ) + db_span.add_event("exception", {"exception.message": unreachable}) + db_span.set_status(Status(StatusCode.ERROR, unreachable)) + with tracer.start_as_current_span("chat claude-haiku") as llm_span: + llm_span.set_attribute("server.address", "api.anthropic.com") + + in_fresh_context(run) + + tenant = {s.name: s for s in dest_exporter.get_finished_spans()} + operator = {s.name: s for s in operator_exporter.get_finished_spans()} + assert set(tenant) == {"postgres get_data", "chat claude-haiku"}, "the tenant keeps the whole tree" + tenant_db = tenant["postgres get_data"] + assert dict(tenant_db.attributes) == { + "db.system.name": "postgresql", + "db.system": "postgresql", + "db.operation.name": "get_data", + "error.type": "PrismaError", + } + assert list(tenant_db.events) == [] + assert tenant_db.status.status_code is StatusCode.ERROR, "the tenant still sees that the call failed" + assert tenant_db.status.description is None + assert "db.internal.example" not in tenant_db.to_json() + assert tenant["chat claude-haiku"].attributes["server.address"] == "api.anthropic.com", ( + "only the operator's datastore is redacted, never the model endpoint" + ) + operator_db = operator["postgres get_data"] + assert operator_db.attributes["server.address"] == "db.internal.example" + assert operator_db.attributes["server.port"] == 15400 + assert operator_db.attributes["db.namespace"] == "litellm" + assert operator_db.attributes["error.message"] == unreachable + assert operator_db.attributes["error"] == unreachable + assert operator_db.status.description == unreachable + assert [event.name for event in operator_db.events] == ["exception"] + + @pytest.mark.parametrize("failure_status", ["guardrail_failed_to_respond", "failure"]) + def test_a_guardrails_failure_text_does_not_ride_along_to_the_tenant(self, failure_status): + dest_exporter, operator_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(operator_exporter)) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) + tracer = get_tracer(provider, "litellm") + unreachable = "Cannot connect to host guardrail.internal.example:9000" + verdict = '{"action": "block", "categories": ["pii"]}' + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("execute_guardrail pii") as down: + down.set_attributes( + { + "litellm.guardrail.name": "pii", + "litellm.guardrail.status": failure_status, + "litellm.guardrail.response": unreachable, + } + ) + with tracer.start_as_current_span("execute_guardrail toxicity") as up: + up.set_attributes( + { + "litellm.guardrail.name": "toxicity", + "litellm.guardrail.status": "guardrail_intervened", + "litellm.guardrail.response": verdict, + } + ) + + in_fresh_context(run) + + tenant = {s.name: s for s in dest_exporter.get_finished_spans()} + operator = {s.name: s for s in operator_exporter.get_finished_spans()} + assert dict(tenant["execute_guardrail pii"].attributes) == { + "litellm.guardrail.name": "pii", + "litellm.guardrail.status": failure_status, + } + assert "guardrail.internal.example" not in tenant["execute_guardrail pii"].to_json() + assert tenant["execute_guardrail toxicity"].attributes["litellm.guardrail.response"] == verdict + assert operator["execute_guardrail pii"].attributes["litellm.guardrail.response"] == unreachable + + def test_the_callers_key_in_the_query_string_does_not_ride_along_to_the_tenant(self): + """A Google AI Studio style request authenticates with ``?key=``, + and the instrumentor stamps the full request URL on the server span. The + tenant keeps the URL up to the query string, and the operator's copy keeps it + whole.""" + dest_exporter, operator_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(operator_exporter)) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) + tracer = get_tracer(provider, "litellm") + path = "/v1beta/models/gemini-2.5-flash:generateContent" + query = "key=sk-another-members-virtual-key&alt=sse" + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + with tracer.start_as_current_span(f"POST {path}") as server_span: + server_span.set_attributes( + { + "http.method": "POST", + "http.route": path, + "http.target": f"{path}?{query}", + "http.url": f"http://proxy.example:4000{path}?{query}", + "url.path": path, + "url.query": query, + "http.status_code": 200, + } + ) + with tracer.start_as_current_span("generate_content gemini-2.5-flash") as llm_span: + llm_span.set_attributes( + { + "gen_ai.operation.name": "generate_content", + "url.full": f"https://generativelanguage.googleapis.com{path}?key=AIza-operator-provider-key", + } + ) + + in_fresh_context(run) + + tenant = {s.name: s for s in dest_exporter.get_finished_spans()} + assert dict(tenant[f"POST {path}"].attributes) == { + "http.method": "POST", + "http.route": path, + "http.target": path, + "http.url": f"http://proxy.example:4000{path}", + "url.path": path, + "http.status_code": 200, + } + assert "sk-another-members-virtual-key" not in tenant[f"POST {path}"].to_json() + assert tenant["generate_content gemini-2.5-flash"].attributes["url.full"] == ( + f"https://generativelanguage.googleapis.com{path}" + ), "the tenant's own span keeps its error text, and still loses a query string" + operator = {s.name: s for s in operator_exporter.get_finished_spans()} + assert operator[f"POST {path}"].attributes["http.url"] == f"http://proxy.example:4000{path}?{query}" + assert operator[f"POST {path}"].attributes["url.query"] == query + assert "AIza-operator-provider-key" in operator["generate_content gemini-2.5-flash"].to_json() + + def test_captured_request_headers_do_not_ride_along_to_the_tenant(self): + """With ``OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST`` set, the + server span carries the caller's bearer token. A team admin's collector must + not receive it, while the operator's own copy keeps it and the tenant keeps the + rest of the span.""" + dest_exporter, operator_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(operator_exporter)) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) + tracer = get_tracer(provider, "litellm") + bearer = "Bearer sk-another-members-virtual-key" + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + with tracer.start_as_current_span("POST /v1/chat/completions") as server_span: + server_span.set_attributes( + { + "http.request.method": "POST", + "http.route": "/v1/chat/completions", + "http.request.header.authorization": (bearer,), + "http.request.header.x_litellm_api_key": (bearer,), + "http.response.header.set_cookie": ("session=abc",), + } + ) + server_span.set_status(Status(StatusCode.ERROR)) + + in_fresh_context(run) + + tenant = dest_exporter.get_finished_spans()[0] + assert dict(tenant.attributes) == {"http.request.method": "POST", "http.route": "/v1/chat/completions"} + assert bearer not in tenant.to_json() + assert tenant.status.status_code is StatusCode.ERROR + operator = operator_exporter.get_finished_spans()[0] + assert operator.attributes["http.request.header.authorization"] == (bearer,) + assert operator.attributes["http.response.header.set_cookie"] == ("session=abc",) + + def test_the_proxys_own_error_text_does_not_ride_along_to_the_tenant(self): + """Postgres failing during auth surfaces as a ``ProxyException`` whose message + quotes the Prisma error, so the auth span and the request root carry the + operator's database endpoint in ``error.message``, in the exception event and + in the status description. None of it is the tenant's, so it all comes off, + while the failure itself (its type, its code, its status) stays. The tenant's + own model call keeps its error text, less the stack trace that walks the + operator's install. The operator's copy keeps everything.""" + dest_exporter, operator_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(operator_exporter)) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) + tracer = get_tracer(provider, "litellm") + unreachable = "Authentication Error, Can't reach database server at db.internal.example:15400" + install = "/srv/litellm/.venv/lib/python3.13/site-packages/opentelemetry/trace/__init__.py" + provider_error = "AnthropicException - invalid x-api-key" + + def fail(span, message: str) -> None: + span.set_attributes( + { + "error.type": "ProxyException", + "error.message": message, + "litellm.provider.error.code": "500", + "litellm.provider.error.stack_trace": f"Traceback\n File {install}\n{message}", + } + ) + span.add_event( + "exception", + {"exception.type": "ProxyException", "exception.message": message, "exception.stacktrace": install}, + ) + span.set_status(Status(StatusCode.ERROR, message)) + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + with tracer.start_as_current_span("auth /v1/chat/completions") as auth: + fail(auth, unreachable) + with tracer.start_as_current_span("chat claude-haiku") as llm: + llm.set_attribute("gen_ai.operation.name", "chat") + fail(llm, provider_error) + fail(root, unreachable) + + in_fresh_context(run) + + tenant = {s.name: s for s in dest_exporter.get_finished_spans()} + operator = {s.name: s for s in operator_exporter.get_finished_spans()} + assert set(tenant) == {"POST /v1/chat/completions", "auth /v1/chat/completions", "chat claude-haiku"} + for name in ("POST /v1/chat/completions", "auth /v1/chat/completions"): + proxy_span = tenant[name] + assert dict(proxy_span.attributes) == {"error.type": "ProxyException", "litellm.provider.error.code": "500"} + assert list(proxy_span.events) == [] + assert proxy_span.status.status_code is StatusCode.ERROR + assert proxy_span.status.description is None + assert "db.internal.example" not in proxy_span.to_json() + assert install not in proxy_span.to_json() + llm_span = tenant["chat claude-haiku"] + assert llm_span.attributes["error.message"] == provider_error, "the tenant's own call keeps its error text" + assert "litellm.provider.error.stack_trace" not in llm_span.attributes + assert llm_span.status.description == provider_error + assert [dict(event.attributes) for event in llm_span.events] == [ + {"exception.type": "ProxyException", "exception.message": provider_error} + ] + assert install not in llm_span.to_json() + for name, message in (("auth /v1/chat/completions", unreachable), ("chat claude-haiku", provider_error)): + assert operator[name].attributes["error.message"] == message + assert install in operator[name].attributes["litellm.provider.error.stack_trace"] + assert operator[name].events[0].attributes["exception.stacktrace"] == install + assert operator[name].status.description == message + + def test_a_tenants_service_name_is_layered_onto_the_operators_resource(self): + """The destination's ``service.name`` replaces the operator's on the tenant's + copy and every other resource attribute travels unchanged. Nothing is detected + afresh per span, so no attribute the operator did not configure appears.""" + dest = InMemorySpanExporter() + provider = TracerProvider( + resource=Resource({"service.name": "litellm-proxy", "deployment.environment.name": "prod"}) + ) + provider.add_span_processor(TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest))) + + def run(): + set_request_destinations( + ( + OtelDestination( + endpoint="http://a.local", + callback_name="langfuse_otel", + resource_attributes={"service.name": "team-checkout"}, + ), + ) + ) + emit(provider) + + in_fresh_context(run) + + (span,) = dest.get_finished_spans() + assert dict(span.resource.attributes) == { + "service.name": "team-checkout", + "deployment.environment.name": "prod", + } + + def test_a_destination_that_cannot_build_a_processor_is_skipped_quietly(self): + """An unbuildable destination must not cost the caller its request.""" + attempts = [] + reached_the_end = [] + + def factory(destination): + attempts.append(destination.endpoint) + + provider = TracerProvider() + provider.add_span_processor(TenantFanOutSpanProcessor(processor_factory=factory)) + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + emit(provider) + reached_the_end.append(True) + + in_fresh_context(run) + + assert attempts == [LANGFUSE_DEST.endpoint] + assert reached_the_end == [True] + + def test_an_unbuildable_destination_leaves_the_span_with_the_operator(self): + """Anchoring the destination is what makes the operator's exporter stand down + for the backend, so a destination nothing can deliver to must never be anchored, + or the span reaches neither account.""" + global_exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(global_exporter), "langfuse_otel")) + provider.add_span_processor(TenantFanOutSpanProcessor(processor_factory=lambda _d: None)) + + def run(): + set_request_destinations(deliverable_destinations((LANGFUSE_DEST,), provider)) + emit(provider) + return request_destinations() + + anchored = in_fresh_context(run) + + assert anchored == () + assert [s.name for s in global_exporter.get_finished_spans()] == ["chat gpt-4"] + + def test_a_buildable_destination_is_still_anchored_and_still_overrides(self): + global_exporter, dest_exporter = InMemorySpanExporter(), InMemorySpanExporter() + provider = wired_provider(dest_exporter, global_exporter) + + def run(): + set_request_destinations(deliverable_destinations((LANGFUSE_DEST,), provider)) + emit(provider) + return request_destinations() + + anchored = in_fresh_context(run) + + assert anchored == (LANGFUSE_DEST,) + assert global_exporter.get_finished_spans() == () + assert [s.name for s in dest_exporter.get_finished_spans()] == ["chat gpt-4"] + + def test_only_the_unbuildable_destination_is_dropped_from_a_mixed_set(self): + dest_exporter = InMemorySpanExporter() + other = LANGFUSE_DEST.model_copy(update={"endpoint": "http://broken.local/otel"}) + fan_out = TenantFanOutSpanProcessor( + processor_factory=lambda d: None if d.endpoint == other.endpoint else SimpleSpanProcessor(dest_exporter) + ) + + assert fan_out.deliverable((other, LANGFUSE_DEST)) == (LANGFUSE_DEST,) + + def test_no_fan_out_means_nothing_is_anchored(self): + """With nothing to carry the spans to the tenant, anchoring would only stop the + operator's exporter from writing them.""" + provider = TracerProvider() + + assert deliverable_destinations((LANGFUSE_DEST,), provider) == () + + def test_a_protocol_with_no_otlp_transport_is_not_deliverable(self): + """An unknown exporter kind falls back to the console exporter, which ignores the + tenant's credentials and prints its spans to the proxy's stdout. Treating that as + deliverable would stand the operator's exporter down for spans nobody stores.""" + typo = LANGFUSE_DEST.model_copy(update={"protocol": "consle"}) + fan_out = TenantFanOutSpanProcessor() + try: + assert fan_out.deliverable((typo, LANGFUSE_DEST)) == (LANGFUSE_DEST,) + finally: + fan_out.shutdown() + + def test_a_closed_fan_out_anchors_nothing(self): + fan_out = TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(InMemorySpanExporter())) + provider = TracerProvider() + provider.add_span_processor(fan_out) + fan_out.shutdown() + + assert deliverable_destinations((LANGFUSE_DEST,), provider) == () + + def test_the_processor_built_to_check_deliverability_is_the_one_that_exports(self): + built = [] + + def factory(_destination): + built.append(SimpleSpanProcessor(InMemorySpanExporter())) + return built[-1] + + provider = TracerProvider() + provider.add_span_processor(TenantFanOutSpanProcessor(processor_factory=factory)) + + def run(): + set_request_destinations(deliverable_destinations((LANGFUSE_DEST,), provider)) + emit(provider) + + in_fresh_context(run) + + assert len(built) == 1 + + def test_one_processor_is_reused_across_spans_of_the_same_destination(self): + built = [] + + def factory(_destination): + processor = SimpleSpanProcessor(InMemorySpanExporter()) + built.append(processor) + return processor + + provider = TracerProvider() + provider.add_span_processor(TenantFanOutSpanProcessor(processor_factory=factory)) + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + emit(provider, "one") + emit(provider, "two") + + in_fresh_context(run) + + assert len(built) == 1 + + +class TestProviderWiring: + def test_build_tracer_provider_only_filters_when_asked(self): + config = OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory", owner=ExporterOwner.LANGFUSE_OTEL)]) + operator = build_tracer_provider(config, tenant_overrides=True) + tenant = build_tracer_provider(config) + + def kinds(provider): + return [type(p).__name__ for p in provider._active_span_processor._span_processors] + + assert "_OverriddenBackendFilter" in kinds(operator) + assert "_OverriddenBackendFilter" not in kinds(tenant), "a per-tenant provider must not filter itself out" + assert "TenantFanOutSpanProcessor" not in kinds(operator), "delivery belongs to the published global alone" + assert "TenantFanOutSpanProcessor" not in kinds(tenant) + + def test_only_the_published_global_provider_delivers_to_tenants(self): + """A second v2 logger's provider never sees the server, auth or database spans, + so fanning out from it would hand the tenant a one-span trace. Publishing is + what picks the one provider the whole request tree passes through.""" + config = OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory", owner=ExporterOwner.ARIZE_AX)]) + published, other = OpenTelemetryV2(config=config, callback_name="arize"), OpenTelemetryV2(config=config) + + publish_global_otel_v2_provider([other], lambda _p: None, registered=published) + + def kinds(logger): + return [type(p).__name__ for p in logger._tracer_provider._active_span_processor._span_processors] + + assert kinds(published).count("TenantFanOutSpanProcessor") == 1 + assert "TenantFanOutSpanProcessor" not in kinds(other) + + @pytest.mark.parametrize("canonical", ["langfuse_otel", "arize"]) + def test_publishing_tells_the_fan_out_about_every_v2_loggers_account(self, monkeypatch, canonical): + monkeypatch.setenv("LITELLM_OTEL_TENANT_DESTINATION_MODE", "additive") + shared = InMemorySpanExporter() + monkeypatch.setattr(otel_providers, "_destination_processor", lambda _d: SimpleSpanProcessor(shared)) + accounts = { + "langfuse_otel": ( + "https://cloud.langfuse.com/api/public/otel/v1/traces", + "authorization=Basic op", + ), + "arize": ( + "https://otlp.arize.com/v1/traces", + "space_id=space-op,api_key=key-op", + ), + } + loggers = { + name: OpenTelemetryV2( + config=OpenTelemetryV2Config( + exporters=(ExporterSpec(kind="otlp_http", endpoint=endpoint, headers=headers),) + ), + callback_name=name, + tracer_provider=TracerProvider(), + ) + for name, (endpoint, headers) in accounts.items() + } + other = "arize" if canonical == "langfuse_otel" else "langfuse_otel" + published = publish_global_otel_v2_provider( + [loggers[other]], + lambda _p: None, + registered=loggers[canonical], + ) + + def destination(name, headers): + return OtelDestination(endpoint=accounts[name][0], headers=headers, callback_name=name) + + def run(destinations): + set_request_destinations(destinations) + emit(published.tracer_provider) + + in_fresh_context(run, (destination(canonical, dict(pair.split("=") for pair in accounts[canonical][1].split(","))),)) + in_fresh_context(run, (destination(other, dict(pair.split("=") for pair in accounts[other][1].split(","))),)) + assert shared.get_finished_spans() == (), "an account the operator already writes to was written twice" + + in_fresh_context(run, (destination(other, {"authorization": "Basic team"}),)) + assert [s.name for s in shared.get_finished_spans()] == ["chat gpt-4"] + + def test_publishing_twice_does_not_double_export(self): + config = OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory", owner=ExporterOwner.ARIZE_AX)]) + logger = OpenTelemetryV2(config=config, callback_name="arize") + + publish_global_otel_v2_provider([], lambda _p: None, registered=logger) + publish_global_otel_v2_provider([], lambda _p: None, registered=logger) + + kinds = [type(p).__name__ for p in logger._tracer_provider._active_span_processor._span_processors] + assert kinds.count("TenantFanOutSpanProcessor") == 1 + + def test_anchoring_reads_the_fan_out_off_the_published_provider_not_the_otel_global(self, monkeypatch): + """``set_tracer_provider`` keeps the first provider it was handed. When + auto-instrumentation or a legacy logger claimed it before the proxy published, + the OTel global carries no fan-out, so reading it there would refuse every + destination the published provider delivers.""" + from litellm.proxy import proxy_server + + config = OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory", owner=ExporterOwner.LANGFUSE_OTEL)]) + logger = OpenTelemetryV2(config=config, callback_name="langfuse_otel") + publish_global_otel_v2_provider([], lambda _p: None, registered=logger) + monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger) + claimed_first = TracerProvider() + + assert fan_out_provider() is logger.tracer_provider + assert deliverable_destinations((LANGFUSE_DEST,), claimed_first) == () + assert deliverable_destinations((LANGFUSE_DEST,), fan_out_provider()) == (LANGFUSE_DEST,) + + def test_a_legacy_v1_logger_holding_the_registered_slot_does_not_hide_the_fan_out(self, monkeypatch): + """The proxy publishes with ``registered=None`` when ``open_telemetry_logger`` + holds a v1 logger, so the fan-out lands on a v2 logger taken from + ``_in_memory_loggers``. Reading the registered slot finds no v2 logger there and + the OTel global belongs to v1, so both detours refuse every destination the + published provider delivers.""" + from litellm.integrations.opentelemetry import OpenTelemetry + from litellm.proxy import proxy_server + + config = OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory", owner=ExporterOwner.LANGFUSE_OTEL)]) + v2 = OpenTelemetryV2(config=config, callback_name="langfuse_otel") + publish_global_otel_v2_provider([v2], lambda _p: None, registered=None) + monkeypatch.setattr(proxy_server, "open_telemetry_logger", OpenTelemetry()) + + assert fan_out_provider() is v2.tracer_provider + assert deliverable_destinations((LANGFUSE_DEST,), fan_out_provider()) == (LANGFUSE_DEST,) + + def test_without_a_publish_anchoring_attaches_fan_out_to_registered_v2_logger(self, monkeypatch): + from litellm.proxy import proxy_server + + config = OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory", owner=ExporterOwner.LANGFUSE_OTEL)]) + logger = OpenTelemetryV2(config=config, callback_name="langfuse_otel") + monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger) + + assert fan_out_provider() is logger.tracer_provider + assert deliverable_destinations((LANGFUSE_DEST,), fan_out_provider()) == (LANGFUSE_DEST,) + + def test_concurrent_anchoring_attaches_exactly_one_fan_out(self): + """Requests race to anchor when the startup publish never ran, and a fan-out + attached twice delivers every tenant span twice.""" + import threading + + from litellm.integrations.otel.plumbing.providers import attach_tenant_fan_out + + class SlowAttachProvider(TracerProvider): + def add_span_processor(self, span_processor): + time.sleep(0.05) + super().add_span_processor(span_processor) + + provider = SlowAttachProvider() + config = OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory", owner=ExporterOwner.LANGFUSE_OTEL)]) + barrier = threading.Barrier(8) + + def anchor(): + barrier.wait(timeout=10) + attach_tenant_fan_out(provider, config) + + threads = [threading.Thread(target=anchor) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + kinds = [type(p).__name__ for p in provider._active_span_processor._span_processors] + assert kinds.count("TenantFanOutSpanProcessor") == 1, f"one fan-out per provider, got {kinds}" + + def test_without_a_publish_anchoring_falls_back_to_the_otel_global(self, monkeypatch): + from opentelemetry import trace + + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "open_telemetry_logger", None) + + assert fan_out_provider() is trace.get_tracer_provider() + + def test_auth_seeds_the_request_with_destinations_the_registered_logger_can_deliver( + self, monkeypatch, allow_test_hosts + ): + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import _seed_request_destinations + + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + config = OpenTelemetryV2Config(exporters=[ExporterSpec(kind="in_memory", owner=ExporterOwner.LANGFUSE_OTEL)]) + logger = OpenTelemetryV2(config=config, callback_name="langfuse_otel") + publish_global_otel_v2_provider([], lambda _p: None, registered=logger) + monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger) + auth = UserAPIKeyAuth( + team_metadata={ + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": "pk-team", + "langfuse_secret_key": "sk-team", + "langfuse_host": "http://team.local", + }, + } + ] + } + ) + expected = resolve_tenant_otel_destinations(auth) + assert expected, "the fixture must resolve to a destination for the test to mean anything" + + def run(): + _seed_request_destinations(auth) + return request_destinations() + + assert deliverable_destinations(expected, TracerProvider()) == () + assert in_fresh_context(run) == expected + + +class TestRouting: + def test_an_overridden_backend_is_not_detached_onto_a_second_provider(self): + config = OpenTelemetryV2Config( + exporters=[ExporterSpec(kind="otlp_http", endpoint="http://op.local", owner=ExporterOwner.LANGFUSE_OTEL)] + ) + cache = TenantTracerCache(config, "langfuse_otel", "litellm") + default = get_tracer(TracerProvider(), "litellm") + params = {"langfuse_public_key": "pk", "langfuse_secret_key": "sk"} + + assert cache.route_for(default, params).detached is True + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + return cache.route_for(default, params) + + route = in_fresh_context(run) + assert route.detached is False + assert route.tracer is default + assert route.provider is None + + def test_an_overridden_backend_does_not_detach_on_a_service_name_either(self): + """A key or team service name is its own reason to build a second provider, so + clearing only the credentials would still take the model call out of the tree.""" + config = OpenTelemetryV2Config( + exporters=[ExporterSpec(kind="otlp_http", endpoint="http://op.local", owner=ExporterOwner.LANGFUSE_OTEL)] + ) + cache = TenantTracerCache(config, "langfuse_otel", "litellm") + default = get_tracer(TracerProvider(), "litellm") + auth_metadata = {"otel_service_name": "team-checkout"} + + assert cache.route_for(default, None, auth_metadata).detached is False + assert cache.route_for(default, None, auth_metadata).tracer is not default + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + return cache.route_for(default, None, auth_metadata) + + route = in_fresh_context(run) + assert route.tracer is default, "the fan-out carries the service name on the destination instead" + assert route.provider is None + + @pytest.mark.parametrize("callback_name", ["arize", None]) + def test_a_service_name_does_not_detach_a_backend_the_destination_does_not_name(self, callback_name): + """The fan-out only sees spans on the published provider, so relabelling this + logger's span onto a second provider would drop the model call out of the + trace another backend's destination receives.""" + config = OpenTelemetryV2Config( + exporters=[ExporterSpec(kind="otlp_http", endpoint="http://op.local", owner=ExporterOwner.ARIZE_AX)] + ) + cache = TenantTracerCache(config, callback_name, "litellm") + default = get_tracer(TracerProvider(), "litellm") + auth_metadata = {"otel_service_name": "team-checkout"} + + relabelled = cache.route_for(default, None, auth_metadata) + assert relabelled.tracer is not default + cache.release(relabelled.provider) + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + return cache.route_for(default, None, auth_metadata) + + route = in_fresh_context(run) + assert route.tracer is default + assert route.detached is False + assert route.provider is None + + @pytest.mark.parametrize( + ("owner", "params", "auth_metadata"), + [ + (ExporterOwner.ARIZE_AX, {"arize_space_key": "space", "arize_api_key": "key"}, {}), + (ExporterOwner.ARIZE_PHOENIX, None, {"phoenix_project_name": "team-project"}), + ], + ) + def test_a_backend_pointed_at_its_own_account_still_routes_next_to_another_backend_destination( + self, owner, params, auth_metadata + ): + """Credentials or a project name the tenant's own account for this backend, which + the other backend's destination cannot stand in for.""" + config = OpenTelemetryV2Config( + exporters=[ExporterSpec(kind="otlp_http", endpoint="http://op.local", owner=owner)] + ) + cache = TenantTracerCache(config, owner.value, "litellm") + default = get_tracer(TracerProvider(), "litellm") + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + return cache.route_for(default, params, {"otel_service_name": "team-checkout", **auth_metadata}) + + route = in_fresh_context(run) + assert route.tracer is not default + assert route.detached is True + cache.release(route.provider) + + +@pytest.mark.usefixtures("allow_test_hosts") +class TestDestinationResolution: + def test_a_langfuse_key_pair_and_host_become_a_destination(self, monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + auth = UserAPIKeyAuth( + team_metadata={ + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": "pk-team", + "langfuse_secret_key": "sk-team", + "langfuse_host": "http://team.local", + }, + } + ] + } + ) + + destinations = resolve_tenant_otel_destinations(auth) + + assert [d.endpoint for d in destinations] == ["http://team.local/api/public/otel"] + assert destinations[0].callback_name == "langfuse_otel" + + def test_a_keys_service_name_outranks_its_teams_on_the_destination(self, monkeypatch): + """The key/team ``otel_service_name`` used to reach the backend through + per-request tracer routing, which an overridden backend skips.""" + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + auth = UserAPIKeyAuth( + metadata={"otel_service_name": "key-svc"}, + team_metadata={ + "otel_service_name": "team-svc", + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": "pk-team", + "langfuse_secret_key": "sk-team", + "langfuse_host": "http://team.local", + }, + } + ], + }, + ) + + destinations = resolve_tenant_otel_destinations(auth) + + assert dict(destinations[0].resource_attributes) == {"service.name": "key-svc"} + + def test_a_team_that_named_no_service_name_gets_no_resource_override(self, monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + auth = UserAPIKeyAuth( + team_metadata={ + "otel_service_name": " ", + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": "pk-team", + "langfuse_secret_key": "sk-team", + "langfuse_host": "http://team.local", + }, + } + ], + } + ) + + destinations = resolve_tenant_otel_destinations(auth) + + assert dict(destinations[0].resource_attributes) == {} + + def test_the_key_wins_over_the_team_for_the_same_backend(self, monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + + def entry(host: str) -> Mapping[str, object]: + return { + "callback_name": "langfuse_otel", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + "langfuse_host": host, + }, + } + + auth = UserAPIKeyAuth( + metadata={"logging": [entry("http://key.local")]}, + team_metadata={"logging": [entry("http://team.local")]}, + ) + + assert [d.endpoint for d in resolve_tenant_otel_destinations(auth)] == ["http://key.local/api/public/otel"] + + def test_nothing_resolves_while_otel_v2_is_off(self, monkeypatch): + monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() + auth = UserAPIKeyAuth( + team_metadata={ + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_type": "success", + "callback_vars": {"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}, + } + ] + } + ) + + assert resolve_tenant_otel_destinations(auth) == () + + def test_a_host_without_its_key_pair_resolves_to_nothing(self): + assert destination_for("langfuse_otel", {"langfuse_host": "http://team.local"}) is None + + def test_a_backend_with_no_dynamic_credentials_has_no_destination(self): + assert "arize_phoenix" not in destination_capable_backends() + assert destination_for("arize_phoenix", {"arize_api_key": "k"}) is None + + def test_the_destination_header_string_survives_the_exporter_round_trip(self): + from litellm.integrations.otel.plumbing.providers import parse_headers + + destination = destination_for( + "langfuse_otel", + {"langfuse_public_key": "pk", "langfuse_secret_key": "sk", "langfuse_host": "http://x"}, + ) + assert parse_headers(destination.header_string())["authorization"] == destination.headers["Authorization"] + + +#: Anything that makes ``OpenTelemetryV2Config`` synthesize a real operator destination. +_OTEL_SHORTHAND_ENV = ( + "OTEL_ENDPOINT", + "OTEL_HEADERS", + "OTEL_EXPORTER", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_EXPORTER_OTLP_HEADERS", + "OTEL_EXPORTER_OTLP_PROTOCOL", +) + + +def credential_less_proxy(monkeypatch) -> None: + """An operator with no Langfuse account and no generic OTLP collector.""" + for name in ("LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY", *_OTEL_SHORTHAND_ENV): + monkeypatch.delenv(name, raising=False) + with pytest.raises(ValueError, match="LANGFUSE_PUBLIC_KEY"): + langfuse_preset() + + +class TestPresetDegradation: + def test_a_credential_less_langfuse_exports_nowhere_instead_of_to_the_console(self, monkeypatch, capfd): + """``_normalize`` folds a console exporter in for an empty list, which would + print every span on a proxy whose teams bring their own credentials.""" + credential_less_proxy(monkeypatch) + + config = langfuse_preset(allow_missing_credentials=True) + provider = build_tracer_provider(config, tenant_overrides=True) + capfd.readouterr() + in_fresh_context(emit, provider) + provider.force_flush() + + assert '"name": "chat gpt-4"' not in capfd.readouterr().out + assert "langfuse" in config.mapper_names + + def test_langfuse_still_raises_for_a_global_callback_with_no_credentials(self, monkeypatch): + monkeypatch.delenv("LANGFUSE_PUBLIC_KEY", raising=False) + monkeypatch.delenv("LANGFUSE_SECRET_KEY", raising=False) + + with pytest.raises(ValueError, match="LANGFUSE_PUBLIC_KEY"): + langfuse_preset() + + def test_a_credential_less_proxy_builds_the_gated_logger_beside_a_v2_carrier(self, monkeypatch): + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + credential_less_proxy(monkeypatch) + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + carrier = build_otel_v2_logger(OpenTelemetryV2Config(exporter="in_memory")) + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + return _maybe_construct_otel_v2("langfuse_otel", [carrier]) + + is_otel_v2_enabled.cache_clear() + logger = in_fresh_context(run) + is_otel_v2_enabled.cache_clear() + + assert logger is not None + assert all(spec.requires_headers and not spec.headers for spec in logger.config.exporters) + + def test_a_credential_less_proxy_with_no_destinations_falls_back_to_the_legacy_path(self, monkeypatch): + """Nothing can use a credential-less langfuse here, so the operator has to get + the same story as before v2: the legacy integration, not a global provider + that exports nowhere.""" + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + credential_less_proxy(monkeypatch) + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + + is_otel_v2_enabled.cache_clear() + logger = in_fresh_context(_maybe_construct_otel_v2, "langfuse_otel", []) + is_otel_v2_enabled.cache_clear() + + assert logger is None + + def test_a_valid_newrelic_base_exporter_survives_without_a_license_key(self, monkeypatch): + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + monkeypatch.delenv("NEW_RELIC_LICENSE_KEY", raising=False) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector.local:4318") + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + + is_otel_v2_enabled.cache_clear() + logger = in_fresh_context(_maybe_construct_otel_v2, "newrelic", []) + is_otel_v2_enabled.cache_clear() + + assert logger is not None + assert [spec.endpoint for spec in logger.config.exporters] == [ + "http://collector.local:4318", + "https://otlp.nr-data.net", + ] + + def test_a_credentialless_newrelic_without_a_base_exporter_falls_back(self, monkeypatch): + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + monkeypatch.delenv("NEW_RELIC_LICENSE_KEY", raising=False) + for name in _OTEL_SHORTHAND_ENV: + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + + is_otel_v2_enabled.cache_clear() + logger = in_fresh_context(_maybe_construct_otel_v2, "newrelic", []) + is_otel_v2_enabled.cache_clear() + + assert logger is None + + def test_an_explicit_console_exporter_keeps_a_credentialless_preset_on_v2(self, monkeypatch, capfd): + """``OTEL_EXPORTER=console`` reads exactly like the placeholder ``_normalize`` + folds in, but the operator asked for it, so a credential-less New Relic keeps + the V2 logger and its spans reach stdout instead of the legacy path.""" + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + monkeypatch.delenv("NEW_RELIC_LICENSE_KEY", raising=False) + for name in _OTEL_SHORTHAND_ENV: + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("OTEL_EXPORTER", "console") + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + + is_otel_v2_enabled.cache_clear() + logger = in_fresh_context(_maybe_construct_otel_v2, "newrelic", []) + is_otel_v2_enabled.cache_clear() + + assert logger is not None + assert logger.config.exporters[0].kind == "console" + assert not logger.config.exporters[0].requires_headers + + def test_a_destination_for_one_backend_does_not_degrade_another(self, monkeypatch): + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + credential_less_proxy(monkeypatch) + monkeypatch.delenv("WANDB_API_KEY", raising=False) + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + return _maybe_construct_otel_v2("weave_otel", []) + + is_otel_v2_enabled.cache_clear() + logger = in_fresh_context(run) + is_otel_v2_enabled.cache_clear() + + assert logger is None + + def test_the_exporter_less_logger_is_not_reused_by_a_request_without_destinations(self, monkeypatch): + """Reusing it would let one team's destination decide how every later request + without one is logged, long after the degrade was justified.""" + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + credential_less_proxy(monkeypatch) + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + loggers = [build_otel_v2_logger(OpenTelemetryV2Config(exporter="in_memory"))] + + def with_destination(): + set_request_destinations((LANGFUSE_DEST,)) + return _maybe_construct_otel_v2("langfuse_otel", loggers) + + is_otel_v2_enabled.cache_clear() + degraded = in_fresh_context(with_destination) + plain = in_fresh_context(_maybe_construct_otel_v2, "langfuse_otel", loggers) + is_otel_v2_enabled.cache_clear() + + assert degraded is not None + assert plain is None + + def test_a_credentialed_logger_is_still_reused_across_requests(self, monkeypatch): + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-1") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-1") + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + loggers = [] + + is_otel_v2_enabled.cache_clear() + first = in_fresh_context(_maybe_construct_otel_v2, "langfuse_otel", loggers) + second = in_fresh_context(_maybe_construct_otel_v2, "langfuse_otel", loggers) + is_otel_v2_enabled.cache_clear() + + assert first is not None + assert second is first + + @staticmethod + def _degraded_langfuse_beside(loggers, monkeypatch): + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + credential_less_proxy(monkeypatch) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector.local:4318") + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + return _maybe_construct_otel_v2("langfuse_otel", loggers) + + is_otel_v2_enabled.cache_clear() + logger = in_fresh_context(run) + is_otel_v2_enabled.cache_clear() + assert logger is not None + return logger + + def test_a_degraded_logger_beside_another_v2_logger_leaves_the_collector_to_it(self, monkeypatch): + """The other logger's provider already exports every span to the operator's + collector, so a second model span from this one would land there twice.""" + collector_logger = build_otel_v2_logger(OpenTelemetryV2Config(exporter="in_memory")) + + logger = self._degraded_langfuse_beside([collector_logger], monkeypatch) + + assert [spec.endpoint for spec in logger.config.exporters] == [None] + assert all(spec.requires_headers and not spec.headers for spec in logger.config.exporters) + + @pytest.mark.parametrize("registered", [(), (CustomLogger(),)]) + def test_a_credential_less_proxy_with_a_destination_but_no_v2_carrier_falls_back(self, monkeypatch, registered): + """Only a V2 logger publishes the provider the fan-out rides on, so a legacy + callback beside this one leaves the destination just as unreachable as no + callback at all, and the operator keeps the pre-V2 story.""" + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + credential_less_proxy(monkeypatch) + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + + def run(): + set_request_destinations((LANGFUSE_DEST,)) + return _maybe_construct_otel_v2("langfuse_otel", list(registered)) + + is_otel_v2_enabled.cache_clear() + logger = in_fresh_context(run) + is_otel_v2_enabled.cache_clear() + + assert logger is None + + def test_a_credentialed_logger_beside_another_v2_logger_keeps_every_exporter(self, monkeypatch): + """Only a degraded preset gives the collector up; an operator who configured + both the backend and the collector still exports to both, as on base.""" + from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 + + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-1") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-1") + monkeypatch.setenv("LANGFUSE_HOST", "https://cloud.langfuse.com") + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector.local:4318") + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + collector_logger = build_otel_v2_logger(OpenTelemetryV2Config(exporter="in_memory")) + + is_otel_v2_enabled.cache_clear() + logger = in_fresh_context(_maybe_construct_otel_v2, "langfuse_otel", [collector_logger]) + is_otel_v2_enabled.cache_clear() + + assert logger is not None + assert [spec.endpoint for spec in logger.config.exporters] == [ + "http://collector.local:4318", + "https://cloud.langfuse.com/api/public/otel", + ] + assert all(spec.headers for spec in logger.config.exporters if spec.requires_headers) + + +class TestContextIsolation: + def test_destinations_do_not_leak_between_requests(self): + def first(): + set_request_destinations((LANGFUSE_DEST,)) + return destination_backends() + + assert in_fresh_context(first) == frozenset({"langfuse_otel"}) + assert in_fresh_context(request_destinations) == () + + +class TestOperatorShorthandSurvivesDegradation: + def test_a_generic_otlp_collector_keeps_receiving_when_langfuse_has_no_credentials(self, monkeypatch): + """Only the stdout placeholder is dropped. An operator who set the standard + OTLP env vars configured a real destination and must keep it.""" + monkeypatch.delenv("LANGFUSE_PUBLIC_KEY", raising=False) + monkeypatch.delenv("LANGFUSE_SECRET_KEY", raising=False) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector.local:4318") + + config = langfuse_preset(allow_missing_credentials=True) + + assert [spec.endpoint for spec in config.exporters] == ["http://collector.local:4318", None] + assert [spec.kind for spec in config.exporters] == ["otlp_http", "console"] + + def test_the_stdout_placeholder_is_still_dropped_when_it_is_the_only_exporter(self, monkeypatch): + credential_less_proxy(monkeypatch) + + config = langfuse_preset(allow_missing_credentials=True) + + assert all(spec.requires_headers and not spec.headers for spec in config.exporters) + + +class TestBackendEndpointParity: + def test_arize_follows_its_own_http_endpoint_instead_of_the_grpc_default(self, monkeypatch): + monkeypatch.delenv("ARIZE_ENDPOINT", raising=False) + monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "https://otlp.arize.com/v1/traces") + + destination = destination_for("arize", {"arize_space_id": "s", "arize_api_key": "k"}) + + assert destination.endpoint == "https://otlp.arize.com/v1/traces" + assert destination.protocol == "otlp_http" + + def test_arize_uses_grpc_when_nothing_is_configured(self, monkeypatch): + monkeypatch.delenv("ARIZE_ENDPOINT", raising=False) + monkeypatch.delenv("ARIZE_HTTP_ENDPOINT", raising=False) + + destination = destination_for("arize", {"arize_space_id": "s", "arize_api_key": "k"}) + + assert destination.endpoint == "https://otlp.arize.com/v1" + assert destination.protocol == "otlp_grpc" + + def test_weave_follows_a_self_hosted_wandb_host(self, monkeypatch): + monkeypatch.setenv("WANDB_HOST", "weave.internal.example") + + destination = destination_for("weave_otel", {"wandb_api_key": "k", "weave_project_id": "e/p"}) + + assert destination.endpoint == "https://weave.internal.example/otel/v1/traces" + + def test_weave_uses_the_cloud_endpoint_without_a_host(self, monkeypatch): + monkeypatch.delenv("WANDB_HOST", raising=False) + + destination = destination_for("weave_otel", {"wandb_api_key": "k", "weave_project_id": "e/p"}) + + assert destination.endpoint == "https://trace.wandb.ai/otel/v1/traces" + + +class TestIncompleteCredentials: + """Half a credential set builds a non-empty but unusable header dict. Accepting it + would suppress the operator's exporter and send the trace where it cannot land.""" + + @pytest.mark.parametrize( + "callback_name,callback_vars", + [ + ("arize", {"arize_api_key": "k"}), + ("arize", {"arize_space_id": "s"}), + ("weave_otel", {"wandb_api_key": "k"}), + ("weave_otel", {"weave_project_id": "e/p"}), + ("langfuse_otel", {"langfuse_public_key": "pk"}), + ], + ) + def test_a_partial_credential_set_resolves_to_nothing(self, callback_name, callback_vars): + assert destination_for(callback_name, callback_vars) is None + + @pytest.mark.parametrize( + "callback_name,callback_vars", + [ + ("arize", {"arize_space_id": "s", "arize_api_key": "k"}), + ("weave_otel", {"wandb_api_key": "k", "weave_project_id": "e/p"}), + ("newrelic", {"newrelic_api_key": "k"}), + ], + ) + def test_a_complete_credential_set_resolves(self, callback_name, callback_vars): + assert destination_for(callback_name, callback_vars) is not None + + +@pytest.mark.usefixtures("allow_test_hosts") +class TestCallbackTypeFilter: + @staticmethod + def _auth(callback_type: str | None) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + team_metadata={ + "logging": [ + { + "callback_name": "langfuse_otel", + "callback_type": callback_type, + "callback_vars": { + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + "langfuse_host": "http://team.local", + }, + } + ] + } + ) + + @pytest.mark.parametrize("callback_type", ["success", "success_and_failure", None]) + def test_an_entry_that_wants_success_traces_gets_the_whole_trace(self, monkeypatch, callback_type): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + + assert resolve_tenant_otel_destinations(self._auth(callback_type)) != () + + def test_a_failure_only_entry_does_not_take_over_the_trace(self, monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() + + assert resolve_tenant_otel_destinations(self._auth("failure")) == () + + +class TestTenantConfigAgreement: + """The destination resolver and ``convert_key_logging_metadata_to_callback`` read + the same stored config, so they must not read it two different ways.""" + + @pytest.fixture(autouse=True) + def _v2_on(self, monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setattr( + litellm, "provider_url_destination_allowed_hosts", ["team.local", "key.local"], raising=False + ) + is_otel_v2_enabled.cache_clear() + yield + is_otel_v2_enabled.cache_clear() + + @staticmethod + def _entry(host, **extra): + return { + "callback_name": "langfuse_otel", + "callback_vars": {"langfuse_public_key": "pk", "langfuse_secret_key": "sk", "langfuse_host": host, **extra}, + } + + def test_a_key_that_disabled_its_callbacks_does_not_fall_back_to_the_team(self): + """Disabling a key's callbacks stores an empty list, which the sibling parser + reads as 'the key configured none'.""" + auth = UserAPIKeyAuth( + metadata={"logging": []}, + team_metadata={"logging": [self._entry("http://team.local")]}, + ) + + assert resolve_tenant_otel_destinations(auth) == () + + def test_two_entries_for_one_backend_merge_their_vars_last_wins(self): + auth = UserAPIKeyAuth( + team_metadata={ + "logging": [ + self._entry("http://team.local"), + {"callback_name": "langfuse_otel", "callback_vars": {"langfuse_host": "http://key.local"}}, + ] + } + ) + + destinations = resolve_tenant_otel_destinations(auth) + + assert [d.endpoint for d in destinations] == ["http://key.local/api/public/otel"] + + def test_a_failure_entry_still_wins_the_merge_next_to_a_success_entry(self): + entries = [ + {**self._entry("http://team.local"), "callback_type": "success"}, + { + **self._entry("http://key.local", langfuse_public_key="pk-failure", langfuse_secret_key="sk-failure"), + "callback_type": "failure", + }, + ] + runtime = reduce( + lambda merged, entry: convert_key_logging_metadata_to_callback(AddTeamCallback(**entry), merged), + entries, + None, + ) + + destinations = resolve_tenant_otel_destinations(UserAPIKeyAuth(team_metadata={"logging": entries})) + + assert runtime.callback_vars["langfuse_host"] == "http://key.local" + assert [d.endpoint for d in destinations] == ["http://key.local/api/public/otel"] + assert destinations[0].headers["Authorization"] == f"Basic {b64encode(b'pk-failure:sk-failure').decode()}" + + @pytest.fixture + def premium(self, monkeypatch): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(litellm, "allow_dynamic_callback_disabling", True) + + @pytest.mark.usefixtures("premium") + def test_a_backend_the_key_disabled_resolves_to_no_destination(self): + """Dispatch skips a callback named in the key's ``litellm_disabled_callbacks``, + so the fan-out must not deliver to it either.""" + auth = UserAPIKeyAuth( + metadata={"litellm_disabled_callbacks": ["Langfuse_OTEL"]}, + team_metadata={"logging": [self._entry("http://team.local")]}, + ) + + assert resolve_tenant_otel_destinations(auth) == () + + @pytest.mark.usefixtures("premium") + @pytest.mark.parametrize( + ("header", "resolved"), + [ + ("langfuse_otel", False), + (" LANGFUSE_OTEL ,arize", False), + ("arize", True), + ], + ) + def test_the_disable_header_wins_over_the_key_list(self, header, resolved): + """Same precedence as dispatch: a header that names other backends re-enables + the one the key stored.""" + auth = UserAPIKeyAuth( + metadata={"litellm_disabled_callbacks": ["langfuse_otel"]}, + team_metadata={"logging": [self._entry("http://team.local")]}, + ) + + destinations = resolve_tenant_otel_destinations(auth, {"x-litellm-disable-callbacks": header}) + + assert bool(destinations) is resolved + + def test_a_non_premium_proxy_ignores_the_disabled_list_like_dispatch_does(self, monkeypatch): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", False) + auth = UserAPIKeyAuth( + metadata={"litellm_disabled_callbacks": ["langfuse_otel"]}, + team_metadata={"logging": [self._entry("http://team.local")]}, + ) + + assert resolve_tenant_otel_destinations(auth, {"x-litellm-disable-callbacks": "langfuse_otel"}) != () + + +class TestEvictionSafety: + class Recording(SimpleSpanProcessor): + def __init__(self): + super().__init__(InMemorySpanExporter()) + self.shutdown_calls = 0 + + def shutdown(self): + self.shutdown_calls += 1 + + def _fan_out(self): + built = [] + + def factory(_destination): + built.append(self.Recording()) + return built[-1] + + return TenantFanOutSpanProcessor(processor_factory=factory), built + + @staticmethod + def _dest(index): + return LANGFUSE_DEST.model_copy(update={"endpoint": f"http://d{index}/otel"}) + + @staticmethod + def _settle(fan_out, processor=None): + """Wait for retirement to clear and, when given, for the drain to run. + + The drain pool is shared and bounded, so a shed processor is closed once a + worker picks it up rather than the moment it is handed over. + """ + for _ in range(500): + if not fan_out._retired and (processor is None or processor.shutdown_calls): + return + time.sleep(0.02) + + def test_a_processor_still_exporting_a_span_is_not_closed_under_it(self): + """``on_end`` holds a processor across the export, so closing an evicted one + there drops the span it is holding.""" + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + fan_out, built = self._fan_out() + held = fan_out._acquire(self._dest(0)) + for index in range(1, _MAX_CACHED_DESTINATION_PROCESSORS + 1): + fan_out._acquire(self._dest(index)) + fan_out._release(built[-1]) + + assert held.shutdown_calls == 0 + assert id(held) in fan_out._retired + + fan_out._release(held) + self._settle(fan_out, held) + + assert held.shutdown_calls == 1 + + def test_a_recently_used_destination_is_not_the_one_evicted(self): + """Without the refresh the cache sheds by insertion order, so the busiest + destination is the one whose exporter is rebuilt on every overflow.""" + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + fan_out, built = self._fan_out() + for index in range(_MAX_CACHED_DESTINATION_PROCESSORS): + fan_out._release(fan_out._acquire(self._dest(index))) + fan_out._release(fan_out._acquire(self._dest(0))) + fan_out._release(fan_out._acquire(self._dest(_MAX_CACHED_DESTINATION_PROCESSORS))) + self._settle(fan_out, built[1]) + + assert built[1].shutdown_calls == 1 + assert built[0].shutdown_calls == 0, "the destination used most recently was the one shed" + + def test_an_idle_evicted_processor_is_closed_off_the_export_path(self): + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + fan_out, built = self._fan_out() + for index in range(_MAX_CACHED_DESTINATION_PROCESSORS + 1): + fan_out._acquire(self._dest(index)) + fan_out._release(built[-1]) + self._settle(fan_out, built[0]) + + assert built[0].shutdown_calls == 1 + assert len(fan_out._processors) == _MAX_CACHED_DESTINATION_PROCESSORS + + def test_a_slow_collector_does_not_hold_up_the_export_path(self): + """``shutdown`` flushes over the network and is reached from ``on_end``, so + closing a shed processor inline lets one unreachable tenant collector stall + every other tenant's spans.""" + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + class Slow(self.Recording): + def shutdown(self): + time.sleep(3) + super().shutdown() + + built = [] + + def factory(_destination): + built.append(Slow()) + return built[-1] + + fan_out = TenantFanOutSpanProcessor(processor_factory=factory) + started = time.monotonic() + for index in range(_MAX_CACHED_DESTINATION_PROCESSORS + 1): + fan_out._acquire(self._dest(index)) + fan_out._release(built[-1]) + + assert time.monotonic() - started < 2 + + def test_shedding_many_processors_does_not_spawn_a_thread_each(self): + """A tenant that cycles its destination config sheds a processor per request, + so a thread per shed processor is a thread per request against a slow + collector.""" + import threading + + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + release = threading.Event() + + class Blocking(self.Recording): + def shutdown(self): + release.wait(timeout=10) + super().shutdown() + + built = [] + + def factory(_destination): + built.append(Blocking()) + return built[-1] + + fan_out = TenantFanOutSpanProcessor(processor_factory=factory) + before = self._drain_workers() + try: + for index in range(_MAX_CACHED_DESTINATION_PROCESSORS + 30): + fan_out._acquire(self._dest(index)) + fan_out._release(built[-1]) + grew = self._drain_workers() - before + assert grew == 0, f"one drain thread per shed processor: {grew} new threads" + finally: + release.set() + self._settle(fan_out, built[0]) + + def test_a_saturated_drain_leaves_new_destinations_with_the_operator(self): + """A shed processor keeps its batch thread until its close returns, and against + a collector that never answers every close waits out the exporter's timeout. + Tenants rotating past the cache cap would otherwise queue one more processor, + and one more thread, per request for as long as the outage lasts.""" + import threading + + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + release = threading.Event() + + class Blocking(self.Recording): + def shutdown(self): + release.wait(timeout=10) + super().shutdown() + + built = [] + + def factory(_destination): + built.append(Blocking()) + return built[-1] + + fan_out = TenantFanOutSpanProcessor(processor_factory=factory, pending_drains=3) + try: + anchored = tuple( + fan_out.deliverable((self._dest(index),)) for index in range(_MAX_CACHED_DESTINATION_PROCESSORS + 40) + ) + + assert len(built) == _MAX_CACHED_DESTINATION_PROCESSORS + 3, "a processor per request during the outage" + assert sum(1 for accepted in anchored if accepted) == len(built), "anchored what it could not build" + assert fan_out.deliverable((self._dest(999),)) == (), "the span would vanish instead of staying with the operator" + finally: + release.set() + for _ in range(500): + if not fan_out._drain.saturated(): + break + time.sleep(0.02) + + assert fan_out.deliverable((self._dest(999),)) == (self._dest(999),), "the fan-out never recovered" + + def test_an_anchored_destination_evicted_under_a_saturated_drain_still_gets_the_span(self): + """``deliverable`` accepted the destination, so the operator's exporter has stood + down for it. Other tenants' auths can then evict it, and the eviction is what + tips the drain into saturation, so refusing the rebuild at ``on_end`` would drop + the span outright.""" + import threading + + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + release = threading.Event() + + class Blocking(self.Recording): + def shutdown(self): + release.wait(timeout=10) + super().shutdown() + + built = [] + + def factory(_destination): + built.append(Blocking()) + return built[-1] + + fan_out = TenantFanOutSpanProcessor( + processor_factory=factory, pending_drains=_MAX_CACHED_DESTINATION_PROCESSORS + 1 + ) + provider = TracerProvider() + provider.add_span_processor(fan_out) + tracer = get_tracer(provider, "litellm") + anchored = self._dest(0) + try: + for index in range(1, _MAX_CACHED_DESTINATION_PROCESSORS + 1): + assert fan_out.deliverable((self._dest(index),)) + assert fan_out.deliverable((anchored,)) == (anchored,) + first = built[-1] + for index in range(_MAX_CACHED_DESTINATION_PROCESSORS + 1, 2 * _MAX_CACHED_DESTINATION_PROCESSORS + 1): + assert fan_out.deliverable((self._dest(index),)) + assert fan_out._drain.saturated(), "the anchored destination's own eviction saturates the drain" + assert first not in fan_out._processors.values(), "the anchored destination was not evicted" + + def run(): + set_request_destinations((anchored,)) + with tracer.start_as_current_span("chat anthropic"): + pass + + before = len(built) + in_fresh_context(run) + assert len(built) == before + 1, "the anchored destination was not rebuilt, so its span went nowhere" + assert [span.name for span in built[-1].span_exporter.get_finished_spans()] == ["chat anthropic"] + assert first.span_exporter.get_finished_spans() == (), "the shed processor was handed out again" + finally: + release.set() + + def _saturated_by_anchoring(self, pending_drains, extra): + """A fan-out whose drain ``extra`` anchorings past the cache cap have saturated. + + Returns it with the processors built, the destinations that anchored, and the + event that lets the blocked closes finish. + """ + import threading + + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + release = threading.Event() + + class Blocking(self.Recording): + def shutdown(self): + release.wait(timeout=10) + super().shutdown() + + built = [] + + def factory(_destination): + built.append(Blocking()) + return built[-1] + + fan_out = TenantFanOutSpanProcessor(processor_factory=factory, pending_drains=pending_drains) + destinations = tuple(self._dest(index) for index in range(_MAX_CACHED_DESTINATION_PROCESSORS + extra)) + anchored = tuple(destination for destination in destinations if fan_out.deliverable((destination,))) + assert fan_out._drain.saturated(), "anchoring past the cap did not saturate the drain" + assert len(anchored) > _MAX_CACHED_DESTINATION_PROCESSORS, "not enough destinations in flight to churn" + return fan_out, built, anchored, release + + def test_anchored_rebuilds_under_a_saturated_drain_do_not_grow_with_the_spans(self): + """Every anchored rebuild past the cap evicts another anchored destination, whose + next span rebuilds it in turn. With more destinations in flight than the cache + holds, each span would then cost one more processor, one more batch thread and + one more close queued behind a collector that never answers.""" + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + fan_out, built, anchored, release = self._saturated_by_anchoring(pending_drains=4, extra=8) + try: + after_anchoring = len(built) + for _ in range(5): + for destination in anchored: + fan_out._release(fan_out._acquire(destination)) + + rebuilt = len(built) - after_anchoring + assert rebuilt == len(anchored) - _MAX_CACHED_DESTINATION_PROCESSORS, ( + f"{rebuilt} rebuilds over 5 rounds of {len(anchored)} anchored destinations: one per evicted one expected" + ) + assert len(fan_out._processors) == len(anchored), "an anchored destination was shed under a saturated drain" + assert all(destination in fan_out.deliverable((destination,)) for destination in anchored) + finally: + release.set() + + def test_the_cache_returns_to_its_cap_once_the_drain_has_room(self): + """Holding above the cap is for the outage only: with the drain caught up, the + entries kept for the destinations in flight are the ones to shed.""" + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + fan_out, built, anchored, release = self._saturated_by_anchoring(pending_drains=4, extra=8) + for destination in anchored: + fan_out._release(fan_out._acquire(destination)) + assert len(fan_out._processors) > _MAX_CACHED_DESTINATION_PROCESSORS + + release.set() + for _ in range(500): + for destination in anchored[-4:]: + fan_out._release(fan_out._acquire(destination)) + if len(fan_out._processors) <= _MAX_CACHED_DESTINATION_PROCESSORS: + break + time.sleep(0.02) + + assert len(fan_out._processors) == _MAX_CACHED_DESTINATION_PROCESSORS, "the cache never came back to its cap" + shed = len(built) - _MAX_CACHED_DESTINATION_PROCESSORS + for _ in range(500): + if sum(processor.shutdown_calls for processor in built) == shed: + break + time.sleep(0.02) + + assert sum(processor.shutdown_calls for processor in built) == shed, "a shed processor was never closed" + + def test_concurrent_eviction_cannot_build_between_retirement_and_drain_submission(self): + """A second request cannot build while the first eviction is being handed to + the drain, or concurrent churn can outrun the pending-drain limit.""" + import threading + + from litellm.integrations.otel.plumbing.providers import ( + _DrainPool, + _MAX_CACHED_DESTINATION_PROCESSORS, + ) + + class GatedDrain(_DrainPool): + def __init__(self): + super().__init__(workers=0) + self.started = threading.Event() + self.release = threading.Event() + + def saturated(self): + return False + + def submit(self, processor): + if not self.started.is_set(): + self.started.set() + self.release.wait(timeout=5) + + built = [] + + def factory(_destination): + built.append(self.Recording()) + return built[-1] + + drain = GatedDrain() + fan_out = TenantFanOutSpanProcessor(processor_factory=factory, drain_pool=drain) + for index in range(_MAX_CACHED_DESTINATION_PROCESSORS): + fan_out._release(fan_out._acquire(self._dest(index))) + + first = threading.Thread(target=lambda: fan_out._release(fan_out._acquire(self._dest(32)))) + first.start() + assert drain.started.wait(timeout=5) + second = threading.Thread(target=lambda: fan_out._release(fan_out._acquire(self._dest(33)))) + second.start() + time.sleep(0.1) + + assert len(built) == _MAX_CACHED_DESTINATION_PROCESSORS + 1 + + drain.release.set() + first.join(timeout=5) + second.join(timeout=5) + assert not first.is_alive() and not second.is_alive() + assert len(built) == _MAX_CACHED_DESTINATION_PROCESSORS + 2 + + def test_drain_workers_are_daemons(self): + """Python joins a ThreadPoolExecutor's workers at interpreter exit, so one + unreachable tenant collector would hold the proxy open for its export + timeout on the way down.""" + import threading + + self._fan_out() + workers = [t for t in threading.enumerate() if t.name.startswith("litellm-otel-destination-drain")] + + assert workers, "no drain worker was started" + assert all(t.daemon for t in workers), "a non-daemon drain worker blocks interpreter exit" + + def test_a_burst_of_first_evictions_starts_one_set_of_drain_workers(self): + """A drain pool built lazily on first use is not built once: several threads + can each finish the build, and every pool but the winner is left with its + workers blocked on a queue nothing will ever feed again.""" + import threading + + from litellm.integrations.otel.plumbing.providers import ( + _DRAIN_WORKERS, + _MAX_CACHED_DESTINATION_PROCESSORS, + ) + + for _ in range(3): + before = self._drain_workers() + fan_out, built = self._fan_out() + for index in range(_MAX_CACHED_DESTINATION_PROCESSORS): + fan_out._release(fan_out._acquire(self._dest(index))) + barrier = threading.Barrier(16) + + def shed(index, fan_out=fan_out, barrier=barrier): + barrier.wait(timeout=10) + fan_out._release(fan_out._acquire(self._dest(index))) + + threads = [ + threading.Thread(target=shed, args=(_MAX_CACHED_DESTINATION_PROCESSORS + index,)) for index in range(16) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + self._settle(fan_out) + + assert self._drain_workers() - before == _DRAIN_WORKERS + + @staticmethod + def _drain_workers(): + import threading + + return len([t for t in threading.enumerate() if t.name.startswith("litellm-otel-destination-drain")]) + + def test_shutdown_does_not_close_a_processor_under_an_in_flight_export(self): + """``on_end`` runs on whichever thread ends a span, so it reaches the fan-out + while the SDK tears the provider down.""" + import threading + + fan_out, _ = self._fan_out() + held = fan_out._acquire(self._dest(0)) + closed = threading.Thread(target=fan_out.shutdown) + closed.start() + try: + time.sleep(0.3) + + assert held.shutdown_calls == 0, "closed a processor with a span still being forwarded" + finally: + fan_out._release(held) + closed.join(timeout=10) + + assert held.shutdown_calls == 1 + + def test_a_closed_fan_out_builds_no_new_processor(self): + """A processor built after shutdown is one nothing will ever close, and it + exports to a tenant on a provider the SDK has already torn down.""" + fan_out, built = self._fan_out() + fan_out.shutdown() + + assert fan_out._acquire(self._dest(0)) is None + assert built == [] + + def test_shutdown_gives_up_on_an_export_that_never_finishes(self): + """The wait is bounded: an exporter stuck on a dead collector must not hold + the proxy open on the way down.""" + import threading + + fan_out = TenantFanOutSpanProcessor(processor_factory=lambda _d: self.Recording(), shutdown_drain_seconds=0.2) + fan_out._acquire(self._dest(0)) + closed = threading.Thread(target=fan_out.shutdown) + closed.start() + closed.join(timeout=5) + + assert not closed.is_alive(), "shutdown blocked on an export that never finished" + + def test_shutdown_retires_the_drain_workers(self): + """A proxy that rebuilds its telemetry builds another fan-out, so workers that + outlive the one that started them are two more threads per reload.""" + from litellm.integrations.otel.plumbing.providers import _DRAIN_WORKERS + + before = self._drain_workers() + fan_out, _ = self._fan_out() + assert self._drain_workers() - before == _DRAIN_WORKERS + + fan_out.shutdown() + for _ in range(500): + if self._drain_workers() == before: + break + time.sleep(0.02) + + assert self._drain_workers() == before, "the drain workers outlived their fan-out" + + def test_a_processor_shed_after_shutdown_is_still_closed(self): + """``close`` retires the workers, so anything handed to the pool afterwards + would sit in a queue nobody reads.""" + fan_out, _ = self._fan_out() + stray = self.Recording() + fan_out.shutdown() + fan_out._drain.submit(stray) + + for _ in range(500): + if stray.shutdown_calls: + break + time.sleep(0.02) + + assert stray.shutdown_calls == 1 + + def test_releasing_a_straggler_after_shutdown_does_not_block_the_span_thread(self): + """The teardown deadline has already expired by then, so closing the straggler + inline would park whichever thread just ended a span on the very flush the + deadline gave up waiting for.""" + import threading + + never = threading.Event() + + class Stuck(self.Recording): + def shutdown(self): + never.wait() + + def factory(_destination): + return Stuck() + + fan_out = TenantFanOutSpanProcessor(processor_factory=factory, shutdown_drain_seconds=0.05) + held = fan_out._acquire(self._dest(0)) + fan_out.shutdown() + + released = threading.Event() + caller = threading.Thread(target=lambda: (fan_out._release(held), released.set()), daemon=True) + caller.start() + came_back = released.wait(timeout=5) + never.set() + + assert came_back, "the thread that ended the span was left holding a stuck teardown" + + def test_shutdown_waits_out_an_export_that_lands_inside_the_bound(self): + """Without the wait the closing is left to a daemon thread, which the + interpreter can retire before it runs, so the last spans never reach the + tenant.""" + import threading + + fan_out, built = self._fan_out() + held = fan_out._acquire(self._dest(0)) + threading.Timer(0.2, lambda: fan_out._release(held)).start() + + fan_out.shutdown() + + assert held.shutdown_calls == 1, "shutdown returned before the export it should have waited out" + + def test_a_straggler_past_the_drain_bound_is_closed_by_its_own_thread(self): + """The wait is bounded so one dead collector cannot hold the proxy open, which + means a processor still exporting when it expires has to be left to the thread + holding it rather than closed under the span it is carrying.""" + built = [] + + def factory(_destination): + built.append(self.Recording()) + return built[-1] + + fan_out = TenantFanOutSpanProcessor(processor_factory=factory, shutdown_drain_seconds=0.05) + held = fan_out._acquire(self._dest(0)) + + fan_out.shutdown() + + assert held.shutdown_calls == 0 + + fan_out._release(held) + self._settle(fan_out, held) + + assert held.shutdown_calls == 1 + + def test_a_processor_built_while_shutdown_waits_is_still_closed(self): + """Shutdown cannot slip between the build and the insert, which would leave a + live exporter, with its batch thread and its connection pool, in a map nothing + will read again.""" + import threading + + built = [] + + def slow(_destination): + time.sleep(0.4) + built.append(self.Recording()) + return built[-1] + + fan_out = TenantFanOutSpanProcessor(processor_factory=slow, shutdown_drain_seconds=0.05) + acquired = [] + caller = threading.Thread(target=lambda: acquired.append(fan_out._acquire(self._dest(0)))) + caller.start() + time.sleep(0.1) + fan_out.shutdown() + caller.join(timeout=10) + + assert acquired == built, "the build shutdown waited out was thrown away" + + fan_out._release(built[0]) + self._settle(fan_out, built[0]) + + assert built[0].shutdown_calls == 1, "the exporter outlived the fan-out" + assert fan_out._processors == {}, "an exporter was left in a cleared cache" + + def test_shutdown_returns_when_a_destination_never_finishes_closing(self): + """Closing an exporter flushes over the network and the SDK joins its own + worker with no timeout, so a tenant collector that answers but never finishes + a response would hold process teardown open for as long as it likes.""" + import threading + + never = threading.Event() + + class Stuck(self.Recording): + def shutdown(self): + never.wait() + + built = [] + + def factory(_destination): + built.append(Stuck()) + return built[-1] + + fan_out = TenantFanOutSpanProcessor(processor_factory=factory, shutdown_drain_seconds=0.3) + fan_out._release(fan_out._acquire(self._dest(0))) + returned = threading.Event() + threading.Thread(target=lambda: (fan_out.shutdown(), returned.set()), daemon=True).start() + + came_back = returned.wait(timeout=8) + never.set() + + assert came_back, "shutdown never returned while a collector held its exporter open" + + def test_a_cold_cache_met_by_a_burst_builds_one_processor_per_destination(self): + """Building outside the cache lock let every thread of the burst construct its + own exporter, each with a batch thread and a connection pool, and shed all but + one into the drain.""" + import threading + + built = [] + + def factory(_destination): + time.sleep(0.01) + built.append(self.Recording()) + return built[-1] + + fan_out = TenantFanOutSpanProcessor(processor_factory=factory) + ready = threading.Barrier(8) + + def acquire(): + ready.wait() + fan_out._release(fan_out._acquire(self._dest(0))) + + callers = [threading.Thread(target=acquire) for _ in range(8)] + for caller in callers: + caller.start() + for caller in callers: + caller.join(timeout=10) + + assert len(built) == 1, f"one destination, {len(built)} exporters built" + + def test_a_submit_racing_close_is_never_stranded_behind_the_sentinels(self): + """A submit that read the closed state and then let ``close`` run queues its + processor after every sentinel, where the workers have already exited.""" + import queue + import threading + + from litellm.integrations.otel.plumbing.providers import _DrainPool + + at_the_put, close_returned = threading.Event(), threading.Event() + + class Gated(queue.Queue): + def put(self, item, *args, **kwargs): + if item is not None: + at_the_put.set() + close_returned.wait(timeout=1) + super().put(item, *args, **kwargs) + + pool = _DrainPool(pending=Gated()) + submitted = self.Recording() + submitter = threading.Thread(target=pool.submit, args=(submitted,)) + submitter.start() + assert at_the_put.wait(timeout=5) + closer = threading.Thread(target=pool.close) + closer.start() + closer.join(timeout=1.5) + close_returned.set() + submitter.join(timeout=5) + closer.join(timeout=5) + for _ in range(250): + if submitted.shutdown_calls: + break + time.sleep(0.02) + + assert submitted.shutdown_calls == 1, "a processor was queued behind the sentinels and never closed" + + def test_a_retired_processor_is_still_closed_after_shutdown(self): + """Eviction and shutdown can both land while a span is being forwarded, and the + evicted processor still has to be closed once that export returns.""" + from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + + fan_out, built = self._fan_out() + held = fan_out._acquire(self._dest(0)) + for index in range(1, _MAX_CACHED_DESTINATION_PROCESSORS + 1): + fan_out._acquire(self._dest(index)) + fan_out._release(built[-1]) + fan_out.shutdown() + + assert held.shutdown_calls == 0 + + fan_out._release(held) + self._settle(fan_out, held) + + assert held.shutdown_calls == 1 + + +class TestCredentialGatedExporters: + def test_layering_a_second_preset_does_not_eat_the_first_gated_exporter(self, monkeypatch): + """``base.Preset`` advertises ``config_overrides`` layering, and the gated spec + is itself a console exporter with no endpoint.""" + credential_less_proxy(monkeypatch) + from litellm.integrations.otel.presets.utils import credential_gated_exporters + + once = credential_gated_exporters((), ExporterOwner.LANGFUSE_OTEL) + twice = credential_gated_exporters(once, ExporterOwner.WEAVE_OTEL) + + assert [spec.owner for spec in twice] == [ExporterOwner.LANGFUSE_OTEL, ExporterOwner.WEAVE_OTEL] + + def test_an_exporter_the_operator_configured_survives(self): + from litellm.integrations.otel.presets.utils import credential_gated_exporters + + operator_console = ExporterSpec(kind="console", use_simple_processor=True) + + kept = credential_gated_exporters((operator_console,), ExporterOwner.LANGFUSE_OTEL) + + assert kept[0] == operator_console + + def test_an_otlp_exporter_on_its_default_endpoint_survives(self): + """``OTEL_EXPORTER=otlp_http`` with no endpoint is a real collector on the SDK's + default port, not the placeholder, so the transport is what tells them apart.""" + from litellm.integrations.otel.presets.utils import credential_gated_exporters + + operator_otlp = ExporterSpec(kind="otlp_http", endpoint=None, headers=None) + + kept = credential_gated_exporters((operator_otlp,), ExporterOwner.LANGFUSE_OTEL) + + assert kept[0] == operator_otlp + + def test_an_in_memory_exporter_the_operator_asked_for_survives(self): + """``OTEL_EXPORTER=in_memory`` stores spans, so it is a destination the operator + chose, not the placeholder that stands in for choosing nothing.""" + from litellm.integrations.otel.presets.utils import credential_gated_exporters + + operator_memory = ExporterSpec(kind="in_memory", endpoint=None, headers=None) + + kept = credential_gated_exporters((operator_memory,), ExporterOwner.LANGFUSE_OTEL) + + assert kept[0] == operator_memory + + def test_the_synthesized_stdout_placeholder_is_dropped(self, monkeypatch): + from litellm.integrations.otel.presets.utils import credential_gated_exporters + + for name in _OTEL_SHORTHAND_ENV: + monkeypatch.delenv(name, raising=False) + placeholder = OpenTelemetryV2Config().exporters[0] + + kept = credential_gated_exporters((placeholder,), ExporterOwner.LANGFUSE_OTEL) + + assert [spec.owner for spec in kept] == [ExporterOwner.LANGFUSE_OTEL] + + def test_a_console_exporter_the_operator_named_survives(self, monkeypatch): + """Same kind, endpoint and headers as the placeholder; only the fact that the + operator set ``OTEL_EXPORTER`` tells them apart.""" + from litellm.integrations.otel.presets.utils import credential_gated_exporters + + for name in _OTEL_SHORTHAND_ENV: + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("OTEL_EXPORTER", "console") + operator_console = OpenTelemetryV2Config().exporters[0] + + kept = credential_gated_exporters((operator_console,), ExporterOwner.LANGFUSE_OTEL) + + assert kept[0] is operator_console + + +class TestTenantHostSsrfGuard: + """Anyone who can mint a key can write ``langfuse_host``, so the host it names has + to be one the operator approved.""" + + @pytest.fixture(autouse=True) + def _guard_on(self, monkeypatch): + from litellm.integrations.otel.presets.destinations import _warn_host_not_allowlisted + + monkeypatch.setattr(litellm, "provider_url_destination_allowed_hosts", [], raising=False) + _warn_host_not_allowlisted.cache_clear() + yield + _warn_host_not_allowlisted.cache_clear() + + @staticmethod + def _langfuse(host: str) -> Mapping[str, str]: + return {"langfuse_public_key": "pk", "langfuse_secret_key": "sk", "langfuse_host": host} + + @pytest.mark.parametrize( + "host", + [ + "http://127.0.0.1:9111", + "http://169.254.169.254", + "http://10.0.0.5:3000", + "https://collector.example.com", + "https://langfuse.corp:99999", + "ftp://collector.example.com", + ], + ) + def test_a_host_the_operator_never_approved_resolves_to_nothing(self, host): + assert destination_for("langfuse_otel", self._langfuse(host)) is None + + def test_userinfo_naming_an_allowlisted_host_does_not_smuggle_a_second_one(self, monkeypatch): + """``https://allowed@10.0.0.5`` reads as the allowlisted host to the eye and + posts to 10.0.0.5 on the wire.""" + monkeypatch.setattr(litellm, "provider_url_destination_allowed_hosts", ["collector.example.com"], raising=False) + + assert destination_for("langfuse_otel", self._langfuse("https://collector.example.com@10.0.0.5")) is None + + def test_a_malformed_host_does_not_take_the_other_backends_with_it(self, monkeypatch): + """``urlparse(...).port`` raises a bare ValueError, which would escape + ``destination_for`` and kill the whole resolution.""" + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("NEW_RELIC_OTEL_ENDPOINT", "https://otlp.nr-data.net") + monkeypatch.setattr(litellm, "provider_url_destination_allowed_hosts", ["collector.example.com"], raising=False) + is_otel_v2_enabled.cache_clear() + auth = UserAPIKeyAuth( + token="hashed", + team_metadata={ + "logging": [ + {"callback_name": "langfuse_otel", "callback_vars": self._langfuse("https://lf.corp:99999")}, + {"callback_name": "newrelic", "callback_vars": {"newrelic_api_key": "nr"}}, + ] + }, + ) + + assert [d.callback_name for d in resolve_tenant_otel_destinations(auth)] == ["newrelic"] + + def test_the_operator_can_allowlist_its_teams_internal_langfuse(self, monkeypatch): + monkeypatch.setattr(litellm, "provider_url_destination_allowed_hosts", ["127.0.0.1:9111"], raising=False) + + destination = destination_for("langfuse_otel", self._langfuse("http://127.0.0.1:9111")) + + assert destination.endpoint == "http://127.0.0.1:9111/api/public/otel" + + def test_the_operators_own_internal_host_is_never_blocked(self, monkeypatch): + """The operator configures ``LANGFUSE_HOST`` themselves, so an internal + collector there is a deployment choice rather than caller-supplied input.""" + monkeypatch.setenv("LANGFUSE_HOST", "http://127.0.0.1:9111") + + destination = destination_for("langfuse_otel", {"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}) + + assert destination.endpoint == "http://127.0.0.1:9111/api/public/otel" + + def test_an_allowlisted_host_is_taken_without_resolving_it(self, monkeypatch): + """The check runs on the asyncio auth path, so it must not block on a name the + caller chose. ``.invalid`` never resolves, and it is still accepted.""" + monkeypatch.setattr(litellm, "provider_url_destination_allowed_hosts", ["lf.invalid"], raising=False) + + destination = destination_for("langfuse_otel", self._langfuse("https://lf.invalid")) + + assert destination.endpoint == "https://lf.invalid/api/public/otel" + + def test_a_rejected_host_is_warned_about_once(self, caplog): + with caplog.at_level("WARNING", logger="LiteLLM"): + for _ in range(3): + destination_for("langfuse_otel", self._langfuse("http://10.0.0.5:3000")) + + assert sum("provider_url_destination_allowed_hosts" in record.message for record in caplog.records) == 1 diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index b735abaf7bf..2869c804c07 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -2041,7 +2041,7 @@ def test_select_global_otel_v2_logger_builds_one_when_none_registered(): assert isinstance(chosen, OpenTelemetryV2) -def test_publish_global_otel_v2_provider_sets_selected_logger_provider(): +def test_publish_global_otel_v2_provider_sets_selected_logger_provider(monkeypatch): """The startup publish must set the OTel global provider to the *selected* logger's provider (the preset logger that owns every exporter), so the FastAPI server span and the gen-ai spans share one provider and one trace. @@ -2051,8 +2051,10 @@ def test_publish_global_otel_v2_provider_sets_selected_logger_provider(): test would otherwise miss: that the published provider is the selected logger's, not some other. """ + from litellm.integrations.otel import logger as otel_logger from litellm.integrations.otel.logger import publish_global_otel_v2_provider + monkeypatch.setattr(otel_logger, "_published_v2_provider", None) cfg = OpenTelemetryV2Config(exporter="in_memory") tp = providers.build_tracer_provider(cfg) preset_logger = OpenTelemetryV2( diff --git a/tests/test_litellm/integrations/pointfive/test_logger.py b/tests/test_litellm/integrations/pointfive/test_logger.py new file mode 100644 index 00000000000..52570751c74 --- /dev/null +++ b/tests/test_litellm/integrations/pointfive/test_logger.py @@ -0,0 +1,759 @@ +import asyncio +import gzip +import json +import logging +from collections.abc import Callable + +import pytest + +from litellm.integrations.pointfive.logger import PointFiveLogger +from litellm.integrations.pointfive.upload_client import PointFiveUploadError +from litellm.types.integrations.pointfive import DEFAULT_API_URL, PointFiveInitParams, PointFiveUploadFailure + +OBJECT_KEY = "some/object.ndjson.gz" + + +class FakeUploadClient: + """Records the objects a flush produced, so tests can read what would have shipped.""" + + def __init__( + self, + outcomes: list[str | PointFiveUploadFailure] | None = None, + ping_failure: PointFiveUploadFailure | None = None, + ) -> None: + self.outcomes = outcomes or [OBJECT_KEY] + self.bodies: list[bytes] = [] + self.on_upload: Callable[[], None] | None = None + self.ping_failure = ping_failure + self.pings = 0 + + async def ping(self) -> PointFiveUploadFailure | None: + self.pings += 1 + return self.ping_failure + + async def upload(self, body: bytes) -> str | PointFiveUploadFailure: + if self.on_upload is not None: + self.on_upload() + self.bodies.append(body) + return self.outcomes.pop(0) if len(self.outcomes) > 1 else self.outcomes[0] + + def records(self) -> list[dict]: + return [json.loads(line) for body in self.bodies for line in gzip.decompress(body).decode().splitlines()] + + +def _logger(upload_client: FakeUploadClient, **params) -> PointFiveLogger: + return PointFiveLogger(params=PointFiveInitParams(**params), upload_client=upload_client) + + +def _event(request_id: str, size: int = 0) -> dict: + return {"standard_logging_object": {"id": request_id, "model": "gpt-4o", "blob": "x" * size}} + + +@pytest.mark.asyncio +async def test_a_flush_ships_one_object_holding_every_buffered_record(): + """One object per flush is the whole point: s3_v2 sends one per request.""" + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=3) + + for request_id in ("a", "b", "c"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + await _settle(logger) + assert len(upload_client.bodies) == 1 + assert [record["id"] for record in upload_client.records()] == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_records_are_held_until_the_batch_is_full(): + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=3) + + await logger.async_log_success_event(_event("a"), None, None, None) + + assert upload_client.bodies == [] + assert len(logger.log_queue) == 1 + + +@pytest.mark.asyncio +async def test_failed_requests_are_logged_too(): + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=1) + + await logger.async_log_failure_event(_event("failed"), None, None, None) + + await _settle(logger) + assert [record["id"] for record in upload_client.records()] == ["failed"] + + +@pytest.mark.asyncio +async def test_an_event_without_a_standard_payload_is_skipped(): + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=1) + + await logger.async_log_success_event({"kwargs": "but no payload"}, None, None, None) + + assert upload_client.bodies == [] + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_a_batch_over_the_byte_cap_ships_as_several_objects(): + """Record count cannot bound an object: an unredacted payload dwarfs a redacted one.""" + cap = 600 + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=4, max_batch_bytes=cap) + + for request_id in ("a", "b", "c", "d"): + await logger.async_log_success_event(_event(request_id, size=200), None, None, None) + + await _settle(logger) + assert len(upload_client.bodies) > 1 + assert [record["id"] for record in upload_client.records()] == ["a", "b", "c", "d"] + assert all(len(gzip.decompress(body)) <= cap for body in upload_client.bodies) + + +@pytest.mark.asyncio +async def test_a_retryable_failure_keeps_the_batch_for_the_next_flush(): + upload_client = FakeUploadClient([PointFiveUploadFailure("upload target is down", retryable=True)]) + logger = _logger(upload_client, batch_size=2) + + for request_id in ("a", "b"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + assert [record["id"] for record in logger.log_queue] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_retryable_failure_surfaces_so_the_base_logger_can_preserve_it(): + upload_client = FakeUploadClient([PointFiveUploadFailure("upload target is down", retryable=True)]) + logger = _logger(upload_client, batch_size=99) + logger.log_queue.append(_event("a")["standard_logging_object"]) + + with pytest.raises(PointFiveUploadError, match="upload target is down"): + await logger.async_send_batch() + + +@pytest.mark.asyncio +async def test_a_rejected_batch_is_dropped_rather_than_blocking_the_queue(): + """Retrying a rejection forever would stall every record queued behind it.""" + upload_client = FakeUploadClient([PointFiveUploadFailure("object too large", retryable=False)]) + logger = _logger(upload_client, batch_size=2) + + for request_id in ("a", "b"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + await _settle(logger) + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_records_already_queued_ship_with_the_event_that_triggers_the_flush(): + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=1) + logger.log_queue.append(_event("mid-flight")["standard_logging_object"]) + + await logger.async_log_success_event(_event("a"), None, None, None) + + await _settle(logger) + assert [record["id"] for record in upload_client.records()] == ["mid-flight", "a"] + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_a_record_that_arrives_mid_flush_is_kept_for_the_next_one(): + """The queue is drained by count, so a record appended mid-upload must survive.""" + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=1) + upload_client.on_upload = lambda: logger.log_queue.append(_event("late")["standard_logging_object"]) + + await logger.async_log_success_event(_event("first"), None, None, None) + + await _settle(logger) + assert [record["id"] for record in upload_client.records()] == ["first"] + assert [record["id"] for record in logger.log_queue] == ["late"] + + +def test_defaults_favour_fewer_larger_uploads_over_freshness(): + upload_client = FakeUploadClient() + + logger = _logger(upload_client) + + assert logger.batch_size == 1_000 + assert logger.flush_interval == 300 + assert logger.max_batch_bytes == 8 * 1024 * 1024 + + +def test_the_default_api_url_is_the_pointfive_ingress(monkeypatch): + """api.pointfive.co is the host the ingress serves; .com does not resolve to it.""" + monkeypatch.setenv("POINTFIVE_API_KEY", "p5tu_env") + + logger = PointFiveLogger() + + assert logger.upload_client.api_url == "https://api.pointfive.co/api/v1/ingestion" + + +def test_the_api_key_can_come_from_the_environment(monkeypatch): + """The proxy ui configures a callback by writing environment variables.""" + monkeypatch.setenv("POINTFIVE_API_KEY", "p5tu_from_env") + + logger = PointFiveLogger() + + assert logger.upload_client.api_key == "p5tu_from_env" + + +def test_the_api_url_can_come_from_the_environment(monkeypatch): + monkeypatch.setenv("POINTFIVE_API_KEY", "p5tu_env") + monkeypatch.setenv("POINTFIVE_API_URL", "https://api.staging.pointfive.co/api/v1/ingestion") + + logger = PointFiveLogger() + + assert logger.upload_client.api_url == "https://api.staging.pointfive.co/api/v1/ingestion" + + +def test_config_yaml_wins_over_the_environment(monkeypatch): + """A value set in config.yaml is explicit, so it outranks whatever the ui left behind.""" + monkeypatch.setenv("POINTFIVE_API_KEY", "p5tu_from_env") + monkeypatch.setenv("POINTFIVE_API_URL", "https://from-env.example/api/v1/ingestion") + + logger = PointFiveLogger( + params=PointFiveInitParams(api_key="p5tu_from_config", api_url="https://from-config.example/api/v1/ingestion") + ) + + assert logger.upload_client.api_key == "p5tu_from_config" + assert logger.upload_client.api_url == "https://from-config.example/api/v1/ingestion" + + +def test_a_missing_api_key_fails_at_startup_not_at_the_first_flush(monkeypatch): + monkeypatch.delenv("POINTFIVE_API_KEY", raising=False) + + with pytest.raises(ValueError, match="api key"): + PointFiveLogger(params=PointFiveInitParams()) + + +def test_an_api_key_can_be_an_environment_reference(monkeypatch): + """config.yaml spells secrets as `os.environ/NAME`, so the plugin must resolve one.""" + monkeypatch.setenv("POINTFIVE_TEST_KEY", "p5tu_from_env") + + logger = PointFiveLogger(params=PointFiveInitParams(api_key="os.environ/POINTFIVE_TEST_KEY")) + + assert logger.upload_client.api_key == "p5tu_from_env" + + +def test_params_are_read_from_litellm_settings(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "pointfive_params", {"api_key": "p5tu_configured", "batch_size": 7}) + + logger = PointFiveLogger() + + assert logger.upload_client.api_key == "p5tu_configured" + assert logger.batch_size == 7 + + +def test_an_out_of_range_setting_is_rejected(): + with pytest.raises(ValueError, match="batch_size"): + PointFiveInitParams(api_key="p5tu_k", batch_size=0) + + +@pytest.mark.asyncio +async def test_an_idle_flush_reports_liveness_instead_of_uploading(): + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=99) + + await logger.flush_queue() + + assert upload_client.pings == 1 + assert upload_client.bodies == [] + + +@pytest.mark.asyncio +async def test_a_flush_with_records_uploads_and_does_not_ping(): + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=1) + + await logger.async_log_success_event(_event("a"), None, None, None) + + await _settle(logger) + assert upload_client.pings == 0 + assert len(upload_client.bodies) == 1 + + +@pytest.mark.asyncio +async def test_a_failed_ping_does_not_raise(): + """Liveness is bookkeeping; a proxy must not see errors from it.""" + upload_client = FakeUploadClient(ping_failure=PointFiveUploadFailure("api down", retryable=True)) + logger = _logger(upload_client, batch_size=99) + + await logger.flush_queue() + + assert upload_client.pings == 1 + + +@pytest.mark.asyncio +async def test_health_check_is_healthy_when_the_api_accepts_the_key(): + upload_client = FakeUploadClient() + + assert await _logger(upload_client).async_health_check() == {"status": "healthy", "error_message": None} + assert upload_client.pings == 1 + + +@pytest.mark.asyncio +async def test_health_check_reports_why_the_api_refused(): + """The ui test button shows this message, so a rejected key has to say so rather than pass.""" + upload_client = FakeUploadClient(ping_failure=PointFiveUploadFailure("key was revoked", retryable=False)) + + outcome = await _logger(upload_client).async_health_check() + + assert outcome == {"status": "unhealthy", "error_message": "key was revoked"} + + +def test_the_client_follows_a_key_and_url_changed_after_startup(monkeypatch): + """ + The proxy ui writes new values into a running proxy's environment. + + Reading them once at construction would leave the logger talking to the old endpoint + until someone restarted the proxy. + """ + monkeypatch.setenv("POINTFIVE_API_KEY", "p5tu_first") + monkeypatch.setenv("POINTFIVE_API_URL", "https://first.example.invalid/api/v1/ingestion") + logger = PointFiveLogger(params=PointFiveInitParams()) + + assert logger.upload_client.api_key == "p5tu_first" + assert logger.upload_client.api_url == "https://first.example.invalid/api/v1/ingestion" + + monkeypatch.setenv("POINTFIVE_API_KEY", "p5tu_second") + monkeypatch.setenv("POINTFIVE_API_URL", "https://second.example.invalid/api/v1/ingestion") + + assert logger.upload_client.api_key == "p5tu_second" + assert logger.upload_client.api_url == "https://second.example.invalid/api/v1/ingestion" + + +def test_a_configured_key_still_wins_over_the_environment(monkeypatch): + monkeypatch.setenv("POINTFIVE_API_KEY", "p5tu_from_env") + logger = PointFiveLogger(params=PointFiveInitParams(api_key="p5tu_from_config")) + + assert logger.upload_client.api_key == "p5tu_from_config" + + +@pytest.mark.asyncio +async def test_health_check_says_so_when_the_key_was_removed(monkeypatch): + monkeypatch.setenv("POINTFIVE_API_KEY", "p5tu_present") + logger = PointFiveLogger(params=PointFiveInitParams()) + monkeypatch.delenv("POINTFIVE_API_KEY") + + outcome = await logger.async_health_check() + + assert outcome["status"] == "unhealthy" + assert "requires an api key" in (outcome["error_message"] or "") + + +def _pending_flush_tasks() -> tuple[asyncio.Task, ...]: + return tuple(task for task in asyncio.all_tasks() if "periodic_flush" in str(task.get_coro())) + + +@pytest.mark.asyncio +async def test_a_one_shot_logger_leaves_no_flush_task_behind(): + """ + A health check builds a logger for a single answer and drops it. + + Without this, every check would leave a flusher running that keeps pinging for the + lifetime of the proxy. + """ + before = _pending_flush_tasks() + + logger = PointFiveLogger(params=PointFiveInitParams(), upload_client=FakeUploadClient(), start_periodic_flush=False) + + assert logger._periodic_flush_task is None + assert _pending_flush_tasks() == before + + +@pytest.mark.asyncio +async def test_the_logger_flushes_periodically_by_default(): + logger = PointFiveLogger(params=PointFiveInitParams(), upload_client=FakeUploadClient()) + + assert logger._periodic_flush_task is not None + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_params_already_built_are_used_as_they_are(monkeypatch): + """config.yaml is validated once into a params object; a second validation would be wasted.""" + import litellm + + monkeypatch.setattr(litellm, "pointfive_params", PointFiveInitParams(max_batch_bytes=4096)) + + logger = PointFiveLogger(upload_client=FakeUploadClient()) + + assert logger.max_batch_bytes == 4096 + + +@pytest.mark.asyncio +async def test_a_dead_flush_task_is_restarted_by_the_next_event(): + """A cancelled or crashed flusher would otherwise leave the queue growing forever.""" + logger = _logger(FakeUploadClient()) + logger._periodic_flush_task.cancel() + await asyncio.sleep(0) # let the cancellation land, so the task reports itself done + + await logger.async_log_success_event(_event("after-cancel"), None, None, None) + + assert logger._periodic_flush_task is not None + assert not logger._periodic_flush_task.done() + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_a_failure_while_queueing_never_breaks_the_request(): + """Logging sits on the request path, so a fault here must not surface to the caller.""" + + class ExplodingQueue(list): + def append(self, _item): + raise RuntimeError("queue is broken") + + upload_client = FakeUploadClient() + logger = _logger(upload_client) + logger.log_queue = ExplodingQueue() + + await logger.async_log_success_event(_event("boom"), None, None, None) + + logger.log_queue = [] + await logger.async_log_success_event(_event("after-the-fault"), None, None, None) + assert [record["id"] for record in logger.log_queue] == ["after-the-fault"] + + +@pytest.mark.asyncio +async def test_a_flush_with_nothing_queued_uploads_nothing(): + upload_client = FakeUploadClient() + logger = _logger(upload_client) + + await logger.async_send_batch() + + assert upload_client.bodies == [] + + +@pytest.mark.asyncio +async def test_the_idle_ping_is_skipped_when_the_key_was_removed(monkeypatch, caplog): + """A key pulled mid-flight must not turn the periodic flush into an exception.""" + monkeypatch.setenv("POINTFIVE_API_KEY", "p5tu_present") + logger = PointFiveLogger(params=PointFiveInitParams()) + monkeypatch.delenv("POINTFIVE_API_KEY") + + with caplog.at_level(logging.WARNING): + await logger.flush_queue() + + assert "liveness ping skipped" in caplog.text + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_a_full_batch_stands_down_while_a_flush_is_already_running(): + """ + Under load every event landing mid-upload also crosses the batch threshold. + + Letting each one flush turns a single burst into a stream of tiny objects, which is + what batching exists to avoid, so a full batch defers to the flush already running. + """ + upload_client = FakeUploadClient() + # No periodic task: this test drives the flushes itself, and the loop's opening cycle + # would otherwise ship the queue it seeds below. + logger = PointFiveLogger( + params=PointFiveInitParams(batch_size=2), + upload_client=upload_client, + start_periodic_flush=False, + ) + release = asyncio.Event() + finish_upload = upload_client.upload + + async def held_upload(body: bytes): + await release.wait() + return await finish_upload(body) + + upload_client.upload = held_upload + logger.log_queue.extend(_event(f"first-{index}")["standard_logging_object"] for index in range(2)) + + flushing = asyncio.create_task(logger.flush_queue()) + await asyncio.sleep(0) + + # Bounded: without the guard these block on the flush lock the held upload owns. + for index in range(6): + await asyncio.wait_for(logger.async_log_success_event(_event(f"mid-{index}"), None, None, None), timeout=2) + + assert upload_client.bodies == [] + + release.set() + await flushing + + assert len(upload_client.bodies) == 1 + assert [record["id"] for record in upload_client.records()] == ["first-0", "first-1"] + assert [record["id"] for record in logger.log_queue] == [f"mid-{index}" for index in range(6)] + + +async def _settle(logger) -> None: + """ + Wait out the flush a full batch schedules, the way the proxy's loop would. + + The upload runs off the request path now, and either the batch task or the periodic + loop can be the one carrying it, so this waits for whichever is in flight to finish. + """ + for _ in range(200): + await asyncio.sleep(0.001) + task = logger._batch_flush_task + if task is not None and not task.done(): + await task + if not logger._flushing: + return + raise AssertionError("the flush never finished") + + +async def _until(done: Callable[[], bool], ticks: int = 400) -> None: + """Wait for a condition the flush path reaches only after gzip finishes on a worker thread.""" + for _ in range(ticks): + if done(): + return + await asyncio.sleep(0.005) + raise AssertionError("condition never became true") + + +@pytest.mark.asyncio +async def test_a_new_logger_announces_itself_without_waiting_for_the_interval(): + """ + Configuring the callback must make the integration connect, with no traffic and no test click. + + The inherited loop sleeps a whole interval before its first flush, which left a freshly + configured proxy silent for five minutes and the integration looking unconfigured. + """ + upload_client = FakeUploadClient() + logger = _logger(upload_client) + + await asyncio.sleep(0) # let the flush task reach its first cycle + + assert upload_client.pings == 1 + assert upload_client.bodies == [] + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_the_first_cycle_ships_records_rather_than_announcing(): + """Announcing is only for an empty queue: records already waiting must go out as an upload.""" + upload_client = FakeUploadClient() + logger = PointFiveLogger( + params=PointFiveInitParams(batch_size=100), + upload_client=upload_client, + start_periodic_flush=False, + ) + logger.log_queue.append(_event("queued-before-start")["standard_logging_object"]) + + logger._periodic_flush_task = logger._start_periodic_flush_task() + await _until(lambda: bool(upload_client.bodies)) + + assert upload_client.pings == 0 + assert [record["id"] for record in upload_client.records()] == ["queued-before-start"] + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_a_failed_request_ships_redacted_when_message_logging_is_off(): + """ + Failure events skip the framework's redaction, so the callback has to redact what it buffers. + + Without this the prompt of every failed request reaches PointFive in full, even though + the integration was configured not to send message content. + """ + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=1, turn_off_message_logging=True) + event = _event("failed-request") + event["standard_logging_object"]["messages"] = [{"role": "user", "content": "my secret prompt"}] + event["standard_logging_object"]["response"] = "the secret answer" + + await logger.async_log_failure_event(event, None, None, None) + await _settle(logger) + + shipped = upload_client.records()[0] + assert "my secret prompt" not in json.dumps(shipped) + assert "the secret answer" not in json.dumps(shipped) + assert event["standard_logging_object"]["messages"][0]["content"] == "my secret prompt" + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_a_dead_loop_does_not_strand_the_flusher(): + """A task whose loop was closed never runs and never reports done, so it must be replaced.""" + logger = PointFiveLogger(params=PointFiveInitParams(), upload_client=FakeUploadClient(), start_periodic_flush=False) + stranded_loop = asyncio.new_event_loop() + forever = asyncio.sleep(3600) + logger._periodic_flush_task = stranded_loop.create_task(forever) + stranded_loop.close() + forever.close() + + await logger.async_log_success_event(_event("after-loop-close"), None, None, None) + + assert logger._periodic_flush_task is not None + assert logger._periodic_flush_task.get_loop() is asyncio.get_running_loop() + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_a_retry_does_not_resend_objects_that_already_landed(): + """ + A failure part way through a multi-object flush used to hand the whole batch back. + + Every record already shipped, and every record already refused for good, went out + again on the next flush, so PointFive received duplicates of both. + """ + upload_client = FakeUploadClient(outcomes=[OBJECT_KEY, PointFiveUploadFailure("service busy", retryable=True)]) + logger = PointFiveLogger( + params=PointFiveInitParams(max_batch_bytes=1), # one record per object + upload_client=upload_client, + start_periodic_flush=False, + ) + logger.log_queue.extend(_event(request_id)["standard_logging_object"] for request_id in ("first", "second")) + + with pytest.raises(PointFiveUploadError): + await logger.async_send_batch() + + assert [record["id"] for record in logger.log_queue] == ["second"] + + +@pytest.mark.asyncio +async def test_the_queue_stops_growing_at_its_cap_without_waiting_for_a_failure(): + """The base class trims only after a failed send, so a proxy that keeps flushing never trims.""" + logger = _logger(FakeUploadClient(), batch_size=10_000) + logger.max_queue_size = 3 + + for request_id in ("a", "b", "c", "d", "e"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + assert [record["id"] for record in logger.log_queue] == ["c", "d", "e"] + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_a_failed_request_honours_the_global_redaction_setting(monkeypatch): + """ + Redaction can be turned on globally or per request, not only on this callback. + + The async failure path hands the payload over untouched, so a tenant could trigger a + provider failure and ship prompts that the operator had already asked to be redacted. + """ + import litellm + + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=1) + event = _event("globally-redacted") + event["standard_logging_object"]["messages"] = [{"role": "user", "content": "my secret prompt"}] + + await logger.async_log_failure_event(event, None, None, None) + + await _settle(logger) + assert "my secret prompt" not in json.dumps(upload_client.records()[0]) + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_excluded_fields_are_dropped_from_a_failed_request(monkeypatch): + """standard_logging_payload_excluded_fields drops a field entirely; failures skipped it too.""" + import litellm + + monkeypatch.setattr(litellm, "standard_logging_payload_excluded_fields", ["messages"]) + upload_client = FakeUploadClient() + logger = _logger(upload_client, batch_size=1, turn_off_message_logging=True) + event = _event("field-excluded") + event["standard_logging_object"]["messages"] = [{"role": "user", "content": "my secret prompt"}] + + await logger.async_log_failure_event(event, None, None, None) + await _settle(logger) + + shipped = upload_client.records()[0] + assert "messages" not in shipped + assert shipped["id"] == "field-excluded" + logger._periodic_flush_task.cancel() + + +def _held_upload(upload_client: FakeUploadClient, release: asyncio.Event) -> None: + finish = upload_client.upload + + async def held(body: bytes): + await release.wait() + return await finish(body) + + upload_client.upload = held + + +@pytest.mark.asyncio +async def test_a_full_batch_does_not_hold_the_request(): + """ + The upload belongs off the request path. + + Awaiting it inline meant a hung PointFive api held the caller's response open for as + long as the attempts and their backoff took. + """ + upload_client = FakeUploadClient() + release = asyncio.Event() + _held_upload(upload_client, release) + logger = _logger(upload_client, batch_size=1) + + await asyncio.wait_for(logger.async_log_success_event(_event("first"), None, None, None), timeout=2) + + assert upload_client.bodies == [] + release.set() + await _settle(logger) + assert [record["id"] for record in upload_client.records()] == ["first"] + logger._periodic_flush_task.cancel() + + +@pytest.mark.asyncio +async def test_records_arriving_during_a_flush_survive_the_queue_cap(): + """ + The flush drains by count, so trimming the front underneath it loses records. + + Records that arrived while the upload was in flight would be deleted by that drain + without ever being sent. + """ + upload_client = FakeUploadClient() + release = asyncio.Event() + _held_upload(upload_client, release) + logger = PointFiveLogger( + params=PointFiveInitParams(batch_size=2), + upload_client=upload_client, + start_periodic_flush=False, + ) + logger.max_queue_size = 2 + logger.log_queue.extend(_event(request_id)["standard_logging_object"] for request_id in ("a", "b")) + + flushing = asyncio.create_task(logger.flush_queue()) + await asyncio.sleep(0.01) + for request_id in ("c", "d", "e"): + await logger.async_log_success_event(_event(request_id), None, None, None) + release.set() + await flushing + + assert [record["id"] for record in upload_client.records()] == ["a", "b"] + assert [record["id"] for record in logger.log_queue] == ["c", "d", "e"] + logger._periodic_flush_task.cancel() + + +def test_an_unset_env_reference_is_never_used_as_the_key(monkeypatch): + """ + A config that names a missing variable has no key, and must say so. + + Falling back to the reference text sent the literal "os.environ/NAME" as the bearer + token, so the callback started and every upload was rejected for the wrong reason. + """ + monkeypatch.delenv("POINTFIVE_API_KEY", raising=False) + monkeypatch.delenv("POINTFIVE_MISSING_KEY", raising=False) + + with pytest.raises(ValueError, match="requires an api key"): + PointFiveLogger( + params=PointFiveInitParams(api_key="os.environ/POINTFIVE_MISSING_KEY"), + start_periodic_flush=False, + ) + + +def test_an_unset_url_reference_falls_back_to_the_public_endpoint(monkeypatch): + """An unresolved url reference must not become the destination the proxy uploads to.""" + from litellm.integrations.pointfive.logger import _resolved_api_url + + monkeypatch.delenv("POINTFIVE_API_URL", raising=False) + monkeypatch.delenv("POINTFIVE_MISSING_URL", raising=False) + + assert _resolved_api_url(PointFiveInitParams(api_url="os.environ/POINTFIVE_MISSING_URL")) == DEFAULT_API_URL diff --git a/tests/test_litellm/integrations/pointfive/test_payload.py b/tests/test_litellm/integrations/pointfive/test_payload.py new file mode 100644 index 00000000000..d310941e396 --- /dev/null +++ b/tests/test_litellm/integrations/pointfive/test_payload.py @@ -0,0 +1,80 @@ +import gzip +import json + +import pytest + +from litellm.integrations.pointfive.payload import chunk_lines, encode_lines, serialize_records + +UNBOUNDED = 10_000_000 + + +def test_each_record_becomes_one_json_line(): + lines = serialize_records([{"id": "a"}, {"id": "b"}, {"id": "c"}]) + + assert len(lines) == 3 + assert [json.loads(line)["id"] for line in lines] == ["a", "b", "c"] + + +def test_non_serializable_values_do_not_raise(): + """An odd payload must not kill the flush.""" + lines = serialize_records([{"id": "a", "when": object()}]) + + assert json.loads(lines[0])["id"] == "a" + + +def test_records_that_fit_stay_in_one_object(): + lines = serialize_records([{"id": f"r{i}"} for i in range(50)]) + + assert chunk_lines(lines, UNBOUNDED) == (lines,) + + +def test_objects_are_capped_by_uncompressed_size(): + lines = serialize_records([{"id": f"r{i}", "blob": "x" * 100} for i in range(10)]) + line_bytes = len(lines[0].encode("utf-8")) + 1 + + chunks = chunk_lines(lines, line_bytes * 3) + + assert [len(chunk) for chunk in chunks] == [3, 3, 3, 1] + + +def test_oversized_single_record_is_sent_alone_not_stalled(): + """A record too big for the cap must still go out, or it blocks everything behind it.""" + lines = serialize_records([{"id": "small"}, {"id": "huge", "blob": "x" * 5000}, {"id": "small2"}]) + + chunks = chunk_lines(lines, 200) + + assert sum(len(chunk) for chunk in chunks) == 3 + huge = [chunk for chunk in chunks if any("huge" in line for line in chunk)] + assert len(huge) == 1 + assert len(huge[0]) == 1 + + +def test_no_records_produces_no_objects(): + assert chunk_lines((), UNBOUNDED) == () + + +def test_every_record_appears_exactly_once(): + lines = serialize_records([{"id": f"r{i}"} for i in range(37)]) + + chunks = chunk_lines(lines, len(lines[0]) * 4) + + assert [line for chunk in chunks for line in chunk] == list(lines) + + +@pytest.mark.asyncio +async def test_encode_lines_round_trips_through_gzip(): + lines = serialize_records([{"id": f"r{i}"} for i in range(5)]) + + encoded = await encode_lines(lines) + + assert encoded[:2] == b"\x1f\x8b" + assert gzip.decompress(encoded).decode("utf-8") == "\n".join(lines) + + +@pytest.mark.asyncio +async def test_encode_lines_compresses_repetitive_records(): + lines = serialize_records([{"id": f"r{i}", "model": "gpt-4o", "cost": 0.01} for i in range(200)]) + + encoded = await encode_lines(lines) + + assert len(encoded) < len(gzip.decompress(encoded)) / 2 diff --git a/tests/test_litellm/integrations/pointfive/test_upload_client.py b/tests/test_litellm/integrations/pointfive/test_upload_client.py new file mode 100644 index 00000000000..50ef085386d --- /dev/null +++ b/tests/test_litellm/integrations/pointfive/test_upload_client.py @@ -0,0 +1,378 @@ +import json +from collections.abc import Sequence + +import httpx +import pytest + +import litellm +from litellm.integrations.pointfive.upload_client import PointFiveUploadClient +from litellm.litellm_core_utils.url_utils import validate_url +from litellm.types.integrations.pointfive import PointFiveUploadFailure + +API_URL = "https://api.pointfive.co/api/v1/ingestion" +UPLOAD_URL = "https://uploads.example.invalid/some/object.ndjson.gz?signature=sig" +OBJECT_KEY = "some/object.ndjson.gz" +BODY = b"gzipped-bytes" + + +def _presigned(status_code: int = 200) -> httpx.Response: + return _response( + status_code, {"uploadUrl": UPLOAD_URL, "objectKey": OBJECT_KEY, "expiresAt": "2026-08-25T14:35:00Z"} + ) + + +def _response(status_code: int, payload: object) -> httpx.Response: + return httpx.Response(status_code, text=json.dumps(payload)) + + +def _refused(status_code: int, error: str) -> httpx.Response: + """The body PointFive sends with every refusal.""" + return _response(status_code, {"success": False, "error": error}) + + +def _accepted() -> httpx.Response: + return httpx.Response(200, text="") + + +def _no_content() -> httpx.Response: + return httpx.Response(204, text="") + + +class FakeHTTPClient: + """ + Stands in for AsyncHTTPHandler, including its habit of raising on error statuses. + + Scripted results are consumed in order, and the last one repeats, so a test that + cares about a single behaviour passes a single result. + """ + + def __init__( + self, + presign: Sequence[httpx.Response | Exception] | None = None, + put: Sequence[httpx.Response | Exception] | None = None, + ) -> None: + self.presign = list(presign) if presign else [_presigned()] # mutable-ok: results are consumed by popping + self.put_results = list(put) if put else [_accepted()] # mutable-ok: results are consumed by popping + self.presign_calls: list[dict] = [] + self.put_calls: list[dict] = [] + + async def post(self, url, json=None, headers=None, **_): + self.presign_calls.append({"url": url, "json": json, "headers": headers or {}}) + return _next_result(self.presign, url) + + async def put(self, url, data=None, headers=None, follow_redirects=None, **_): + self.put_calls.append( + {"url": url, "data": data, "headers": headers or {}, "follow_redirects": follow_redirects} + ) + return _next_result(self.put_results, url) + + +def _next_result(results: list, url: str) -> httpx.Response: + result = results.pop(0) if len(results) > 1 else results[0] + if isinstance(result, Exception): + raise result + if result.status_code >= 300: + request = httpx.Request("POST", url) + raise httpx.HTTPStatusError( + "boom", + request=request, + response=httpx.Response(result.status_code, text=result.text, headers=result.headers), + ) + return result + + +async def _no_backoff(_seconds: float) -> None: + return None + + +def _trusting_validator(url: str) -> tuple[str, str]: + """Stands in for validate_url so the fixture hosts need no DNS; the SSRF tests use the real one.""" + return url, httpx.URL(url).host + + +def _client( + http_client: FakeHTTPClient, + max_retries: int = 3, + api_url: str = API_URL, + validate_upload_url=_trusting_validator, +) -> PointFiveUploadClient: + return PointFiveUploadClient( + api_key="p5tu_testkey", + api_url=api_url, + http_client=http_client, + max_retries=max_retries, + sleep=_no_backoff, + validate_upload_url=validate_upload_url, + ) + + +def _presigned_for(upload_url: str) -> httpx.Response: + return _response(200, {"uploadUrl": upload_url, "objectKey": OBJECT_KEY, "expiresAt": "2026-08-25T14:35:00Z"}) + + +@pytest.mark.asyncio +async def test_uploads_the_body_to_the_url_the_api_returned(): + http_client = FakeHTTPClient() + + outcome = await _client(http_client).upload(BODY) + + assert outcome == OBJECT_KEY + assert http_client.put_calls[0]["url"] == UPLOAD_URL + assert http_client.put_calls[0]["data"] == BODY + + +@pytest.mark.asyncio +async def test_presign_request_is_authenticated_and_sized(): + http_client = FakeHTTPClient() + + await _client(http_client).upload(BODY) + + call = http_client.presign_calls[0] + assert call["url"] == "https://api.pointfive.co/api/v1/ingestion/upload-url" + assert call["headers"]["Authorization"] == "Bearer p5tu_testkey" + assert call["json"] == {"kind": "LITELLM", "byteCount": len(BODY)} + + +@pytest.mark.asyncio +async def test_a_trailing_slash_on_the_api_url_is_tolerated(): + """A pasted URL often ends in a slash; it must not produce a double slash in the path.""" + http_client = FakeHTTPClient() + + await _client(http_client, api_url=API_URL + "/").upload(BODY) + + assert http_client.presign_calls[0]["url"] == "https://api.pointfive.co/api/v1/ingestion/upload-url" + + +@pytest.mark.asyncio +async def test_no_bearer_token_is_sent_to_the_presigned_url(): + """The URL carries its own authorization, so the api key must not travel with it.""" + http_client = FakeHTTPClient() + + await _client(http_client).upload(BODY) + + assert "Authorization" not in http_client.put_calls[0]["headers"] + + +@pytest.mark.asyncio +async def test_the_upload_pins_the_host_and_never_follows_a_redirect(): + http_client = FakeHTTPClient() + + await _client(http_client).upload(BODY) + + call = http_client.put_calls[0] + assert call["headers"]["Host"] == "uploads.example.invalid" + assert call["follow_redirects"] is False + + +@pytest.mark.asyncio +async def test_a_redirected_upload_is_refused_rather_than_followed(): + """A presigned URL never redirects legitimately; following one is how a bad endpoint reaches inside.""" + http_client = FakeHTTPClient(put=[httpx.Response(301, headers={"location": "http://169.254.169.254/"})]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure( + "presigned upload redirected with 301, refusing to follow", retryable=False + ) + assert len(http_client.put_calls) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "upload_url", + [ + "http://169.254.169.254/latest/meta-data", + "https://10.0.0.7/internal/bucket/object", + "http://127.0.0.1:9000/bucket/object", + ], +) +async def test_an_upload_url_inside_the_network_is_refused_before_any_bytes_leave(upload_url): + http_client = FakeHTTPClient(presign=[_presigned_for(upload_url)]) + + outcome = await _client(http_client, validate_upload_url=validate_url).upload(BODY) + + assert isinstance(outcome, PointFiveUploadFailure) + assert not outcome.retryable + assert outcome.detail.startswith("presigned upload url refused: ") + assert http_client.put_calls == [] + + +@pytest.mark.asyncio +async def test_an_operator_can_switch_destination_validation_off(monkeypatch): + """litellm.user_url_validation is the proxy-wide switch every SSRF guard honours.""" + monkeypatch.setattr(litellm, "user_url_validation", False) + http_client = FakeHTTPClient(presign=[_presigned_for("http://10.0.0.7/bucket/object")]) + + outcome = await _client(http_client, validate_upload_url=validate_url).upload(BODY) + + assert outcome == OBJECT_KEY + assert http_client.put_calls[0]["url"] == "http://10.0.0.7/bucket/object" + assert "Host" not in http_client.put_calls[0]["headers"] + + +@pytest.mark.asyncio +async def test_the_object_is_declared_as_gzipped_ndjson(): + http_client = FakeHTTPClient() + + await _client(http_client).upload(BODY) + + assert http_client.put_calls[0]["headers"]["Content-Encoding"] == "gzip" + assert http_client.put_calls[0]["headers"]["Content-Type"] == "application/x-ndjson" + + +@pytest.mark.asyncio +async def test_each_retry_presigns_again(): + """A retry must never reuse a URL that was consumed or has expired.""" + http_client = FakeHTTPClient(put=[httpx.Response(503), _accepted()]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == OBJECT_KEY + assert len(http_client.presign_calls) == 2 + assert len(http_client.put_calls) == 2 + + +@pytest.mark.asyncio +async def test_retryable_upload_failure_gives_up_after_max_retries(): + http_client = FakeHTTPClient(put=[httpx.Response(503)]) + + outcome = await _client(http_client, max_retries=2).upload(BODY) + + assert outcome == PointFiveUploadFailure("presigned upload returned 503, gave up after 2 attempts", retryable=True) + assert len(http_client.put_calls) == 2 + + +@pytest.mark.asyncio +async def test_rejected_upload_is_not_retried(): + http_client = FakeHTTPClient(put=[httpx.Response(403)]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("presigned upload returned 403", retryable=False) + assert len(http_client.put_calls) == 1 + + +@pytest.mark.asyncio +async def test_bad_api_key_is_not_retried(): + http_client = FakeHTTPClient(presign=[httpx.Response(401)]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("pointfive api returned 401", retryable=False) + assert http_client.put_calls == [] + + +@pytest.mark.asyncio +async def test_the_reason_for_a_refusal_is_surfaced(): + """A 403 means the key no longer maps to an integration; the operator needs to read why.""" + http_client = FakeHTTPClient(presign=[_refused(403, "no integration accepts uploads from this api key")]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure( + "pointfive api returned 403, no integration accepts uploads from this api key", retryable=False + ) + assert http_client.put_calls == [] + + +@pytest.mark.asyncio +async def test_api_server_error_is_retried(): + http_client = FakeHTTPClient(presign=[httpx.Response(503), _presigned()]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == OBJECT_KEY + assert len(http_client.presign_calls) == 2 + + +@pytest.mark.asyncio +async def test_too_many_requests_is_retried(): + http_client = FakeHTTPClient(presign=[httpx.Response(429), _presigned()]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == OBJECT_KEY + assert len(http_client.presign_calls) == 2 + + +@pytest.mark.asyncio +async def test_unreachable_api_is_retried_then_reported_as_retryable(): + http_client = FakeHTTPClient(presign=(ConnectionError("down"),)) + + outcome = await _client(http_client, max_retries=2).upload(BODY) + + assert isinstance(outcome, PointFiveUploadFailure) + assert outcome.retryable + assert "unreachable" in outcome.detail + assert len(http_client.presign_calls) == 2 + + +@pytest.mark.asyncio +async def test_malformed_api_body_is_not_retried(): + http_client = FakeHTTPClient(presign=[_response(200, {"objectKey": "k"})]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False) + assert http_client.put_calls == [] + + +@pytest.mark.asyncio +async def test_a_body_that_is_not_json_is_reported_as_unreadable(): + http_client = FakeHTTPClient(presign=(httpx.Response(200, text="gateway"),)) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False) + assert http_client.put_calls == [] + + +@pytest.mark.asyncio +async def test_ping_reports_a_live_shipper(): + http_client = FakeHTTPClient(presign=(_no_content(),)) + + failure = await _client(http_client).ping() + + assert failure is None + assert http_client.presign_calls[0]["url"] == "https://api.pointfive.co/api/v1/ingestion/ping" + assert http_client.presign_calls[0]["json"] == {"kind": "LITELLM"} + + +@pytest.mark.asyncio +async def test_ping_surfaces_a_revoked_key(): + http_client = FakeHTTPClient(presign=(_refused(403, "no integration accepts uploads from this api key"),)) + + failure = await _client(http_client).ping() + + assert failure is not None + assert not failure.retryable + assert "no integration accepts uploads from this api key" in failure.detail + + +@pytest.mark.asyncio +async def test_ping_surfaces_an_unreachable_api(): + http_client = FakeHTTPClient(presign=(ConnectionError("down"),)) + + failure = await _client(http_client).ping() + + assert failure is not None + assert failure.retryable + + +@pytest.mark.asyncio +async def test_a_transport_fault_on_the_upload_itself_is_retryable(): + http_client = FakeHTTPClient(put=(ConnectionError("reset"),)) + + outcome = await _client(http_client, max_retries=1).upload(BODY) + + assert isinstance(outcome, PointFiveUploadFailure) + assert outcome.retryable + assert "presigned upload unreachable" in outcome.detail + + +@pytest.mark.asyncio +async def test_a_client_that_may_not_try_at_all_says_so(): + """max_upload_retries is validated as >= 1, so this guards the loop against a future zero.""" + outcome = await _client(FakeHTTPClient(), max_retries=0).upload(BODY) + + assert outcome == PointFiveUploadFailure("max_upload_retries must be at least 1", retryable=False) diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index de8b654987b..6b7780acd20 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1777,6 +1777,224 @@ class TestEnableAnthropicPromptCaching: assert messages == before +class TestClaudeCodeOneShotAutoCaching: + BILLING_TEXT = "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli; cc_is_subagent=true;" + BILLING_SYSTEM = [{"type": "text", "text": BILLING_TEXT}] + MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "unique fetched document"}]}] + + @staticmethod + def _kwargs(configured=None): + kwargs = { + "litellm_metadata": {}, + "proxy_server_request": { + "headers": { + "user-agent": "claude-cli/2.1.263 (external, cli)", + "x-app": "cli-bg", + } + }, + } + if configured is not None: + kwargs["cache_control_injection_points"] = configured + return kwargs + + @pytest.mark.parametrize( + "system", + [ + BILLING_TEXT, + BILLING_SYSTEM, + [*BILLING_SYSTEM, {"type": "text", "text": " "}], + [ + *BILLING_SYSTEM, + {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli;"}, + ], + ], + ids=["string", "text_block", "whitespace_block", "multiple_billing_blocks"], + ) + @pytest.mark.parametrize("tools", [None, []], ids=["absent_tools", "empty_tools"]) + def test_skips_defaults_and_attribution_for_one_shot_subagent(self, monkeypatch, system, tools): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages = copy.deepcopy(self.MESSAGES) + kwargs = self._kwargs() + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + messages, + copy.deepcopy(system), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=tools, + ) + + assert result_messages == self.MESSAGES + assert result_system == system + assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"] + + def test_user_agent_header_lookup_is_case_insensitive(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs() + user_agent = kwargs["proxy_server_request"]["headers"].pop("user-agent") + kwargs["proxy_server_request"]["headers"]["User-Agent"] = user_agent + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(self.MESSAGES), + copy.deepcopy(self.BILLING_SYSTEM), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert result_messages == self.MESSAGES + assert result_system == self.BILLING_SYSTEM + + def test_router_affinity_skips_string_billing_system(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages = copy.deepcopy(self.MESSAGES) + kwargs = self._kwargs() + kwargs["system"] = self.BILLING_TEXT + + result = AnthropicCacheControlHook.messages_with_default_injections( + messages=messages, + models=("claude-sonnet-4-5",), + request_kwargs=kwargs, + ) + + assert result == messages + + @pytest.mark.parametrize( + "headers,system", + [ + ("not-a-mapping", BILLING_SYSTEM), + ( + {"user-agent": "claude-cli/2.1.263 (external, cli)"}, + [{"type": "text", "text": "x-anthropic-billing-header: malformed"}], + ), + ({"user-agent": "claude-cli/2.1.263 (external, cli)"}, None), + ({"user-agent": "claude-cli/2.1.263 (external, cli)"}, ["not-a-mapping"]), + ( + {"user-agent": "claude-cli/2.1.263 (external, cli)"}, + [{"type": "image", "text": BILLING_TEXT}], + ), + ], + ids=["malformed_headers", "malformed_billing", "missing_system", "malformed_block", "non_text_block"], + ) + def test_malformed_untrusted_context_keeps_defaults(self, monkeypatch, headers, system): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + points = AnthropicCacheControlHook.get_default_injection_points( + messages=copy.deepcopy(self.MESSAGES), + system=copy.deepcopy(system), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + request_kwargs={"proxy_server_request": {"headers": headers}}, + ) + + assert len(points) == 2 + + def test_message_without_role_keeps_defaults(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + points = AnthropicCacheControlHook.get_default_injection_points( + messages=[{"content": "missing role"}], + system=copy.deepcopy(self.BILLING_SYSTEM), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + request_kwargs=self._kwargs(), + ) + + assert len(points) == 2 + + @pytest.mark.parametrize( + "messages,system,tools", + [ + ( + MESSAGES, + BILLING_SYSTEM, + [{"name": "WebFetch", "description": "fetch", "input_schema": {"type": "object"}}], + ), + (MESSAGES, [*BILLING_SYSTEM, {"type": "text", "text": "Explore the repository"}], None), + ( + [ + {"role": "user", "content": "first turn"}, + {"role": "assistant", "content": "reply"}, + *MESSAGES, + ], + BILLING_SYSTEM, + None, + ), + ], + ids=["tools", "real_system", "history"], + ) + def test_keeps_defaults_for_reusable_subagents(self, monkeypatch, messages, system, tools): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs() + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(messages), + copy.deepcopy(system), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=copy.deepcopy(tools), + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2 + assert kwargs["litellm_metadata"]["litellm_gateway_injected_cache"] == "" + + @pytest.mark.parametrize( + "user_agent,system", + [ + ("anthropic-sdk-python/0.75.0", BILLING_SYSTEM), + ( + "claude-cli/2.1.263 (external, cli)", + [ + { + "type": "text", + "text": f"{BILLING_TEXT}\nadditional system instructions", + } + ], + ), + ( + "claude-cli/2.1.263 (external, cli)", + [ + { + "type": "text", + "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=false;", + } + ], + ), + ], + ids=["different_client", "appended_instructions", "not_a_subagent"], + ) + def test_ambiguous_or_unmatched_signals_fail_open(self, monkeypatch, user_agent, system): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs() + kwargs["proxy_server_request"]["headers"]["user-agent"] = user_agent + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(self.MESSAGES), + copy.deepcopy(system), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2 + + def test_explicit_injection_points_remain_authoritative(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs([{"location": "message", "role": "user"}]) + + result_messages, _ = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(self.MESSAGES), + copy.deepcopy(self.BILLING_SYSTEM), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert result_messages[0]["content"][-1]["cache_control"] == {"type": "ephemeral"} + + class TestPerKeyEnablePromptCaching: """Per-request enable_prompt_caching override (stamped from key metadata) with the global flag off.""" @@ -1977,6 +2195,25 @@ class TestConfiguredInjectionPointsStandDown: _, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) assert result_sys == [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}] + @pytest.mark.parametrize( + "configured", + [None, CONFIGURED], + ids=["automatic_defaults", "configured_points"], + ) + def test_v1_messages_stands_down_for_root_cache_control(self, monkeypatch, configured): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + root_cache_control = {"type": "ephemeral"} + kwargs = {"cache_control": root_cache_control, "litellm_metadata": {}} + if configured is not None: + kwargs["cache_control_injection_points"] = copy.deepcopy(configured) + + result_messages, result_system = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) + + assert result_messages == self.V1_MESSAGES + assert result_system == "sys" + assert kwargs["cache_control"] is root_cache_control + assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"] + def test_v1_messages_reentry_flow_preserves_tool_config_remainder(self): """The advisor interceptor re-enters anthropic_messages() with the outer request's kwargs and post-injection messages. The first pass applies the diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 038197c06c5..ecb7afd4ffb 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -867,6 +867,45 @@ async def test_azure_sentinel_concurrent_threshold_sends_collapse_into_one_attem assert getattr(logger, queue_attr) == [] +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_batch_size_bounds_every_request_under_concurrent_events( + queue_attr, send_method, build_payloads +): + """Lowering batch_size is the documented way to stay under the ingestion cap, so no request may + carry more than batch_size records even when events keep landing while a send is on the wire, + and every one of those records still has to arrive exactly once.""" + logger = _build_logger(batch_size=5) + records = build_payloads(40) + + attempts = [] + first_send_started = asyncio.Event() + release_first_send = asyncio.Event() + + async def _on_ingest(data): + attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))]) + if len(attempts) == 1: + first_send_started.set() + await release_first_send.wait() + return _accepted() + + _install_ingestion(logger, _on_ingest) + + sends = [asyncio.create_task(_log(logger, queue_attr, record)) for record in records] + await asyncio.wait_for(first_send_started.wait(), timeout=10) + + assert attempts == [[record["id"] for record in records[:5]]] + assert getattr(logger, queue_attr) == records[5:] + + release_first_send.set() + await asyncio.wait_for(asyncio.gather(*sends), timeout=10) + await logger.flush_queue() + + assert max(len(attempt) for attempt in attempts) <= 5 + assert [record_id for attempt in attempts for record_id in attempt] == [record["id"] for record in records] + assert getattr(logger, queue_attr) == [] + + @pytest.mark.asyncio @pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) async def test_azure_sentinel_requeues_a_cancelled_send( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 885bd1d4d72..ddc8439a83a 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,5 +1,5 @@ import asyncio -from typing import TYPE_CHECKING, Literal, Optional +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional from unittest.mock import AsyncMock import pytest @@ -10,6 +10,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import CallTypes, UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail if TYPE_CHECKING: @@ -1841,12 +1842,14 @@ class _ApplyStyleGuardrail(CustomGuardrail): self.block = block self.apply_called = False self.seen_texts = None + self.seen_request_data = None async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): from fastapi import HTTPException self.apply_called = True self.seen_texts = inputs.get("texts") + self.seen_request_data = request_data if self.block: raise HTTPException(status_code=400, detail={"error": "Violated moderation policy"}) return inputs @@ -2378,11 +2381,108 @@ def _logged_call(messages: list | str) -> tuple[dict, object]: return kwargs, response +class _NativeApplyGuardrail(_InheritedApplyGuardrail): + use_native_lifecycle_hooks: ClassVar[bool] = True + + +@pytest.mark.parametrize("guardrail_type", (CustomGuardrail, _NativeApplyGuardrail, _InheritedApplyGuardrail)) +@pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.logging_only, + "logging_only", + [GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ["pre_call", "logging_only"], + Mode(tags={"audit": "logging_only"}, default="pre_call"), + Mode(tags={"audit": ["pre_call", "logging_only"]}), + Mode(tags={"enforce": "pre_call"}, default="logging_only"), + Mode(tags={}, default=["pre_call", "logging_only"]), + ), +) +def test_logging_only_requires_framework_support_or_explicit_declaration( + guardrail_type: type[CustomGuardrail], + event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, +) -> None: + supported: Final = [GuardrailEventHooks.pre_call] + if guardrail_type is _InheritedApplyGuardrail: + guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) + assert guardrail.event_hook == event_hook + assert supported == [GuardrailEventHooks.pre_call] + else: + with pytest.raises(ValueError, match=r"logging_only.*not in the supported event hooks"): + guardrail_type(event_hook=event_hook, supported_event_hooks=supported) + + explicitly_supported: Final = guardrail_type( + event_hook=event_hook, + supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ) + assert explicitly_supported.event_hook == event_hook + + +@pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.post_call, + "post_call", + [GuardrailEventHooks.logging_only, GuardrailEventHooks.post_call], + ["logging_only", "post_call"], + Mode(tags={"enforce": "post_call"}, default="logging_only"), + Mode(tags={"enforce": ["logging_only", "post_call"]}), + Mode(tags={"audit": "logging_only"}, default="post_call"), + Mode(tags={}, default=["logging_only", "post_call"]), + ), +) +def test_framework_logging_only_does_not_allow_other_unsupported_modes( + event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, +) -> None: + with pytest.raises(ValueError, match=r"post_call.*not in the supported event hooks"): + _InheritedApplyGuardrail(event_hook=event_hook, supported_event_hooks=[GuardrailEventHooks.pre_call]) + + class TestLoggingOnlyApplyGuardrail: """LIT-4876 regression: a guardrail in mode logging_only that implements only apply_guardrail must still run against the logged request and response and record guardrail_information, instead of inheriting the CustomLogger no-op.""" + @pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.logging_only, + "logging_only", + [GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ["pre_call", "logging_only"], + Mode(tags={"audit": "logging_only"}, default="pre_call"), + Mode(tags={"audit": ["pre_call", "logging_only"]}), + Mode(tags={"enforce": "pre_call"}, default="logging_only"), + Mode(tags={}, default=["pre_call", "logging_only"]), + ), + ) + @pytest.mark.asyncio + async def test_content_filter_accepts_logging_only_and_records_detection( + self, event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode + ) -> None: + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks + + guardrail: Final = ContentFilterGuardrail( + guardrail_name="content-review", + event_hook=event_hook, + default_on=True, + blocked_words=[BlockedWord(keyword="hello", action=ContentFilterAction.BLOCK)], + ) + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert out_response is response + assert out_kwargs["messages"] == kwargs["messages"] + assert ( + out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"] + == "guardrail_intervened" + ) + @pytest.mark.asyncio async def test_runs_apply_guardrail_observe_only_and_records_verdict(self): guardrail = _ApplyOnlyObserver() @@ -2548,6 +2648,91 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: which starved every later callback in litellm.callbacks (notably the lazily-appended VectorStorePreCallHook that attaches provider_specific_fields["search_results"]).""" + @pytest.mark.asyncio + async def test_apply_guardrail_retains_request_identity(self) -> None: + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail: Final = _ApplyStyleGuardrail(block=False) + guardrail.event_hook = GuardrailEventHooks.post_call + request_data: Final = {"guardrails": ["apply-style-guardrail"]} + response: Final = ModelResponse(choices=[Choices(message=Message(content="review me"))]) + + await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, response=response, call_type=CallTypes.acompletion + ) + + assert guardrail.seen_request_data is request_data + assert guardrail.seen_texts == ["review me"] + assert "guardrail_to_apply" not in request_data + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", (None, CallTypes.acompletion)) + async def test_apply_guardrail_masks_response_and_records_metadata(self, call_type: CallTypes | None) -> None: + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail: Final = ContentFilterGuardrail( + guardrail_name="response-filter", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="secret", action=ContentFilterAction.MASK)], + ) + request_data: Final = {"guardrails": ["response-filter"]} + response: Final = ModelResponse(choices=[Choices(message=Message(content="a secret"))]) + + result: Final = await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, response=response, call_type=call_type + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == f"a {guardrail.keyword_redaction_tag}" + entries: Final = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "response-filter" + assert entries[0]["guardrail_mode"] == "post_call" + assert "guardrail_to_apply" not in request_data + + @pytest.mark.asyncio + @pytest.mark.parametrize("error_type", (None, RuntimeError, asyncio.CancelledError)) + async def test_dispatch_cleans_up_request_on_every_exit(self, error_type: type[BaseException] | None) -> None: + from contextlib import nullcontext + + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.utils import LLMResponseTypes, ModelResponse + + error: Final = error_type("dispatch interrupted") if error_type is not None else None + + class Dispatch(CustomLogger): + request_data: dict[str, object] | None = None + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> LLMResponseTypes: + self.request_data = data + if error is not None: + raise error + return response + + dispatch: Final = Dispatch() + + class Guardrail(_ApplyStyleGuardrail): + def _deployment_hook_target(self) -> CustomLogger: + return dispatch + + guardrail: Final = Guardrail(block=False) + guardrail.event_hook = GuardrailEventHooks.post_call + request_data: Final = {"guardrails": ["apply-style-guardrail"]} + with pytest.raises(error_type) if error_type is not None else nullcontext(): + await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, response=ModelResponse(), call_type=CallTypes.acompletion + ) + + assert dispatch.request_data is request_data + assert "guardrail_to_apply" not in request_data + @pytest.mark.asyncio async def test_returns_none_when_request_has_no_guardrails(self): from litellm.types.utils import ModelResponse @@ -2610,3 +2795,37 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: ) assert result is replacement + + @pytest.mark.asyncio + async def test_apply_guardrail_interface_modifies_deployment_response(self): + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import ModelResponse + + class ReplacingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + assert input_type == "response" + return {**inputs, "texts": ["filtered response"]} + + guardrail = ReplacingGuardrail( + guardrail_name="test-guardrail", + event_hook=GuardrailEventHooks.post_call, + ) + response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "original response"}}]) + request_data = {"guardrails": ["test-guardrail"]} + + result = await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, + response=response, + call_type=CallTypes.acompletion, + ) + + assert result is response + assert response.choices[0].message.content == "filtered response" + assert "guardrail_to_apply" not in request_data + assert len(_guardrail_entries(request_data)) == 1 diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index d36878e455f..87e76499b84 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1341,6 +1341,257 @@ def _emit(logger: LangFuseLogger, *, metadata=None, headers=None): ) +@pytest.mark.parametrize("level", ["DEFAULT", "ERROR"]) +@pytest.mark.parametrize( + "headers,metadata,expected_id", + [ + ({"x-litellm-session-id": "session-7125"}, {}, "call"), + ({"X-Claude-Code-Session-Id": "session-7125"}, {}, "call"), + ({"x-session-id": "session-7125"}, {}, "call"), + ({"session-id": "session-7125", "user-agent": "codex_cli_rs/1.0"}, {}, "call"), + ({"thread-id": "session-7125", "user-agent": "codex-tui"}, {}, "call"), + ({"session_id": "session-7125", "user-agent": "Codex 1.0"}, {}, "call"), + ({"conversation_id": "session-7125", "user-agent": "codex_vscode/1.0"}, {}, "call"), + ({"x-litellm-session-id": "short"}, {}, "call"), + ({"x-litellm-trace-id": "session-7125"}, {}, "session-7125"), + ( + {"X-LiteLLM-Trace-Id": "session-7125", "x-litellm-session-id": "session-7125"}, + {}, + "session-7125", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_trace_id": "session-7125"}, + {}, + "session-7125", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_trace_id": "explicit-trace"}, + {}, + "explicit-trace", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_existing_trace_id": "existing-trace"}, + {}, + "existing-trace", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + {"x-litellm-session-id": "short", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + {"X-Claude-Code-Session-Id": "session-7125", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + {"x-session-id": "session-7125", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + { + "session-id": "session-7125", + "user-agent": "codex_cli_rs/1.0", + "langfuse_session_id": "custom-session", + }, + {}, + "call", + ), + ( + { + "x-litellm-session-id": "session-7125", + "langfuse_session_id": "custom-session", + "x-litellm-trace-id": "explicit-trace", + }, + {}, + "explicit-trace", + ), + ( + { + "x-litellm-session-id": "session-7125", + "langfuse_session_id": "custom-session", + "langfuse_trace_id": "explicit-trace", + }, + {}, + "explicit-trace", + ), + ( + { + "x-litellm-session-id": "session-7125", + "langfuse_session_id": "custom-session", + "langfuse_existing_trace_id": "existing-trace", + }, + {}, + "existing-trace", + ), + ({}, {"trace_id": "session-7125", "session_id": "session-7125"}, "session-7125"), + ({}, {"trace_id": "explicit-trace", "session_id": "session-7125"}, "explicit-trace"), + ( + {"x-vendor-session-id": "short"}, + {"trace_id": "short", "session_id": "short"}, + "short", + ), + ( + {"x-session-id": "invalid value"}, + {"trace_id": "invalid value", "session_id": "invalid value"}, + "invalid value", + ), + ( + {"session-id": "session-7125", "user-agent": "codexfoo/1.0"}, + {"trace_id": "session-7125", "session_id": "session-7125"}, + "session-7125", + ), + ( + {"x-vendor-session-id": "short"}, + {"trace_id": "session-7125", "session_id": "session-7125"}, + "session-7125", + ), + ({}, {}, "call"), + ], +) +def test_session_header_trace_provenance(headers, metadata, expected_id, level): + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + clean_headers, + redact_credential_headers, + ) + + logger: Final = _steering_logger() + for turn in range(2): + call_id = f"call-{turn}" + request_headers = Headers(headers) + data = LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=request_headers, data={"metadata": dict(metadata)}, _metadata_variable_name="metadata" + ) + original_metadata = dict(data["metadata"]) + now = datetime.datetime.now() + result = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_call_id": call_id, + "litellm_trace_id": data.get("litellm_trace_id"), + "litellm_params": { + "metadata": data["metadata"], + "proxy_server_request": {"headers": redact_credential_headers(clean_headers(request_headers))}, + }, + "messages": [{"role": "user", "content": f"turn {turn}"}], + "optional_params": {}, + }, + response_obj=( + None + if level == "ERROR" + else litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]) + ), + start_time=now, + end_time=now, + level=level, + status_message="provider error" if level == "ERROR" else None, + ) + trace_params = logger.Langfuse.trace.call_args.kwargs + assert trace_params["id"] == (call_id if expected_id == "call" else expected_id) + assert result["trace_id"] == trace_params["id"] + if expected_id != "existing-trace": + assert trace_params["session_id"] == headers.get("langfuse_session_id", original_metadata.get("session_id")) + steering = {key[len("langfuse_") :]: value for key, value in headers.items() if key.startswith("langfuse_")} + assert data["metadata"] == {**original_metadata, **steering} + + +def test_session_header_trace_without_call_id_keeps_session_alias(): + logger: Final = _steering_logger() + now: Final = datetime.datetime.now() + + result: Final = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_call_id": "", + "litellm_params": { + "metadata": {"trace_id": "session-7125", "session_id": "session-7125"}, + "proxy_server_request": {"headers": {"x-litellm-session-id": "session-7125"}}, + }, + "messages": [{"role": "user", "content": "no call id"}], + "optional_params": {}, + }, + response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]), + start_time=now, + end_time=now, + ) + + assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125" + assert result["trace_id"] == "session-7125" + + +def test_every_proxy_session_header_shape_is_classified_as_a_session_alias(): + """The classifier must cover every header shape the proxy turns into a chain id.""" + from litellm.integrations.langfuse.langfuse import _is_session_header_trace + from litellm.proxy.litellm_pre_call_utils import ( + _CODEX_SESSION_ID_HEADERS, + get_chain_id_from_headers, + ) + + session: Final = "session-7125-abcdef" + session_shapes: Final = ( + {"x-litellm-session-id": session}, + {"X-Claude-Code-Session-Id": session}, + {"x-session-id": session}, + *({header: session, "user-agent": "codex_cli_rs/1.0"} for header in _CODEX_SESSION_ID_HEADERS), + ) + for headers in session_shapes: + assert get_chain_id_from_headers(dict(headers)) == session, headers + assert _is_session_header_trace(session, session, {"headers": headers}) is True, headers + + explicit_trace: Final = {"x-litellm-trace-id": session, "x-litellm-session-id": session} + assert get_chain_id_from_headers(dict(explicit_trace)) == session + assert _is_session_header_trace(session, session, {"headers": explicit_trace}) is False + + +@pytest.mark.parametrize( + "proxy_server_request", + [None, {}, {"headers": None}], + ids=["no-proxy-request", "no-headers-key", "null-headers"], +) +def test_sdk_caller_without_request_headers_keeps_its_trace(proxy_server_request): + """A direct SDK caller has no request headers, so a session-shaped trace id stays the caller's.""" + logger: Final = _steering_logger() + now: Final = datetime.datetime.now() + + result: Final = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_call_id": "call-0", + "litellm_params": { + "metadata": {"trace_id": "session-7125", "session_id": "session-7125"}, + "proxy_server_request": proxy_server_request, + }, + "messages": [{"role": "user", "content": "sdk turn"}], + "optional_params": {}, + }, + response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]), + start_time=now, + end_time=now, + ) + + assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125" + assert result["trace_id"] == "session-7125" + + +def test_session_header_classifier_survives_non_string_header_keys(): + """A non-string header key must not cost the caller its whole trace.""" + from litellm.integrations.langfuse.langfuse import _is_session_header_trace + + session: Final = "session-7125-abcdef" + headers: Final = {7: "numeric key", "x-litellm-session-id": session} + assert _is_session_header_trace(session, session, {"headers": headers}) is True + assert _is_session_header_trace(session, session, {"headers": {7: "numeric key"}}) is False + + def test_mask_input_header_false_keeps_the_prompt(): logger = _steering_logger() diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/test_litellm/integrations/test_mlflow.py index 61010f8531c..f828c34a9ff 100644 --- a/tests/test_litellm/integrations/test_mlflow.py +++ b/tests/test_litellm/integrations/test_mlflow.py @@ -195,3 +195,70 @@ def test_mlflow_stream_handler_uses_async_complete_response(): is final_response ) assert "abc123" not in mlflow_logger._stream_id_to_span + + +def test_mlflow_stream_handler_pops_span_when_end_raises(): + modules = _mock_mlflow_modules() + with patch.dict("sys.modules", modules): + from litellm.integrations.mlflow import MlflowLogger + + mlflow_logger = MlflowLogger() + mlflow_logger._start_span_or_trace = MagicMock(return_value="mock_span") + mlflow_logger._end_span_or_trace = MagicMock( + side_effect=TypeError("unexpected keyword argument 'trace_id'") + ) + mlflow_logger._extract_and_set_chat_attributes = MagicMock() + + response_obj = MagicMock() + response_obj.choices = [] + + kwargs = { + "litellm_call_id": "leak123", + "complete_streaming_response": MagicMock(), + } + + with pytest.raises(TypeError): + mlflow_logger._handle_stream_event( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.utcnow(), + end_time=datetime.utcnow(), + ) + + assert "leak123" not in mlflow_logger._stream_id_to_span + + +class _Mlflow2StyleClient: + """Mimics the mlflow 2.x client signatures, which have no trace_id kwarg.""" + + def __init__(self): + self.ended_traces = [] + self.ended_spans = [] + + def end_trace(self, request_id, outputs=None, attributes=None, status="OK", end_time_ns=None): + self.ended_traces.append(request_id) + + def end_span(self, request_id, span_id, outputs=None, attributes=None, status="OK", end_time_ns=None): + self.ended_spans.append((request_id, span_id)) + + +def test_mlflow_end_span_or_trace_works_with_mlflow_2x_client(): + modules = _mock_mlflow_modules() + with patch.dict("sys.modules", modules): + from litellm.integrations.mlflow import MlflowLogger + + mlflow_logger = MlflowLogger() + client = _Mlflow2StyleClient() + mlflow_logger._client = client + + root_span = MagicMock(parent_id=None, request_id="req-1") + mlflow_logger._end_span_or_trace( + span=root_span, outputs="out", end_time_ns=1, status="OK" + ) + assert client.ended_traces == ["req-1"] + + child_span = MagicMock(parent_id="parent-1", request_id="req-2", span_id="span-2") + mlflow_logger._end_span_or_trace( + span=child_span, outputs="out", end_time_ns=1, status="OK" + ) + assert client.ended_spans == [("req-2", "span-2")] diff --git a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py index ea661d2ea78..004ac4dbffb 100644 --- a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py +++ b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py @@ -84,6 +84,7 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent(): "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None ): logger = PrometheusLogger() + logger._emit_input_sequence_length_label = False logger.litellm_proxy_total_requests_metric = MagicMock() logger.get_labels_for_metric = MagicMock( return_value=["client_ip", "user_agent"] diff --git a/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py b/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py new file mode 100644 index 00000000000..bc922061544 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py @@ -0,0 +1,428 @@ +import asyncio +import datetime +from collections.abc import Mapping +from copy import deepcopy +from typing import Final, cast + +import pytest +from prometheus_client import REGISTRY +from prometheus_client.samples import Sample + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import ( + PrometheusMetricLabels, + UserAPIKeyLabelNames, + UserAPIKeyLabelValues, + get_input_sequence_length_bucket, +) +from litellm.types.utils import StandardLoggingPayload + +LATENCY_METRICS: Final = ( + "litellm_llm_api_latency_metric", + "litellm_llm_api_time_to_first_token_metric", + "litellm_request_total_latency_metric", +) +FLAG: Final = "prometheus_emit_input_sequence_length_label" + + +def _clear_prometheus_registry() -> None: + for collector in tuple(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage] # test registry reset + REGISTRY.unregister(collector) + + +@pytest.fixture(autouse=True) +def isolated_registry(monkeypatch: pytest.MonkeyPatch): + _clear_prometheus_registry() + monkeypatch.setattr(litellm, FLAG, False) + yield + _clear_prometheus_registry() + + +@pytest.mark.parametrize("metric", LATENCY_METRICS) +def test_input_sequence_length_label_is_opt_in(monkeypatch: pytest.MonkeyPatch, metric: str): + assert UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value not in PrometheusMetricLabels.get_labels(metric) + + monkeypatch.setattr(litellm, FLAG, True) + assert UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value in PrometheusMetricLabels.get_labels(metric) + + +def test_input_sequence_length_label_stays_off_non_latency_metrics(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, FLAG, True) + assert UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value not in PrometheusMetricLabels.get_labels( + "litellm_proxy_total_requests_metric" + ) + + +@pytest.mark.parametrize( + "prompt_tokens, expected", + [ + (None, "unknown"), + (0, "0-1k"), + (999, "0-1k"), + (1_000, "1k-4k"), + (3_999, "1k-4k"), + (4_000, "4k-16k"), + (15_999, "4k-16k"), + (16_000, "16k-64k"), + (63_999, "16k-64k"), + (64_000, "64k+"), + (10_000_000, "64k+"), + (-1, "unknown"), + ], +) +def test_input_sequence_length_bucket_boundaries(prompt_tokens: int | None, expected: str): + assert get_input_sequence_length_bucket(prompt_tokens) == expected + + +def test_user_api_key_label_values_carries_input_sequence_length(): + values: Final = UserAPIKeyLabelValues(input_sequence_length="4k-16k") + + assert values.input_sequence_length == "4k-16k" + assert values.model_dump()["input_sequence_length"] == "4k-16k" + + +def _assert_latency_metrics(expected: str | None, stream: bool = True, queue_time: float = 0) -> None: + samples: Final = tuple(sample for metric in REGISTRY.collect() for sample in metric.samples) + for metric, duration in zip(LATENCY_METRICS, (2, 1, 3 + queue_time)): + counts: Final = tuple(sample for sample in samples if sample.name == f"{metric}_count") + sums: Final = tuple(sample for sample in samples if sample.name == f"{metric}_sum") + buckets: Final = tuple(sample for sample in samples if sample.name == f"{metric}_bucket") + if not stream and metric == "litellm_llm_api_time_to_first_token_metric": + assert not counts and not sums and not buckets + continue + assert len(counts) == len(sums) == 1 + assert counts[0].value == 1 + assert sums[0].value == pytest.approx(duration) + assert buckets and any(sample.labels["le"] == "+Inf" for sample in buckets) + assert all(sample.value == int(float(sample.labels["le"]) >= duration) for sample in buckets) + assert all(sample.labels.get("input_sequence_length") == expected for sample in (*counts, *sums, *buckets)) + + +def _non_target_samples() -> tuple[Sample, ...]: + return tuple( + sample + for metric in REGISTRY.collect() + if metric.name not in LATENCY_METRICS + for sample in metric.samples + if "input_sequence_length" in sample.labels and not sample.name.endswith("_created") + ) + + +def _standard_logging_payload(now: datetime.datetime, prompt_tokens: int) -> StandardLoggingPayload: + return cast( + StandardLoggingPayload, + { + "id": "t", + "call_type": "completion", + "response_cost": 0.001, + "status": "success", + "total_tokens": prompt_tokens + 20, + "prompt_tokens": prompt_tokens, + "completion_tokens": 20, + "startTime": now - datetime.timedelta(seconds=3), + "endTime": now, + "completionStartTime": now - datetime.timedelta(seconds=1), + "model": "gpt-4o-mini", + "model_id": "model-123", + "model_group": "gpt-4o-mini", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "request_tags": [], + "stream": True, + "metadata": { + "user_api_key_hash": "h", + "user_api_key_alias": "a", + "user_api_key_team_id": "t", + "user_api_key_team_alias": "ta", + "user_api_key_user_id": "u", + "user_api_key_user_email": "e@x.com", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, + "requester_metadata": None, + "user_api_key_end_user_id": None, + "usage_object": None, + }, + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + }, + ) + + +def _success_kwargs( + now: datetime.datetime, prompt_tokens: int, requester_metadata: Mapping[str, object] | None = None +) -> Mapping[str, object]: + payload: Final = _standard_logging_payload(now, prompt_tokens) + return { + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {}}, + "standard_logging_object": { + **payload, + "metadata": {**payload["metadata"], "requester_metadata": requester_metadata}, + }, + "stream": True, + "start_time": now - datetime.timedelta(seconds=3), + "api_call_start_time": now - datetime.timedelta(seconds=2), + "completion_start_time": now - datetime.timedelta(seconds=1), + "end_time": now, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("flag_at_request_time", (True, False)) +async def test_logger_emits_bucket_from_its_startup_label_set( + monkeypatch: pytest.MonkeyPatch, flag_at_request_time: bool +): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + monkeypatch.setattr(litellm, FLAG, flag_at_request_time) + + await logger.async_log_success_event(dict(_success_kwargs(now, prompt_tokens=4_000)), None, now, now) + + _assert_latency_metrics("4k-16k") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "combined_usage", "expected"), + ( + ({"id": "moderation", "results": []}, None, "unknown"), + ({"usage": None}, None, "unknown"), + ({"usage": {}}, None, "unknown"), + ({"usage": {"completion_tokens": 3}}, None, "unknown"), + ({"usage": {"total_tokens": 5}}, None, "unknown"), + ({"usage": {"prompt_tokens": 0}}, None, "0-1k"), + ({"usage": {"prompt_tokens": 4_000}}, None, "4k-16k"), + ({"usage": {"input_tokens": 0, "output_tokens": 3, "total_tokens": 3}}, None, "0-1k"), + ({"usage": {"input_tokens": 4_000, "output_tokens": 3, "total_tokens": 4_003}}, None, "4k-16k"), + (litellm.ModelResponse(usage=litellm.Usage(prompt_tokens=0)), None, "0-1k"), + (None, litellm.Usage(prompt_tokens=0), "0-1k"), + ), +) +@pytest.mark.parametrize("include_usage_metadata", (True, False)) +async def test_logger_distinguishes_missing_usage_from_reported_zero( + monkeypatch: pytest.MonkeyPatch, + response: object, + combined_usage: object, + expected: str, + include_usage_metadata: bool, +): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + usage: Final = StandardLoggingPayloadSetup.get_usage_as_dict( + response_obj=response if isinstance(response, dict) else None + ) + + payload: Final = _standard_logging_payload(now, usage.get("prompt_tokens", 0)) + await logger.async_log_success_event( + { + **_success_kwargs(now, prompt_tokens=usage.get("prompt_tokens", 0)), + "combined_usage_object": combined_usage, + "standard_logging_object": { + **payload, + "metadata": {**payload["metadata"], "usage_object": usage if include_usage_metadata else None}, + }, + }, + response, + now, + now, + ) + + _assert_latency_metrics(expected) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("total_tokens", (None, 0, 5_000)) +@pytest.mark.parametrize("prompt_tokens", (0, 4_000)) +async def test_upstream_total_only_usage_has_unknown_input_length( + monkeypatch: pytest.MonkeyPatch, total_tokens: int | None, prompt_tokens: int +): + import httpx + + from litellm.litellm_core_utils.litellm_logging import Logging, StandardLoggingPayloadSetup + from litellm.proxy.pass_through_endpoints.upstream_usage_headers import apply_upstream_reported_usage + + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + logging_obj: Final = Logging( + model="gpt-4o-mini", + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=now, + litellm_call_id="test-call-id", + function_id="1", + ) + headers: Final = httpx.Headers( + { + "x-litellm-response-cost": "0.001", + **({"x-litellm-total-tokens": str(total_tokens)} if total_tokens is not None else {}), + } + ) + reported: Final = apply_upstream_reported_usage(logging_obj=logging_obj, headers=headers) + assert reported is not None + combined_usage: Final = logging_obj.model_call_details.get("combined_usage_object") + response: Final = {"usage": {"prompt_tokens": prompt_tokens}} + usage: Final = StandardLoggingPayloadSetup.get_usage_as_dict(response, combined_usage) + payload: Final = _standard_logging_payload(now, usage.get("prompt_tokens", 0)) + + await logger.async_log_success_event( + { + **logging_obj.model_call_details, + **_success_kwargs(now, usage.get("prompt_tokens", 0)), + "standard_logging_object": {**payload, "metadata": {**payload["metadata"], "usage_object": usage}}, + }, + response, + now, + now, + ) + + _assert_latency_metrics("unknown" if total_tokens is not None else get_input_sequence_length_bucket(prompt_tokens)) + + +@pytest.mark.asyncio +async def test_logger_built_with_flag_off_emits_no_bucket_label(monkeypatch: pytest.MonkeyPatch): + now: Final = datetime.datetime.now() + logger: Final = PrometheusLogger() + monkeypatch.setattr(litellm, FLAG, True) + + await logger.async_log_success_event(dict(_success_kwargs(now, prompt_tokens=4_000)), None, now, now) + + _assert_latency_metrics(None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("flag_at_startup", (True, False)) +@pytest.mark.parametrize("stream", (True, False)) +@pytest.mark.parametrize( + "metadata", + ( + None, + {}, + {"input_sequence_length": None}, + {"input_sequence_length": False}, + {"input_sequence_length": True}, + {"input_sequence_length": 0}, + {"input_sequence_length": []}, + {"input_sequence_length": {}}, + {"input_sequence_length": ""}, + {"input_sequence_length": "from-metadata"}, + ), +) +async def test_custom_input_length_label_is_scoped_to_target_histograms( + monkeypatch: pytest.MonkeyPatch, flag_at_startup: bool, stream: bool, metadata: Mapping[str, object] | None +): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, "custom_prometheus_metadata_labels", ["input_sequence_length"]) + kwargs: Final = { + **_success_kwargs(now, prompt_tokens=4_000, requester_metadata=metadata), + "stream": stream, + "litellm_params": {"metadata": {"queue_time_seconds": 0.25}}, + } + original_kwargs: Final = deepcopy(kwargs) + baseline_logger: Final = PrometheusLogger() + await baseline_logger.async_log_success_event(kwargs, None, now, now) + baseline_samples: Final = _non_target_samples() + _clear_prometheus_registry() + monkeypatch.setattr(litellm, FLAG, flag_at_startup) + logger: Final = PrometheusLogger() + monkeypatch.setattr(litellm, FLAG, not flag_at_startup) + + await logger.async_log_success_event(kwargs, None, now, now) + + assert kwargs == original_kwargs + custom_value: Final = (metadata or {}).get("input_sequence_length") + expected: Final = custom_value if isinstance(custom_value, str) else ("4k-16k" if flag_at_startup else "None") + _assert_latency_metrics(expected, stream=stream, queue_time=0.25) + assert all(logger.get_labels_for_metric(metric).count("input_sequence_length") == 1 for metric in LATENCY_METRICS) + non_target_samples: Final = _non_target_samples() + assert { + "litellm_requests_metric_total", + "litellm_spend_metric_total", + "litellm_total_tokens_metric_total", + "litellm_request_queue_time_seconds_count", + "litellm_deployment_success_responses_total", + }.issubset({sample.name for sample in non_target_samples}) + assert non_target_samples == baseline_samples + queue_sum: Final = tuple( + sample for sample in non_target_samples if sample.name == "litellm_request_queue_time_seconds_sum" + ) + assert len(queue_sum) == 1 and queue_sum[0].value == 0.25 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", (True, False)) +async def test_concurrent_requests_keep_independent_buckets(monkeypatch: pytest.MonkeyPatch, enabled: bool): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, enabled) + logger: Final = PrometheusLogger() + cases: Final = ( + (None, "unknown"), + (0, "0-1k"), + (1_000, "1k-4k"), + (4_000, "4k-16k"), + (16_000, "16k-64k"), + (64_000, "64k+"), + ) + calls: Final = tuple( + ( + {**_success_kwargs(now, prompt_tokens=tokens or 0), "stream": stream}, + {"usage": {"prompt_tokens": tokens}} if tokens is not None else None, + ) + for tokens, _ in cases + for stream in (True, False) + for _ in range(2) + ) + original_calls: Final = deepcopy(calls) + + await asyncio.gather(*(logger.async_log_success_event(kwargs, response, now, now) for kwargs, response in calls)) + + assert calls == original_calls + samples: Final = tuple(sample for metric in REGISTRY.collect() for sample in metric.samples) + for metric, duration in zip(LATENCY_METRICS, (2, 1, 3)): + expected_count: Final = 2 if metric == "litellm_llm_api_time_to_first_token_metric" else 4 + counts: Final = tuple(sample for sample in samples if sample.name == f"{metric}_count") + sums: Final = tuple(sample for sample in samples if sample.name == f"{metric}_sum") + buckets: Final = tuple(sample for sample in samples if sample.name == f"{metric}_bucket") + expected: Final = ( + {bucket: expected_count for _, bucket in cases} if enabled else {None: expected_count * len(cases)} + ) + assert len(counts) == len(sums) == len(expected) + assert {sample.labels.get("input_sequence_length"): sample.value for sample in counts} == expected + assert {sample.labels.get("input_sequence_length"): sample.value for sample in sums} == { + bucket: count * duration for bucket, count in expected.items() + } + assert sum(sample.value for sample in buckets if sample.labels["le"] == "+Inf") == expected_count * len(cases) + assert all( + sample.value + == expected[sample.labels.get("input_sequence_length")] * int(float(sample.labels["le"]) >= duration) + for sample in buckets + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", (True, False)) +async def test_failed_request_does_not_observe_latency(monkeypatch: pytest.MonkeyPatch, enabled: bool): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, enabled) + monkeypatch.setattr(litellm, "custom_prometheus_metadata_labels", ["input_sequence_length"]) + logger: Final = PrometheusLogger() + kwargs: Final = { + **_success_kwargs(now, prompt_tokens=4_000), + "standard_logging_object": {**_standard_logging_payload(now, 4_000), "status": "failure"}, + "exception": RuntimeError("upstream request failed"), + } + + await logger.async_log_failure_event(kwargs, None, now, now) + + samples: Final = tuple(sample for metric in REGISTRY.collect() for sample in metric.samples) + assert not any(sample.name.startswith(LATENCY_METRICS) for sample in samples) + for metric in ("litellm_llm_api_failed_requests_metric_total", "litellm_deployment_failure_responses_total"): + counts: Final = tuple(sample for sample in samples if sample.name == metric) + assert len(counts) == 1 + assert counts[0].value == 1 + assert counts[0].labels["input_sequence_length"] == "None" diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index a037284d7c1..08d37297ab1 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1,10 +1,19 @@ import asyncio +import re +import sys +import textwrap +import uuid +from contextlib import asynccontextmanager from datetime import datetime -from unittest.mock import MagicMock, patch +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, call, patch +import httpx import pytest from litellm.integrations.s3_v2 import S3Logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.integrations.s3_v2 import s3BatchLoggingElement from litellm.types.utils import StandardLoggingPayload @@ -21,9 +30,7 @@ class TestS3V2UnitTests: source_code = inspect.getsource(s3_v2) # Verify that json.dumps is not used directly in the code - assert ( - "json.dumps(" not in source_code - ), "S3 v2 should not use json.dumps directly" + assert "json.dumps(" not in source_code, "S3 v2 should not use json.dumps directly" @patch("asyncio.create_task") @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") @@ -86,12 +93,8 @@ class TestS3V2UnitTests: call_args_minio = s3_logger_minio.async_httpx_client.put.call_args assert call_args_minio is not None url_minio = call_args_minio[0][0] - expected_minio_url = ( - "https://minio.example.com:9000/litellm-logs/2025-09-14/test-key.json" - ) - assert ( - url_minio == expected_minio_url - ), f"Expected MinIO URL {expected_minio_url}, got {url_minio}" + expected_minio_url = "https://minio.example.com:9000/litellm-logs/2025-09-14/test-key.json" + assert url_minio == expected_minio_url, f"Expected MinIO URL {expected_minio_url}, got {url_minio}" # Test 3: Custom endpoint without bucket name (should fall back to default) s3_logger_no_bucket = S3Logger( @@ -136,12 +139,8 @@ class TestS3V2UnitTests: call_args_sync = mock_sync_client.put.call_args assert call_args_sync is not None url_sync = call_args_sync[0][0] - expected_sync_url = ( - "https://custom.s3.endpoint.com/sync-bucket/2025-09-14/test-key.json" - ) - assert ( - url_sync == expected_sync_url - ), f"Expected sync URL {expected_sync_url}, got {url_sync}" + expected_sync_url = "https://custom.s3.endpoint.com/sync-bucket/2025-09-14/test-key.json" + assert url_sync == expected_sync_url, f"Expected sync URL {expected_sync_url}, got {url_sync}" # Test 5: Download method with custom endpoint s3_logger_download = S3Logger( @@ -158,19 +157,15 @@ class TestS3V2UnitTests: s3_logger_download.async_httpx_client = AsyncMock() s3_logger_download.async_httpx_client.get.return_value = mock_download_response - result = asyncio.run( - s3_logger_download._download_object_from_s3( - "2025-09-14/download-test-key.json" - ) - ) + result = asyncio.run(s3_logger_download._download_object_from_s3("2025-09-14/download-test-key.json")) call_args_download = s3_logger_download.async_httpx_client.get.call_args assert call_args_download is not None url_download = call_args_download[0][0] expected_download_url = "https://download.s3.endpoint.com/download-bucket/2025-09-14/download-test-key.json" - assert ( - url_download == expected_download_url - ), f"Expected download URL {expected_download_url}, got {url_download}" + assert url_download == expected_download_url, ( + f"Expected download URL {expected_download_url}, got {url_download}" + ) assert result == {"downloaded": "data"} @@ -216,12 +211,8 @@ class TestS3V2UnitTests: call_args = s3_logger_virtual.async_httpx_client.put.call_args assert call_args is not None url = call_args[0][0] - expected_url = ( - "https://test-bucket.s3.custom-endpoint.com/2025-09-14/test-key.json" - ) - assert ( - url == expected_url - ), f"Expected virtual-hosted-style URL {expected_url}, got {url}" + expected_url = "https://test-bucket.s3.custom-endpoint.com/2025-09-14/test-key.json" + assert url == expected_url, f"Expected virtual-hosted-style URL {expected_url}, got {url}" # Test 2: Path-style (default behavior with s3_use_virtual_hosted_style=False) s3_logger_path = S3Logger( @@ -241,12 +232,8 @@ class TestS3V2UnitTests: call_args_path = s3_logger_path.async_httpx_client.put.call_args assert call_args_path is not None url_path = call_args_path[0][0] - expected_path_url = ( - "https://s3.custom-endpoint.com/test-bucket/2025-09-14/test-key.json" - ) - assert ( - url_path == expected_path_url - ), f"Expected path-style URL {expected_path_url}, got {url_path}" + expected_path_url = "https://s3.custom-endpoint.com/test-bucket/2025-09-14/test-key.json" + assert url_path == expected_path_url, f"Expected path-style URL {expected_path_url}, got {url_path}" # Test 3: Virtual-hosted-style with http protocol s3_logger_http = S3Logger( @@ -266,12 +253,10 @@ class TestS3V2UnitTests: call_args_http = s3_logger_http.async_httpx_client.put.call_args assert call_args_http is not None url_http = call_args_http[0][0] - expected_http_url = ( - "http://http-bucket.minio.local:9000/2025-09-14/test-key.json" + expected_http_url = "http://http-bucket.minio.local:9000/2025-09-14/test-key.json" + assert url_http == expected_http_url, ( + f"Expected virtual-hosted-style URL with http {expected_http_url}, got {url_http}" ) - assert ( - url_http == expected_http_url - ), f"Expected virtual-hosted-style URL with http {expected_http_url}, got {url_http}" # Test 4: Sync upload method with virtual-hosted-style s3_logger_sync_virtual = S3Logger( @@ -295,12 +280,10 @@ class TestS3V2UnitTests: call_args_sync = mock_sync_client.put.call_args assert call_args_sync is not None url_sync = call_args_sync[0][0] - expected_sync_url = ( - "https://sync-bucket.storage.example.com/2025-09-14/test-key.json" + expected_sync_url = "https://sync-bucket.storage.example.com/2025-09-14/test-key.json" + assert url_sync == expected_sync_url, ( + f"Expected virtual-hosted-style sync URL {expected_sync_url}, got {url_sync}" ) - assert ( - url_sync == expected_sync_url - ), f"Expected virtual-hosted-style sync URL {expected_sync_url}, got {url_sync}" # Test 5: Download method with virtual-hosted-style s3_logger_download_virtual = S3Logger( @@ -316,34 +299,27 @@ class TestS3V2UnitTests: mock_download_response.status_code = 200 mock_download_response.json = MagicMock(return_value={"downloaded": "data"}) s3_logger_download_virtual.async_httpx_client = AsyncMock() - s3_logger_download_virtual.async_httpx_client.get.return_value = ( - mock_download_response - ) + s3_logger_download_virtual.async_httpx_client.get.return_value = mock_download_response - result = asyncio.run( - s3_logger_download_virtual._download_object_from_s3( - "2025-09-14/download-test-key.json" - ) - ) + result = asyncio.run(s3_logger_download_virtual._download_object_from_s3("2025-09-14/download-test-key.json")) call_args_download = s3_logger_download_virtual.async_httpx_client.get.call_args assert call_args_download is not None url_download = call_args_download[0][0] expected_download_url = "https://download-bucket.download.endpoint.com/2025-09-14/download-test-key.json" - assert ( - url_download == expected_download_url - ), f"Expected virtual-hosted-style download URL {expected_download_url}, got {url_download}" + assert url_download == expected_download_url, ( + f"Expected virtual-hosted-style download URL {expected_download_url}, got {url_download}" + ) assert result == {"downloaded": "data"} @patch("asyncio.create_task") @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") - def test_s3_v2_put_url_encodes_spaces_in_object_key( - self, mock_periodic_flush, mock_create_task - ): - import requests + def test_s3_v2_put_url_encodes_spaces_in_object_key(self, mock_periodic_flush, mock_create_task): from unittest.mock import AsyncMock + import requests + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement mock_periodic_flush.return_value = None @@ -487,9 +463,7 @@ async def test_async_upload_exhausts_retries_on_persistent_503(): # All 3 attempts return 503 response_503 = MagicMock() response_503.status_code = 503 - response_503.raise_for_status = MagicMock( - side_effect=Exception("503 Service Unavailable") - ) + response_503.raise_for_status = MagicMock(side_effect=Exception("503 Service Unavailable")) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock(return_value=response_503) @@ -528,12 +502,12 @@ async def test_async_upload_no_retry_on_4xx(): s3_object_download_filename="test-no-retry.json", ) - response_403 = MagicMock() - response_403.status_code = 403 - response_403.raise_for_status = MagicMock(side_effect=Exception("403 Forbidden")) + response_400 = MagicMock() + response_400.status_code = 400 + response_400.raise_for_status = MagicMock(side_effect=Exception("400 Bad Request")) logger.async_httpx_client = AsyncMock() - logger.async_httpx_client.put = AsyncMock(return_value=response_403) + logger.async_httpx_client.put = AsyncMock(return_value=response_400) with patch.object(logger, "handle_callback_failure") as mock_failure: await logger.async_upload_data_to_s3(test_element) @@ -543,6 +517,190 @@ async def test_async_upload_no_retry_on_4xx(): mock_failure.assert_called_once_with(callback_name="S3Logger") +_SIGV4_ACCESS_KEY = re.compile(r"Credential=(AKIA\d+)/") + + +@pytest.fixture +def rotating_profile(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> str: + """ + A real botocore profile whose credential_process hands out a new key generation on every call and + expires inside the advisory refresh window, so RefreshableCredentials re-runs it on every property read. + """ + counter = tmp_path / "generation" + script = tmp_path / "rotate_credentials.py" + script.write_text( + textwrap.dedent( + f""" + import json, sys + from datetime import datetime, timedelta, timezone + from pathlib import Path + + counter = Path({str(counter)!r}) + generation = int(counter.read_text()) if counter.exists() else 0 + counter.write_text(str(generation + 1)) + expiry = (datetime.now(timezone.utc) + timedelta(minutes=12)).strftime("%Y-%m-%dT%H:%M:%SZ") + json.dump( + {{ + "Version": 1, + "AccessKeyId": f"AKIA{{generation}}", + "SecretAccessKey": f"secret-{{generation}}", + "SessionToken": f"token-{{generation}}", + "Expiration": expiry, + }}, + sys.stdout, + ) + """ + ) + ) + profile = f"rotating-{uuid.uuid4().hex}" + (tmp_path / "config").write_text(f"[profile {profile}]\ncredential_process = {sys.executable} {script}\n") + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) + return profile + + +def _generation(request: httpx.Request) -> tuple[str, str]: + """(access key generation, session token generation) SigV4 baked into one request.""" + access_key = _SIGV4_ACCESS_KEY.search(request.headers["Authorization"]) + assert access_key is not None + return access_key.group(1).removeprefix("AKIA"), request.headers["X-Amz-Security-Token"].removeprefix("token-") + + +@asynccontextmanager +async def _s3_logger_on_production_handler(profile: str, statuses: list[int]): + """ + S3Logger wired to the real AsyncHTTPHandler over an httpx MockTransport that answers with the given + statuses in order, so the handler's own raise_for_status behaviour is exercised end to end. + """ + requests: list[httpx.Request] = [] + replies = iter(statuses) + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(next(replies), request=request, text="SignatureDoesNotMatch") + + handler = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_region_name="us-east-1", + s3_aws_profile_name=profile, + s3_flush_interval=3600, + ) + logger.async_httpx_client = handler + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + yield logger, requests, mock_sleep + await handler.client.aclose() + + +@pytest.mark.asyncio +async def test_async_upload_signs_with_one_frozen_credential_snapshot(rotating_profile: str, caplog): + """ + RefreshableCredentials refreshes on every property read once inside the advisory window, so signing + off the live object would mix the access key of one generation with the token of the next. + """ + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-frozen.json", + payload={"test": "frozen"}, + s3_object_download_filename="test-frozen.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [200]) as (logger, requests, _): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 1 + access_key_generation, token_generation = _generation(requests[0]) + assert access_key_generation == token_generation + assert "Error uploading to s3" not in caplog.text + + +@pytest.mark.asyncio +async def test_async_upload_retries_403_with_fresh_credentials_and_signature(rotating_profile: str, caplog): + """ + A 403 (SignatureDoesNotMatch after an IMDS rotation) must be retried, and the retry must fetch + credentials again and carry a signature computed from that newer generation. + """ + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-403.json", + payload={"test": "403"}, + s3_object_download_filename="test-403.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [403, 200]) as (logger, requests, mock_sleep): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 2 + first_key, first_token = _generation(requests[0]) + second_key, second_token = _generation(requests[1]) + assert first_key == first_token + assert second_key == second_token + assert int(second_key) > int(first_key) + assert requests[1].headers["Authorization"] != requests[0].headers["Authorization"] + mock_sleep.assert_awaited_once_with(1) + assert "Error uploading to s3" not in caplog.text + + +@pytest.mark.asyncio +async def test_async_upload_exhausts_403_retries_through_production_http_handler(rotating_profile: str, caplog): + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-403-exhausted.json", + payload={"test": "403-exhausted"}, + s3_object_download_filename="test-403-exhausted.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [403, 403, 403]) as (logger, requests, mock_sleep): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 3 + assert mock_sleep.await_args_list == [call(1), call(2)] + assert "Error uploading to s3" in caplog.text + + +@pytest.mark.asyncio +async def test_async_upload_does_not_retry_404_through_production_http_handler(rotating_profile: str, caplog): + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-404.json", + payload={"test": "404"}, + s3_object_download_filename="test-404.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [404]) as (logger, requests, mock_sleep): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 1 + mock_sleep.assert_not_awaited() + assert "Error uploading to s3" in caplog.text + + +def test_sync_upload_retries_403_with_fresh_signature(rotating_profile: str, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_PROFILE", rotating_profile) + logger = S3Logger(s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_flush_interval=3600) + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-sync-403.json", + payload={"test": "sync-403"}, + s3_object_download_filename="test-sync-403.json", + ) + requests: list[httpx.Request] = [] + replies = iter([403, 200]) + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(next(replies), request=request) + + handler = HTTPHandler() + handler.client = httpx.Client(transport=httpx.MockTransport(respond)) + with ( + patch( # test-quality-ok: sync upload builds its HTTPHandler per call, there is no injection seam for it + "litellm.integrations.s3_v2._get_httpx_client", return_value=handler + ), + patch("time.sleep") as mock_sleep, + ): + logger.upload_data_to_s3(test_element) + + assert len(requests) == 2 + first_key, first_token = _generation(requests[0]) + second_key, second_token = _generation(requests[1]) + assert first_key == first_token + assert second_key == second_token + assert int(second_key) > int(first_key) + mock_sleep.assert_called_once_with(1) + + def test_sync_upload_retries_on_s3_503(): """ Test that the sync upload_data_to_s3 retries on transient S3 503. @@ -626,9 +784,7 @@ async def test_async_log_event_skips_when_standard_logging_object_missing(): # Nothing should have been queued (catches the case where code falls # through without returning and appends None to the queue) - assert ( - len(logger.log_queue) == 0 - ), "log_queue should be empty when standard_logging_object is missing" + assert len(logger.log_queue) == 0, "log_queue should be empty when standard_logging_object is missing" @pytest.mark.asyncio @@ -767,20 +923,18 @@ async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch): litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, # This should NOT be ignored - "s3_use_ssl": False, # This should also NOT be ignored - }, + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, # This should NOT be ignored + "s3_use_ssl": False, # This should also NOT be ignored + }, ) with patch("asyncio.create_task"): - with patch( - "litellm.integrations.s3_v2.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client: mock_client = AsyncMock() mock_get_client.return_value = mock_client @@ -788,22 +942,16 @@ async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch): logger = S3Logger() # Verify s3_verify is False, not None - assert ( - logger.s3_verify is False - ), f"Expected s3_verify=False, got {logger.s3_verify}" - assert ( - logger.s3_use_ssl is False - ), f"Expected s3_use_ssl=False, got {logger.s3_use_ssl}" + assert logger.s3_verify is False, f"Expected s3_verify=False, got {logger.s3_verify}" + assert logger.s3_use_ssl is False, f"Expected s3_use_ssl=False, got {logger.s3_use_ssl}" # Verify that get_async_httpx_client was called with ssl_verify=False mock_get_client.assert_called_once() call_kwargs = mock_get_client.call_args.kwargs - assert ( - "params" in call_kwargs - ), "params should be passed to get_async_httpx_client" - assert call_kwargs["params"] == { - "ssl_verify": False - }, f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}" + assert "params" in call_kwargs, "params should be passed to get_async_httpx_client" + assert call_kwargs["params"] == {"ssl_verify": False}, ( + f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}" + ) @pytest.mark.asyncio @@ -820,17 +968,15 @@ async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch): litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_aws_access_key_id": "test-key", - "s3_aws_secret_access_key": "test-secret", - "s3_region_name": "us-east-1", - }, + "s3_bucket_name": "test-bucket", + "s3_aws_access_key_id": "test-key", + "s3_aws_secret_access_key": "test-secret", + "s3_region_name": "us-east-1", + }, ) with patch("asyncio.create_task"): - with patch( - "litellm.integrations.s3_v2.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client: mock_client = AsyncMock() mock_get_client.return_value = mock_client @@ -838,9 +984,7 @@ async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch): logger = S3Logger() # Verify s3_verify is None (default) - assert ( - logger.s3_verify is None - ), f"Expected s3_verify=None, got {logger.s3_verify}" + assert logger.s3_verify is None, f"Expected s3_verify=None, got {logger.s3_verify}" # Verify that get_async_httpx_client was called mock_get_client.assert_called_once() @@ -868,13 +1012,13 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatc litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - }, + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, ) with patch("asyncio.create_task"): @@ -890,9 +1034,7 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatc httpx_client = logger.async_httpx_client.client # Check the _verify attribute (httpx internal) if hasattr(httpx_client, "_verify"): - assert ( - httpx_client._verify is False - ), f"Expected httpx client _verify=False, got {httpx_client._verify}" + assert httpx_client._verify is False, f"Expected httpx client _verify=False, got {httpx_client._verify}" @pytest.mark.asyncio @@ -910,13 +1052,13 @@ async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch): litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - }, + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, ) with patch("asyncio.create_task"): @@ -948,9 +1090,9 @@ async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch): if hasattr(logger.async_httpx_client, "client"): httpx_client = logger.async_httpx_client.client if hasattr(httpx_client, "_verify"): - assert ( - httpx_client._verify is False - ), f"Expected async httpx client _verify=False, got {httpx_client._verify}" + assert httpx_client._verify is False, ( + f"Expected async httpx client _verify=False, got {httpx_client._verify}" + ) @pytest.mark.asyncio @@ -1017,9 +1159,7 @@ def patch_asyncio_create_task(): (True, True, None, None, ""), ], ) -def test_s3_object_key_prefix_combinations( - use_team_prefix, use_key_prefix, team_alias, key_alias, expected_prefix -): +def test_s3_object_key_prefix_combinations(use_team_prefix, use_key_prefix, team_alias, key_alias, expected_prefix): """ Validate correct S3 prefix composition for team alias + key alias combinations. """ @@ -1490,9 +1630,7 @@ def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch): logger = S3Logger(s3_callback_params_override=override) assert logger.s3_bucket_name == "resolved-bucket" assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - assert ( - litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - ) + assert litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" def test_s3_callback_params_override_none_falls_back_to_global(monkeypatch): @@ -1520,9 +1658,7 @@ def _expected_content_md5(payload: dict) -> str: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps json_string = safe_dumps(payload) - return base64.b64encode( - hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest() - ).decode() + return base64.b64encode(hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest()).decode() def _require_non_security_md5(monkeypatch): @@ -1658,9 +1794,9 @@ def test_s3_server_side_encryption_read_from_callback_params(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @@ -1789,10 +1925,10 @@ def test_s3_sse_kms_key_id_read_from_callback_params(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, ) logger = S3Logger() assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") @@ -1863,10 +1999,10 @@ def test_kms_key_id_dropped_when_algorithm_is_not_kms(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "AES256", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "AES256", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "AES256" @@ -1884,10 +2020,10 @@ def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(monkeypatch litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": True, - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": True, + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @@ -1902,10 +2038,10 @@ def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": 12345, - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": 12345, + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @@ -2045,6 +2181,7 @@ async def test_download_signs_object_key_with_space_the_way_s3_does(): headers=call.kwargs["headers"], ) + _RESERVED_CHAR_KEYS = ( "2026-08-21/time-05-29-36_resp_bGl0ZWxsbTpjdXN0b20=.json", "session=logs/2026-08-21/time-05-29-36_abc.json", diff --git a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py index ae5cffd8ab0..f1f9f7c3f3f 100644 --- a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py +++ b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -11,7 +11,16 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i ProxyServerRuntime, VectorStorePreCallHook, ) -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse +from litellm.types.utils import ( + CallTypes, + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) from litellm.types.vector_stores import ( VectorStoreResultContent, VectorStoreSearchResponse, @@ -36,6 +45,18 @@ def _search_response(text: str) -> VectorStoreSearchResponse: ) +def _first_message(response: ModelResponse) -> Message: + choice = response.choices[0] + assert isinstance(choice, Choices) + return choice.message + + +@dataclass(frozen=True) +class ExplodingRegistry: + async def pop_vector_stores_to_run_with_db_fallback(self, **kwargs: object) -> list[LiteLLM_ManagedVectorStore]: + raise RuntimeError("the registry blew up") + + @dataclass class RecordingRouter: failing_vector_store_ids: frozenset[str] = frozenset() @@ -285,3 +306,264 @@ def test_the_default_runtime_follows_the_proxy_globals(monkeypatch: pytest.Monke assert runtime.llm_router() is router assert runtime.prisma_client() is prisma + + +@pytest.mark.asyncio +async def test_a_failing_vector_store_is_reported_back_to_the_caller( + registry_with: RegisterStores, +) -> None: + """Regression (LIT-6809): a silently dropped store left the caller with an un-augmented answer and no signal.""" + registry_with("vs-broken", "vs-healthy") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))) + ), + ["vs-broken", "vs-healthy"], + logging_obj, + ) + + response = ModelResponse(choices=[Choices(message=Message(content="an answer"))]) + await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.acompletion, + ) + + provider_specific_fields = _first_message(response).provider_specific_fields or {} + assert provider_specific_fields["vector_store_search_failures"] == ( + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + }, + ) + assert len(provider_specific_fields["search_results"]) == 1 + + +@pytest.mark.asyncio +async def test_a_healthy_vector_store_alone_reports_no_failures(registry_with: RegisterStores) -> None: + registry_with("vs-healthy") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=RecordingRouter())), + ["vs-healthy"], + logging_obj, + ) + + response = ModelResponse(choices=[Choices(message=Message(content="an answer"))]) + await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.acompletion, + ) + + assert "vector_store_search_failures" not in (_first_message(response).provider_specific_fields or {}) + + +@pytest.mark.asyncio +async def test_a_failing_vector_store_is_reported_on_the_responses_api_response( + registry_with: RegisterStores, +) -> None: + """Regression (LIT-6809): /v1/responses answered 200 with no sign the knowledge base was missing.""" + registry_with("vs-broken") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))) + ), + ["vs-broken"], + logging_obj, + ) + + response = ResponsesAPIResponse(id="resp-lit6809", created_at=0, output=[]) + await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.aresponses, + ) + + assert response.model_dump()["vector_store_search_failures"] == [ + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + } + ] + + +@pytest.mark.asyncio +async def test_a_healthy_vector_store_leaves_the_responses_api_response_alone(registry_with: RegisterStores) -> None: + registry_with("vs-healthy") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=RecordingRouter())), + ["vs-healthy"], + logging_obj, + ) + + response = ResponsesAPIResponse(id="resp-lit6809", created_at=0, output=[]) + await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.aresponses, + ) + + assert "vector_store_search_failures" not in response.model_dump() + + +@pytest.mark.asyncio +async def test_a_failing_vector_store_is_reported_on_the_streaming_chunk(registry_with: RegisterStores) -> None: + registry_with("vs-broken") + logging_obj = FakeLoggingObj({}) + + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))) + ), + ["vs-broken"], + logging_obj, + ) + + chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="an answer"))]) + await VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=None) + ).async_post_call_streaming_deployment_hook( + request_data=logging_obj.model_call_details, + response_chunk=chunk, + call_type=CallTypes.acompletion, + ) + + assert (chunk.choices[0].delta.provider_specific_fields or {})["vector_store_search_failures"] == ( + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + }, + ) + + +@pytest.mark.asyncio +async def test_error_mode_fails_the_request_instead_of_answering_without_the_knowledge_base( + registry_with: RegisterStores, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression (LIT-6809): opting in must turn an ungrounded answer into a 400 the caller can act on.""" + registry_with("vs-broken", "vs-healthy") + monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "error") + + with pytest.raises(litellm.VectorStoreSearchError) as raised: + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime( + router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"})) + ) + ), + ["vs-broken", "vs-healthy"], + FakeLoggingObj({}), + ) + + assert raised.value.status_code == 400 + assert raised.value.failures == ( + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + }, + ) + assert "vs-broken: litellm.BadRequestError: no healthy deployments for vs-broken" in raised.value.message + + +@pytest.mark.asyncio +async def test_a_misspelled_failure_mode_annotates_instead_of_erroring_the_request( + registry_with: RegisterStores, + monkeypatch: pytest.MonkeyPatch, + warnings: list[logging.LogRecord], +) -> None: + """Regression (LIT-6809): litellm_settings takes any value, so a typo must not become a 500.""" + registry_with("vs-broken") + monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "erorr") + + logging_obj = FakeLoggingObj({}) + _, messages, _ = await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))) + ), + ["vs-broken"], + logging_obj, + ) + + assert messages[0]["content"] == "what is litellm?" + assert logging_obj.model_call_details["vector_store_search_failures"] == ( + { + "vector_store_id": "vs-broken", + "custom_llm_provider": "bedrock", + "error": "litellm.BadRequestError: no healthy deployments for vs-broken", + }, + ) + assert any("erorr" in record.getMessage() for record in warnings) + + +@pytest.mark.asyncio +async def test_error_mode_leaves_a_fully_healthy_request_alone( + registry_with: RegisterStores, + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry_with("vs-healthy") + monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "error") + + _, messages, _ = await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=RecordingRouter())), + ["vs-healthy"], + FakeLoggingObj({}), + ) + + assert messages[0]["content"] == "Context:\n\ncontext from vs-healthy\n\n" + + +@pytest.mark.asyncio +async def test_error_mode_does_not_swallow_the_raise_in_the_hooks_own_catch_all( + registry_with: RegisterStores, + monkeypatch: pytest.MonkeyPatch, + warnings: list[logging.LogRecord], +) -> None: + """Regression (LIT-6809): the catch-all around the hook must not turn the opted-in failure back into a 200.""" + registry_with("vs-broken") + monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "error") + + with pytest.raises(litellm.VectorStoreSearchError): + await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime( + router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"})) + ) + ), + ["vs-broken"], + FakeLoggingObj({}), + ) + + assert [record.levelname for record in warnings] == ["WARNING"] + + +@pytest.mark.asyncio +async def test_a_crash_outside_the_search_names_the_requested_vector_stores( + monkeypatch: pytest.MonkeyPatch, + warnings: list[logging.LogRecord], +) -> None: + """Regression (LIT-6809): the catch-all logged no store id, so an operator could not tell which store broke.""" + monkeypatch.setattr(litellm, "vector_store_registry", ExplodingRegistry()) + + _, messages, _ = await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)), + ["vs-one", "vs-two"], + FakeLoggingObj({}), + ) + + assert messages == [{"role": "user", "content": "what is litellm?"}] + assert [record.getMessage() for record in warnings] == [ + "Error in VectorStorePreCallHook for vector_store_ids=('vs-one', 'vs-two'): the registry blew up" + ] diff --git a/tests/test_litellm/litellm_core_utils/event_loop_lag.py b/tests/test_litellm/litellm_core_utils/event_loop_lag.py new file mode 100644 index 00000000000..1cac0365547 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/event_loop_lag.py @@ -0,0 +1,40 @@ +import asyncio +import time +from collections.abc import Awaitable, Callable +from typing import Final, TypeVar + +import litellm + +T = TypeVar("T") + + +def warm_tokenizer(model: str) -> None: + litellm.token_counter(model=model, text="load the tokenizer before anything is timed") + + +async def loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: + async def wake_lag() -> float: + started: Final = time.perf_counter() + await asyncio.sleep(0.001) + return time.perf_counter() - started - 0.001 + + return tuple([await wake_lag() for _ in iter(until.is_set, True)]) + + +async def timed_with_loop_lags(run: Callable[[], Awaitable[T]]) -> tuple[T, float, tuple[float, ...]]: + finished: Final = asyncio.Event() + + async def timed() -> tuple[T, float]: + await asyncio.sleep(0) + started: Final = time.perf_counter() + try: + return await run(), time.perf_counter() - started + finally: + finished.set() + + (result, took), lags = await asyncio.gather(timed(), loop_wake_lags(finished)) + return result, took, lags + + +def assert_loop_stayed_free(took: float, lags: tuple[float, ...]) -> None: + assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index df680b7cb0e..cbe6fe198c9 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1,9 +1,11 @@ import json +from datetime import datetime, timezone import pytest from fastapi.testclient import TestClient import litellm +from litellm._internal_context import pinned_billing_time from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) @@ -27,10 +29,10 @@ from litellm.types.utils import ( ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( + BilledTokenRates, CostCalculatorUtils, PromptTokensDetailsResult, TokenRates, - TokenTypeCostBreakdown, _calculate_input_cost, _get_token_base_cost, _is_off_peak, @@ -38,6 +40,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( apply_off_peak_pricing, calculate_cache_writing_cost, generic_cost_per_token, + get_billed_token_rates, get_token_type_cost_breakdown, ) from litellm.types.utils import CacheCreationTokenDetails, Usage @@ -49,6 +52,44 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) +@pytest.mark.parametrize("prompt_tokens", [100, 200000, 200001]) +@pytest.mark.parametrize("read_rate", [None, 0.0, 0.25e-6]) +@pytest.mark.parametrize("service_tier", [None, "priority"]) +def test_missing_cache_read_policy_preserves_billing(prompt_tokens, read_rate, service_tier): + info = { + "input_cost_per_token": 3e-6, + "input_cost_per_token_priority": 4e-6, + "input_cost_per_token_above_200k_tokens": 6e-6, + "input_cost_per_token_above_200k_tokens_priority": 8e-6, + "output_cost_per_token": 1e-6, + "cache_read_input_token_cost": read_rate, + } + usage = Usage(prompt_tokens=prompt_tokens, prompt_tokens_details={"cached_tokens": 100}) + billed = _get_token_base_cost(info, usage, service_tier=service_tier) + savings = _get_token_base_cost(info, usage, service_tier=service_tier, missing_cache_read_uses_input=True) + prompt_cost, _ = generic_cost_per_token("policy-fixture", usage, "openai", service_tier=service_tier, model_info=info) + assert billed[4] == pytest.approx(read_rate or 0.0) + assert savings[:4] == billed[:4] + assert savings[4] == pytest.approx(billed[0] if read_rate is None else read_rate) + assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4]) + + +def test_missing_cache_read_uses_off_peak_input_rate(): + from datetime import datetime, timezone + + info = { + "input_cost_per_token": 3e-6, + "off_peak_pricing": {"hours_utc": "00:00-23:59", "input_cost_per_token": 5e-6}, + } + when = datetime(2026, 9, 7, 12, tzinfo=timezone.utc) + billed = _get_token_base_cost(info, Usage(prompt_tokens=100), current_time=when) + savings = _get_token_base_cost( + info, Usage(prompt_tokens=100), current_time=when, missing_cache_read_uses_input=True + ) + assert billed[4] == 0.0 + assert savings[0] == savings[4] == 5e-6 + + def test_reasoning_tokens_no_price_set(_local_model_cost_map): # Use o1 - o1-mini was deprecated/renamed; o1 has same reasoning-token semantics # (no separate output_cost_per_reasoning_token, so all completion tokens use output_cost_per_token) @@ -1674,7 +1715,7 @@ def test_generic_cost_per_token_gpt55(_local_model_cost_map): def test_generic_cost_per_token_gpt55_pro(_local_model_cost_map): - """gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input.""" + """gpt-5.5-pro: responses-only model, $30/1M input, $180/1M output, no cached input rate published.""" model = "gpt-5.5-pro" custom_llm_provider = "openai" @@ -1683,7 +1724,7 @@ def test_generic_cost_per_token_gpt55_pro(_local_model_cost_map): # Sanity-check the map values match OpenAI's published pricing. assert model_cost_map["input_cost_per_token"] == 3e-5 assert model_cost_map["output_cost_per_token"] == 1.8e-4 - assert model_cost_map["cache_read_input_token_cost"] == 3e-6 + assert "cache_read_input_token_cost" not in model_cost_map assert model_cost_map["litellm_provider"] == "openai" # gpt-5.5-pro is a responses-only model (no /v1/chat/completions endpoint). assert model_cost_map["mode"] == "responses" @@ -2008,7 +2049,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 +2064,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 +2072,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 +2091,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 +2101,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", [ @@ -2599,6 +2661,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): "image_count": 0, "video_length_seconds": 0.0, "audio_length_seconds": 0.0, + "query_count": 0, } model_info: ModelInfo = {} @@ -3180,6 +3243,37 @@ def test_image_count_prevents_text_tokens_fallback(_local_model_cost_map): assert completion_cost == 0.0 +def test_query_count_bills_input_cost_per_query(_local_model_cost_map): + usage = Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + prompt_tokens_details=PromptTokensDetailsWrapper(query_count=3, image_count=1), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model="us.twelvelabs.marengo-embed-3-0-v1:0", + usage=usage, + custom_llm_provider="bedrock", + ) + + assert prompt_cost == pytest.approx(3 * 7e-05 + 1e-04) + assert completion_cost == 0.0 + + +def test_query_count_is_free_without_a_per_query_price(_local_model_cost_map): + usage = Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + prompt_tokens_details=PromptTokensDetailsWrapper(query_count=1), + ) + + prompt_cost, _ = generic_cost_per_token(model="text-embedding-3-small", usage=usage, custom_llm_provider="openai") + + assert prompt_cost == 0.0 + + # --------------------------------------------------------------------------- # Data-residency (OpenAI regional processing) tests # --------------------------------------------------------------------------- @@ -3815,6 +3909,200 @@ def test_token_type_cost_breakdown_reconciles_with_generic_total(_local_model_co assert text_input_cost + breakdown.cache_read_cost == pytest.approx(prompt_cost) +def _custom_priced_usage() -> Usage: + return Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800, cache_creation_tokens=100), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=200), + ) + + +def test_token_type_cost_breakdown_prices_custom_pricing_from_its_flat_rates(): + """ + A custom-priced deployment, usually absent from the cost map, used to get zero cache and + reasoning lines while its total already billed cache tokens at the custom cache rates. + The lines must come from the same flat rates: a configured cache rate, else the input + rate for cache tokens and the output rate for reasoning tokens. + """ + from litellm.types.utils import CostPerToken + + breakdown = get_token_type_cost_breakdown( + model="openai/onprem-model", + custom_llm_provider="openai", + usage=_custom_priced_usage(), + custom_cost_per_token=CostPerToken( + input_cost_per_token=1e-6, output_cost_per_token=2e-6, cache_read_input_token_cost=1e-7 + ), + ) + + assert breakdown.cache_read_cost == pytest.approx(800 * 1e-7) + assert breakdown.cache_creation_cost == pytest.approx(100 * 1e-6) + assert breakdown.reasoning_cost == pytest.approx(200 * 2e-6) + + +def test_token_type_cost_breakdown_reconciles_with_custom_pricing_totals(): + from litellm.cost_calculator import cost_per_token + from litellm.types.utils import CostPerToken + + usage = _custom_priced_usage() + custom_cost_per_token = CostPerToken( + input_cost_per_token=1e-6, + output_cost_per_token=2e-6, + cache_read_input_token_cost=1e-7, + cache_creation_input_token_cost=1.25e-6, + ) + + prompt_cost, completion_cost = cost_per_token( + model="openai/onprem-model", + custom_llm_provider="openai", + prompt_tokens=1000, + completion_tokens=500, + usage_object=usage, + custom_cost_per_token=custom_cost_per_token, + ) + breakdown = get_token_type_cost_breakdown( + model="openai/onprem-model", + custom_llm_provider="openai", + usage=usage, + custom_cost_per_token=custom_cost_per_token, + ) + + assert 100 * 1e-6 + breakdown.cache_read_cost + breakdown.cache_creation_cost == pytest.approx(prompt_cost) + assert 300 * 2e-6 + breakdown.reasoning_cost == pytest.approx(completion_cost) + + +def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + "tiered-cache-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + usage = Usage( + prompt_tokens=250_000, + completion_tokens=1_000, + total_tokens=251_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200_000, cache_creation_tokens=10_000), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=200), + ) + + rates = get_billed_token_rates(model="tiered-cache-model", custom_llm_provider="openai", usage=usage) + breakdown = get_token_type_cost_breakdown(model="tiered-cache-model", custom_llm_provider="openai", usage=usage) + + assert rates == BilledTokenRates( + input_cost_per_token=6e-6, + output_cost_per_token=3e-5, + cache_read_input_token_cost=6e-7, + cache_creation_input_token_cost=7.5e-6, + cache_creation_input_token_cost_above_1hr=0.0, + output_cost_per_reasoning_token=3e-5, + ) + assert breakdown.cache_read_cost == pytest.approx(200_000 * rates.cache_read_input_token_cost) + assert breakdown.cache_creation_cost == pytest.approx(10_000 * rates.cache_creation_input_token_cost) + assert breakdown.reasoning_cost == pytest.approx(200 * rates.output_cost_per_reasoning_token) + + +def test_a_pinned_billing_time_prices_the_totals_and_the_reported_rates_at_one_moment(monkeypatch): + """Totals and reported rates resolve off-peak pricing on separate paths that each read the + clock, so a window opening between the two reads used to leave them describing one request + at two different prices. Pinned, both must answer for the pinned moment.""" + monkeypatch.setitem( + litellm.model_cost, + "off-peak-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "off_peak_pricing": { + "hours_utc": "02:00-03:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 5e-6, + }, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + with pinned_billing_time(datetime(2026, 1, 1, 2, 30, tzinfo=timezone.utc)): + off_peak_prompt_cost, off_peak_completion_cost = generic_cost_per_token( + model="off-peak-model", usage=usage, custom_llm_provider="openai" + ) + off_peak_rates = get_billed_token_rates(model="off-peak-model", custom_llm_provider="openai", usage=usage) + with pinned_billing_time(datetime(2026, 1, 1, 12, 30, tzinfo=timezone.utc)): + peak_prompt_cost, peak_completion_cost = generic_cost_per_token( + model="off-peak-model", usage=usage, custom_llm_provider="openai" + ) + peak_rates = get_billed_token_rates(model="off-peak-model", custom_llm_provider="openai", usage=usage) + + assert off_peak_rates.input_cost_per_token == pytest.approx(1e-6) + assert peak_rates.input_cost_per_token == pytest.approx(3e-6) + assert off_peak_prompt_cost == pytest.approx(1000 * off_peak_rates.input_cost_per_token) + assert off_peak_completion_cost == pytest.approx(500 * off_peak_rates.output_cost_per_token) + assert peak_prompt_cost == pytest.approx(1000 * peak_rates.input_cost_per_token) + assert peak_completion_cost == pytest.approx(500 * peak_rates.output_cost_per_token) + + +def test_the_token_type_breakdown_carries_the_rates_it_billed_at(monkeypatch): + """Callers that report both the lines and the rates read the rates off the breakdown rather than + resolving them a second time, so the breakdown has to hand back exactly what it billed at.""" + monkeypatch.setitem( + litellm.model_cost, + "xai/tiered-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "litellm_provider": "xai", + "mode": "chat", + }, + ) + usage = Usage( + prompt_tokens=200_000, + completion_tokens=1_000, + total_tokens=201_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100_000), + ) + + breakdown = get_token_type_cost_breakdown(model="xai/tiered-model", custom_llm_provider="xai", usage=usage) + + assert breakdown.rates == get_billed_token_rates( + model="xai/tiered-model", custom_llm_provider="xai", usage=usage + ) + assert breakdown.rates.cache_read_input_token_cost == pytest.approx(6e-7) + assert breakdown.cache_read_cost == pytest.approx(100_000 * breakdown.rates.cache_read_input_token_cost) + + +def test_the_token_type_breakdown_reports_no_rates_for_an_unpriced_model(): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + + breakdown = get_token_type_cost_breakdown( + model="no-such-model-anywhere", custom_llm_provider="openai", usage=usage + ) + + assert breakdown.rates is None + + +def test_billed_token_rates_are_none_for_an_unpriced_model(): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + + assert get_billed_token_rates(model="no-such-model-anywhere", custom_llm_provider="openai", usage=usage) is None + + def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost_map): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) @@ -3822,9 +4110,7 @@ def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost model="gpt-4o", custom_llm_provider="openai", usage=usage ) - assert breakdown == TokenTypeCostBreakdown( - reasoning_cost=0.0, cache_read_cost=0.0, cache_creation_cost=0.0 - ) + assert (breakdown.reasoning_cost, breakdown.cache_read_cost, breakdown.cache_creation_cost) == (0.0, 0.0, 0.0) @pytest.mark.parametrize( @@ -3896,9 +4182,7 @@ def test_token_type_cost_breakdown_handles_unknown_model_gracefully(): completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=5), ), ) - assert breakdown == TokenTypeCostBreakdown( - reasoning_cost=0.0, cache_read_cost=0.0, cache_creation_cost=0.0 - ) + assert (breakdown.reasoning_cost, breakdown.cache_read_cost, breakdown.cache_creation_cost) == (0.0, 0.0, 0.0) def test_token_type_cost_breakdown_applies_regional_uplift(_local_model_cost_map): @@ -4766,3 +5050,100 @@ def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, r ) assert cost == expected_cost + + +def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_local_model_cost_map: None) -> None: + """ + Realtime usage (OpenAI and Azure) reports output_tokens == text_tokens + audio_tokens with + reasoning_tokens already counted inside text_tokens, so reasoning must not be billed on top. + """ + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=346, + completion_tokens=29, + total_tokens=375, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=152, image_tokens=194, audio_tokens=0, cached_tokens=128 + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=29, audio_tokens=0, reasoning_tokens=19 + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert completion_cost == pytest.approx(29 * info["output_cost_per_token"]) + assert completion_cost - breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"]) + assert prompt_cost == pytest.approx( + 24 * info["input_cost_per_token"] + + 128 * info["cache_read_input_token_cost"] + + 194 * info["input_cost_per_image_token"] + ) + + +def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens( + _local_model_cost_map: None, +) -> None: + """Providers whose text_tokens exclude reasoning (text + reasoning == completion) stay billed in full.""" + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=100, + completion_tokens=44, + total_tokens=144, + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=25, audio_tokens=0, reasoning_tokens=19 + ), + ) + + _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert completion_cost == pytest.approx(44 * info["output_cost_per_token"]) + + +def test_generic_cost_per_token_strips_only_the_reasoning_share_when_text_over_reports( + _local_model_cost_map: None, +) -> None: + """Text over-reported past the reasoning share keeps its extra tokens billed; only the nested reasoning is netted out.""" + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=120, + completion_tokens=100, + total_tokens=220, + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=100, audio_tokens=70, reasoning_tokens=10), + ) + + _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"]) + assert completion_cost == pytest.approx( + 100 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"] + ) + + +def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output(_local_model_cost_map: None) -> None: + """Audio-output realtime usage nests reasoning inside text_tokens next to audio_tokens; text is billed net of it.""" + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=120, + completion_tokens=100, + total_tokens=220, + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=30, audio_tokens=70, reasoning_tokens=20), + ) + + _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert breakdown.reasoning_cost == pytest.approx(20 * info["output_cost_per_token"]) + assert completion_cost == pytest.approx( + 30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"] + ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 6cc3dcceebc..bbb7b5f9c35 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,5 +1,6 @@ -import os +import json from collections.abc import Mapping, Sequence +from pathlib import Path import pytest @@ -11,8 +12,6 @@ from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebS from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams - - def test_web_search_cost_low(): web_search_options = WebSearchOptions(search_context_size="low") model_info = litellm.get_model_info("gpt-4o-search-preview") @@ -683,12 +682,13 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model( def _openai_responses_with_web_search_calls(model, num_calls): - from litellm.types.llms.openai import ResponsesAPIResponse from openai.types.responses.response_function_web_search import ( ActionSearch, ResponseFunctionWebSearch, ) + from litellm.types.llms.openai import ResponsesAPIResponse + output = [ ResponseFunctionWebSearch( id=f"ws_{i}", @@ -859,11 +859,62 @@ def test_dated_search_preview_entries_carry_search_pricing(local_model_cost_map) custom_llm_provider="openai", standard_built_in_tools_params=None, ) - assert cost == pytest.approx(0.035), ( - f"dated search-preview id must bill the $0.035 search fee, got ${cost}" + assert cost == pytest.approx(0.025), ( + f"dated search-preview id must bill the $0.025 search fee, got ${cost}" ) +@pytest.mark.parametrize( + "web_search_options", + [ + None, + WebSearchOptions(search_context_size="low"), + WebSearchOptions(search_context_size="medium"), + WebSearchOptions(search_context_size="high"), + ], +) +def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias( + web_search_options: WebSearchOptions | None, local_model_cost_map: None +) -> None: + alias_info = litellm.get_model_info("gpt-4o-mini") + snapshot_info = litellm.get_model_info("gpt-4o-mini-2024-07-18") + + assert not snapshot_info["supports_web_search"] + assert not alias_info["supports_web_search"] + + snapshot_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search( + web_search_options=web_search_options, model_info=snapshot_info + ) + alias_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search( + web_search_options=web_search_options, model_info=alias_info + ) + + assert snapshot_cost == alias_cost == 0.025 + + +def test_gpt_4o_mini_web_search_price_matches_in_both_cost_maps(): + repo_root = Path(__file__).parents[4] + cost_maps = tuple( + json.loads((repo_root / path).read_text(encoding="utf-8")) + for path in ( + "model_prices_and_context_window.json", + "litellm/model_prices_and_context_window_backup.json", + ) + ) + canonical, backup = cost_maps + expected_search_price = { + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025, + "search_context_size_high": 0.025, + } + for model_name in ("gpt-4o-mini", "gpt-4o-mini-2024-07-18"): + canonical_entry = canonical[model_name] + backup_entry = backup[model_name] + assert canonical_entry["search_context_cost_per_query"] == expected_search_price + assert backup_entry["search_context_cost_per_query"] == expected_search_price + assert canonical_entry == backup_entry + + # Note: File search integration test removed due to complex annotation detection logic # The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index 304d732c518..8e46ae21de6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -1,4 +1,6 @@ +from typing import Final +import pytest from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -99,3 +101,97 @@ def test_handle_invalid_parallel_tool_calls_skips_custom_tool_calls(): ) result = _handle_invalid_parallel_tool_calls([custom_tool_call, function_tool_call]) assert result == [custom_tool_call, function_tool_call] + + +def test_convert_empty_choices_response() -> None: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + ) + + resp: Final = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, + "vertex_ai_safety_results": ["blocked"], + } + result: Final = convert_to_model_response_object( + response_object=resp, + model_response_object=ModelResponse(), + response_type="completion", + ) + assert result.choices == [] + assert getattr(result, "vertex_ai_safety_results") == ["blocked"] + + sync_stream: Final = list(convert_to_streaming_response(response_object=resp)) + assert len(sync_stream) == 1 + assert sync_stream[0].choices == [] + + +@pytest.mark.asyncio +async def test_convert_empty_choices_response_async() -> None: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response_async, + ) + + resp: Final = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, + } + async_chunks: Final = [chunk async for chunk in convert_to_streaming_response_async(response_object=resp)] + assert len(async_chunks) == 1 + assert async_chunks[0].choices == [] + + +def test_convert_missing_choices_raises_api_error() -> None: + from litellm.exceptions import APIError + + resp: Final = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + } + with pytest.raises(APIError) as exc_info: + convert_to_model_response_object( + response_object=resp, + model_response_object=ModelResponse(), + response_type="completion", + ) + assert "no 'choices'" in str(exc_info.value) + + +@pytest.mark.parametrize(("choices", "type_name"), [({}, "dict"), ("", "str"), (None, "NoneType"), (0, "int")]) +@pytest.mark.asyncio +async def test_convert_non_list_choices_raises_api_error(choices: object, type_name: str) -> None: + from litellm.exceptions import APIError + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, + ) + + resp: Final = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + "choices": choices, + } + expected: Final = f"'choices' that is not a list \\({type_name}\\)" + with pytest.raises(APIError, match=expected): + convert_to_model_response_object( + response_object=resp, + model_response_object=ModelResponse(), + response_type="completion", + ) + with pytest.raises(APIError, match=expected): + list(convert_to_streaming_response(response_object=resp)) + with pytest.raises(APIError, match=expected): + async for _ in convert_to_streaming_response_async(response_object=resp): + pass diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index c037f928593..b5890d1a5b0 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1,23 +1,33 @@ +import copy import functools import json import os +import sys +from typing import Final from unittest.mock import MagicMock, patch import pytest - from litellm.litellm_core_utils.prompt_templates.common_utils import ( + ENCRYPTED_REASONING_SIGNATURE_PREFIX, TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, add_system_prompt_to_messages, + encrypted_content_from_signature, + encrypted_reasoning_signature, get_file_ids_from_messages, get_format_from_file_id, handle_any_messages_to_chat_completion_str_messages_conversion, hoist_images_from_tool_messages, + is_encrypted_reasoning_block, + responses_reasoning_items_from_thinking_blocks, split_concatenated_json_objects, + strip_encrypted_reasoning_from_messages, update_messages_with_model_file_ids, ) +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + def test_get_format_from_file_id(): unified_file_id = "litellm_proxy:application/pdf;unified_id,cbbe3534-8bf8-4386-af00-f5f6b7e370bf" @@ -1435,7 +1445,7 @@ class TestFlattenTopLevelSchemaCombinators: assert schema == snapshot -class TestToolWithFlattenedParameters: +class TestToolWithSanitizedParameters: def _anyof_tool(self): return { "type": "function", @@ -1462,11 +1472,12 @@ class TestToolWithFlattenedParameters: def test_flattens_anyof_parameters_into_new_tool(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( - tool_with_flattened_parameters, + flatten_combinators_and_drop_non_python_regex_patterns, + tool_with_sanitized_parameters, ) tool = self._anyof_tool() - result = tool_with_flattened_parameters(tool) + result = tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) assert result is not tool parameters = result["function"]["parameters"] @@ -1479,7 +1490,8 @@ class TestToolWithFlattenedParameters: def test_clean_parameters_return_the_same_tool_object(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( - tool_with_flattened_parameters, + flatten_combinators_and_drop_non_python_regex_patterns, + tool_with_sanitized_parameters, ) tool = { @@ -1490,7 +1502,23 @@ class TestToolWithFlattenedParameters: }, } - assert tool_with_flattened_parameters(tool) is tool + assert tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) is tool + + def test_pattern_only_sanitizer_drops_the_regex_and_keeps_the_union(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + tool_with_sanitized_parameters, + ) + + tool = self._anyof_tool() + tool["function"]["parameters"]["properties"]["id"]["pattern"] = _ARTIFACT_FIELD_PATTERN + + result = tool_with_sanitized_parameters(tool, drop_non_python_regex_patterns) + + parameters = result["function"]["parameters"] + assert parameters["properties"]["id"] == {"type": "string"} + assert parameters["anyOf"] == self._anyof_tool()["function"]["parameters"]["anyOf"] + assert tool["function"]["parameters"]["properties"]["id"]["pattern"] == _ARTIFACT_FIELD_PATTERN @pytest.mark.parametrize( "tool", @@ -1503,10 +1531,127 @@ class TestToolWithFlattenedParameters: ) def test_non_dict_function_or_parameters_return_the_same_tool_object(self, tool): from litellm.litellm_core_utils.prompt_templates.common_utils import ( - tool_with_flattened_parameters, + flatten_combinators_and_drop_non_python_regex_patterns, + tool_with_sanitized_parameters, ) - assert tool_with_flattened_parameters(tool) is tool + assert tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) is tool + + +class TestDropNonPythonRegexPatterns: + """Claude Code's Artifact tool declares ECMA-262 ``\\p{..}`` escapes that OpenAI's + validator, which compiles ``pattern`` values and ``patternProperties`` keys with + Python ``re``, refuses as "not a 'regex'".""" + + def _schema(self, pattern): + return { + "type": "object", + "properties": { + "field": {"type": "string", "pattern": pattern}, + "writes": { + "type": "array", + "items": {"properties": {"doc_id": {"type": "string", "pattern": pattern}}}, + }, + "query": {"anyOf": [{"type": "string", "pattern": pattern}, {"type": "null"}]}, + "pair": {"type": "array", "prefixItems": [{"type": "string", "pattern": pattern}]}, + "extra": {"type": "object", "additionalProperties": {"type": "string", "pattern": pattern}}, + "tagged": { + "type": "object", + "patternProperties": {pattern: {"type": "string"}, "^x_": {"type": "integer"}}, + }, + }, + "$defs": {"segment": {"type": "string", "pattern": pattern}}, + "required": ["field"], + } + + def test_drops_every_regex_python_re_rejects_from_every_schema_position(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = self._schema(_ARTIFACT_FIELD_PATTERN) + + result = drop_non_python_regex_patterns(schema) + + assert '"pattern"' not in json.dumps(result) + properties = result["properties"] + assert properties["field"] == {"type": "string"} + assert properties["writes"]["items"]["properties"]["doc_id"] == {"type": "string"} + assert properties["query"]["anyOf"] == [{"type": "string"}, {"type": "null"}] + assert properties["pair"]["prefixItems"] == [{"type": "string"}] + assert properties["extra"]["additionalProperties"] == {"type": "string"} + assert properties["tagged"]["patternProperties"] == {"^x_": {"type": "integer"}} + assert result["$defs"]["segment"] == {"type": "string"} + assert result["required"] == ["field"] + assert schema == self._schema(_ARTIFACT_FIELD_PATTERN) + + def test_keeps_regexes_python_re_compiles_and_returns_the_same_object(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = self._schema(r'^(?!__.*__$)[^"\\./[\]]{1,200}$') + + assert drop_non_python_regex_patterns(schema) is schema + + def test_pattern_keys_inside_data_positions_are_not_regexes(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = { + "type": "object", + "properties": { + "pattern": {"type": "string"}, + "template": {"type": "object", "default": {"pattern": _ARTIFACT_FIELD_PATTERN}}, + "samples": {"type": "array", "examples": [{"pattern": _ARTIFACT_FIELD_PATTERN}]}, + "fixed": {"const": {"pattern": _ARTIFACT_FIELD_PATTERN}}, + "vendor": {"type": "string", "x-litellm": {"pattern": _ARTIFACT_FIELD_PATTERN}}, + }, + "required": ["pattern"], + } + + assert drop_non_python_regex_patterns(schema) is schema + + def test_regex_nested_past_what_python_re_can_parse_is_dropped_not_raised(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = { + "type": "object", + "properties": {"deep": {"type": "string", "pattern": "(" * 2000 + "a" + ")" * 2000}}, + } + + assert drop_non_python_regex_patterns(schema)["properties"]["deep"] == {"type": "string"} + + def test_walks_schemas_deeper_than_the_interpreter_recursion_limit(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + depth = sys.getrecursionlimit() + leaf = {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN} + schema = functools.reduce( + lambda inner, _: {"type": "object", "properties": {"child": inner}}, range(depth), leaf + ) + + result = drop_non_python_regex_patterns(schema) + + assert functools.reduce(lambda node, _: node["properties"]["child"], range(depth), result) == {"type": "string"} + assert functools.reduce(lambda node, _: node["properties"]["child"], range(depth), schema) is leaf + + def test_leaves_levels_past_the_json_nesting_limit_alone(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + leaf = {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN} + schema = functools.reduce( + lambda inner, _: {"type": "object", "properties": {"child": inner}}, range(1100), leaf + ) + + assert drop_non_python_regex_patterns(schema) is schema class TestRequestContainsImageContent: @@ -1554,3 +1699,117 @@ class TestRequestContainsImageContent: for _ in range(50): nested = {"type": "tool_result", "content": [nested]} assert request_contains_image_content([{"role": "user", "content": [nested]}]) is False + + +class TestEncryptedReasoningReplay: + """Regression for https://github.com/BerriAI/litellm/issues/40288.""" + + def test_signature_round_trips_the_encrypted_content(self): + assert encrypted_content_from_signature(encrypted_reasoning_signature("gAAAA_bytes")) == "gAAAA_bytes" + + @pytest.mark.parametrize( + "signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7] + ) + def test_anything_else_is_not_encrypted_content(self, signature): + assert encrypted_content_from_signature(signature) is None + + def test_encrypted_thinking_block_replays_its_own_item(self): + items = responses_reasoning_items_from_thinking_blocks( + [{"type": "thinking", "thinking": "Plan.", "signature": encrypted_reasoning_signature("gAAAA_1")}] + ) + assert items == ( + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Plan."}], + "encrypted_content": "gAAAA_1", + }, + ) + + def test_encrypted_redacted_block_replays_with_an_empty_summary(self): + items = responses_reasoning_items_from_thinking_blocks( + [{"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_1")}] + ) + assert items == ({"type": "reasoning", "summary": [], "encrypted_content": "gAAAA_1"},) + + def test_plain_blocks_collapse_into_one_summary_item_around_encrypted_ones(self): + items = responses_reasoning_items_from_thinking_blocks( + [ + {"type": "thinking", "thinking": "A.", "signature": None}, + {"type": "thinking", "thinking": "B.", "signature": ""}, + {"type": "thinking", "thinking": "C.", "signature": encrypted_reasoning_signature("gAAAA_c")}, + {"type": "redacted_thinking", "data": "anthropic-minted-opaque-data"}, + {"type": "thinking", "thinking": "D."}, + ] + ) + assert items == ( + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}], + }, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "C."}], "encrypted_content": "gAAAA_c"}, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "D."}]}, + ) + assert all("id" not in item for item in items) + + def test_blocks_without_text_or_encrypted_content_produce_nothing(self): + assert responses_reasoning_items_from_thinking_blocks([{"type": "thinking", "thinking": ""}]) == () + assert responses_reasoning_items_from_thinking_blocks([]) == () + + @pytest.mark.parametrize( + ("block", "expected"), + [ + ({"type": "thinking", "thinking": "x", "signature": encrypted_reasoning_signature("g")}, True), + ({"type": "redacted_thinking", "data": encrypted_reasoning_signature("g")}, True), + ({"type": "thinking", "thinking": "x", "signature": ENCRYPTED_REASONING_SIGNATURE_PREFIX}, True), + ({"type": "redacted_thinking", "data": ENCRYPTED_REASONING_SIGNATURE_PREFIX}, True), + ({"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}, False), + ({"type": "redacted_thinking", "data": "EmwKAhgBEgy"}, False), + ({"type": "text", "text": encrypted_reasoning_signature("g")}, False), + ("not a block", False), + ], + ) + def test_is_encrypted_reasoning_block(self, block, expected): + assert is_encrypted_reasoning_block(block) is expected + + def test_strip_drops_every_bridge_tagged_block_and_leaves_no_unsigned_thinking_behind(self): + assistant_content = [ + {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, + {"type": "thinking", "thinking": "packed by the bridge", "signature": encrypted_reasoning_signature("g1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("g2")}, + {"type": "thinking", "thinking": "", "signature": encrypted_reasoning_signature("g3")}, + {"type": "text", "text": "answer"}, + ] + messages = [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": assistant_content}, + {"role": "user", "content": [{"type": "text", "text": "follow-up"}]}, + ] + + strip_encrypted_reasoning_from_messages(messages) + + assert messages[1]["content"] is assistant_content + assert assistant_content == [ + {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, + {"type": "text", "text": "answer"}, + ] + assert all(block["signature"] for block in assistant_content if block["type"] == "thinking") + assert messages[0] == {"role": "user", "content": "question"} + assert messages[2] == {"role": "user", "content": [{"type": "text", "text": "follow-up"}]} + + @pytest.mark.parametrize( + "messages", + [ + "not a list", + None, + [{"role": "user", "content": None}], + [{"role": "user", "content": "plain string"}], + ["not a message"], + [{"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}]}], + ], + ) + def test_strip_leaves_history_without_bridge_reasoning_untouched(self, messages): + before = copy.deepcopy(messages) + + strip_encrypted_reasoning_from_messages(messages) + + assert messages == before diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index dd2d45f00c6..66d10fd1407 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -191,8 +191,16 @@ def test_bedrock_converse_assistant_with_empty_thinking_block_and_tool_calls(): {"type": "thinking", "thinking": "oss reasoning", "signature": None}, {"type": "thinking", "thinking": "oss reasoning", "signature": ""}, {"type": "thinking", "thinking": "oss reasoning"}, + {"type": "thinking", "thinking": "openai reasoning", "signature": "litellm_encrypted_reasoning:gAAAA"}, + {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:gAAAA"}, + ], + ids=[ + "null_signature", + "empty_signature", + "missing_signature", + "encrypted_reasoning_signature", + "encrypted_reasoning_redacted_data", ], - ids=["null_signature", "empty_signature", "missing_signature"], ) def test_anthropic_messages_pt_drops_unsignable_thinking_block(thinking_block): """Open-source reasoning models (DeepSeek-R1, Qwen, etc.) emit thinking blocks @@ -219,7 +227,7 @@ def test_anthropic_messages_pt_drops_unsignable_thinking_block(thinking_block): assistant = next(m for m in result if m["role"] == "assistant") content = assistant["content"] assert all( - block.get("type") != "thinking" for block in content + block.get("type") not in ("thinking", "redacted_thinking") for block in content ), f"unsignable thinking block must be dropped, got {content!r}" assert any( block.get("type") == "text" and block.get("text") == "2+2 equals 4." diff --git a/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py b/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py new file mode 100644 index 00000000000..42ca91bfd8f --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py @@ -0,0 +1,112 @@ +""" +Tests for route -> CallTypes resolution (api_route_to_call_types). + +Regression coverage for the guardrail route table bugs: +- placeholder segments with a literal suffix ({model}:generateContent) never matched +- the /v1beta generateContent routes were missing from the table +- /llm_passthrough listed the sync call type first, resolving consumers that + take call_types[0] to a handler-less type +""" + +from litellm.litellm_core_utils.api_route_to_call_types import ( + get_call_types_for_route, +) +from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes + + +class TestGenerateContentRouteResolution: + def test_bare_generate_content_route_resolves(self): + call_types = get_call_types_for_route("/models/gemini-2.5-flash:generateContent") + assert call_types is not None + assert list(call_types) == [CallTypes.agenerate_content, CallTypes.generate_content] + + def test_v1beta_generate_content_route_resolves(self): + call_types = get_call_types_for_route("/v1beta/models/gemini-2.5-flash:generateContent") + assert call_types is not None + assert list(call_types) == [CallTypes.agenerate_content, CallTypes.generate_content] + + def test_bare_stream_generate_content_route_resolves(self): + call_types = get_call_types_for_route("/models/gemini-2.5-flash:streamGenerateContent") + assert call_types is not None + assert list(call_types) == [ + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream, + ] + + def test_v1beta_stream_generate_content_route_resolves(self): + call_types = get_call_types_for_route("/v1beta/models/gemini-2.5-flash:streamGenerateContent") + assert call_types is not None + assert list(call_types) == [ + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream, + ] + + def test_slash_containing_model_name_resolves(self): + call_types = get_call_types_for_route("/v1beta/models/gemini/gemini-2.5-flash:generateContent") + assert call_types is not None + assert list(call_types) == [CallTypes.agenerate_content, CallTypes.generate_content] + + def test_empty_model_name_does_not_match(self): + assert get_call_types_for_route("/models/:generateContent") is None + + def test_unrelated_model_action_does_not_match(self): + assert get_call_types_for_route("/models/gemini-2.5-flash:countTokens") is None + + +class TestPassthroughRouteOrdering: + def test_llm_passthrough_lists_async_call_type_first(self): + call_types = get_call_types_for_route("/llm_passthrough") + assert call_types is not None + assert list(call_types) == [ + CallTypes.allm_passthrough_route, + CallTypes.llm_passthrough_route, + ] + + def test_v1_llm_passthrough_lists_async_call_type_first(self): + call_types = get_call_types_for_route("/v1/llm_passthrough") + assert call_types is not None + assert list(call_types) == [ + CallTypes.allm_passthrough_route, + CallTypes.llm_passthrough_route, + ] + + +class TestFirstCallTypeHasTranslationHandler: + def test_first_call_type_is_translatable_whenever_any_is(self): + """ + Consumers (unified guardrail post-call and streaming resolution) take + call_types[0]. A route whose first call type lacks a guardrail + translation handler while a later one has it silently skips guardrail + scanning, so the table must list a handler-backed call type first. + """ + from litellm.llms import load_guardrail_translation_mappings + + mappings = load_guardrail_translation_mappings() + misordered = { + route: [call_type.value for call_type in call_types] + for route, call_types in API_ROUTE_TO_CALL_TYPES.items() + if call_types + and call_types[0] not in mappings + and any(call_type in mappings for call_type in call_types) + } + assert misordered == {} + + +class TestExistingRouteResolutionUnchanged: + def test_exact_route_still_resolves(self): + call_types = get_call_types_for_route("/chat/completions") + assert call_types is not None + assert CallTypes.acompletion in call_types + + def test_single_segment_placeholder_still_resolves(self): + call_types = get_call_types_for_route("/a2a/my-agent/message/send") + assert call_types is not None + assert list(call_types) == [CallTypes.asend_message, CallTypes.send_message] + + def test_longer_route_does_not_collapse_into_bare_placeholder_pattern(self): + call_types = get_call_types_for_route("/responses/resp_123/input_items") + assert call_types is not None + assert list(call_types) == [CallTypes.alist_input_items] + + def test_unknown_route_returns_none(self): + assert get_call_types_for_route("/not/a/real/route") is None diff --git a/tests/test_litellm/litellm_core_utils/test_classifier_logging.py b/tests/test_litellm/litellm_core_utils/test_classifier_logging.py new file mode 100644 index 00000000000..ccf14ff06c7 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_classifier_logging.py @@ -0,0 +1,51 @@ +from typing import Final + +import pytest + +from litellm.litellm_core_utils.classifier_logging import classifier_input_snapshot, masked_originating_request + + +@pytest.mark.parametrize("encoded", [False, True]) +def test_classifier_snapshot_preserves_provider_shape_and_is_independent(encoded: bool) -> None: + import json + + provider_body: Final = {"system": [{"text": "rubric"}], "messages": [{"role": "user", "content": "ask"}]} + snapshot: Final = classifier_input_snapshot(json.dumps(provider_body) if encoded else provider_body) + assert snapshot == provider_body + provider_body["messages"][0]["content"] = "later mutation" + assert snapshot == {"system": [{"text": "rubric"}], "messages": [{"role": "user", "content": "ask"}]} + + +def test_originating_snapshot_masks_nested_credentials_without_altering_source() -> None: + body: Final = { + "model": "router", + "input": [{"type": "message", "role": "user", "content": "source-only"}], + "api_key": "short", + "metadata": {"nested": [{"Authorization": "Bearer secret", "access_token": 123}]}, + } + snapshot: Final = masked_originating_request({"proxy_server_request": {"body": body}}) + assert snapshot is not None + assert snapshot["model"] == "router" + assert snapshot["input"] == body["input"] + assert snapshot["api_key"] == "REDACTED" + assert snapshot["metadata"] == {"nested": [{"Authorization": "REDACTED", "access_token": "REDACTED"}]} + assert body["api_key"] == "short" + assert body["metadata"]["nested"][0]["Authorization"] == "Bearer secret" + + +@pytest.mark.parametrize("header", ["Cookie", "cookie", "COOKIE", "sEt-CoOkIe"]) +def test_originating_snapshot_redacts_cookie_headers_shared_with_caller_metadata(header: str) -> None: + headers: Final = {header: "session=synthetic-session-credential", "content-type": "application/json"} + body: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {"headers": headers}} + snapshot: Final = masked_originating_request({"proxy_server_request": {"body": body, "headers": headers}}) + assert snapshot == { + "messages": [{"role": "user", "content": "hello"}], + "metadata": {"headers": {header: "REDACTED", "content-type": "application/json"}}, + } + assert headers[header] == "session=synthetic-session-credential" + assert body["metadata"]["headers"][header] == "session=synthetic-session-credential" + + +@pytest.mark.parametrize("value", [None, "not-json", [], {"messages": object()}]) +def test_invalid_provider_payload_is_not_reported_as_captured(value: object) -> None: + assert classifier_input_snapshot(value) is None diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 93cb01e1969..e937be47441 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -1,11 +1,16 @@ """Tests for litellm_core_utils.core_helpers module.""" +import logging + import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, + drop_params_env_flag, + drop_params_flag, get_or_create_metadata_bucket, map_finish_reason, + normalize_drop_params, reconstruct_model_name, redact_nested_match_and_regex_keys, ) @@ -257,6 +262,73 @@ class TestRedactNestedMatchAndRegexKeys: assert redact_nested_match_and_regex_keys("plain") == "plain" +@pytest.mark.parametrize( + "value, expected", + [ + (True, True), + (False, False), + ("true", True), + ("True", True), + (" TRUE ", True), + ("false", False), + ("False", False), + ("yes", True), + ("off", False), + ("1", True), + (1, True), + (0, False), + (None, None), + ("", None), + ("os.environ/DROP_PARAMS", None), + ("v2:gcm:not-a-flag", None), + (2, None), + ], +) +def test_normalize_drop_params(value, expected): + assert normalize_drop_params(value) is expected + + +@pytest.mark.parametrize("value, expected", [("true", True), ("off", False), (None, False)]) +def test_drop_params_flag_returns_a_bool_without_a_warning(value, expected, caplog): + with caplog.at_level(logging.WARNING, logger="drop-params-test"): + assert drop_params_flag(value, "LITELLM_DROP_PARAMS", logging.getLogger("drop-params-test")) is expected + assert caplog.text == "" + + +@pytest.mark.parametrize("value", ["temperature", "ture", 2]) +def test_drop_params_flag_treats_non_flag_values_as_off_with_a_warning(value, caplog): + with caplog.at_level(logging.WARNING, logger="drop-params-test"): + assert drop_params_flag(value, "LITELLM_DROP_PARAMS", logging.getLogger("drop-params-test")) is False + assert f"LITELLM_DROP_PARAMS={value!r} is not a flag value, treating it as off" in caplog.text + + +@pytest.mark.parametrize( + "environ, expected", + [ + ({}, False), + ({"LITELLM_DROP_PARAMS": ""}, False), + ({"LITELLM_DROP_PARAMS": " "}, False), + ({"LITELLM_DROP_PARAMS": "true"}, True), + ({"LITELLM_DROP_PARAMS": " False "}, False), + ({"LITELLM_DROP_PARAMS": "0"}, False), + ], +) +def test_drop_params_env_flag_reads_a_flag_without_a_warning(environ, expected, caplog): + with caplog.at_level(logging.WARNING, logger="drop-params-test"): + assert drop_params_env_flag(environ, logging.getLogger("drop-params-test")) is expected + assert caplog.text == "" + + +@pytest.mark.parametrize("configured", ["temperature", "temperature,top_p", "enabled"]) +def test_drop_params_env_flag_keeps_a_non_flag_value_on_with_a_warning(configured, caplog): + with caplog.at_level(logging.WARNING, logger="drop-params-test"): + assert drop_params_env_flag({"LITELLM_DROP_PARAMS": configured}, logging.getLogger("drop-params-test")) is True + assert ( + f"LITELLM_DROP_PARAMS={configured!r} is not a flag value, treating it as on. Set it to true or false" + in caplog.text + ) + + class TestIsExpectedClientError: def test_status_ranges(self): from litellm.litellm_core_utils.core_helpers import is_expected_client_error diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 1778eca25ef..42d3df76902 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -9,9 +9,11 @@ import litellm from litellm.litellm_core_utils.exception_mapping_utils import ( ExceptionCheckers, _get_body_error_code, + _get_response_headers, exception_type, extract_and_raise_litellm_exception, ) +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.openai.common_utils import OpenAIError from litellm.types.utils import LlmProviders @@ -1254,3 +1256,156 @@ def test_handle_error_marks_only_a_status_code_it_never_received(): raise handler._handle_error(e=upstream, provider_config=None) assert received.value.status_code == 500 assert received.value.status_code_is_synthesized is False + + +def test_bedrock_500_preserves_provider_response_headers(): + """A Bedrock 5xx must keep x-amzn-RequestId so AWS support can trace it (LIT-5428).""" + provider_response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-map-500"}, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=500, + message=provider_response.text, + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-map-500" + + +@pytest.mark.parametrize( + "custom_llm_provider, status_code, provider_message, expected_exception", + [ + ( + "bedrock_mantle", + 400, + ( + '{"error":{"code":"validation_error",' + '"message":"prompt tokens (1055489) exceed model maximum (1050000) for openai.gpt-5.6-sol",' + '"param":null,"type":"invalid_request_error"}}' + ), + litellm.ContextWindowExceededError, + ), + ( + "bedrock", + 400, + '{"message":"Input is too long for requested model."}', + litellm.ContextWindowExceededError, + ), + ( + "bedrock", + 400, + '{"message":"Could not process image"}', + litellm.InternalServerError, + ), + ], +) +def test_bedrock_classified_errors_preserve_provider_response_headers( + custom_llm_provider, status_code, provider_message, expected_exception +): + """Branches that classify a Bedrock error by its text must keep x-amzn-RequestId (LIT-5428).""" + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-classified"}, + text=provider_message, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message=provider_message, + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(expected_exception) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider=custom_llm_provider, + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-classified" + + +@pytest.mark.parametrize( + "status_code, provider_message", + [ + (504, '{"message":"Gateway timeout"}'), + (408, '{"message":"Bedrock did not answer in time"}'), + (408, '{"message":"Connect timeout on endpoint URL"}'), + ], +) +def test_bedrock_timeout_mapping_preserves_provider_headers(status_code, provider_message): + """A mapped bedrock timeout keeps the upstream response, like every other mapped bedrock error. + + The proxy prefixes those headers on the way out, while retry and cooldown + logic still reads the raw retry-after off the response. + """ + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-timeout", "set-cookie": "session=attacker"}, + text=provider_message, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message=provider_message, + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.Timeout) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-timeout" + assert exc_info.value.headers is None + + +@pytest.mark.parametrize("status_code", [504, 408]) +def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code): + """Cooldown and retry timing read retry-after through _get_response_headers.""" + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-retry-after", "retry-after": "7"}, + text='{"message":"Bedrock did not answer in time"}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message='{"message":"Bedrock did not answer in time"}', + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.Timeout) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + exception_headers = _get_response_headers(original_exception=exc_info.value) + assert exception_headers is not None + assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7 diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index b1e8163b91d..b0220a36054 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -571,3 +571,109 @@ def test_shipped_mid_conversation_gate_on_bedrock_ids(shipped_cost_map): ): matched = match_capability_generalizations(unflagged) assert matched is None or not matched.get("supports_mid_conversation_system"), unflagged + + +def test_shipped_rules_flag_unmapped_wandb_ids_as_reasoning(shipped_cost_map): + """W&B ships reasoning models faster than the registry names them, so an unmapped + wandb id resolves as reasoning-capable and its reasoning_effort survives instead of + being dropped. The rule carries no mode and no pricing, so cost stays on the standard + unpriced behavior and the deployment does not read as catalog-mapped.""" + model = "wandb/zai-org/GLM-6-Turbo" + assert model not in litellm.model_cost + + info = litellm.get_model_info(model, custom_llm_provider="wandb") + assert info["litellm_provider"] == "wandb" + assert info["supports_reasoning"] is True + assert info.get("mode") is None + assert not info.get("input_cost_per_token") + assert not info.get("output_cost_per_token") + + assert litellm.supports_reasoning(model="zai-org/GLM-6-Turbo", custom_llm_provider="wandb") is True + + +def test_shipped_wandb_rule_loses_to_mapped_non_reasoning_entries(shipped_cost_map): + """The whole point of a fallback is that it only fills gaps. A wandb model the map + describes as non-reasoning must stay non-reasoning, otherwise the rule silently + re-introduces the blanket supports_reasoning it exists to avoid.""" + for model in ( + "meta-llama/Llama-3.1-8B-Instruct", + "microsoft/Phi-4-mini-instruct", + "moonshotai/Kimi-K2-Instruct", + "Qwen/Qwen3-Coder-480B-A35B-Instruct", + ): + assert f"wandb/{model}" in litellm.model_cost, model + assert litellm.supports_reasoning(model=model, custom_llm_provider="wandb") is False, model + + +def test_shipped_wandb_rule_is_anchored_to_the_wandb_namespace(shipped_cost_map): + """``^wandb/`` is anchored, so it cannot leak onto another provider's ids.""" + assert match_capability_generalizations("wandb/some-new-model") == {"supports_reasoning": True} + for foreign in ("openai/some-new-model", "notwandb/some-new-model", "together_ai/wandb/some-new-model"): + matched = match_capability_generalizations(foreign) + assert matched is None or not matched.get("supports_reasoning"), foreign + + +def test_shipped_wandb_rule_keeps_reasoning_effort_on_an_unmapped_model(shipped_cost_map): + """End to end through the provider config: the gate WandbConfig applies reads the + rule, so reasoning_effort is advertised and survives get_optional_params rather than + raising UnsupportedParamsError.""" + model = "zai-org/GLM-6-Turbo" + assert f"wandb/{model}" not in litellm.model_cost + + supported = litellm.get_supported_openai_params(model=f"wandb/{model}") + assert supported is not None + assert "reasoning_effort" in supported + + optional_params = litellm.utils.get_optional_params( + model=model, + custom_llm_provider="wandb", + reasoning_effort="medium", + drop_params=False, + ) + assert optional_params["reasoning_effort"] == "medium" + + +def test_router_registration_does_not_shadow_shipped_rules(shipped_cost_map): + """Regression: Router writes every configured deployment into ``litellm.model_cost``, + and an exact entry ends the lookup ladder before the rules are consulted. Registering + an unmapped model has to carry the rule defaults forward, or configuring a model on a + proxy silently strips the capabilities the same model resolves to off-proxy.""" + from litellm import Router + + unmapped_wandb = "wandb/zai-org/GLM-6-Turbo" + unmapped_claude = "anthropic/claude-opus-9" + assert unmapped_wandb not in litellm.model_cost + assert unmapped_claude not in litellm.model_cost + + Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": name, "api_key": "fake"}} + for name in (unmapped_wandb, unmapped_claude) + ] + ) + + assert unmapped_wandb in litellm.model_cost + assert unmapped_claude in litellm.model_cost + assert litellm.supports_reasoning(model="zai-org/GLM-6-Turbo", custom_llm_provider="wandb") is True + assert litellm.supports_reasoning(model="claude-opus-9", custom_llm_provider="anthropic") is True + + +def test_deployment_model_info_beats_the_seeded_rule_defaults(shipped_cost_map): + """Seeding a registration from the rules is a floor, not an override: an explicit + model_info on the deployment still wins, so a non-reasoning model can be configured + under a reasoning-first namespace.""" + from litellm import Router + + model = "wandb/some-org/NoThink-1" + Router( + model_list=[ + { + "model_name": model, + "litellm_params": {"model": model, "api_key": "fake"}, + "model_info": {"supports_reasoning": False}, + } + ] + ) + + assert litellm.model_cost[model]["supports_reasoning"] is False + assert litellm.supports_reasoning(model="some-org/NoThink-1", custom_llm_provider="wandb") is False diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index 956da571d43..f026ff57719 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -217,30 +217,9 @@ class TestMetadataFallsBackToLitellmMetadata: assert litellm_metadata == {"trace_id": "trace-1"} -class TestRustOptIn: - """`rust: true` is a litellm param, so it has to reach `litellm_params`. - - `all_litellm_params` keeps it out of the provider body; without it also - being carried into `litellm_params` the chat completions handlers cannot - see the opt-in and the Rust path is silently never taken. - """ - - def test_rust_is_an_optional_kwargs_key(self): - assert "rust" in _OPTIONAL_KWARGS_KEYS - - def test_rust_is_forwarded_from_completion_kwargs(self): - from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS - - assert "rust" in FORWARDED_KWARGS_KEYS - - def test_rust_survives_into_litellm_params(self): - params = get_litellm_params(rust=True) - assert params["rust"] is True - - def test_rust_is_absent_when_the_deployment_did_not_set_it(self): - assert "rust" not in get_litellm_params() - - def test_rust_stays_out_of_the_provider_body(self): - from litellm.types.utils import all_litellm_params - - assert "rust" in all_litellm_params +@pytest.mark.parametrize( + "value, expected", + [("true", True), ("false", False), (" TRUE ", True), (True, True), (None, None), ("os.environ/DROP_PARAMS", None)], +) +def test_drop_params_strings_reach_litellm_params_as_flags(value, expected): + assert get_litellm_params(drop_params=value)["drop_params"] is expected diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 18185126775..c509c8399c9 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -6,6 +6,8 @@ count actual model entries, not reserved meta keys) and the extraction of the import json import os +import sys +import threading import pytest @@ -20,21 +22,29 @@ from litellm.litellm_core_utils.get_model_cost_map import ( GetModelCostMap, _count_model_entries, _finalize_model_cost_map, + get_model_cost_map_provenance, + git_blob_id, ) def _load_root_cost_map() -> dict: - path = os.path.join( - os.path.dirname(__file__), "../../../model_prices_and_context_window.json" - ) + path = os.path.join(os.path.dirname(__file__), "../../../model_prices_and_context_window.json") with open(path) as f: return json.load(f) +def _bundled_blob_id() -> str: + path = os.path.join(os.path.dirname(__file__), "../../../litellm/model_prices_and_context_window_backup.json") + with open(path, "rb") as f: + return git_blob_id(f.read()) + + +def test_git_blob_id_is_what_git_hash_object_prints(): + assert git_blob_id(b'{"gpt-5.4-mini": {"mode": "chat"}}\n') == "18b9a8381e13a3b38a2128f184f631f95829e987" + + def _make_models(n: int) -> dict: - return { - f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n) - } + return {f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n)} def test_count_model_entries_excludes_reserved_keys(): @@ -117,9 +127,7 @@ def test_finalize_pops_key_and_installs_rules(): def test_finalize_with_no_block_clears_rules(): previous = list(get_fallback_generalization_rules()) try: - set_fallback_generalizations( - [{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}] - ) + set_fallback_generalizations([{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}]) _finalize_model_cost_map(_make_models(2)) assert match_capability_generalizations("x-1") is None finally: @@ -298,17 +306,14 @@ def test_openrouter_catalog_costs_match_live_headline_rates(cost_map: dict): assert entry["output_cost_per_token"] != stale_out, model -def test_get_model_cost_map_stamps_loaded_at(monkeypatch): +def test_get_model_cost_map_stamps_loaded_at(): """The load time feeds each pod's reload-due decision; a load that does not stamp it would make manual reload requests race the proxy's startup""" from datetime import datetime, timezone from litellm.litellm_core_utils import get_model_cost_map as module - monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None) - client, _calls = _mock_client( - [httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client - ) + client, _calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) before = datetime.now(timezone.utc) module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client) @@ -317,12 +322,14 @@ def test_get_model_cost_map_stamps_loaded_at(monkeypatch): assert loaded_at is not None assert before <= loaded_at <= datetime.now(timezone.utc) + # --------------------------------------------------------------------------- # refetch_model_cost_map: retry/backoff behavior for runtime reloads # --------------------------------------------------------------------------- import functools import random +from datetime import datetime, timezone import httpx @@ -382,9 +389,7 @@ async def test_refetch_retries_429_honoring_retry_after(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloaded) assert len(result.model_cost_map) > 100 assert calls["count"] == 3 @@ -396,9 +401,7 @@ async def test_refetch_gives_up_after_max_attempts_with_exponential_backoff(): """All 429 without Retry-After: exponential backoff waits, then a failure value.""" client, calls = _mock_client([httpx.Response(429)]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "429" in result.reason assert "after 3 attempts" in result.reason @@ -418,9 +421,7 @@ async def test_refetch_caps_retry_after_wait(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloaded) assert sleeper.waits == [30.0] @@ -435,9 +436,7 @@ async def test_refetch_retries_transport_errors(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloaded) assert calls["count"] == 2 assert len(sleeper.waits) == 1 @@ -448,9 +447,7 @@ async def test_refetch_non_retryable_status_fails_immediately(): """A 404 is permanent: one attempt, no sleeps, failure value.""" client, calls = _mock_client([httpx.Response(404)]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "404" in result.reason assert calls["count"] == 1 @@ -461,9 +458,7 @@ async def test_refetch_non_retryable_status_fails_immediately(): async def test_refetch_invalid_json_fails_immediately(): client, calls = _mock_client([httpx.Response(200, content=b"not json")]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "invalid JSON" in result.reason assert calls["count"] == 1 @@ -475,9 +470,7 @@ async def test_refetch_shrunk_map_fails_integrity_not_swapped_in(): """A drastically shrunk upstream file is rejected instead of being adopted.""" tiny = json.dumps(_make_models(60)).encode() client, _calls = _mock_client([httpx.Response(200, content=tiny)]) - result = await refetch_model_cost_map( - url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "integrity validation" in result.reason @@ -500,6 +493,67 @@ async def test_refetch_respects_local_env_override(monkeypatch): assert len(result.model_cost_map) > 100 +@pytest.mark.asyncio +async def test_refetch_records_the_blob_id_of_the_bytes_served_and_the_fetch_etag(): + body = _real_map_bytes() + client, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"abc123"'}, content=body)]) + + result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) + + assert isinstance(result, ModelCostMapReloaded) + assert result.revision == git_blob_id(body) + assert result.etag == 'W/"abc123"' + assert get_model_cost_map_provenance() == {"source_revision": git_blob_id(body), "etag": 'W/"abc123"'} + + +@pytest.mark.asyncio +async def test_refetch_revision_follows_the_bytes_not_the_url(): + edited = json.loads(_real_map_bytes()) + edited["gpt-5.4-mini"]["input_cost_per_token"] = 0.5 + client, _ = _mock_client( + [httpx.Response(200, content=_real_map_bytes()), httpx.Response(200, content=json.dumps(edited).encode())] + ) + + first = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) + second = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) + + assert isinstance(first, ModelCostMapReloaded) and isinstance(second, ModelCostMapReloaded) + assert first.revision != second.revision + assert get_model_cost_map_provenance()["source_revision"] == second.revision + + +@pytest.mark.asyncio +async def test_refetch_local_override_reports_the_bundled_blob_id_without_an_etag(monkeypatch): + remote, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"remote"'}, content=_real_map_bytes())]) + await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=remote) + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0)) + + assert isinstance(result, ModelCostMapReloaded) + assert result.revision == _bundled_blob_id() + assert get_model_cost_map_provenance() == {"source_revision": _bundled_blob_id(), "etag": None} + + +@pytest.mark.asyncio +async def test_refetch_stamps_loaded_at_on_remote_and_local_reloads(monkeypatch): + from litellm.litellm_core_utils import get_model_cost_map as module + + client, _ = _mock_client([httpx.Response(200, content=_real_map_bytes())]) + before_remote = datetime.now(timezone.utc) + await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) + remote_loaded_at = module.get_model_cost_map_loaded_at() + assert remote_loaded_at is not None + assert before_remote <= remote_loaded_at <= datetime.now(timezone.utc) + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + before_local = datetime.now(timezone.utc) + await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0)) + local_loaded_at = module.get_model_cost_map_loaded_at() + assert local_loaded_at is not None + assert before_local <= local_loaded_at <= datetime.now(timezone.utc) + + # --------------------------------------------------------------------------- # get_model_cost_map: the boot-time load retries transient failures like a reload does # --------------------------------------------------------------------------- @@ -513,69 +567,125 @@ from litellm.litellm_core_utils.get_model_cost_map import ( class _SyncSleepRecorder: """Injected in place of time.sleep so the boot path's waits are asserted without delay.""" - def __init__(self): + def __init__(self, block=False): self.waits = [] + self.block = block + self.started = threading.Event() + self.release = threading.Event() def __call__(self, seconds: float) -> None: + if self.block: + self.started.set() + self.release.wait(timeout=10) self.waits.append(seconds) -def test_boot_load_retries_transient_failures_instead_of_falling_back(): - """A refused connection then a 503 at pod boot used to pin the process to the bundled - backup for its lifetime; both are transient and must be retried before giving up.""" +def _retry_threads(): + return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-retry"] + + +def test_boot_load_success_does_not_start_background_retry(): + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + + cost_map = get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + ) + assert calls["count"] == 1 + assert sleeper.waits == [] + assert _retry_threads() == [] + assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + assert get_model_cost_map_source_info()["source"] == "remote" + + +def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(monkeypatch): + import litellm + from litellm import utils as litellm_utils + from litellm.litellm_core_utils import get_model_cost_map as module + + original_model_cost = litellm.model_cost + monkeypatch.setattr(litellm, "model_cost", dict(original_model_cost)) + for name, provider_models in tuple(vars(litellm).items()): + if name.endswith("_models") and isinstance(provider_models, set): + monkeypatch.setattr(litellm, name, set(provider_models)) + monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider)) + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + source_info = module._cost_map_source_info + for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): + monkeypatch.setattr(source_info, name, getattr(source_info, name)) + + remote_map = _load_root_cost_map() + remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} client, calls = _mock_client( [ httpx.ConnectError("connection refused"), - httpx.Response(503), - httpx.Response(200, content=_real_map_bytes()), + httpx.Response(200, content=json.dumps(remote_map).encode()), ], client_cls=httpx.Client, ) - sleeper = _SyncSleepRecorder() + sleeper = _SyncSleepRecorder(block=True) + litellm.register_model({"my-runtime-model": {"litellm_provider": "custom", "max_input_tokens": 4321}}) - cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - - assert calls["count"] == 3 - assert len(sleeper.waits) == 2 - assert 2.0 <= sleeper.waits[0] < 3.0 - assert 4.0 <= sleeper.waits[1] < 5.0 - source = get_model_cost_map_source_info() - assert source["source"] == "remote" - assert source["fallback_reason"] is None - assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} - - -def test_boot_load_honors_retry_after_then_falls_back_after_max_attempts(): - """An outage longer than the retry budget still ends on the bundled backup, and the - recorded fallback reason says how many attempts were spent so operators can tell.""" - client, calls = _mock_client( - [httpx.Response(429, headers={"Retry-After": "7"})], client_cls=httpx.Client + cost_map = get_model_cost_map( + url=_URL, + max_attempts=3, + sleep=sleeper, + rng=random.Random(0), + client=client, ) - sleeper = _SyncSleepRecorder() - cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - - assert calls["count"] == 3 - assert sleeper.waits == [7.0, 7.0] - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert "after 3 attempts" in source["fallback_reason"] - assert len(cost_map) > 100 + assert calls["count"] == 1 + assert sleeper.waits == [] + assert sleeper.started.wait(timeout=10) + threads = _retry_threads() + try: + assert len(threads) == 1 + assert "claude-remote-only-test" not in cost_map + assert cost_map.keys() == _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()).keys() + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"].startswith("Remote fetch failed:") + sleeper.release.set() + for thread in threads: + thread.join(timeout=10) + assert all(not thread.is_alive() for thread in threads) + assert sleeper.waits and 2.0 <= sleeper.waits[0] < 3.0 + assert calls["count"] == 2 + assert "claude-remote-only-test" in litellm.model_cost + assert "claude-remote-only-test" in litellm.anthropic_models + assert "my-runtime-model" in litellm.model_cost + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None + finally: + sleeper.release.set() + for thread in _retry_threads(): + thread.join(timeout=10) -def test_boot_load_does_not_retry_permanent_failures(): - """A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup.""" +def test_boot_load_does_not_retry_non_retryable_failure(): client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) sleeper = _SyncSleepRecorder() - get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + ) assert calls["count"] == 1 assert sleeper.waits == [] - assert get_model_cost_map_source_info()["source"] == "local" - - get_model_cost_map(url="not a url", sleep=sleeper, rng=random.Random(0)) - assert sleeper.waits == [] - assert get_model_cost_map_source_info()["source"] == "local" + assert _retry_threads() == [] + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] is not None def test_boot_load_respects_local_env_override(monkeypatch): @@ -592,3 +702,77 @@ def test_boot_load_respects_local_env_override(monkeypatch): ) assert len(cost_map) > 100 assert get_model_cost_map_source_info()["is_env_forced"] is True + + +def test_boot_load_records_the_blob_id_of_the_bytes_served_and_the_fetch_etag(): + body = _real_map_bytes() + client, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=body)], client_cls=httpx.Client) + + get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=client) + + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["etag"] == 'W/"boot"' + assert source["source_revision"] == git_blob_id(body) + assert source["loaded_at"] is not None + + +def test_boot_load_fallback_to_the_backup_reports_its_blob_id_and_drops_the_remote_etag(): + remote, _ = _mock_client( + [httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=_real_map_bytes())], client_cls=httpx.Client + ) + get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=remote) + failing, _ = _mock_client([httpx.Response(404)], client_cls=httpx.Client) + + get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=failing) + + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["etag"] is None + assert source["source_revision"] == _bundled_blob_id() + + +def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rejected_fetch(): + remote, _ = _mock_client( + [httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=_real_map_bytes())], client_cls=httpx.Client + ) + get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=remote) + shrunk_body = b'{"gpt-5.4-mini": {"mode": "chat", "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}}' + shrunk, _ = _mock_client( + [httpx.Response(200, headers={"ETag": 'W/"shrunk"'}, content=shrunk_body)], client_cls=httpx.Client + ) + + get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=shrunk) + + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] == "Remote data failed integrity validation" + assert source["etag"] is None + assert source["source_revision"] == _bundled_blob_id() + assert source["source_revision"] != git_blob_id(shrunk_body) + + +@pytest.mark.parametrize( + ("argv0", "request_count"), + [ + ("/some/venv/bin/lite", 0), + ("/some/venv/bin/lite.exe", 0), + ("/some/venv/bin/python", 1), + ], +) +def test_boot_load_skips_remote_fetch_for_cli_processes( + monkeypatch: pytest.MonkeyPatch, argv0: str, request_count: int +) -> None: + monkeypatch.setattr(sys, "argv", [argv0, "--version"]) + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + + cost_map = get_model_cost_map(url=_URL, client=client) + + assert calls["count"] == request_count + assert cost_map + source = get_model_cost_map_source_info() + if request_count == 0: + assert source["source"] == "local" + else: + assert source["source"] == "remote" diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index 893472d63ae..8fa4bd6c14d 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -1,3 +1,7 @@ +import asyncio +import copy +import time +import uuid from unittest.mock import patch import pytest @@ -7,9 +11,13 @@ import litellm from litellm import constants from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.image_handling import ( + MAX_CONCURRENT_REMOTE_MEDIA_FETCHES, + RemoteMedia, async_convert_url_to_base64, + async_inline_remote_media, convert_url_to_base64, ) +from litellm.litellm_core_utils.url_utils import SSRFError @pytest.fixture(autouse=True) @@ -107,9 +115,7 @@ class StreamingLargeImageClient: request=Request("GET", url), ) # Mock the iter_bytes method to return our generator - response.iter_bytes = lambda chunk_size=8192: generate_chunks( - size_bytes, chunk_size - ) + response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) return response @@ -207,9 +213,7 @@ def test_streaming_download_handles_petabyte_file(monkeypatch): """ # Simulate a 1 petabyte file (1,000,000 GB) # Without streaming protection, this would cause OOM or hang indefinitely - client = StreamingLargeImageClient( - size_mb=1_000_000_000, include_content_length=False - ) + client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) monkeypatch.setattr(litellm, "module_level_client", client) with pytest.raises(litellm.ImageFetchError) as excinfo: @@ -268,3 +272,259 @@ def test_image_size_limit_disabled(monkeypatch): assert "Image URL download is disabled" in str(excinfo.value) assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value) + + +async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + messages = [ + {"role": "system", "content": "be terse"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + {"type": "image_url", "image_url": {"url": image_url, "detail": "low"}}, + {"type": "image_url", "image_url": image_url}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + {"type": "file", "file": {"file_id": pdf_url}}, + {"type": "file", "file": {"file_id": image_url, "format": "image/png"}}, + {"type": "document", "source": {"type": "url", "url": pdf_url}, "title": "the doc"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + {"type": "document", "source": {"type": "file", "file_id": "file_abc"}}, + ], + }, + ] + snapshot = copy.deepcopy(messages) + + inlined = await async_inline_remote_media(messages) + + data_url = async_only_image_fetch.data_url + base64_png = async_only_image_fetch.base64_png + assert inlined[0] == {"role": "system", "content": "be terse"} + assert inlined[1]["content"] == [ + {"type": "text", "text": "what is this?"}, + {"type": "image_url", "image_url": {"url": data_url, "detail": "low"}}, + {"type": "image_url", "image_url": data_url}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + {"type": "file", "file": {"format": "application/pdf", "file_data": data_url}}, + {"type": "file", "file": {"format": "image/png", "file_data": data_url}}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": base64_png}, + "title": "the doc", + }, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64_png}}, + {"type": "document", "source": {"type": "file", "file_id": "file_abc"}}, + ] + assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url]) + assert messages == snapshot + + +async def test_async_inline_remote_media_inlines_only_the_parts_the_predicate_accepts(async_only_image_fetch): + files_api_prefix = "https://generativelanguage.googleapis.com/v1beta/files/" + files_api_pdf = f"{files_api_prefix}{uuid.uuid4().hex}" + hinted_image = f"https://img.example/{uuid.uuid4()}.png" + plain_image = f"https://img.example/{uuid.uuid4()}.png" + hinted_document = f"https://docs.example/{uuid.uuid4()}.pdf" + seen = [] + + def inline_unhinted_outside_files_api(media: RemoteMedia) -> bool: + seen.append(media) + return not media.url.startswith(files_api_prefix) and "format" not in media.fields + + messages = [ + { + "role": "user", + "content": [ + {"type": "file", "file": {"file_id": files_api_pdf}}, + {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": plain_image}}, + {"type": "image_url", "image_url": plain_image}, + {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}}, + ], + } + ] + snapshot = copy.deepcopy(messages) + + inlined = await async_inline_remote_media(messages, should_inline=inline_unhinted_outside_files_api) + + assert inlined[0]["content"] == [ + {"type": "file", "file": {"file_id": files_api_pdf}}, + {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}}, + {"type": "image_url", "image_url": async_only_image_fetch.data_url}, + {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}}, + ] + assert async_only_image_fetch.fetched == [plain_image] + assert [(media.url, dict(media.fields)) for media in seen[:5]] == [ + (files_api_pdf, {"file_id": files_api_pdf}), + (hinted_image, {"url": hinted_image, "format": "image/png"}), + (plain_image, {"url": plain_image}), + (plain_image, {}), + (hinted_document, {"type": "url", "url": hinted_document, "format": "application/pdf"}), + ] + assert messages == snapshot + + +async def test_async_inline_remote_media_inlines_a_shared_url_only_where_the_predicate_accepts_it( + async_only_image_fetch, +): + shared = f"https://img.example/{uuid.uuid4()}.png" + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": shared}}, + ], + } + ] + + inlined = await async_inline_remote_media(messages, should_inline=lambda media: "format" not in media.fields) + + assert inlined[0]["content"] == [ + {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}}, + ] + assert async_only_image_fetch.fetched == [shared] + + +async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fails(monkeypatch): + missing = f"http://img.example/{uuid.uuid4()}-missing.png" + slow = f"http://img.example/{uuid.uuid4()}-slow.png" + slow_fetch_outcomes = [] + + async def serve(client, url, **kwargs): + if url == missing: + return Response(404, request=Request("GET", url)) + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + slow_fetch_outcomes.append("cancelled") + raise + slow_fetch_outcomes.append("finished") + return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve) + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": missing}}, + {"type": "image_url", "image_url": {"url": slow}}, + ], + } + ] + started = time.perf_counter() + + with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): + await async_inline_remote_media(messages) + + assert slow_fetch_outcomes == ["cancelled"] + assert time.perf_counter() - started < 1 + + +_SSRF_VERDICTS = ( + SSRFError( + "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " + "add the host to `user_url_allowed_hosts` in general_settings." + ), + SSRFError("DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known"), + SSRFError("No addresses found for 'internal.example'"), +) + + +def _assert_verdict_free_messages(messages, url): + assert len(messages) == len(_SSRF_VERDICTS) + assert len(set(messages)) == 1, "a caller must not be able to tell a blocked host from one that does not resolve" + message = messages[0] + assert "The proxy could not resolve this host or its URL policy rejected it" in message + assert "user_url_allowed_hosts" in message + assert url in message + assert "10.0.0.8" not in message + assert "DNS" not in message + assert "No addresses" not in message + + +async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): + attempts = [] + messages = [] + url = f"http://internal.example/{uuid.uuid4()}.png" + + for verdict in _SSRF_VERDICTS: + + async def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise verdict + + monkeypatch.setattr(image_handling, "async_safe_get", block) + with pytest.raises(litellm.ImageFetchError) as raised: + await async_convert_url_to_base64(url) + messages.append(raised.value.message) + + assert attempts == [url] * len(_SSRF_VERDICTS) + _assert_verdict_free_messages(messages, url) + + +def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): + attempts = [] + messages = [] + url = f"http://internal.example/{uuid.uuid4()}.png" + + for verdict in _SSRF_VERDICTS: + + def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise verdict + + monkeypatch.setattr(image_handling, "safe_get", block) + with pytest.raises(litellm.ImageFetchError) as raised: + convert_url_to_base64(url) + messages.append(raised.value.message) + + assert attempts == [url] * len(_SSRF_VERDICTS) + _assert_verdict_free_messages(messages, url) + + +async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): + in_flight = {"now": 0, "peak": 0} + + async def serve_png_slowly(client, url, **kwargs): + in_flight["now"] += 1 + in_flight["peak"] = max(in_flight["peak"], in_flight["now"]) + await asyncio.sleep(0.01) + in_flight["now"] -= 1 + return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve_png_slowly) + urls = [f"https://img.example/{uuid.uuid4()}.png" for _ in range(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + 5)] + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}} for url in urls]}] + + inlined = await async_inline_remote_media(messages) + + assert in_flight["peak"] == MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + assert all(part["image_url"]["url"].startswith("data:image/png;base64,") for part in inlined[0]["content"]) + + +async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch): + messages = [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}], + }, + ] + + assert await async_inline_remote_media(messages) is messages + assert async_only_image_fetch.fetched == [] + + +async def test_async_inline_remote_media_raises_image_fetch_error_when_the_fetch_fails(monkeypatch): + async def serve_404(client, url, **kwargs): + return Response(404, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve_404) + url = f"http://img.example/{uuid.uuid4()}.png" + + with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): + await async_inline_remote_media([{"role": "user", "content": [{"type": "image_url", "image_url": url}]}]) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0f44cbbb2a7..2f6339dcdbb 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,7 +1,10 @@ +import asyncio import contextlib +import datetime import os import sys -import asyncio +from collections.abc import Callable +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -20,7 +23,13 @@ from litellm.litellm_core_utils.litellm_logging import ( _get_status_fields, set_callbacks, ) -from litellm.types.utils import ModelResponse, TextCompletionResponse +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.utils import ( + CallTypes, + LiteLLMRealtimeStreamLoggingObject, + ModelResponse, + TextCompletionResponse, +) @pytest.fixture @@ -632,6 +641,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.""" @@ -4391,6 +4480,39 @@ def test_handle_anthropic_messages_response_logging_translates_bare_responses_ap assert result.usage.total_tokens == 18 # type: ignore[attr-defined] +def test_handle_anthropic_messages_response_logging_keeps_the_served_response_id(): + from openai.types.responses import ResponseOutputMessage, ResponseOutputText + + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + served_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-1", response_id="resp_upstream" + ) + logging_obj = _anthropic_messages_logging_obj() + result = logging_obj._handle_anthropic_messages_response_logging( + result=ResponsesAPIResponse( + id=served_id, + created_at=1700000000, + output=[ + ResponseOutputMessage( + id="msg-1", + type="message", + role="assistant", + status="completed", + content=[ResponseOutputText(annotations=[], text="hi", type="output_text")], + ) + ], + usage=ResponseAPIUsage(input_tokens=2, output_tokens=1, total_tokens=3), + service_tier="flex", + ) + ) + + assert isinstance(result, ModelResponse) + assert result.id == served_id, "the spend log row must keep the id the caller was served" + assert result.service_tier == "flex" + + def test_handle_anthropic_messages_response_logging_passes_model_response_through(): """Anthropic-native path already yields a ModelResponse; it must be returned unchanged.""" logging_obj = _anthropic_messages_logging_obj() @@ -5953,15 +6075,17 @@ def test_prompt_hooks_skip_prompt_managers_when_no_prompt_id(logging_obj, tmp_pa for hook in [cb for cb in litellm.callbacks if isinstance(cb, VectorStorePreCallHook)]: litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, hook) def test_newrelic_dispatch_prefers_otel_v2_when_flag_on(monkeypatch): - """With LITELLM_OTEL_V2 on, the "newrelic" callback builds the OTel v2 - logger (per-team credential routing); with the flag off (default) it keeps - the legacy agent-based logger, so existing deployments are untouched.""" + """With LITELLM_OTEL_V2 on and operator credentials present, the "newrelic" + callback builds the OTel v2 logger (per-team credential routing); with the + flag off (default) it keeps the legacy agent-based logger, so existing + deployments are untouched.""" from litellm.integrations.otel.logger import OpenTelemetryV2 from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.litellm_core_utils import litellm_logging as logging_module logging_module._in_memory_loggers.clear() monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("NEW_RELIC_LICENSE_KEY", "test-license-key") is_otel_v2_enabled.cache_clear() try: v2_logger = logging_module._init_custom_logger_compatible_class( @@ -6016,6 +6140,7 @@ def test_get_custom_logger_compatible_class_finds_v2_newrelic(monkeypatch): logging_module._in_memory_loggers.clear() monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("NEW_RELIC_LICENSE_KEY", "test-license-key") is_otel_v2_enabled.cache_clear() try: created = logging_module._init_custom_logger_compatible_class( @@ -6278,6 +6403,90 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == "" +def _responses_ws_logging_obj() -> LitellmLogging: + return LitellmLogging( + model="gpt-4o", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="responses-ws-usage-test", + function_id="responses-ws-usage-test", + ) + + +def test_normalize_logging_result_extracts_usage_for_responses_websocket(monkeypatch): + """LIT-6512: native /v1/responses WebSocket sessions logged $0 spend because the usage + carried by stored response.completed events was never extracted. The session must cost + exactly what the same usage costs over HTTP /v1/responses, discounts included.""" + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.5}) + logging_obj = _responses_ws_logging_obj() + events = [ + {"type": "response.created", "response": {}}, + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140}}, + }, + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 60, "output_tokens": 10, "total_tokens": 70}}, + }, + ] + + normalized = logging_obj.normalize_logging_result(result=events) + + assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject) + assert normalized.usage.prompt_tokens == 160 + assert normalized.usage.completion_tokens == 50 + + ws_cost = litellm.completion_cost( + completion_response=normalized, + model="gpt-4o", + call_type=CallTypes.aresponses_websocket.value, + custom_llm_provider="openai", + ) + http_cost = litellm.completion_cost( + completion_response=ResponsesAPIResponse( + id="resp-6512", + created_at=1700000000, + output=[], + usage=ResponseAPIUsage(input_tokens=160, output_tokens=50, total_tokens=210), + ), + model="gpt-4o", + call_type=CallTypes.aresponses.value, + custom_llm_provider="openai", + ) + assert ws_cost > 0 + assert ws_cost == http_cost + + +def test_normalize_logging_result_bills_incomplete_responses_websocket_turns(): + """LIT-6512: a turn cut short by max_output_tokens ends in response.incomplete, which + OpenAI bills, so its usage counts toward the session like a completed turn.""" + events = [ + { + "type": "response.created", + "response": {"usage": {"input_tokens": 999, "output_tokens": 999, "total_tokens": 1998}}, + }, + { + "type": "response.incomplete", + "response": {"usage": {"input_tokens": 15, "output_tokens": 16, "total_tokens": 31}}, + }, + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 40, "output_tokens": 4, "total_tokens": 44}}, + }, + {"type": "response.failed", "response": {"usage": None}}, + ] + + normalized = _responses_ws_logging_obj().normalize_logging_result(result=events) + + assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject) + assert normalized.usage.prompt_tokens == 55 + assert normalized.usage.completion_tokens == 20 + assert normalized.usage.total_tokens == 75 + + def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj): """LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs).""" @@ -6326,6 +6535,113 @@ def test_get_standard_logging_object_payload_survives_logging_obj_without_timing assert payload["hidden_params"]["litellm_overhead_time_ms"] is None +@pytest.mark.parametrize( + ("header_name", "header_source"), + ( + ("x-amzn-RequestId", "response"), + ("x-request-id", "response"), + ("request-id", "response"), + ("x-ms-request-id", "response"), + ("apim-request-id", "response"), + ("x-goog-request-id", "response"), + ("cf-ray", "response"), + ("X-Request-Id", "litellm_response_headers"), + ("X-MS-Request-ID", "headers"), + ), +) +def test_failure_standard_logging_payload_captures_provider_request_id( + logging_obj: LitellmLogging, + header_name: str, + header_source: Literal["response", "litellm_response_headers", "headers"], +): + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + request_id = "provider-request-123" + request = httpx.Request("POST", "https://provider.example/v1/chat/completions") + response = httpx.Response(429, headers={header_name: request_id}, request=request) + provider_error = httpx.HTTPStatusError("provider error", request=request, response=response) + if header_source == "litellm_response_headers": + response.headers.clear() + provider_error.litellm_response_headers = {header_name: request_id} + elif header_source == "headers": + response.headers.clear() + provider_error.headers = {header_name: request_id} + now = datetime.datetime.now() + + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "call-1", "model": "test-model", "messages": []}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="failure", + original_exception=provider_error, + ) + + assert payload is not None + assert payload["error_information"] is not None + assert payload["error_information"]["error_provider_request_id"] == request_id + + +def test_get_error_information_ignores_unsupported_headers() -> None: + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + request = httpx.Request("POST", "https://provider.example/v1/chat/completions") + response = httpx.Response(429, headers={"retry-after": "3"}, request=request) + provider_error = httpx.HTTPStatusError("provider error", request=request, response=response) + + error_information = StandardLoggingPayloadSetup.get_error_information(provider_error) + + assert error_information["error_provider_request_id"] is None + + +def test_get_error_information_uses_header_precedence_and_fallback() -> None: + from litellm.exceptions import RateLimitError + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + request = httpx.Request("POST", "https://provider.example/v1/chat/completions") + response = httpx.Response( + 429, + headers={"x-request-id": "response-id", "x-amzn-requestid": "amazon-id"}, + request=request, + ) + provider_error = RateLimitError( + message="provider error", + llm_provider="test-provider", + model="test-model", + response=response, + headers={"retry-after": "3"}, + ) + + error_information = StandardLoggingPayloadSetup.get_error_information(provider_error) + + assert error_information["error_provider_request_id"] == "amazon-id" + + +def test_get_error_information_ignores_malformed_headers() -> None: + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + provider_error = Exception("provider error") + provider_error.headers = [("x-request-id", "provider-request-123")] + + error_information = StandardLoggingPayloadSetup.get_error_information(provider_error) + + assert error_information["error_provider_request_id"] is None + + +def test_get_provider_request_id_ignores_header_lookup_errors() -> None: + from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id + + class HeaderLookupError(Exception): + @property + def response(self) -> object: + raise RuntimeError("headers unavailable") + + assert _get_provider_request_id(HeaderLookupError("provider error")) is None + + def test_get_standard_logging_object_payload_failure_status_keeps_overhead_none(logging_obj): """A post_call guardrail can fail the request after the upstream call succeeded; the failure payload keeps litellm_overhead_time_ms None, matching responses that carry their own _hidden_params.""" @@ -6486,3 +6802,205 @@ def test_get_error_information_redacts_provider_key_from_upstream_url(): assert "REDACTED" in result["traceback"] assert "REDACTED" in result["error_message"] assert result["error_code"] == "400" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ["openai", "azure", "anthropic", "bedrock", "responses"]) +async def test_classifier_audit_matches_provider_transport(provider: str) -> None: + import json + + from openai import AsyncAzureOpenAI, AsyncOpenAI + + from litellm.litellm_core_utils.classifier_logging import classifier_input_snapshot + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + outbound: Final = asyncio.Queue() + logs: Final = asyncio.Queue() + + def respond(request: httpx.Request) -> httpx.Response: + outbound.put_nowait(json.loads(request.content)) + content: Final = '{"tier":"SIMPLE"}' + if provider == "responses": + from litellm.responses.main import mock_responses_api_response + + return httpx.Response(200, json=mock_responses_api_response(content).model_dump()) + if provider == "anthropic": + return httpx.Response(200, json={ + "id": "msg-audit", "type": "message", "role": "assistant", "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": content}], "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }) + if provider == "bedrock": + return httpx.Response(200, json={ + "output": {"message": {"role": "assistant", "content": [{"text": content}]}}, + "stopReason": "end_turn", "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + "metrics": {"latencyMs": 1}, + }) + return httpx.Response(200, json={ + "id": "chatcmpl-audit", "object": "chat.completion", "created": 0, "model": "gpt-5.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }) + + async def capture(kwargs, response_obj, start_time, end_time): + logs.put_nowait(kwargs["standard_logging_object"]) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + handler: Final = AsyncHTTPHandler() + await handler.close() + handler.client = http_client + client: Final = ( + AsyncAzureOpenAI( + api_key="transport-only", azure_endpoint="https://azure.invalid", + api_version="2025-04-01-preview", http_client=http_client, + ) + if provider == "azure" else AsyncOpenAI(api_key="transport-only", http_client=http_client) + if provider == "openai" else handler + ) + model: Final = { + "openai": "openai/gpt-5.6", + "azure": "azure/gpt-5.6", + "anthropic": "anthropic/claude-haiku-4-5", + "bedrock": "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0", + "responses": "openai/gpt-5.6", + }[provider] + + async def run(marker: str) -> None: + if provider == "responses": + await litellm.aresponses( + model=model, api_key="transport-only", client=client, max_output_tokens=128, + instructions="classifier-rubric", input=marker, + metadata={"internal_call_origin": "autorouter_classifier"}, + proxy_server_request={"body": {}, "originating_request_masked": {"input": f"source-only-{marker}"}}, + success_callback=[capture], num_retries=0, + ) + return + await litellm.acompletion( + model=model, api_key="transport-only", client=client, max_tokens=128, + aws_access_key_id="transport-only", aws_secret_access_key="transport-only", aws_region_name="us-east-1", + messages=[{"role": "system", "content": "classifier-rubric"}, {"role": "user", "content": marker}], + metadata={"internal_call_origin": "autorouter_classifier"}, + proxy_server_request={"body": {}, "originating_request_masked": {"input": f"source-only-{marker}"}}, + success_callback=[capture], num_retries=0, + **({"api_base": "https://azure.invalid", "api_version": "2025-04-01-preview"} if provider == "azure" else {}), + **({"extra_body": {"audit_context": "provider-extra"}, "extra_headers": {"X-Audit": "header-only-secret"}} + if provider in ("openai", "azure") else {}), + ) + + await asyncio.gather(run("request-one"), run("request-two")) + requests: Final = await asyncio.wait_for(asyncio.gather(outbound.get(), outbound.get()), timeout=10) + payloads: Final = await asyncio.wait_for(asyncio.gather(logs.get(), logs.get()), timeout=10) + for payload in payloads: + snapshot: Final = payload["classifier_input"] + assert snapshot in requests + assert "source-only" not in json.dumps(snapshot) + assert "classifier-rubric" in json.dumps(snapshot) + assert "transport-only" not in json.dumps(snapshot) + assert "header-only-secret" not in json.dumps(snapshot) + assert "SIMPLE" in json.dumps(payload["response"]) + marker: Final = "request-one" if "request-one" in json.dumps(snapshot) else "request-two" + assert payload["originating_request_masked"] == {"input": f"source-only-{marker}"} + assert classifier_input_snapshot(snapshot) is not None + if provider not in ("openai", "azure", "responses"): + assert all("system" in request for request in requests) + + +@pytest.mark.parametrize("redaction", ["none", "global", "request", "header"]) +@pytest.mark.parametrize("status", ["success", "failure"]) +@pytest.mark.parametrize("call_type", ["completion", "acompletion", "responses", "aresponses"]) +def test_classifier_audit_obeys_message_logging_before_payload_emission(logging_obj, monkeypatch, redaction, status, call_type): + from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload + + monkeypatch.setattr(litellm, "turn_off_message_logging", redaction == "global") + params: Final = { + "metadata": {"internal_call_origin": "autorouter_classifier", **( + {"headers": {"x-litellm-enable-message-redaction": "true"}} if redaction == "header" else {} + )}, + "proxy_server_request": {"body": {}, "originating_request_masked": {"input": "source-only"}}, + } + logging_obj.call_type = call_type + logging_obj.model_call_details["litellm_params"] = params + logging_obj.model_call_details["standard_callback_dynamic_params"] = ( + {"turn_off_message_logging": True} if redaction == "request" else {} + ) + logging_obj.pre_call( + input=[], api_key=None, additional_args={"complete_input_dict": {"system": "rubric", "messages": []}} + ) + now: Final = datetime.datetime.now() + payload: Final = get_standard_logging_object_payload( + kwargs={**logging_obj.model_call_details, "call_type": call_type}, init_response_obj={}, + start_time=now, end_time=now, logging_obj=logging_obj, status=status, + ) + assert payload is not None + if redaction == "none": + assert payload["classifier_input"] == {"system": "rubric", "messages": []} + assert payload["originating_request_masked"] == {"input": "source-only"} + else: + assert "classifier_input" not in payload + assert "originating_request_masked" not in payload + + +@pytest.mark.parametrize("call_type,origin", [("completion", None), ("aembedding", "autorouter_classifier")]) +def test_classifier_audit_is_not_added_to_other_calls(logging_obj, call_type, origin): + logging_obj.call_type = call_type + logging_obj.model_call_details["litellm_params"] = {"metadata": {"internal_call_origin": origin}} + logging_obj.pre_call(input=[], api_key=None, additional_args={"complete_input_dict": {"input": "embedding"}}) + assert logging_obj.classifier_input is None + + +def _run_while_a_thread_grows(target: dict, read: Callable[[], None], reads: int) -> None: + import itertools + import threading + + stop: Final = threading.Event() + + def grow() -> None: + for counter in itertools.count(): + if stop.is_set(): + return + key: Final = f"late_{counter % 64}" + if key in target: + del target[key] + else: + target[key] = counter + + writer: Final = threading.Thread(target=grow, daemon=True) + previous_interval: Final = sys.getswitchinterval() + sys.setswitchinterval(1e-6) + writer.start() + try: + for _ in range(reads): + read() + finally: + stop.set() + writer.join(timeout=5) + sys.setswitchinterval(previous_interval) + + +def test_merge_litellm_metadata_survives_a_thread_growing_metadata_mid_merge(): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + metadata: Final = {f"key_{i}": i for i in range(2000)} + litellm_params: Final = {"metadata": metadata, "litellm_metadata": {"model_group": "gpt"}} + + def read() -> None: + merged: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) + assert merged["key_1999"] == 1999 + assert merged["model_group"] == "gpt" + + _run_while_a_thread_grows(metadata, read, reads=300) + + +def test_get_additional_headers_survives_a_thread_growing_headers_mid_copy(): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + headers: Final = {f"llm_provider-x-custom-{i}": str(i) for i in range(2000)} + headers["x-ratelimit-remaining-requests"] = "7" + + def read() -> None: + copied: Final = StandardLoggingPayloadSetup.get_additional_headers(headers) + assert copied is not None + assert copied["x_ratelimit_remaining_requests"] == 7 + assert copied["llm_provider-x-custom-1999"] == "1999" + + _run_while_a_thread_grows(headers, read, reads=300) diff --git a/tests/test_litellm/litellm_core_utils/test_private_json.py b/tests/test_litellm/litellm_core_utils/test_private_json.py index cedff61959f..3c9f49607f9 100644 --- a/tests/test_litellm/litellm_core_utils/test_private_json.py +++ b/tests/test_litellm/litellm_core_utils/test_private_json.py @@ -4,7 +4,11 @@ import stat import pytest -from litellm.litellm_core_utils.private_json import overwrite_private_json, write_private_json +from litellm.litellm_core_utils.private_json import ( + overwrite_private_json, + write_private_bytes, + write_private_json, +) class TestOverwritePrivateJson: @@ -35,3 +39,32 @@ class TestOverwritePrivateJson: overwrite_private_json(str(path), {"user_id": "u-1"}) assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + +class TestWritePrivateBytes: + def test_replaces_the_file_in_one_step_so_a_reader_holding_the_old_one_keeps_it_whole(self, tmp_path): + path = tmp_path / "script.py" + write_private_bytes(str(path), b"print('one')\n" * 200) + before = path.stat().st_ino + + with path.open("rb") as reader: + write_private_bytes(str(path), b"print('two')\n") + assert reader.read() == b"print('one')\n" * 200 + + assert path.read_bytes() == b"print('two')\n" + assert path.stat().st_ino != before + assert [child.name for child in tmp_path.iterdir()] == ["script.py"] + + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") + def test_lands_owner_only_and_a_refused_stage_leaves_the_previous_file_untouched(self, tmp_path): + path = tmp_path / "script.py" + write_private_bytes(str(path), b"first") + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + tmp_path.chmod(0o500) + try: + with pytest.raises(PermissionError): + write_private_bytes(str(path), b"second") + finally: + tmp_path.chmod(0o700) + assert path.read_bytes() == b"first" diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 3be0bae4120..584a3ac471c 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -6,6 +6,7 @@ but litellm_params["litellm_metadata"] is None. """ import threading +from typing import Final from types import SimpleNamespace import pytest @@ -16,6 +17,7 @@ from litellm.litellm_core_utils.redact_messages import ( _redact_responses_api_output, perform_redaction, redact_streaming_responses_for_custom_logger, + redacted_standard_logging_payload, should_redact_message_logging, ) from litellm.responses.main import mock_responses_api_response @@ -858,3 +860,104 @@ class TestRedactStreamingResponsesForCustomLogger: assert result_details is model_call_details assert response_obj.choices[0].message.content == "secret content" + + +@pytest.mark.parametrize("callback_only", [False, True]) +def test_classifier_audit_redaction_removes_both_fields_and_source_carrier(callback_only: bool) -> None: + audit: Final = {"classifier_input": {"system": "private rubric"}, "originating_request_masked": {"input": "private source"}} + standard_payload: Final = { + **audit, + "messages": [{"role": "user", "content": "private prompt"}], + "response": {"choices": [{"message": {"content": "private answer"}}]}, + "model": "classifier", + } + details: Final = { + "standard_logging_object": standard_payload, + "litellm_params": {"proxy_server_request": {"body": {}, "originating_request_masked": audit["originating_request_masked"]}}, + } + logger: Final = CustomLogger() + logger.turn_off_message_logging = True + if callback_only: + redacted: Final = logger.redact_standard_logging_payload_from_model_call_details(details) + assert "classifier_input" not in redacted["standard_logging_object"] + assert "originating_request_masked" not in redacted["standard_logging_object"] + assert "originating_request_masked" not in redacted["litellm_params"]["proxy_server_request"] + assert details["standard_logging_object"]["classifier_input"] == audit["classifier_input"] + assert details["litellm_params"]["proxy_server_request"]["originating_request_masked"] == audit["originating_request_masked"] + else: + perform_redaction(details, result=None) + assert "classifier_input" not in details["standard_logging_object"] + assert "originating_request_masked" not in details["standard_logging_object"] + assert "originating_request_masked" not in details["litellm_params"]["proxy_server_request"] + + assert standard_payload["classifier_input"] == audit["classifier_input"] + assert standard_payload["originating_request_masked"] == audit["originating_request_masked"] + assert standard_payload["messages"][0]["content"] == "private prompt" + assert standard_payload["response"]["choices"][0]["message"]["content"] == "private answer" + + +@pytest.mark.parametrize("excluded", [False, True]) +def test_classifier_callback_redaction_preserves_exclusions(monkeypatch: pytest.MonkeyPatch, excluded: bool) -> None: + monkeypatch.setattr(litellm, "standard_logging_payload_excluded_fields", ["messages", "response"] if excluded else []) + payload: Final = { + "classifier_input": {"system": "private rubric"}, + "originating_request_masked": {"input": "private source"}, + "messages": [{"role": "user", "content": "private prompt"}], + "response": {"choices": [{"message": {"content": "private answer"}}]}, + "model": "classifier", + } + logger: Final = CustomLogger() + logger.turn_off_message_logging = True + redacted: Final = logger.redact_standard_logging_payload_from_model_call_details({"standard_logging_object": payload}) + stored: Final = redacted["standard_logging_object"] + assert "classifier_input" not in stored + assert "originating_request_masked" not in stored + assert stored["model"] == "classifier" + assert ("messages" not in stored) is excluded + assert ("response" not in stored) is excluded + if not excluded: + assert stored["messages"][0]["content"] == "redacted-by-litellm" + assert stored["response"]["choices"][0]["message"]["content"] == "redacted-by-litellm" + + failure_payload: Final = redacted_standard_logging_payload(payload) + assert "classifier_input" not in failure_payload + assert "originating_request_masked" not in failure_payload + assert failure_payload["messages"][0]["content"] == "redacted-by-litellm" + assert failure_payload["response"]["choices"][0]["message"]["content"] == "redacted-by-litellm" + assert payload["classifier_input"] == {"system": "private rubric"} + assert payload["response"]["choices"][0]["message"]["content"] == "private answer" + + +class _SelfRedactingLogger(CustomLogger): + def redacts_messages_itself(self) -> bool: + return True + + +@pytest.mark.parametrize("logger", [CustomLogger(), _SelfRedactingLogger()], ids=["default", "redacts_itself"]) +def test_field_exclusion_alone_leaves_messages_and_responses_intact(monkeypatch: pytest.MonkeyPatch, logger: CustomLogger) -> None: + monkeypatch.setattr(litellm, "standard_logging_payload_excluded_fields", ["model"]) + payload: Final = { + "messages": [{"role": "user", "content": "private prompt"}], + "response": {"choices": [{"message": {"content": "private answer"}}]}, + "model": "classifier", + } + stored: Final = logger.redact_standard_logging_payload_from_model_call_details({"standard_logging_object": payload})[ + "standard_logging_object" + ] + assert stored == {"messages": payload["messages"], "response": payload["response"]} + + +def test_a_callback_that_redacts_itself_keeps_its_messages_but_not_the_classifier_audit() -> None: + payload: Final = { + "classifier_input": {"system": "private rubric"}, + "messages": [{"role": "user", "content": "private prompt"}], + "response": {"choices": [{"message": {"content": "private answer"}}]}, + } + logger: Final = _SelfRedactingLogger() + logger.turn_off_message_logging = True + stored: Final = logger.redact_standard_logging_payload_from_model_call_details({"standard_logging_object": payload})[ + "standard_logging_object" + ] + assert "classifier_input" not in stored + assert stored["messages"] == payload["messages"] + assert stored["response"] == payload["response"] diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py index c8b6772d965..1551c3fd6e6 100644 --- a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py +++ b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py @@ -358,7 +358,12 @@ def test_redact_credentials_in_payload_leaves_no_fragment_of_the_secret(): "azure_ad_token": fake_token, "aws_secret_access_key": "fake-aws-secret-0000", "vertex_credentials": {"private_key": "fake-pem"}, - "extra_headers": {"Authorization": "Bearer fake-bearer-0000", "x-request-id": "abc123"}, + "extra_headers": { + "Authorization": "Bearer fake-bearer-0000", + "Cookie": "session=fake-session", + "Set-Cookie": "session=fake-session; HttpOnly", + "x-request-id": "abc123", + }, "model": "gpt-4o-mini", "max_tokens": 17, "temperature": 0.25, @@ -373,6 +378,8 @@ def test_redact_credentials_in_payload_leaves_no_fragment_of_the_secret(): assert "fake-bearer-0000" not in str(result) assert result["api_key"] == "REDACTED" assert result["extra_headers"]["Authorization"] == "REDACTED" + assert result["extra_headers"]["Cookie"] == "REDACTED" + assert result["extra_headers"]["Set-Cookie"] == "REDACTED" assert result["extra_headers"]["x-request-id"] == "abc123" assert result["model"] == "gpt-4o-mini" assert result["max_tokens"] == 17 diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index bacbcbf132b..efe4209c1c9 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1,4 +1,6 @@ import json +from collections.abc import Mapping, Sequence +from typing import Final import pytest @@ -1476,3 +1478,128 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate( assert usage.completion_tokens == 100 assert usage.completion_tokens_details.reasoning_tokens == expected_reasoning_tokens assert usage.completion_tokens_details.text_tokens == expected_text_tokens + + +def _openai_chunk( + choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None +) -> dict[str, object]: + base: Final = { + "id": "chatcmpl-lit6552", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.4-mini", + "choices": list(choices), + } + return base if usage is None else {**base, "usage": dict(usage)} + + +@pytest.mark.parametrize( + "chunks", + [ + pytest.param([_openai_chunk(choices=[]), _openai_chunk(choices=[])], id="all_empty_choices_dicts"), + pytest.param( + [ModelResponseStream(model="gpt-5.4-mini", choices=[]) for _ in range(2)], + id="all_empty_choices_objects", + ), + ], +) +def test_stream_chunk_builder_survives_all_empty_choices(chunks: Sequence[object]) -> None: + response: Final = stream_chunk_builder(chunks=list(chunks)) + + assert response is not None + assert response.choices[0].message.role == "assistant" + assert response.choices[0].finish_reason == "stop" + + +def test_stream_chunk_builder_keeps_usage_from_usage_only_frames() -> None: + usage_frame: Final = _openai_chunk( + choices=[], usage={"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10} + ) + + response: Final = stream_chunk_builder(chunks=[usage_frame]) + + assert response is not None + assert response.choices[0].message.role == "assistant" + assert response.usage.prompt_tokens == 10 + assert response.usage.total_tokens == 10 + + +@pytest.mark.parametrize( + "delta", + [pytest.param({"content": "Hi"}, id="delta_without_role"), pytest.param({}, id="empty_delta")], +) +def test_stream_chunk_builder_defaults_role_when_delta_omits_it(delta: Mapping[str, str]) -> None: + chunks: Final = [ + _openai_chunk(choices=[{"index": 0, "delta": dict(delta), "finish_reason": None}]), + _openai_chunk(choices=[{"index": 0, "delta": {"content": "!"}, "finish_reason": "stop"}]), + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert response.choices[0].message.role == "assistant" + assert response.choices[0].message.content == delta.get("content", "") + "!" + assert response.choices[0].finish_reason == "stop" + + +def test_stream_chunk_builder_reads_role_from_first_frame_with_choices() -> None: + chunks: Final = [ + _openai_chunk(choices=[]), + _openai_chunk(choices=[{"index": 0, "delta": {"role": "user", "content": "Hi"}, "finish_reason": None}]), + _openai_chunk(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]), + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert response.choices[0].message.role == "user" + assert response.choices[0].message.content == "Hi" + + +def _fail_prompt_token_count() -> int: + raise AssertionError("prompt tokens must come from the usage chunk, not the tokenizer") + + +def test_calculate_usage_reads_prompt_tokens_from_mock_stream_usage_chunk_without_tokenizer_fallback() -> None: + from litellm.utils import mock_completion_streaming_obj + + chunks: Final = list( + mock_completion_streaming_obj( + ModelResponseStream(model="gpt-5.4-mini"), + mock_response="ok", + model="gpt-5.4-mini", + prompt_tokens=51234, + ) + ) + assert chunks[-1].choices == [] + + usage: Final = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="gpt-5.4-mini", + completion_output="ok", + count_prompt_tokens=_fail_prompt_token_count, + ) + + assert usage.prompt_tokens == 51234 + assert usage.completion_tokens == chunks[-1].usage.completion_tokens + assert usage.total_tokens == 51234 + usage.completion_tokens + + +def test_calculate_usage_falls_back_to_prompt_counter_when_mock_stream_has_no_admission_count() -> None: + from litellm.utils import mock_completion_streaming_obj + + chunks: Final = list( + mock_completion_streaming_obj( + ModelResponseStream(model="gpt-5.4-mini"), mock_response="ok", model="gpt-5.4-mini" + ) + ) + assert all(chunk.choices for chunk in chunks) + + usage: Final = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="gpt-5.4-mini", + completion_output="ok", + count_prompt_tokens=lambda: 77, + ) + + assert usage.prompt_tokens == 77 diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 0aa73833677..37e2031fdf4 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -6,7 +6,7 @@ import pytest import asyncio import traceback -from typing import Optional +from typing import Final, Optional import litellm from litellm import verbose_logger @@ -2633,6 +2633,48 @@ def test_dispatch_cached_response_extracts_delta( assert initialized_custom_stream_wrapper.response_id == "chatcmpl-cache-1" +def test_dispatch_cached_response_without_choices_is_an_empty_chunk( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """A cached completion with no choices replays as an empty, unfinished chunk + instead of raising IndexError on choices[0].""" + initialized_custom_stream_wrapper.custom_llm_provider = "cached_response" + chunk: Final = ModelResponseStream(id="chatcmpl-cache-empty", choices=[]) + + result, model_response, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, chunk + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] is None + assert initialized_custom_stream_wrapper.received_finish_reason is None + assert model_response.id == "chatcmpl-cache-empty" + + +@pytest.mark.asyncio +async def test_cached_response_without_choices_streams_a_single_stop_chunk( + logging_obj: Logging, +): + """A stream cache hit on a completion stored with choices == [] ends with one + finish_reason=stop chunk, the same shape the live empty stream produced.""" + + async def cached_chunks(): + yield ModelResponseStream(id="chatcmpl-cache-empty", choices=[]) + + wrapper: Final = CustomStreamWrapper( + completion_stream=cached_chunks(), + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="cached_response", + ) + + chunks: Final = tuple([chunk async for chunk in wrapper]) + + assert len(chunks) == 1 + assert tuple(choice.finish_reason for chunk in chunks for choice in chunk.choices) == ("stop",) + assert all(choice.delta.content in (None, "") for chunk in chunks for choice in chunk.choices) + + def test_dispatch_vertex_ai_legacy_text_and_finish_reason( initialized_custom_stream_wrapper: CustomStreamWrapper, ): diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 4694fa8fbed..60f25c48443 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1,10 +1,16 @@ #### What this tests #### # This tests litellm.token_counter.token_counter() function +import asyncio +import base64 import importlib +import threading import time import traceback +from concurrent.futures import Future, wait +from typing import Final from unittest.mock import MagicMock +import anyio.to_thread import pytest import tiktoken @@ -14,9 +20,23 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.litellm_core_utils.token_counter import _get_tiktoken_count_function +from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.token_counter import ( + _get_exact_count_function, + _get_extrapolating_count_function, + _get_tiktoken_count_function, + calculate_img_tokens, + high_detail_image_token_upper_bound, + offload_token_count, +) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text +from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, +) from tests.test_litellm.litellm_core_utils.messages_with_counts import ( MESSAGES_TEXT, MESSAGES_WITH_IMAGES, @@ -120,6 +140,135 @@ def test_valid_chunk_size_config_is_honoured(monkeypatch): importlib.reload(litellm.constants) +async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free(): + warm_tokenizer("claude-fable-5") + + tokens, took, lags = await timed_with_loop_lags( + lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) + ) + + assert tokens > 0 + assert_loop_stayed_free(took, lags) + + +@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) +def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int): + count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk)) + front_heavy: Final = "a" * 1_000 + "b" * 4_000 + exact: Final = 1_000 + len(front_heavy) + + estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy) + + assert abs(estimate - exact) <= exact // 100 + assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars + + +def test_count_at_or_below_the_cap_is_exact(): + count_exactly: Final = MagicMock(side_effect=len) + + assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000 + assert count_exactly.call_args_list == [(("a" * 5_000,),)] + + +class _SlowEncoder: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self.in_flight = 0 + self.peak_in_flight = 0 + + def encode_batch_fast(self, texts: list[str]) -> list[list[int]]: + with self._lock: + self.in_flight += 1 + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + time.sleep(0.1) + with self._lock: + self.in_flight -= 1 + return [[0] * len(text) for text in texts] + + +@pytest.mark.asyncio +async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool(): + encoder: Final = _SlowEncoder() + count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder}) + shared_pool: Final = anyio.to_thread.current_default_thread_limiter() + burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + + async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]: + if counting.done(): + return () + await asyncio.sleep(0.01) + return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting)) + + counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst)))) + borrowed: Final = await shared_pool_borrowed_until_done(counting) + + assert await counting == [3] * burst + assert len(borrowed) > 1 and max(borrowed) == 0 + assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + + +def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None: + def slow_count(counted: str) -> int: + time.sleep(0.1) + return len(counted) + + result.set_result(asyncio.run(offload_token_count(slow_count)(text))) + + +def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process(): + loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + results: Final = tuple(Future[int]() for _ in range(loops)) + threads: Final = tuple( + threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True) + for size, result in enumerate(results, start=1) + ) + for thread in threads: + thread.start() + + _, pending = wait(results, timeout=5) + + assert not pending + assert tuple(result.result() for result in results) == tuple(range(1, loops + 1)) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("8", 8), ("0", 4), ("not-an-int", 4)], +) +def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS") + importlib.reload(litellm.constants) + + +def test_token_counter_applies_the_default_cap(): + max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS + prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] + over_the_cap: Final = prose + "a" * 200_000 + exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap) + + estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap) + + assert estimate != exact + assert abs(estimate - exact) <= exact // 100 + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)], +) +def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS") + importlib.reload(litellm.constants) + + def test_token_counter_with_prefix(): messages = [ {"role": "user", "content": "Who won the world cup in 2022?"}, @@ -1412,3 +1561,18 @@ def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): assert _count_user_content([prompt, named]) == _count_user_content( [prompt, {"type": "text", "text": "report.pdf"}] ) + + +def _png_data_url(width: int, height: int) -> str: + ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big") + return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode() + + +@pytest.mark.parametrize(("width", "height"), [(1, 1), (768, 768), (2000, 768), (768, 2000), (4096, 4096), (8000, 3072)]) +def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, height: int) -> None: + assert calculate_img_tokens(_png_data_url(width, height), mode="high") <= high_detail_image_token_upper_bound() + + +def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None: + assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound() + assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index aaaa43a0dc4..fccdc1a2a0a 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -1,5 +1,9 @@ +import asyncio import socket +import threading +import time +import httpx import pytest import litellm @@ -535,3 +539,34 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames(): detail = str(exc.value) assert "attacker.example.com" not in detail assert "api.internal-corp.example" not in detail + + +async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch): + loop_thread = threading.current_thread() + resolver_threads = [] + + def slow_getaddrinfo(host, port, *args, **kwargs): + resolver_threads.append(threading.current_thread()) + time.sleep(0.4) + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))] + + monkeypatch.setattr(url_utils.socket, "getaddrinfo", slow_getaddrinfo) + + class FakeClient: + async def get(self, url, **kwargs): + return httpx.Response(200, request=httpx.Request("GET", url)) + + ticks = [time.perf_counter()] + + async def heartbeat(): + while True: + await asyncio.sleep(0.01) + ticks.append(time.perf_counter()) + + beating = asyncio.create_task(heartbeat()) + response = await url_utils.async_safe_get(FakeClient(), "https://img.example/a.png") + beating.cancel() + + assert response.status_code == 200 + assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads) + assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.2 diff --git a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py index ca25ee80c23..d24e2b58db8 100644 --- a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py +++ b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py @@ -1,6 +1,5 @@ -import litellm from litellm import LlmProviders from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -46,14 +45,6 @@ def test_xai_openai_compatible_provider_info(): assert dynamic_api_key == "api-key" -def test_xai_get_model_info_uses_xai_pricing_metadata(): - model_info = litellm.get_model_info("xai/grok-3-mini") - - assert model_info["litellm_provider"] == "xai" - assert model_info["key"] == "xai/grok-3-mini" - assert model_info["mode"] == "chat" - - def test_xai_validate_environment_reads_api_key(monkeypatch): monkeypatch.setenv("XAI_API_KEY", "api-key") diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 3044a321aa6..9fe56f4dc65 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -264,6 +264,231 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: # Should return the responses unchanged assert result == responses_so_far + @staticmethod + def _ended_sse_chunks() -> list: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + @staticmethod + def _masking_guardrail() -> CustomGuardrail: + class MaskWorld(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs.get("texts", [])]} + + return MaskWorld(guardrail_name="test") + + @staticmethod + def _delta_texts(chunks: list) -> list: + texts = [] + for chunk in chunks: + for line in chunk.decode().split("\n"): + if not line.startswith("data:"): + continue + data = json.loads(line[len("data:") :].strip()) + if data.get("type") == "content_block_delta": + texts.append(data["delta"]["text"]) + return texts + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_text_back_into_sse_chunks(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert self._delta_texts(chunks) == ["hello [MASKED]", ""] + raw = b"".join(chunks).decode() + assert "event: message_start" in raw and "event: message_stop" in raw + assert '"stop_reason": "end_turn"' in raw + + @staticmethod + def _ended_tool_use_sse_chunks() -> list: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + @staticmethod + def _argument_masking_guardrail() -> CustomGuardrail: + class MaskArguments(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call.function.arguments = '{"fruit": "[MASKED]"}' + return inputs + + return MaskArguments(guardrail_name="test") + + @staticmethod + def _partial_jsons(chunks: list) -> list: + return [ + json.loads(line[len("data:") :].strip())["delta"]["partial_json"] + for chunk in chunks + for line in chunk.decode().split("\n") + if line.startswith("data:") and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta" + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_use_input_back_into_sse_chunks(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', "", ""] + raw = b"".join(chunks).decode() + assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw + assert '"stop_reason": "tool_use"' in raw + assert "persim" not in raw + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_use_name_back_into_sse_chunks(self): + class RenameTool(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call.function.name = "lookup_fruit_reviewed" + return inputs + + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=RenameTool(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + raw = b"".join(chunks).decode() + assert '"name": "lookup_fruit_reviewed"' in raw and '"id": "toolu_1"' in raw + assert '"name": "lookup_fruit"' not in raw + assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"} + + @pytest.mark.asyncio + async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + original = [bytes(chunk) for chunk in chunks] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert chunks == original + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = AnthropicMessagesHandler() + server_tool_use = [ + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ] + tool_use = self._ended_tool_use_sse_chunks() + chunks = ( + tool_use[:1] + + [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in server_tool_use] + + [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]] + ) + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks() + original = [bytes(chunk) for chunk in chunks] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert chunks == original + + @pytest.mark.asyncio + async def test_unended_stream_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks()[:-2] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_unended_stream_without_rewrite_is_released_with_delivery_expected(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks()[:-2] + original = [bytes(chunk) for chunk in chunks] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert chunks == original + + @pytest.mark.asyncio + async def test_unended_stream_rewrite_without_delivery_expected_does_not_raise(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks()[:-2] + original = [bytes(chunk) for chunk in chunks] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert result is chunks + assert chunks == original + class TestAnthropicMessagesHandlerInputProcessing: """Test input processing preserves litellm_metadata for dynamic guardrails.""" @@ -2045,3 +2270,29 @@ class TestAnthropicMessagesHandlerStreamingScanKey: assert open_key == StreamingScanKey(texts=("hi",)) assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] assert ended_key != open_key + + +class TestAnthropicMessagesHandlerPostCallHookResponse: + def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self): + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + assembled = ModelResponse( + id="msg_1", + model="claude", + choices=[Choices(message=Message(role="assistant", content="hello world"), finish_reason="stop")], + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + ) + + hook_response = AnthropicMessagesHandler().post_call_hook_response(assembled) + + assert hook_response["type"] == "message" + assert hook_response["role"] == "assistant" + assert hook_response["content"] == [{"type": "text", "text": "hello world"}] + assert hook_response["stop_reason"] == "end_turn" + assert hook_response["usage"]["input_tokens"] == 1 + assert hook_response["usage"]["output_tokens"] == 2 + + def test_anything_else_reaches_the_hook_untouched(self): + native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]} + + assert AnthropicMessagesHandler().post_call_hook_response(native) is native diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index 043537f8c1f..18309595414 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -9,7 +9,8 @@ import pytest import litellm from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm._uuid import uuid +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, @@ -85,6 +86,54 @@ def test_anthropic_completion_does_not_send_deployment_default_limits(): assert "default_api_key_tpm_limit" not in request_body +async def test_anthropic_async_completion_inlines_http_images_off_the_event_loop(async_only_image_fetch): + http_image_url = f"http://img.example/{uuid.uuid4()}.png" + https_image_url = f"https://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="anthropic/claude-sonnet-4-6", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": http_image_url}}, + {"type": "image_url", "image_url": {"url": https_image_url}}, + ], + } + ], + api_key="test-key", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [http_image_url] + sources = [part["source"] for part in captured["body"]["messages"][0]["content"] if part["type"] == "image"] + assert sources == [ + {"type": "base64", "media_type": "image/png", "data": async_only_image_fetch.base64_png}, + {"type": "url", "url": https_image_url}, + ] + + def test_redacted_thinking_content_block_delta(): chunk = { "type": "content_block_start", @@ -2256,7 +2305,7 @@ class TestRustChatCompletionsHook: def _reset_bridge(self, monkeypatch): from litellm.rust_bridge import chat_completions as bridge - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") bridge.set_rust_chat_completions( chat_completions=None, achat_completions=None, decline=None ) @@ -2282,7 +2331,7 @@ class TestRustChatCompletionsHook: "logging_obj": MagicMock(), "optional_params": {"max_tokens": 16}, "timeout": 30.0, - "litellm_params": {"rust": True}, + "litellm_params": {}, "acompletion": False, "headers": {}, "client": None, @@ -2366,7 +2415,8 @@ class TestRustChatCompletionsHook: ) assert seen["call"][0]["optional_params"]["max_tokens"] == 7 - def test_without_the_opt_in_the_core_is_never_consulted(self): + def test_without_the_opt_in_the_core_is_never_consulted(self, monkeypatch): + monkeypatch.setenv("LITELLM_RUST", "0") from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -2587,6 +2637,7 @@ class TestRustChatCompletionsHook: def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch): """The suppression must not swallow the log on the ordinary path.""" + monkeypatch.setenv("LITELLM_RUST", "0") from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion from litellm.llms.anthropic.chat.transformation import AnthropicConfig diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index ea3b19fba2b..f6ee1cd71c0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1,4 +1,5 @@ import base64 +import json from typing import Any, Final, cast import pytest @@ -9,6 +10,7 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_PLACEHOLDER, + encrypted_reasoning_signature, ) from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, @@ -40,6 +42,66 @@ from litellm.types.utils import ( ) +def test_translate_openai_response_to_anthropic_empty_choices() -> None: + response: Final = ModelResponse( + id="chatcmpl-empty", + model="gemini-3.5-flash", + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10), + ) + + result: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "end_turn" + assert result["usage"]["input_tokens"] == 10 + + +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( @@ -362,6 +424,43 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks(): assert result[1]["tool_calls"][0]["id"] == "toolu_01234" +def test_translate_anthropic_messages_to_openai_drops_bridge_encrypted_reasoning_blocks(): + """A session that moves from an OpenAI reasoning model to a chat provider replays reasoning only OpenAI can read. + + Gemini rejects the whole request when such a block reaches it as a thought_signature, so the + adapter drops those blocks and keeps the provider-signed ones. + """ + + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[{"type": "text", "text": "Who drinks water?"}], + ), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + {"type": "text", "text": "The Norwegian."}, + ], + ), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "thinking", "thinking": "native", "signature": "EqQBCkYIAxgCIkA_signed"}, + {"type": "text", "text": "Still the Norwegian."}, + ], + ), + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=anthropic_messages) + + assert [m["role"] for m in result] == ["user", "assistant", "assistant"] + assert not result[1].get("thinking_blocks") + assert result[1]["content"] == "The Norwegian." + assert [b["signature"] for b in result[2]["thinking_blocks"]] == ["EqQBCkYIAxgCIkA_signed"] + + def test_translate_anthropic_messages_to_openai_sets_reasoning_content(): """Reasoning-aware chat providers read reasoning_content, so thinking text must land there. @@ -679,9 +778,14 @@ def test_translate_anthropic_to_openai_orders_top_level_and_midturn_system(): ] -def _translate_with_metadata( - model: str, metadata: dict[str, str], custom_llm_provider: str | None -) -> dict[str, Any]: +def _claude_code_user_id(session_id: str) -> str: + return json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": session_id}) + + +CLAUDE_CODE_USER_ID: Final = _claude_code_user_id("session-abc") + + +def _translate_with_metadata(model: str, metadata: dict[str, str], custom_llm_provider: str | None) -> dict[str, Any]: openai_request, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( anthropic_message_request={ "model": model, @@ -694,23 +798,51 @@ def _translate_with_metadata( return cast(dict[str, Any], openai_request) -def test_translate_anthropic_to_openai_maps_user_id_to_prompt_cache_key_for_openai(): - openai_request = _translate_with_metadata("openai/gpt-5.6-luna", {"user_id": "session-abc"}, "openai") - assert openai_request["user"] == "session-abc" +def test_translate_anthropic_to_openai_maps_claude_code_session_id_to_prompt_cache_key_for_openai(): + openai_request = _translate_with_metadata("openai/gpt-5.6-luna", {"user_id": CLAUDE_CODE_USER_ID}, "openai") + assert openai_request["user"] == CLAUDE_CODE_USER_ID assert openai_request["prompt_cache_key"] == "session-abc" -def test_translate_anthropic_to_openai_truncates_prompt_cache_key_but_keeps_full_user(): - long_id = "".join(str(i % 10) for i in range(100)) - openai_request = _translate_with_metadata("openai/gpt-5.6-luna", {"user_id": long_id}, "openai") - assert openai_request["user"] == long_id - assert openai_request["prompt_cache_key"] == long_id[:64] - assert len(openai_request["prompt_cache_key"]) == 64 +def test_translate_anthropic_to_openai_gives_each_claude_code_session_its_own_prompt_cache_key(): + """BerriAI/litellm#39145: the first 64 chars of Claude Code's user_id are the per-install device_id.""" + keys = tuple( + _translate_with_metadata("openai/gpt-5.6-luna", {"user_id": _claude_code_user_id(session_id)}, "openai")[ + "prompt_cache_key" + ] + for session_id in ("11111111-1111-1111-1111-111111111111", "22222222-2222-2222-2222-222222222222") + ) + assert keys == ("11111111-1111-1111-1111-111111111111", "22222222-2222-2222-2222-222222222222") + + +def test_translate_anthropic_to_openai_truncates_long_session_id_to_openai_limit(): + long_session_id = "".join(str(i % 10) for i in range(100)) + openai_request = _translate_with_metadata( + "openai/gpt-5.6-luna", {"user_id": _claude_code_user_id(long_session_id)}, "openai" + ) + assert openai_request["prompt_cache_key"] == long_session_id[:64] + + +@pytest.mark.parametrize( + "user_id", + [ + "alice", + "".join(str(i % 10) for i in range(100)), + json.dumps({"device_id": "d" * 64, "account_uuid": ""}), + json.dumps({"session_id": ""}), + json.dumps({"session_id": 123}), + "{not json", + ], +) +def test_translate_anthropic_to_openai_keeps_plain_user_id_off_prompt_cache_key(user_id: str): + openai_request = _translate_with_metadata("openai/gpt-5.6-luna", {"user_id": user_id}, "openai") + assert openai_request["user"] == user_id + assert "prompt_cache_key" not in openai_request @pytest.mark.parametrize("model", ["azure/my-gpt-5-deployment", "my-gpt-5-deployment"]) def test_translate_anthropic_to_openai_sets_prompt_cache_key_for_azure(model: str): - openai_request = _translate_with_metadata(model, {"user_id": "session-abc"}, "azure") + openai_request = _translate_with_metadata(model, {"user_id": CLAUDE_CODE_USER_ID}, "azure") assert openai_request["prompt_cache_key"] == "session-abc" @@ -727,8 +859,8 @@ def test_translate_anthropic_to_openai_sets_prompt_cache_key_for_azure(model: st def test_translate_anthropic_to_openai_skips_prompt_cache_key_when_provider_lacks_it( model: str, custom_llm_provider: str ): - openai_request = _translate_with_metadata(model, {"user_id": "session-abc"}, custom_llm_provider) - assert openai_request["user"] == "session-abc" + openai_request = _translate_with_metadata(model, {"user_id": CLAUDE_CODE_USER_ID}, custom_llm_provider) + assert openai_request["user"] == CLAUDE_CODE_USER_ID assert "prompt_cache_key" not in openai_request @@ -736,14 +868,14 @@ def test_translate_anthropic_to_openai_skips_prompt_cache_key_for_chained_litell assert "prompt_cache_key" in litellm.get_supported_openai_params( model="xai", custom_llm_provider="litellm_proxy" ) - openai_request = _translate_with_metadata("litellm_proxy/xai", {"user_id": "session-abc"}, "litellm_proxy") - assert openai_request["user"] == "session-abc" + openai_request = _translate_with_metadata("litellm_proxy/xai", {"user_id": CLAUDE_CODE_USER_ID}, "litellm_proxy") + assert openai_request["user"] == CLAUDE_CODE_USER_ID assert "prompt_cache_key" not in openai_request def test_translate_anthropic_to_openai_skips_prompt_cache_key_without_provider(): - openai_request = _translate_with_metadata("openai/gpt-5.6-luna", {"user_id": "session-abc"}, None) - assert openai_request["user"] == "session-abc" + openai_request = _translate_with_metadata("openai/gpt-5.6-luna", {"user_id": CLAUDE_CODE_USER_ID}, None) + assert openai_request["user"] == CLAUDE_CODE_USER_ID assert "prompt_cache_key" not in openai_request @@ -798,6 +930,7 @@ def test_translate_openai_content_to_anthropic_empty_function_arguments(): assert ( result[0]["input"] == {} ), "Empty function arguments should result in empty dict" + assert "provider_specific_fields" not in result[0] def test_translate_openai_content_to_anthropic_text_and_tool_calls(): @@ -843,6 +976,11 @@ def test_translate_openai_content_to_anthropic_strips_gemini_thought_from_tool_c base = "call_3e9417b7925e49aca9a71dc1885e" sig = "CiIBDDnWx+/a==" combined = f"{base}{THOUGHT_SIGNATURE_SEPARATOR}{sig}" + function = Function( + name="get_weather", + arguments='{"location": "Boston"}', + ) + function.provider_specific_fields = {"thought_signature": sig} openai_choices = [ Choices( message=Message( @@ -852,10 +990,7 @@ def test_translate_openai_content_to_anthropic_strips_gemini_thought_from_tool_c ChatCompletionAssistantToolCall( id=combined, type="function", - function=Function( - name="get_weather", - arguments='{"location": "Boston"}', - ), + function=function, ) ], ) @@ -871,6 +1006,7 @@ def test_translate_openai_content_to_anthropic_strips_gemini_thought_from_tool_c assert THOUGHT_SIGNATURE_SEPARATOR not in result[0]["id"] assert result[0]["name"] == "get_weather" assert result[0]["input"] == {"location": "Boston"} + assert result[0]["provider_specific_fields"] == {"signature": sig} def test_translate_openai_content_to_anthropic_sanitizes_colon_dot_tool_call_ids(): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py index f48d51dbe1e..7dc7507120f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py @@ -1,3 +1,4 @@ +import json import os import sys @@ -10,6 +11,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( ) MESSAGES = [{"role": "user", "content": "hello"}] +CLAUDE_CODE_USER_ID = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "session-abc"}) def _prepare(model: str, extra_kwargs: dict[str, object], thinking: dict[str, object] | None = None): @@ -17,7 +19,7 @@ def _prepare(model: str, extra_kwargs: dict[str, object], thinking: dict[str, ob max_tokens=1024, messages=MESSAGES, model=model, - metadata={"user_id": "session-abc"}, + metadata={"user_id": CLAUDE_CODE_USER_ID}, thinking=thinking, extra_kwargs=extra_kwargs, ) @@ -26,7 +28,7 @@ def _prepare(model: str, extra_kwargs: dict[str, object], thinking: dict[str, ob def test_prepare_completion_kwargs_derives_prompt_cache_key_for_openai_provider(): completion_kwargs = _prepare("openai/gpt-5.6-luna", {"custom_llm_provider": "openai"}) - assert completion_kwargs["user"] == "session-abc" + assert completion_kwargs["user"] == CLAUDE_CODE_USER_ID assert completion_kwargs["prompt_cache_key"] == "session-abc" @@ -35,7 +37,7 @@ def test_prepare_completion_kwargs_prefers_explicit_prompt_cache_key_over_derive "openai/gpt-5.6-luna", {"custom_llm_provider": "openai", "prompt_cache_key": "explicit-key"}, ) - assert completion_kwargs["user"] == "session-abc" + assert completion_kwargs["user"] == CLAUDE_CODE_USER_ID assert completion_kwargs["prompt_cache_key"] == "explicit-key" @@ -50,13 +52,13 @@ def test_prepare_completion_kwargs_skips_prompt_cache_key_without_provider_suppo model: str, extra_kwargs: dict[str, object] ): completion_kwargs = _prepare(model, extra_kwargs) - assert completion_kwargs["user"] == "session-abc" + assert completion_kwargs["user"] == CLAUDE_CODE_USER_ID assert "prompt_cache_key" not in completion_kwargs def test_prepare_completion_kwargs_skips_prompt_cache_key_for_chained_litellm_proxy(): completion_kwargs = _prepare("litellm_proxy/xai", {"custom_llm_provider": "litellm_proxy"}) - assert completion_kwargs["user"] == "session-abc" + assert completion_kwargs["user"] == CLAUDE_CODE_USER_ID assert "prompt_cache_key" not in completion_kwargs diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py index 56b754c3476..af7befecc33 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py @@ -82,3 +82,22 @@ class TestTheNormalizedTierIsTheTierSent: self, local_model_cost_map, model, provider, effort, expected ): assert _reasoning_effort_sent(model, provider, effort) == expected + + @pytest.mark.parametrize( + "model, provider", + [ + ("gpt-6-astra", "azure_ai"), + ("azure_ai/gpt-6-astra", "azure_ai"), + ("gpt-6-astra", "azure"), + ("us/gpt-6-astra", "azure"), + ], + ) + def test_an_azure_hosted_astra_deployment_drops_to_the_tier_it_accepts( + self, local_model_cost_map, model, provider + ): + """The deployment answers ``max`` with a 400 naming ``none`` through ``xhigh``, so the rows + say so and the adapter sends the tier below instead of the rejected one.""" + assert _reasoning_effort_sent(model, provider, "max") == "xhigh" + + def test_the_openai_hosted_twin_still_sends_max(self, local_model_cost_map): + assert _reasoning_effort_sent("gpt-6-astra", "openai", "max") == "max" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 17d42f55ae0..fdd08eaa182 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -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"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index e819433c269..01f7a2fb7ab 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1072,6 +1072,54 @@ async def test_messages_strips_provider_prefix_exactly_once(requested_model, exp assert captured["url"] == expected_url +@pytest.mark.asyncio +async def test_native_messages_strips_replayed_provider_specific_fields_from_wire(): + captured = {} + + async def fake_send(self, request, **kwargs): + captured["body"] = json.loads(request.content) + raise httpx.ConnectError("cut at the wire", request=request) + + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "get_weather", + "input": {"city": "Paris"}, + "provider_specific_fields": {"signature": "sig_abc"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": "Sunny", + } + ], + }, + ] + + with ( + patch.object(httpx.AsyncClient, "send", fake_send), + pytest.raises(litellm.exceptions.InternalServerError), + ): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=messages, + model="anthropic/claude-haiku-4-5-20251001", + api_key="test-api-key", + ) + + assert "provider_specific_fields" in messages[0]["content"][0] + assert "provider_specific_fields" not in captured["body"]["messages"][0]["content"][0] + + @pytest.mark.asyncio @pytest.mark.parametrize( "requested_model, expected_reported_model", diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py new file mode 100644 index 00000000000..c64e9d392e5 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py @@ -0,0 +1,50 @@ +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, +) +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) + + +def _transform(messages): + return AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-sonnet-4-5", + messages=messages, + anthropic_messages_optional_request_params={"max_tokens": 1024}, + litellm_params={}, + headers={}, + ) + + +def test_reasoning_replayed_from_the_responses_bridge_never_reaches_anthropic(): + """Claude Code resumed on a Claude model echoes the thinking blocks a gpt turn produced.""" + messages = [ + {"role": "user", "content": "Solve it."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + {"type": "text", "text": "The answer."}, + ], + }, + {"role": "user", "content": "And the next one?"}, + ] + request = _transform(messages) + assert request["messages"][1]["content"] == [{"type": "text", "text": "The answer."}] + assert len(messages[1]["content"]) == 3 + + +def test_anthropic_signed_thinking_blocks_are_forwarded_untouched(): + messages = [ + {"role": "user", "content": "Solve it."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."}, + ], + }, + ] + assert _transform(messages)["messages"] == messages diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py index 3383813245a..b66075f691b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py @@ -16,6 +16,7 @@ from litellm.llms.anthropic.experimental_pass_through.responses_adapters.handler ) MESSAGES = [{"role": "user", "content": "hello"}] +CLAUDE_CODE_USER_ID = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "session-abc"}) RESPONSES_SSE_BODY = ( b"event: response.created\n" @@ -30,7 +31,19 @@ RESPONSES_SSE_BODY = ( ) -def test_build_responses_kwargs_derives_prompt_cache_key_from_user_id(): +def test_build_responses_kwargs_derives_prompt_cache_key_from_claude_code_session_id(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + metadata={"user_id": CLAUDE_CODE_USER_ID}, + extra_kwargs={"custom_llm_provider": "openai"}, + ) + assert responses_kwargs["user"] == CLAUDE_CODE_USER_ID[:64] + assert responses_kwargs["prompt_cache_key"] == "session-abc" + + +def test_build_responses_kwargs_sets_no_prompt_cache_key_for_plain_user_id(): responses_kwargs = _build_responses_kwargs( max_tokens=1024, messages=MESSAGES, @@ -39,7 +52,7 @@ def test_build_responses_kwargs_derives_prompt_cache_key_from_user_id(): extra_kwargs={"custom_llm_provider": "openai"}, ) assert responses_kwargs["user"] == "session-abc" - assert responses_kwargs["prompt_cache_key"] == "session-abc" + assert "prompt_cache_key" not in responses_kwargs def test_build_responses_kwargs_prefers_explicit_prompt_cache_key_over_derived(): @@ -47,13 +60,46 @@ def test_build_responses_kwargs_prefers_explicit_prompt_cache_key_over_derived() max_tokens=1024, messages=MESSAGES, model="openai/gpt-5.6-luna", - metadata={"user_id": "session-abc"}, + metadata={"user_id": CLAUDE_CODE_USER_ID}, extra_kwargs={"custom_llm_provider": "openai", "prompt_cache_key": "explicit-key"}, ) - assert responses_kwargs["user"] == "session-abc" + assert responses_kwargs["user"] == CLAUDE_CODE_USER_ID[:64] assert responses_kwargs["prompt_cache_key"] == "explicit-key" +def test_build_responses_kwargs_asks_openai_for_encrypted_reasoning_without_thinking(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + extra_kwargs={"custom_llm_provider": "openai"}, + ) + assert responses_kwargs["include"] == ["reasoning.encrypted_content"] + assert "reasoning" not in responses_kwargs + + +def test_build_responses_kwargs_skips_include_for_a_responses_provider_that_rejects_it(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="perplexity/sonar", + thinking={"type": "enabled", "budget_tokens": 4096}, + extra_kwargs={"custom_llm_provider": "perplexity"}, + ) + assert "include" not in responses_kwargs + assert "reasoning" in responses_kwargs + + +def test_build_responses_kwargs_keeps_the_deployment_include_next_to_encrypted_reasoning(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + extra_kwargs={"custom_llm_provider": "openai", "include": ["file_search_call.results"]}, + ) + assert responses_kwargs["include"] == ["reasoning.encrypted_content", "file_search_call.results"] + + def test_build_responses_kwargs_without_metadata_sets_no_prompt_cache_key(): responses_kwargs = _build_responses_kwargs( max_tokens=1024, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index aebbed88c70..bfe2d6b7cea 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -10,6 +10,9 @@ from types import SimpleNamespace sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, +) from litellm.llms.anthropic.experimental_pass_through.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, ) @@ -34,6 +37,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 +66,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 @@ -97,7 +117,7 @@ class TestReasoningItemWithoutSummaryText: """ @staticmethod - def _gpt_turn(reasoning_summary_deltas: list) -> list: + def _gpt_turn(reasoning_summary_deltas: list, encrypted_content: str | None = None) -> list: return [ {"type": "response.created"}, {"type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}}, @@ -105,7 +125,10 @@ class TestReasoningItemWithoutSummaryText: {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": delta} for delta in reasoning_summary_deltas ), - {"type": "response.output_item.done", "item": {"type": "reasoning", "id": "rs_1"}}, + { + "type": "response.output_item.done", + "item": {"type": "reasoning", "id": "rs_1", "encrypted_content": encrypted_content}, + }, {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}, {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hello"}, {"type": "response.output_item.done", "item": {"type": "message", "id": "msg_1"}}, @@ -154,6 +177,70 @@ class TestReasoningItemWithoutSummaryText: assert not [c for c in chunks if c.get("delta", {}).get("type") == "signature_delta"] +_ENCRYPTED_REASONING = "gAAAAABp_encrypted_reasoning_bytes_only_openai_can_read" + + +class TestEncryptedReasoningIsStreamedForReplay: + """Regression for https://github.com/BerriAI/litellm/issues/40288. + + The client echoes a thinking block's signature (or a redacted block's data) back on the + next turn, so the item's ``encrypted_content`` has to reach it through one of those. + """ + + def test_encrypted_content_is_streamed_as_the_signature_before_the_block_closes(self): + chunks = _drain_async( + TestReasoningItemWithoutSummaryText._gpt_turn( + reasoning_summary_deltas=["Weighing options"], encrypted_content=_ENCRYPTED_REASONING + ) + ) + + assert [(c["type"], c.get("index"), c.get("delta", {}).get("type")) for c in chunks[1:5]] == [ + ("content_block_start", 0, None), + ("content_block_delta", 0, "thinking_delta"), + ("content_block_delta", 0, "signature_delta"), + ("content_block_stop", 0, None), + ] + assert chunks[3]["delta"]["signature"] == encrypted_reasoning_signature(_ENCRYPTED_REASONING) + + def test_reasoning_without_summary_streams_a_redacted_thinking_block(self): + chunks = _drain_async( + TestReasoningItemWithoutSummaryText._gpt_turn( + reasoning_summary_deltas=[], encrypted_content=_ENCRYPTED_REASONING + ) + ) + + assert [(c["type"], c.get("index")) for c in chunks[1:]] == [ + ("content_block_start", 0), + ("content_block_stop", 0), + ("content_block_start", 1), + ("content_block_delta", 1), + ("content_block_stop", 1), + ] + assert chunks[1]["content_block"] == { + "type": "redacted_thinking", + "data": encrypted_reasoning_signature(_ENCRYPTED_REASONING), + } + + def test_summary_parts_are_separated_inside_the_one_thinking_block(self): + """Two summary parts read as two paragraphs, not as one run-on sentence.""" + events = [ + {"type": "response.created"}, + {"type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}}, + {"type": "response.reasoning_summary_part.added", "item_id": "rs_1", "summary_index": 0}, + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": "First."}, + {"type": "response.reasoning_summary_part.added", "item_id": "rs_1", "summary_index": 1}, + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": "Second."}, + {"type": "response.output_item.done", "item": {"type": "reasoning", "id": "rs_1"}}, + ] + chunks = _process_all(events) + + thinking = "".join( + c["delta"]["thinking"] for c in chunks if c.get("delta", {}).get("type") == "thinking_delta" + ) + assert thinking == "First.\n\nSecond." + assert [c["type"] for c in chunks].count("content_block_start") == 1 + + class TestToolUseBlockClosedExactlyOnce: """Regression for https://github.com/BerriAI/litellm/issues/37273. @@ -308,3 +395,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"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 5ecf604f096..4ad559aa547 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -19,6 +19,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, + encrypted_reasoning_signature, ) from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, @@ -147,9 +148,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 @@ -568,6 +567,66 @@ class TestTranslateMessagesToResponsesInput: result = _translate_messages(messages) assert "id" not in result[0] + def test_thinking_block_with_encrypted_signature_replays_the_encrypted_content(self): + """Regression for https://github.com/BerriAI/litellm/issues/40288 (inbound fault site).""" + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "Private reasoning.", + "signature": encrypted_reasoning_signature("gAAAA_turn_one"), + } + ], + } + ] + result = _translate_messages(messages) + assert result == [ + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Private reasoning."}], + "encrypted_content": "gAAAA_turn_one", + } + ] + + def test_redacted_thinking_with_encrypted_data_replays_the_encrypted_content(self): + messages = [ + { + "role": "assistant", + "content": [{"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_turn_one")}], + } + ] + result = _translate_messages(messages) + assert result == [{"type": "reasoning", "summary": [], "encrypted_content": "gAAAA_turn_one"}] + + def test_each_encrypted_thinking_block_stays_its_own_reasoning_item(self): + """Two upstream items must not be merged into one, or the encrypted content of one is lost.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "First.", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "thinking", "thinking": "Second.", "signature": encrypted_reasoning_signature("gAAAA_2")}, + ], + } + ] + result = _translate_messages(messages) + assert [item["encrypted_content"] for item in result] == ["gAAAA_1", "gAAAA_2"] + + def test_anthropic_signed_thinking_block_replays_as_a_summary_only_item(self): + """A real Anthropic signature is opaque here, so it never masquerades as encrypted content.""" + messages = [ + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "Private reasoning.", "signature": "ErcBCkgIValid"}], + } + ] + result = _translate_messages(messages) + assert result == [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "Private reasoning."}]} + ] + def test_consecutive_thinking_blocks_become_one_reasoning_item(self): """Summary parts of one upstream reasoning item are regrouped into that item.""" messages = [ @@ -1104,6 +1163,23 @@ class TestTranslateRequestBroaderCoverage: kwargs = _ADAPTER.translate_request(req) assert "reasoning" not in kwargs + def test_thinking_asks_for_the_encrypted_reasoning(self): + """The documented way to get reasoning that survives store=false is to ask for it.""" + req = _make_request(thinking={"type": "enabled", "budget_tokens": 12000}) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["include"] == ["reasoning.encrypted_content"] + + def test_encrypted_reasoning_is_asked_for_without_a_thinking_block(self): + """A reasoning model reasons whether or not the client sent `thinking`, so the replay needs it either way.""" + kwargs = _ADAPTER.translate_request(_make_request()) + assert kwargs["include"] == ["reasoning.encrypted_content"] + + def test_encrypted_reasoning_is_not_asked_for_when_the_provider_rejects_include(self): + req = _make_request(thinking={"type": "enabled", "budget_tokens": 12000}) + kwargs = _ADAPTER.translate_request(req, include_encrypted_reasoning=False) + assert kwargs["reasoning"] == {"effort": "high"} + assert "include" not in kwargs + def test_metadata_user_id_mapped_to_user(self): req = _make_request(metadata={"user_id": "user-42"}) kwargs = _ADAPTER.translate_request(req) @@ -1115,17 +1191,34 @@ class TestTranslateRequestBroaderCoverage: kwargs = _ADAPTER.translate_request(req) assert len(kwargs["user"]) == 64 - def test_metadata_user_id_mapped_to_prompt_cache_key(self): - req = _make_request(metadata={"user_id": "user-42"}) + def test_metadata_claude_code_session_id_mapped_to_prompt_cache_key(self): + user_id = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "session-42"}) + req = _make_request(metadata={"user_id": user_id}) kwargs = _ADAPTER.translate_request(req) - assert kwargs["prompt_cache_key"] == "user-42" + assert kwargs["user"] == user_id[:64] + assert kwargs["prompt_cache_key"] == "session-42" - def test_metadata_user_id_prompt_cache_key_truncated_to_first_64_chars(self): - long_id = "".join(str(i % 10) for i in range(100)) - req = _make_request(metadata={"user_id": long_id}) + def test_metadata_claude_code_sessions_get_distinct_prompt_cache_keys(self): + """BerriAI/litellm#39145: the first 64 chars of Claude Code's user_id are the per-install device_id.""" + keys = tuple( + _ADAPTER.translate_request( + _make_request( + metadata={"user_id": json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": sid})} + ) + )["prompt_cache_key"] + for sid in ("11111111-1111-1111-1111-111111111111", "22222222-2222-2222-2222-222222222222") + ) + assert keys == ("11111111-1111-1111-1111-111111111111", "22222222-2222-2222-2222-222222222222") + + @pytest.mark.parametrize( + "user_id", + ["user-42", "".join(str(i % 10) for i in range(100)), json.dumps({"device_id": "d" * 64}), "{not json"], + ) + def test_metadata_plain_user_id_sets_no_prompt_cache_key(self, user_id: str): + req = _make_request(metadata={"user_id": user_id}) kwargs = _ADAPTER.translate_request(req) - assert kwargs["prompt_cache_key"] == long_id[:64] - assert len(kwargs["prompt_cache_key"]) == 64 + assert kwargs["user"] == user_id[:64] + assert "prompt_cache_key" not in kwargs def test_metadata_empty_user_id_sets_no_prompt_cache_key(self): req = _make_request(metadata={"user_id": ""}) @@ -1207,6 +1300,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] @@ -1219,7 +1324,9 @@ def _make_function_call_item(call_id: str, name: str, arguments: str) -> MagicMo return item -def _make_reasoning_item(summaries: List[str], item_id: str = "rs_test_1") -> MagicMock: +def _make_reasoning_item( + summaries: List[str], item_id: str = "rs_test_1", encrypted_content: str | None = None +) -> MagicMock: """Build a mock ResponseReasoningItem.""" from openai.types.responses import ResponseReasoningItem # type: ignore[import] @@ -1232,9 +1339,13 @@ def _make_reasoning_item(summaries: List[str], item_id: str = "rs_test_1") -> Ma item = MagicMock(spec=ResponseReasoningItem) item.id = item_id item.summary = summary_mocks + item.encrypted_content = encrypted_content return item +_ENCRYPTED_REASONING = "gAAAAABp_encrypted_reasoning_bytes_only_openai_can_read" + + class TestTranslateResponse: """Responses API -> AnthropicMessagesResponse conversion.""" @@ -1265,6 +1376,7 @@ class TestTranslateResponse: assert block["id"] == "call_99" assert block["name"] == "get_weather" assert block["input"] == {"city": "NYC"} + assert "provider_specific_fields" not in block def test_function_call_sets_stop_reason_tool_use(self): """Presence of a function_call sets stop_reason to 'tool_use'.""" @@ -1279,6 +1391,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( @@ -1323,7 +1470,81 @@ class TestTranslateResponse: reasoning = _make_reasoning_item(["Part one.", "Part two."], item_id="rs_abc123") response = _make_mock_response(output=[reasoning]) result: Any = _ADAPTER.translate_response(response) - assert [block["signature"] for block in result["content"]] == [None, None] + assert [block["signature"] for block in result["content"]] == [None] + assert "rs_abc123" not in json.dumps(result["content"]) + + def test_summary_parts_join_into_one_thinking_block(self): + """One reasoning item is one block, so its signature is echoed back exactly once.""" + reasoning = _make_reasoning_item(["Part one.", "Part two."]) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert [block["thinking"] for block in result["content"]] == ["Part one.\n\nPart two."] + + def test_encrypted_content_rides_the_thinking_signature(self): + """Regression for https://github.com/BerriAI/litellm/issues/40288 (outbound fault site).""" + reasoning = _make_reasoning_item(["Part one."], item_id="rs_abc123", encrypted_content=_ENCRYPTED_REASONING) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [ + { + "type": "thinking", + "thinking": "Part one.", + "signature": encrypted_reasoning_signature(_ENCRYPTED_REASONING), + } + ] + + def test_reasoning_without_summary_becomes_redacted_thinking(self): + """With summaries off the encrypted reasoning still has to reach the client to be replayed.""" + reasoning = _make_reasoning_item([], encrypted_content=_ENCRYPTED_REASONING) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature(_ENCRYPTED_REASONING)} + ] + + def test_dict_reasoning_item_carries_its_encrypted_content(self): + response = _make_mock_response( + output=[ + { + "type": "reasoning", + "id": "rs_dict_1", + "encrypted_content": _ENCRYPTED_REASONING, + "summary": [{"type": "summary_text", "text": "Weighing the options."}], + } + ] + ) + result: Any = _ADAPTER.translate_response(response) + assert result["content"][0]["signature"] == encrypted_reasoning_signature(_ENCRYPTED_REASONING) + + def test_reasoning_item_round_trip_is_byte_stable(self): + """Regression for https://github.com/BerriAI/litellm/issues/40288. + + The reasoning item the next turn replays must be the one OpenAI produced, with its + encrypted reasoning intact, and identical on every later turn so the prompt cache + prefix keeps matching. + """ + reasoning = _make_reasoning_item(["Part one.", "Part two."], encrypted_content=_ENCRYPTED_REASONING) + turn: Any = _ADAPTER.translate_response(_make_mock_response(output=[reasoning])) + history = [{"role": "assistant", "content": turn["content"]}] + + replayed_items = [_translate_messages(history) for _ in range(2)] + + assert replayed_items[0] == replayed_items[1] + assert replayed_items[0] == [ + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Part one.\n\nPart two."}], + "encrypted_content": _ENCRYPTED_REASONING, + } + ] + + def test_redacted_reasoning_round_trip_replays_the_encrypted_content(self): + reasoning = _make_reasoning_item([], encrypted_content=_ENCRYPTED_REASONING) + turn: Any = _ADAPTER.translate_response(_make_mock_response(output=[reasoning])) + + replayed = _translate_messages([{"role": "assistant", "content": turn["content"]}]) + + assert replayed == [{"type": "reasoning", "summary": [], "encrypted_content": _ENCRYPTED_REASONING}] def test_dict_reasoning_item_becomes_thinking_block(self): """A reasoning item arriving as a plain dict is kept, not dropped.""" @@ -1337,18 +1558,28 @@ 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): + @pytest.mark.parametrize( + ("summaries", "encrypted_content"), + [ + (["Part one."], None), + (["Part one."], _ENCRYPTED_REASONING), + ([], _ENCRYPTED_REASONING), + ], + ids=["unsigned_thinking", "encrypted_thinking", "encrypted_redacted_thinking"], + ) + def test_thinking_blocks_are_dropped_when_replayed_to_anthropic(self, summaries, encrypted_content): """Replaying this turn to an Anthropic model must not send a signature it cannot verify.""" from litellm.litellm_core_utils.prompt_templates.factory import ( _drop_unsignable_thinking_blocks, ) - response = _make_mock_response(output=[_make_reasoning_item(["Part one."], item_id="rs_abc123")]) + response = _make_mock_response( + output=[_make_reasoning_item(summaries, item_id="rs_abc123", encrypted_content=encrypted_content)] + ) result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 1 assert _drop_unsignable_thinking_blocks(result["content"]) == [] def test_usage_mapped_correctly(self): @@ -1447,6 +1678,7 @@ class TestTranslateResponse: assert result["content"][0]["type"] == "tool_use" assert result["content"][0]["name"] == "search" assert result["content"][0]["input"] == {"query": "cats"} + assert "provider_specific_fields" not in result["content"][0] assert result["stop_reason"] == "tool_use" def test_mixed_reasoning_text_and_tool_use(self): @@ -1481,9 +1713,7 @@ class TestToolResultImages: }, { "role": "user", - "content": [ - {"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content} - ], + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}], }, ] @@ -1630,9 +1860,7 @@ class TestToolResultDocuments: }, { "role": "user", - "content": [ - {"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content} - ], + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}], }, ] @@ -1665,9 +1893,7 @@ class TestToolResultDocuments: def test_document_title_becomes_filename(self): output = self._tool_output(self._translate([self._base64_document(title="quarterly-report.pdf")])) - assert output == [ - {"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI} - ] + assert output == [{"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI}] def test_url_document_becomes_file_url_part(self): output = self._tool_output( @@ -1776,9 +2002,7 @@ class TestUserContentDocuments: def test_document_title_becomes_filename(self): content = self._user_content(self._translate([self._base64_document(title="quarterly-report.pdf")])) - assert content == [ - {"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI} - ] + assert content == [{"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI}] def test_url_document_becomes_file_url_part(self): content = self._user_content( @@ -1808,9 +2032,7 @@ class TestUserContentDocuments: assert content == [{"type": "input_text", "text": "still here"}] def test_document_breakpoint_rides_on_the_file_part(self): - content = self._user_content( - self._translate([self._base64_document(prompt_cache_breakpoint=self.EXPLICIT)]) - ) + content = self._user_content(self._translate([self._base64_document(prompt_cache_breakpoint=self.EXPLICIT)])) assert content == [ { "type": "input_file", @@ -1857,7 +2079,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"] == [ diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 794613942a1..e1b39c4ba13 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -26,6 +26,33 @@ FAKE_REGULAR_KEY = "sk-ant-api03-regular-key-for-testing-123456789" FAKE_AUTH_TOKEN = "sk-ant-aut01-fake-auth-token-for-testing-123456789" +@pytest.mark.parametrize( + "messages,system,expected", + [ + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_is_subagent=true;", True), + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: =junk; cc_is_subagent=true;", False), + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_version=; cc_is_subagent=true;", False), + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: malformed", False), + ([{"content": "missing role"}], "x-anthropic-billing-header: cc_is_subagent=true;", False), + (["not-a-mapping"], "x-anthropic-billing-header: cc_is_subagent=true;", False), + ([{"role": "user", "content": "hi"}], ["not-a-mapping"], False), + ([{"role": "user", "content": "hi"}], None, False), + ], +) +def test_is_claude_code_one_shot_subagent_request(messages, system, expected): + from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request + + assert ( + is_claude_code_one_shot_subagent_request( + messages=messages, + system=system, + tools=None, + user_agent="claude-cli/2.1.263 (external, cli)", + ) + is expected + ) + + class TestOptionallyHandleAnthropicOAuth: """Tests for optionally_handle_anthropic_oauth function.""" @@ -1517,6 +1544,71 @@ class TestAnthropicThinkingSignatureSelfHeal: out = strip_empty_content_blocks_from_anthropic_messages(msgs) assert [b["type"] for b in out[0]["content"]] == ["thinking"] + def test_strip_keeps_encrypted_reasoning_blocks_for_the_responses_bridge(self): + """The /v1/messages handler runs this before dispatch, so the bridge must still see the replay.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, + ) + from litellm.llms.anthropic.common_utils import ( + strip_empty_content_blocks_from_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + {"type": "text", "text": "The answer."}, + ], + } + ] + assert strip_empty_content_blocks_from_anthropic_messages(msgs) == msgs + + def test_strip_encrypted_reasoning_drops_only_the_bridge_tagged_blocks(self): + """A session resumed on an Anthropic model replays reasoning only OpenAI can verify.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, + ) + from litellm.llms.anthropic.common_utils import ( + strip_encrypted_reasoning_blocks_from_anthropic_messages, + ) + + msgs = [ + {"role": "user", "content": "Solve it."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_3")}, + {"type": "thinking", "thinking": "native", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."}, + ], + }, + ] + out = strip_encrypted_reasoning_blocks_from_anthropic_messages(msgs) + assert [m["role"] for m in out] == ["user", "assistant"] + assert [b["type"] for b in out[1]["content"]] == ["thinking", "redacted_thinking", "text"] + assert out[1]["content"][0]["signature"] == "EqQBCkYIAxgCIkA_anthropic_signed" + assert len(msgs[1]["content"]) == 2 + assert len(msgs[2]["content"]) == 4 + + def test_strip_encrypted_reasoning_leaves_malformed_messages_for_the_provider_to_reject(self): + """A bare string in messages must reach Anthropic as a 400, not die in the stripper as a 500.""" + from litellm.llms.anthropic.common_utils import ( + strip_encrypted_reasoning_blocks_from_anthropic_messages, + ) + + msgs = ["hi", {"role": "user", "content": "hello"}] + assert strip_encrypted_reasoning_blocks_from_anthropic_messages(msgs) == msgs + def test_strip_empty_text_blocks_treats_null_text_as_empty(self): from litellm.llms.anthropic.common_utils import ( strip_empty_content_blocks_from_anthropic_messages, @@ -2175,3 +2267,25 @@ def test_create_anthropic_model_list_response_empty(): assert response["has_more"] is False assert response["first_id"] is None assert response["last_id"] is None + + +def test_create_anthropic_model_list_response_lists_ids_as_told(): + """listed_ids renames an entry for the caller while display_name and every other field stay keyed to the served + id, and the envelope's first/last ids follow the renamed entries.""" + from litellm.llms.anthropic.common_utils import ( + create_anthropic_model_list_response, + ) + + response = create_anthropic_model_list_response( + [ + {"id": "gpt-4o", "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": 1000000}, + {"id": "claude-haiku-4-5", "object": "model", "created": 0, "owned_by": "openai"}, + ], + display_names={"gpt-4o": "GPT 4o"}, + listed_ids={"gpt-4o": "claude-router-gpt-4o[1m]"}, + ) + + gpt, haiku = response["data"] + assert (gpt["id"], gpt["display_name"], gpt["max_input_tokens"]) == ("claude-router-gpt-4o[1m]", "GPT 4o", 1000000) + assert (haiku["id"], haiku["display_name"]) == ("claude-haiku-4-5", "claude-haiku-4-5") + assert (response["first_id"], response["last_id"]) == ("claude-router-gpt-4o[1m]", "claude-haiku-4-5") diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py index 4e6b9ed0188..bc6cb0c0fed 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -198,6 +198,9 @@ def test_azure_gpt_5_takes_the_reasoning_path() -> None: assert "reasoning_effort" in supported +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + + class TestAzureToolSchemaCombinatorFlattening: """ Regression tests for LIT-6510: Azure's chat completions validator rejects @@ -259,6 +262,26 @@ class TestAzureToolSchemaCombinatorFlattening: self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) assert tool == self._anyof_tool() + def test_transform_request_drops_non_python_regex_pattern(self): + tool = { + "type": "function", + "function": { + "name": "Artifact", + "parameters": { + "type": "object", + "properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}, + }, + }, + } + + request = self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) + + assert request["tools"][0]["function"]["parameters"] == { + "type": "object", + "properties": {"field": {"type": "string"}}, + } + assert tool["function"]["parameters"]["properties"]["field"]["pattern"] == _ARTIFACT_FIELD_PATTERN + def test_clean_object_schema_passes_through_as_same_object(self): tool = { "type": "function", diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 29b74c2ee4a..c7e86616ee2 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -1,11 +1,20 @@ import json +from datetime import datetime from unittest.mock import MagicMock import httpx +import pytest - -from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig -from litellm.types.utils import ModelResponse +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound +from litellm.llms.azure.passthrough.transformation import ( + AzurePassthroughConfig, + azure_router_model_in_endpoint, + foreign_azure_deployment, +) +from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse +from litellm.types.utils import EmbeddingResponse, ModelResponse def _azure_chat_completion_body(): @@ -73,22 +82,408 @@ def test_azure_passthrough_logging_non_streaming_response_chat_completions(): assert result.usage.total_tokens == 18 -def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_returns_none(): - """ - Endpoints other than chat/completions (responses, messages, images) fall - through to None — matches base-class behavior and Bedrock's "unknown - endpoint" handling. Not a regression; just scoping. - """ - config = AzurePassthroughConfig() - logging_obj = MagicMock() - - result = config.logging_non_streaming_response( - model="gpt-4.1-mini", +def _relay_logging_obj(model: str) -> Logging: + logging_obj = Logging( + model=model, + messages=[], + stream=False, + call_type="allm_passthrough_route", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + model=model, + litellm_params={"api_base": "https://my-resource.openai.azure.com", "custom_llm_provider": "azure"}, + optional_params={}, custom_llm_provider="azure", - httpx_response=_make_httpx_response(_azure_chat_completion_body()), + ) + return logging_obj + + +def _relay_logging_result(model: str, endpoint: str, body, status_code: int = 200): + logging_obj = _relay_logging_obj(model) + response = httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request( + "POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview" + ), + ) + result = AzurePassthroughConfig().logging_non_streaming_response( + model=model, + custom_llm_provider="azure", + httpx_response=response, request_data={}, logging_obj=logging_obj, - endpoint="openai/responses", + endpoint=endpoint, + ) + return result, logging_obj + + +EMBEDDINGS_BODY = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1000, "total_tokens": 1000}, +} + +RESPONSES_BODY = { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4.1-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, +} + + +def test_azure_passthrough_embeddings_relay_is_costed_per_input_token(): + result, logging_obj = _relay_logging_result( + "text-embedding-3-small", "openai/deployments/text-embedding-3-small/embeddings", EMBEDDINGS_BODY + ) + per_token = litellm.get_model_info("azure/text-embedding-3-small")["input_cost_per_token"] + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(1000 * per_token) + + +def test_azure_passthrough_responses_relay_is_costed_per_token(): + result, logging_obj = _relay_logging_result("gpt-4.1-mini", "openai/responses", RESPONSES_BODY) + info = litellm.get_model_info("azure/gpt-4.1-mini") + + assert isinstance(result, ResponsesAPIResponse) + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=result) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +def test_azure_passthrough_failed_embeddings_relay_is_not_costed(): + result, logging_obj = _relay_logging_result( + "text-embedding-3-small", + "openai/deployments/text-embedding-3-small/embeddings", + {"error": {"code": "429", "message": "rate limited"}}, + status_code=429, ) assert result is None + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_returns_none(): + result, logging_obj = _relay_logging_result( + "gpt-4o-mini-tts", "openai/deployments/gpt-4o-mini-tts/audio/speech", {"audio": "..."} + ) + + assert result is None + assert logging_obj.call_type == "allm_passthrough_route" + + +def _sse_line(payload: dict) -> str: + return "data: " + json.dumps(payload) + + +def _azure_chat_completion_chunks() -> list[str]: + head = {"id": "chatcmpl-abc123", "object": "chat.completion.chunk", "created": 1700000000, "model": "gpt-4.1-mini"} + return [ + _sse_line( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}], + } + ), + _sse_line( + {**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]} + ), + _sse_line({**head, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}), + _sse_line({**head, "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}}), + "data: [DONE]", + ] + + +def test_azure_passthrough_streaming_chat_chunks_build_the_complete_response(): + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_chat_completion_chunks(), + litellm_logging_obj=MagicMock(), + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Hello! How can I assist?" + assert response.usage.prompt_tokens == 10 + assert response.usage.completion_tokens == 8 + + +def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_from_the_relayed_request(): + messages = [{"role": "user", "content": "Say hi in three words"}] + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}} + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk], + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Hello! How can I assist?" + assert response.usage.prompt_tokens > 0 + assert response.usage.prompt_tokens == litellm.token_counter(model="gpt-4.1-mini", messages=messages) + assert response.usage.completion_tokens > 0 + + +def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_without_fetching_the_image(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": "http://127.0.0.1:9/doc.png", "detail": "high"}}, + ], + } + ] + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}} + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk], + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + text_only_messages = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] + assert isinstance(response, ModelResponse) + assert response.usage.prompt_tokens == ( + litellm.token_counter(model="gpt-4.1-mini", messages=text_only_messages) + high_detail_image_token_upper_bound() + ) + + +def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_chat_completion_chunks(), + litellm_logging_obj=MagicMock(), + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/embeddings", + ) + + assert response is None + + +def _azure_responses_stream_chunks(terminal_event: str | None = "response.completed") -> list[str]: + in_progress = {**RESPONSES_BODY, "status": "in_progress", "output": [], "usage": None} + events = [ + ("response.created", {"type": "response.created", "sequence_number": 0, "response": in_progress}), + ( + "response.output_text.delta", + {"type": "response.output_text.delta", "sequence_number": 1, "item_id": "msg_1", "delta": "hi"}, + ), + ] + ( + [(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] + if terminal_event + else [] + ) + return [line for name, payload in events for line in (f"event: {name}", _sse_line(payload))] + + +def test_azure_passthrough_streaming_responses_chunks_are_costed_per_token(): + logging_obj = _relay_logging_obj("gpt-4.1-mini") + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_responses_stream_chunks(), + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/responses", + ) + info = litellm.get_model_info("azure/gpt-4.1-mini") + + assert isinstance(response, ResponseCompletedEvent) + assert response.response.usage.input_tokens == 1000 + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +def test_azure_passthrough_streaming_responses_without_a_terminal_event_are_not_costed(): + logging_obj = _relay_logging_obj("gpt-4.1-mini") + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_responses_stream_chunks(terminal_event=None), + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/responses", + ) + + assert response is None + assert logging_obj.call_type == "allm_passthrough_route" + + +def _complete_url(request_query_params: dict, litellm_params: dict) -> httpx.URL: + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params=request_query_params, + litellm_params=litellm_params, + ) + return url + + +def test_azure_passthrough_url_forwards_the_callers_api_version(): + url = _complete_url(request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={}) + + assert url.path == "/openai/deployments/gpt-4.1-mini/chat/completions" + assert url.params["api-version"] == "2025-04-01-preview" + + +def test_azure_passthrough_url_prefers_the_callers_api_version_over_the_deployments(): + url = _complete_url( + request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={"api_version": "2024-10-21"} + ) + + assert url.params["api-version"] == "2025-04-01-preview" + + +def test_azure_passthrough_url_fills_in_the_deployments_api_version_when_the_caller_sends_none(): + url = _complete_url(request_query_params={}, litellm_params={"api_version": "2024-10-21"}) + + assert url.params["api-version"] == "2024-10-21" + + +FULL_URL_API_BASE = ( + "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" +) + + +def _full_url_complete_url(request_query_params: dict) -> httpx.URL: + url, _ = AzurePassthroughConfig().get_complete_url( + api_base=FULL_URL_API_BASE, + api_key="key", + model="gpt-4.1-mini", + endpoint="chat/completions", + request_query_params=request_query_params, + litellm_params={}, + ) + return url + + +def test_azure_passthrough_url_prefers_the_callers_api_version_over_a_full_url_api_bases(): + url = _full_url_complete_url(request_query_params={"api-version": "2025-04-01-preview"}) + + assert str(url) == ( + "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions" + "?api-version=2025-04-01-preview" + ) + + +def test_azure_passthrough_url_keeps_a_full_url_api_bases_api_version_when_the_caller_sends_none(): + url = _full_url_complete_url(request_query_params={}) + + assert url.params["api-version"] == "2024-10-21" + + +def test_azure_passthrough_url_strips_the_leading_router_model_segment(): + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="gpt-4.1-mini/openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={}, + ) + + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) + + +def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment(): + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="gpt/openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "gpt"}}, + ) + + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) + + +@pytest.mark.parametrize( + "request_data, expected", + [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], +) +def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_data, expected): + assert ( + AzurePassthroughConfig().is_streaming_request( + endpoint="openai/deployments/x/chat/completions", request_data=request_data + ) + is expected + ) + + +@pytest.mark.parametrize( + "endpoint, expected", + [ + ("gpt/openai/deployments/gpt/chat/completions", None), + ("openai/deployments/gpt/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), + ("gpt/openai/deployments/GPT-5.4-MINI/chat/completions", None), + ("gpt/openai/deployments/Gpt/chat/completions", "Gpt"), + ("gpt/models/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), + ("gpt/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/victim/gpt/chat/completions", "victim"), + ("gpt/openai/deployments/GPT-5.4/chat/completions", "GPT-5.4"), + ], +) +def test_foreign_azure_deployment_names_a_segment_outside_the_group(endpoint, expected): + assert foreign_azure_deployment(endpoint, "gpt", lambda: frozenset({"gpt-5.4-mini"})) == expected + + +def test_foreign_azure_deployment_skips_the_router_when_the_segment_is_the_group_itself(): + def served_models(): + raise AssertionError("the router must not be consulted for the group's own name") + + assert foreign_azure_deployment("gpt/openai/deployments/gpt/chat/completions", "gpt", served_models) is None + + +@pytest.mark.parametrize( + "endpoint, expected", + [ + ("other-group/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/gpt/chat/completions", "gpt"), + ("openai/deployments/my-azure-deployment/chat/completions", None), + ("gpt", None), + ], +) +def test_azure_router_model_in_endpoint_picks_the_first_router_model_segment(endpoint, expected): + assert azure_router_model_in_endpoint(endpoint, frozenset({"gpt", "other-group"})) == expected diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/test_litellm/llms/azure/response/test_azure_transformation.py index 0cac2705ab0..726c9f65681 100644 --- a/tests/test_litellm/llms/azure/response/test_azure_transformation.py +++ b/tests/test_litellm/llms/azure/response/test_azure_transformation.py @@ -1,4 +1,5 @@ from copy import deepcopy +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -243,6 +244,9 @@ def test_provider_config_manager_o_series_selection(): assert not isinstance(default_config, AzureOpenAIOSeriesResponsesAPIConfig) +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + + class TestAzureResponsesAPIConfig: def setup_method(self): self.config = AzureOpenAIResponsesAPIConfig() @@ -599,6 +603,31 @@ class TestAzureResponsesAPIConfig: assert result["tools"][0] is tool assert "anyOf" in result["tools"][0]["parameters"] + def test_azure_drops_non_python_regex_pattern_while_keeping_gpt5_combinators(self): + tool = { + "type": "function", + "name": "Artifact", + "parameters": { + "type": "object", + "anyOf": [{"properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}}], + "properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}, + }, + } + + result = self.config.transform_responses_api_request( + model="my-eastus-deployment", + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(model_info={"base_model": "azure/gpt-5.4-mini"}), + headers={}, + ) + + assert result["tools"][0]["parameters"] == { + "type": "object", + "anyOf": [{"properties": {"field": {"type": "string"}}}], + "properties": {"field": {"type": "string"}}, + } + def test_azure_keeps_combinators_for_unrecognized_deployment_without_base_model(self): tool = self._anyof_tool() diff --git a/tests/test_litellm/llms/azure/test_audio_transcriptions.py b/tests/test_litellm/llms/azure/test_audio_transcriptions.py new file mode 100644 index 00000000000..cd5fcbd85a9 --- /dev/null +++ b/tests/test_litellm/llms/azure/test_audio_transcriptions.py @@ -0,0 +1,61 @@ +import json +from pathlib import Path +from typing import Final + +import httpx +import pytest +from openai import AzureOpenAI + +import litellm +from litellm.cost_calculator import completion_cost +from litellm.litellm_core_utils.audio_utils.utils import calculate_request_duration + +AUDIO_FILE: Final = Path(__file__).parents[3] / "gettysburg.wav" +WHISPER_COST_PER_SECOND: Final = 0.0001 + + +def _transcription_client() -> AzureOpenAI: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"text": "Four score and seven years ago"}) + + return AzureOpenAI( + api_key="test-key", + api_version="2024-06-01", + azure_endpoint="https://example.cognitiveservices.azure.com", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + + +def test_azure_ai_transcription_is_priced_at_the_azure_ai_entry(): + with AUDIO_FILE.open("rb") as audio: + response = litellm.transcription( + model="azure_ai/whisper", + file=audio, + api_base="https://example.cognitiveservices.azure.com", + api_key="test-key", + api_version="2024-06-01", + client=_transcription_client(), + ) + with AUDIO_FILE.open("rb") as audio: + duration = calculate_request_duration(audio) + + assert duration is not None and duration > 0 + assert response._hidden_params["custom_llm_provider"] == "azure_ai" + assert completion_cost(completion_response=response, call_type="transcription") == pytest.approx( + WHISPER_COST_PER_SECOND * duration + ) + + +def test_azure_transcription_keeps_the_azure_provider(): + with AUDIO_FILE.open("rb") as audio: + response = litellm.transcription( + model="azure/whisper-1", + file=audio, + api_base="https://example.openai.azure.com", + api_key="test-key", + api_version="2024-06-01", + client=_transcription_client(), + ) + + assert response._hidden_params["custom_llm_provider"] == "azure" + assert json.loads(response.model_dump_json())["text"] == "Four score and seven years ago" diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index f000abb4c9a..c959c201ccb 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -385,6 +385,58 @@ def test_select_azure_base_url_called(setup_mocks): setup_mocks["select_url"].assert_called_once() +def test_initialize_defaults_max_retries_to_litellm_default(setup_mocks): + result = BaseAzureLLM().initialize_azure_sdk_client( + litellm_params={}, + api_key="test-api-key", + api_base="https://test.openai.azure.com", + model_name="gpt-4", + api_version="2023-06-01", + is_async=False, + ) + + assert result["max_retries"] == litellm.constants.DEFAULT_MAX_RETRIES + + +@pytest.mark.parametrize( + "configured, expected", + [(0, 0), (5, 5), (None, litellm.constants.DEFAULT_MAX_RETRIES)], +) +def test_initialize_honors_explicit_max_retries(setup_mocks, configured, expected): + result = BaseAzureLLM().initialize_azure_sdk_client( + litellm_params={"max_retries": configured}, + api_key="test-api-key", + api_base="https://test.openai.azure.com", + model_name="gpt-4", + api_version="2023-06-01", + is_async=False, + ) + + assert result["max_retries"] == expected + + +def test_default_max_retries_env_var_reaches_azure_sdk_client(): + import subprocess + import sys + + code = ( + "from litellm.llms.azure.common_utils import BaseAzureLLM\n" + "client = BaseAzureLLM().get_azure_openai_client(" + "api_key='test-api-key', api_base='https://test.openai.azure.com', api_version='2024-02-01'," + " client=None, _is_async=True, litellm_params={}, model='gpt-4')\n" + "print(client.max_retries)" + ) + completed = subprocess.run( + [sys.executable, "-c", code], + env={**os.environ, "DEFAULT_MAX_RETRIES": "0"}, + capture_output=True, + text=True, + check=True, + ) + + assert completed.stdout.strip() == "0" + + @pytest.mark.parametrize( "call_type", [ diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 33fbb4e8fc7..f8cc0b5071e 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -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. diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py index 284a912d9a4..75e046825a3 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -6,6 +6,7 @@ import pytest import litellm +from litellm.images.utils import ImageEditRequestUtils from litellm.llms.azure_ai.image_edit import ( AzureFoundryMAIImageEditConfig, get_azure_ai_image_edit_config, @@ -70,44 +71,48 @@ class TestAzureMAIImageEdit: assert "/mai/v1/images/edits" in url assert "api-version=preview" in url - def test_map_openai_params_keeps_size(self): - config = AzureFoundryMAIImageEditConfig() - optional_params = config.map_openai_params( - image_edit_optional_params={"size": "1792x1024", "n": 1}, + def test_get_optional_params_image_edit_size_raises_400(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError, match="size") as exc_info: + ImageEditRequestUtils.get_optional_params_image_edit( + model="MAI-Image-2.5", + image_edit_provider_config=AzureFoundryMAIImageEditConfig(), + image_edit_optional_params={"size": "1024x1024", "n": 1}, + ) + assert exc_info.value.status_code == 400 + + def test_get_optional_params_image_edit_size_dropped_with_drop_params(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + optional_params = ImageEditRequestUtils.get_optional_params_image_edit( model="MAI-Image-2.5", + image_edit_provider_config=AzureFoundryMAIImageEditConfig(), + image_edit_optional_params={"size": "1024x1024", "n": 1}, drop_params=True, ) - assert optional_params["size"] == "1792x1024" + assert "size" not in optional_params assert optional_params["n"] == 1 - assert "width" not in optional_params - assert "height" not in optional_params - def test_map_openai_params_defaults_size(self): - config = AzureFoundryMAIImageEditConfig() - optional_params = config.map_openai_params( - image_edit_optional_params={}, + def test_get_optional_params_image_edit_without_size_forwards_nothing_extra(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + optional_params = ImageEditRequestUtils.get_optional_params_image_edit( model="MAI-Image-2.5", - drop_params=True, + image_edit_provider_config=AzureFoundryMAIImageEditConfig(), + image_edit_optional_params={}, ) - assert optional_params["size"] == "1024x1024" + assert optional_params == {} - def test_map_openai_params_unsupported_size_raises(self): - config = AzureFoundryMAIImageEditConfig() - with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): - config.map_openai_params( - image_edit_optional_params={"size": "auto"}, - model="MAI-Image-2.5", - drop_params=True, - ) - - def test_map_openai_params_invalid_size_format_raises(self): - config = AzureFoundryMAIImageEditConfig() - with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): - config.map_openai_params( - image_edit_optional_params={"size": "1024xabc"}, - model="MAI-Image-2.5", - drop_params=True, + def test_image_edit_size_surfaces_as_400(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.image_edit( + model="azure_ai/MAI-Image-2.5", + image=io.BytesIO(b"fake-image-bytes"), + prompt="Turn this into a studio product shot", + size="1024x1024", + api_key="test-key", + api_base="https://my-resource.services.ai.azure.com", ) + assert exc_info.value.status_code == 400 def test_transform_image_edit_request_uses_image_field(self): config = AzureFoundryMAIImageEditConfig() @@ -117,14 +122,14 @@ class TestAzureMAIImageEdit: model="MAI-Image-2.5", prompt="Turn this into a studio product shot", image=image_bytes, - image_edit_optional_request_params={"size": "1024x1024", "n": 1}, + image_edit_optional_request_params={"n": 1}, litellm_params={}, headers={}, ) assert data["model"] == "MAI-Image-2.5" assert data["prompt"] == "Turn this into a studio product shot" - assert data["size"] == "1024x1024" + assert "size" not in data assert data["n"] == 1 assert len(files) == 1 assert files[0][0] == "image" diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index 2a44e77ce09..55656b97c57 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -1,11 +1,10 @@ -import os from unittest.mock import MagicMock import httpx import pytest - import litellm +from litellm.exceptions import UnsupportedParamsError from litellm.llms.azure.azure import AzureChatCompletion from litellm.llms.azure.image_generation import get_azure_image_generation_config from litellm.llms.azure.image_generation.http_utils import ( @@ -30,48 +29,21 @@ from litellm.utils import get_optional_params_image_gen class TestAzureMAIImageGeneration: def test_is_mai_model(self): assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5") - assert AzureFoundryMAIImageGenerationConfig.is_mai_model( - "azure_ai/MAI-Image-2.5" - ) + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("azure_ai/MAI-Image-2.5") assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5-Flash") assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2e") assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1") - def test_mai_flash_and_2e_model_pricing_in_cost_map(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - - flash_info = litellm.get_model_info( - model="azure_ai/MAI-Image-2.5-Flash", - custom_llm_provider="azure_ai", - ) - assert flash_info["input_cost_per_token"] == 1.75e-06 - assert flash_info["input_cost_per_image_token"] == 1.75e-06 - assert flash_info["output_cost_per_image_token"] == 3.3e-05 - - image_2e_info = litellm.get_model_info( - model="azure_ai/MAI-Image-2e", - custom_llm_provider="azure_ai", - ) - assert image_2e_info["input_cost_per_token"] == 5e-06 - assert image_2e_info["output_cost_per_image_token"] == 1.95e-05 - def test_get_mai_image_generation_url(self): url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( api_base="https://my-resource.services.ai.azure.com", api_version="preview", ) - assert ( - url - == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" - ) + assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" def test_get_mai_image_generation_url_preserves_full_path(self): - api = ( - "https://my-resource.services.ai.azure.com/mai/v1/images/generations" - "?api-version=preview" - ) + api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( api_base=api, api_version="preview", @@ -83,10 +55,7 @@ class TestAzureMAIImageGeneration: api_base="https://my-resource.services.ai.azure.com/mai/v1", api_version="preview", ) - assert ( - url - == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" - ) + assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" def test_get_azure_ai_image_generation_config_returns_mai(self): config = get_azure_ai_image_generation_config("MAI-Image-2.5") @@ -124,13 +93,13 @@ class TestAzureMAIImageGeneration: config = AzureFoundryMAIImageGenerationConfig() optional_params = get_optional_params_image_gen( model="MAI-Image-2.5", - size="1792x1024", + size="1024x1024", n=1, custom_llm_provider="azure_ai", provider_config=config, drop_params=True, ) - assert optional_params["width"] == 1792 + assert optional_params["width"] == 1024 assert optional_params["height"] == 1024 assert "size" not in optional_params @@ -147,10 +116,7 @@ class TestAzureMAIImageGeneration: assert "api-version=preview" in url def test_mai_json_body_keeps_model(self): - api = ( - "https://my-resource.services.ai.azure.com/mai/v1/images/generations" - "?api-version=preview" - ) + api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" data = { "model": "MAI-Image-2.5", "prompt": "A photograph of a red fox", @@ -196,7 +162,7 @@ class TestAzureMAIImageGeneration: def test_map_openai_params_unsupported_size_raises(self): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): + with pytest.raises(UnsupportedParamsError, match="Unsupported size value: 'auto'"): config.map_openai_params( non_default_params={"size": "auto"}, optional_params={}, @@ -206,7 +172,7 @@ class TestAzureMAIImageGeneration: def test_map_openai_params_invalid_custom_size_raises(self): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): + with pytest.raises(UnsupportedParamsError, match="Invalid size format: '1024xabc'"): config.map_openai_params( non_default_params={"size": "1024xabc"}, optional_params={}, @@ -214,9 +180,138 @@ class TestAzureMAIImageGeneration: drop_params=True, ) + @pytest.mark.parametrize("size", ["512x512", "256x256", "700x1400"]) + def test_map_openai_params_size_below_minimum_dimension_raises(self, size): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="at least 768 pixels"): + config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + @pytest.mark.parametrize("size", ["1792x1024", "1024x1792"]) + def test_map_openai_params_size_over_total_pixel_budget_raises(self, size): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"): + config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + @pytest.mark.parametrize("size", ["1032x1024", "1376x768"]) + def test_map_openai_params_size_at_live_pixel_cap_passes_through(self, size): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert optional_params["width"] * optional_params["height"] == 1_056_768 + + def test_map_openai_params_size_one_pixel_over_live_cap_raises(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"): + config.map_openai_params( + non_default_params={"size": "1033x1024"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + + def test_map_openai_params_explicit_width_height_not_range_checked(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"width": 1792, "height": 1024}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1792 + assert optional_params["height"] == 1024 + + @pytest.mark.parametrize("n", [2, 4, "2", 0, -1]) + def test_map_openai_params_n_other_than_one_raises(self, n): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"): + config.map_openai_params( + non_default_params={"n": n}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + + def test_map_openai_params_non_numeric_n_raises_400(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="not a whole number of images") as exc_info: + config.map_openai_params( + non_default_params={"n": "abc"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert exc_info.value.status_code == 400 + + def test_get_optional_params_image_gen_global_drop_params_drops_multi_image_n(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", True) + optional_params = get_optional_params_image_gen( + model="MAI-Image-2.5", + n=4, + custom_llm_provider="azure_ai", + provider_config=AzureFoundryMAIImageGenerationConfig(), + ) + assert "n" not in optional_params + assert optional_params["width"] == 1024 + + def test_get_optional_params_image_gen_without_any_drop_params_still_raises(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"): + get_optional_params_image_gen( + model="MAI-Image-2.5", + n=4, + custom_llm_provider="azure_ai", + provider_config=AzureFoundryMAIImageGenerationConfig(), + ) + + def test_map_openai_params_multi_image_n_dropped_with_drop_params(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"n": 4}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert "n" not in optional_params + + def test_map_openai_params_single_image_n_still_passes_through(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"n": 1}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert optional_params["n"] == 1 + + @pytest.mark.parametrize("params", [{"n": 2}, {"n": "abc"}, {"size": "512x512"}, {"size": "1792x1024"}]) + def test_image_generation_rejected_params_surface_as_400(self, params): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.image_generation( + model="azure_ai/MAI-Image-2.5", + prompt="A photograph of a red fox", + api_key="test-key", + api_base="https://my-resource.services.ai.azure.com", + **params, + ) + assert exc_info.value.status_code == 400 + def test_map_openai_params_unsupported_param_raises(self): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="Parameter quality is not supported"): + with pytest.raises(UnsupportedParamsError, match="Parameter quality is not supported"): config.map_openai_params( non_default_params={"quality": "hd"}, optional_params={}, @@ -363,16 +458,12 @@ class TestAzureMAIImageGeneration: litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure_ai/MAI-Image-2.5" model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") - image_response = ImageResponse( - data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] - ) + image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]) cost = azure_ai_image_cost_calculator( model=model, image_response=image_response, ) - assert ( - cost == len(image_response.data or []) * model_info["output_cost_per_image"] - ) + assert cost == len(image_response.data or []) * model_info["output_cost_per_image"] assert cost > 0 diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py new file mode 100644 index 00000000000..f00698a6624 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -0,0 +1,619 @@ +import json +from datetime import datetime +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.types.rerank import RerankResponse +from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, ModelResponse +from litellm.utils import ProviderConfigManager + +FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" +RESPONSES_COMPLETED_EVENT = { + "type": "response.completed", + "sequence_number": 2, + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + }, +} + + +class _SpendProbe(CustomLogger): + logged_call_type: str | None = None + logged_cost: float | None = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_call_type = kwargs["call_type"] + self.logged_cost = kwargs["response_cost"] + + +@pytest.fixture(autouse=True) +def clear_azure_ai_env(monkeypatch): + for env_var in ("AZURE_AI_API_BASE", "AZURE_AI_API_KEY", "AZURE_AD_TOKEN", "AZURE_API_KEY"): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setattr(litellm, "api_key", None) + + +def test_provider_config_manager_resolves_azure_ai_passthrough_config(): + config = ProviderConfigManager.get_provider_passthrough_config( + model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI + ) + + assert isinstance(config, AzureAIPassthroughConfig) + + +def test_router_model_prefix_is_stripped_and_native_path_kept_verbatim(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key=None, + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert base == FOUNDRY_BASE + + +def test_model_group_prefix_is_stripped_when_router_metadata_names_it(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key=None, + model="Cohere-parse-v5", + endpoint="/parse-alias/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={"litellm_metadata": {"model_group": "parse-alias"}}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + + +def test_model_inside_the_path_stays_and_query_params_are_forwarded(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/", + api_key=None, + model="gpt-5.4-mini", + endpoint="openai/deployments/gpt-5.4-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/openai/deployments/gpt-5.4-mini/chat/completions?api-version=2024-10-21" + + +def test_api_base_that_already_ends_in_models_is_cut_back_to_the_foundry_root(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models", + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params={"api-version": "2024-05-01-preview"}, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + assert base == FOUNDRY_BASE + + +def test_full_url_api_base_that_already_ends_with_the_native_path_is_not_doubled(): + model_router_url = ( + "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + ) + + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{model_router_url}?api-version=2025-01-01-preview", + api_key="key", + model="model_router/model-router", + endpoint="model-router/chat/completions", + request_query_params=None, + litellm_params={"litellm_metadata": {"model_group": "model-router"}}, + ) + + assert str(url) == f"{model_router_url}?api-version=2025-01-01-preview" + assert base == "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router" + + +@pytest.mark.parametrize("relayed_deployment", ["gpt-4o", "GPT-4o"]) +def test_deployment_root_api_base_is_not_repeated_when_the_relay_carries_the_deployment_path(relayed_deployment): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com/openai/deployments/gpt-4o", + api_key="key", + model="gpt-4o", + endpoint=f"aoai-gpt-4o/openai/deployments/{relayed_deployment}/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "aoai-gpt-4o"}}, + ) + + assert str(url) == ( + f"https://my-resource.openai.azure.com/openai/deployments/{relayed_deployment}/chat/completions" + "?api-version=2024-10-21" + ) + assert base == "https://my-resource.openai.azure.com" + + +def test_deployment_named_like_the_first_native_segment_keeps_its_deployment_root(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com/openai/deployments/chat", + api_key="key", + model="chat", + endpoint="aoai-chat/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "aoai-chat"}}, + ) + + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/chat/chat/completions?api-version=2024-10-21" + assert base == "https://my-resource.openai.azure.com/openai/deployments/chat" + + +def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models", + api_key="key", + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + + +def test_deployment_api_version_fills_in_when_the_caller_sends_none(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params=None, + litellm_params={"api_version": "2024-05-01-preview"}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + + +def test_callers_api_version_beats_the_deployments(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params={"api-version": "2025-04-01-preview"}, + litellm_params={"api_version": "2024-05-01-preview"}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2025-04-01-preview" + + +def test_api_version_on_the_configured_api_base_is_the_last_fallback(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview", + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + + +def test_missing_api_base_raises_instead_of_building_a_relative_url(): + with pytest.raises(ValueError, match="AZURE_AI_API_BASE"): + AzureAIPassthroughConfig().get_complete_url( + api_base=None, + api_key=None, + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + +def _auth_headers(api_key: str | None, api_base: str, litellm_params: dict | None = None) -> dict: + return AzureAIPassthroughConfig().validate_environment( + headers={"content-type": "application/json"}, + model="Cohere-parse-v5", + messages=[], + optional_params={}, + litellm_params=litellm_params or {}, + api_key=api_key, + api_base=api_base, + ) + + +def test_foundry_host_gets_the_api_key_header(): + headers = _auth_headers(api_key="deployment-key", api_base=FOUNDRY_BASE) + + assert headers == {"content-type": "application/json", "api-key": "deployment-key"} + + +def test_serverless_host_gets_a_bearer_token(): + headers = _auth_headers(api_key="deployment-key", api_base="https://cohere-parse.eastus.models.ai.azure.com") + + assert headers["Authorization"] == "Bearer deployment-key" + assert "api-key" not in headers + + +def test_entra_token_is_used_when_the_deployment_has_no_api_key(): + headers = _auth_headers(api_key=None, api_base=FOUNDRY_BASE, litellm_params={"azure_ad_token": "entra-token"}) + + assert headers["Authorization"] == "Bearer entra-token" + + +def test_no_credentials_at_all_raises(): + with pytest.raises(ValueError, match="Missing Azure AI credentials"): + _auth_headers(api_key=None, api_base=FOUNDRY_BASE) + + +@pytest.mark.parametrize( + "request_data, expected", + [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], +) +def test_is_streaming_request_reads_the_stream_flag(request_data, expected): + assert ( + AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) + is expected + ) + + +def _chat_completion_response() -> httpx.Response: + body = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, + } + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", f"{FOUNDRY_BASE}/models/chat/completions"), + ) + + +def test_chat_completions_relay_yields_a_model_response_for_cost_tracking(): + result = AzureAIPassthroughConfig().logging_non_streaming_response( + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + httpx_response=_chat_completion_response(), + request_data={"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "hi"}]}, + logging_obj=MagicMock(), + endpoint="models/chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "hi" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 8 + + +def _non_chat_logging_result(content: bytes, content_type: str): + parse_response = httpx.Response( + status_code=200, + headers={"content-type": content_type}, + content=content, + request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), + ) + return AzureAIPassthroughConfig().logging_non_streaming_response( + model="Cohere-parse-v5", + custom_llm_provider="azure_ai", + httpx_response=parse_response, + request_data={"model": "Cohere-parse-v5"}, + logging_obj=MagicMock(), + endpoint="providers/cohere/v2/parse", + ) + + +def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text(): + assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"} + + +def _relay_logging_obj( + model: str, + api_base: str, + stream: bool = False, + callbacks: list[CustomLogger] | None = None, + endpoint: str = "", +) -> Logging: + logging_obj = Logging( + model=model, + messages=[], + stream=stream, + call_type="allm_passthrough_route", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + dynamic_async_success_callbacks=callbacks, + ) + logging_obj.update_environment_variables( + model=model, + litellm_params={"api_base": api_base, "custom_llm_provider": "azure_ai"}, + optional_params={}, + custom_llm_provider="azure_ai", + endpoint=endpoint, + ) + return logging_obj + + +def _relay_logging_result( + config: AzureAIPassthroughConfig, + model: str, + native_path: str, + body, + api_base: str = FOUNDRY_BASE, + status_code: int = 200, +): + relayed_url = f"{FOUNDRY_BASE}/{native_path}?api-version=2024-05-01-preview" + logging_obj = _relay_logging_obj(model, api_base) + response = httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", relayed_url), + ) + result = config.logging_non_streaming_response( + model=model, + custom_llm_provider="azure_ai", + httpx_response=response, + request_data={"model": model}, + logging_obj=logging_obj, + endpoint=f"{model}/{native_path}", + ) + return result, logging_obj + + +MISTRAL_OCR_BODY = { + "pages": [{"index": 0, "markdown": "page one"}, {"index": 1, "markdown": "page two"}], + "model": "mistral-document-ai-2512", + "usage_info": {"pages_processed": 2, "doc_size_bytes": 4321}, +} + + +def test_mistral_document_ai_relay_is_costed_per_page(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY + ) + per_page = litellm.get_model_info("azure_ai/mistral-document-ai-2512")["ocr_cost_per_page"] + + assert isinstance(result, OCRResponse) + assert result.usage_info.pages_processed == 2 + assert per_page > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(2 * per_page) + + +def test_ocr_route_under_a_models_api_base_is_still_recognised(): + result, _ = _relay_logging_result( + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + MISTRAL_OCR_BODY, + api_base=f"{FOUNDRY_BASE}/models", + ) + + assert isinstance(result, OCRResponse) + + +def test_relay_to_a_non_ocr_route_keeps_the_passthrough_object_and_call_type(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "models/info", {"name": "mistral-document-ai-2512"} + ) + + assert result == {"response": {"name": "mistral-document-ai-2512"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +COHERE_PARSE_BODY = {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 3}}} + + +def test_cohere_parse_relay_is_costed_per_billed_page(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "Cohere-parse-v5", "providers/cohere/v2/parse", COHERE_PARSE_BODY + ) + per_page = litellm.get_model_info("azure_ai/Cohere-parse-v5")["ocr_cost_per_page"] + + assert isinstance(result, OCRResponse) + assert result.usage_info.pages_processed == 3 + assert logging_obj.call_type == "aocr" + assert per_page > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(3 * per_page) + + +def test_deployment_without_an_ocr_config_is_never_costed_as_ocr(): + config = AzureAIPassthroughConfig(ocr_config_for=lambda model: None) + result, logging_obj = _relay_logging_result( + config, "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY + ) + + assert result == {"response": MISTRAL_OCR_BODY} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_accepted_ocr_job_without_a_result_body_is_not_costed(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + {"status": "running"}, + status_code=202, + ) + + assert result == {"response": {"status": "running"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_unparseable_ocr_body_falls_back_to_the_passthrough_object(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + ["not", "an", "ocr", "body"], + ) + + assert result == {"response": '["not", "an", "ocr", "body"]'} + assert logging_obj.call_type == "allm_passthrough_route" + + +EMBEDDINGS_BODY = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "embed-v-4-0", + "usage": {"prompt_tokens": 1200, "total_tokens": 1200}, +} + +RERANK_BODY = { + "id": "rerank-1", + "results": [{"index": 1, "relevance_score": 0.9}, {"index": 0, "relevance_score": 0.2}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 2}}, +} + +IMAGE_BODY = {"created": 1, "data": [{"b64_json": "AAAA"}]} + + +def test_foundry_embeddings_relay_is_costed_per_input_token(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "embed-v-4-0", "models/embeddings", EMBEDDINGS_BODY + ) + per_token = litellm.get_model_info("azure_ai/embed-v-4-0")["input_cost_per_token"] + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(1200 * per_token) + + +def test_cohere_rerank_relay_is_costed_per_search_unit(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "cohere-rerank-v4.0-fast", "providers/cohere/v2/rerank", RERANK_BODY + ) + per_query = litellm.get_model_info("azure_ai/cohere-rerank-v4.0-fast")["input_cost_per_query"] + + assert isinstance(result, RerankResponse) + assert logging_obj.call_type == "arerank" + assert per_query > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(2 * per_query) + + +def test_image_generation_relay_is_costed_per_image(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "FLUX.2-pro", "openai/deployments/FLUX.2-pro/images/generations", IMAGE_BODY + ) + per_image = litellm.get_model_info("azure_ai/FLUX.2-pro")["output_cost_per_image"] + + assert isinstance(result, ImageResponse) + assert logging_obj.call_type == "aimage_generation" + assert per_image > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) + + +def test_flux_2_relay_through_the_provider_route_is_costed_per_image(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "FLUX.2-pro", "providers/blackforestlabs/v1/flux-2-pro", IMAGE_BODY + ) + per_image = litellm.get_model_info("azure_ai/FLUX.2-pro")["output_cost_per_image"] + + assert isinstance(result, ImageResponse) + assert logging_obj.call_type == "aimage_generation" + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) + + +def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), + "cohere-rerank-v4.0-fast", + "providers/cohere/v2/rerank", + {"message": "invalid request"}, + status_code=400, + ) + + assert result == {"response": {"message": "invalid request"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_streaming_chat_completion_chunks_are_costed_like_azure(): + head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} + chunks = [ + "data: " + + json.dumps( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + } + ), + "data: " + + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), + "data: [DONE]", + ] + + response = AzureAIPassthroughConfig().handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=MagicMock(), + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + endpoint="chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "hi" + assert response.usage.total_tokens == 4 + + +def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure(): + logging_obj = _relay_logging_obj("gpt-5.4-mini", FOUNDRY_BASE) + + response = AzureAIPassthroughConfig().handle_logging_collected_chunks( + all_chunks=["event: response.completed", "data: " + json.dumps(RESPONSES_COMPLETED_EVENT)], + litellm_logging_obj=logging_obj, + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + endpoint="gpt/openai/responses", + ) + info = litellm.get_model_info("azure_ai/gpt-5.4-mini") + + assert response is not None + assert response.response.usage.output_tokens == 100 + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +async def test_streaming_responses_relay_flush_reaches_the_success_callbacks_with_a_price(): + probe = _SpendProbe() + logging_obj = _relay_logging_obj( + "gpt-5.4-mini", FOUNDRY_BASE, stream=True, callbacks=[probe], endpoint="gpt/openai/responses" + ) + stream = "event: response.completed\ndata: " + json.dumps(RESPONSES_COMPLETED_EVENT) + "\n\n" + + collector = AzureAIPassthroughConfig().create_stream_collector( + model="gpt-5.4-mini", custom_llm_provider="azure_ai", endpoint="gpt/openai/responses" + ) + collector.add(stream.encode()) + await logging_obj.async_flush_passthrough_collected_chunks(collector=collector) + info = litellm.get_model_info("azure_ai/gpt-5.4-mini") + + assert probe.logged_call_type == "allm_passthrough_route" + assert probe.logged_cost == pytest.approx(1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"]) diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py index 9612d97d946..a43fc3332af 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py @@ -2,20 +2,25 @@ Test Azure AI cost calculator, especially Model Router flat cost. """ +from datetime import datetime +from typing import Final + import pytest +import litellm +from litellm.cost_calculator import completion_cost +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure_ai.cost_calculator import ( - _is_azure_model_router, + calculate_azure_model_router_flat_cost, cost_per_token, + is_azure_model_router, ) -from litellm.types.utils import Usage +from litellm.types.utils import Choices, Message, ModelResponse, Usage from litellm.utils import get_model_info # Get the flat cost from model_prices_and_context_window.json _model_info = get_model_info(model="model_router", custom_llm_provider="azure_ai") -AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS = ( - _model_info.get("input_cost_per_token", 0) * 1_000_000 -) +AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS = _model_info.get("input_cost_per_token", 0) * 1_000_000 class TestAzureModelRouterDetection: @@ -49,7 +54,7 @@ class TestAzureModelRouterDetection: ) def test_is_azure_model_router(self, model: str, expected: bool): """Test Azure Model Router detection.""" - assert _is_azure_model_router(model) == expected + assert is_azure_model_router(model) == expected class TestAzureModelRouterPrefix: @@ -80,108 +85,60 @@ class TestAzureModelRouterPrefix: assert result == expected +ROUTER_FEE_PER_TOKEN: Final = AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 +ROUTED_MODEL: Final = "gpt-4.1-nano-2025-04-14" +ROUTED_USAGE: Final = Usage(prompt_tokens=5000, completion_tokens=2000, total_tokens=7000) +ROUTED_FEE: Final = 5000 * ROUTER_FEE_PER_TOKEN + + +def _router_logging(request_model: str) -> Logging: + return Logging( + model=request_model, + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="test-123", + function_id="test-function", + ) + + +def _azure_ai_response(response_model: str, litellm_model_name: str | None = None) -> ModelResponse: + response: Final = ModelResponse( + id="test-123", + choices=[Choices(finish_reason="stop", index=0, message=Message(role="assistant", content="Hello"))], + created=1234567890, + model=response_model, + object="chat.completion", + usage=ROUTED_USAGE, + ) + response._hidden_params = ( + {"custom_llm_provider": "azure_ai"} + if litellm_model_name is None + else {"custom_llm_provider": "azure_ai", "litellm_model_name": litellm_model_name} + ) + return response + + +def _routed_model_cost() -> tuple[float, float]: + routed_info: Final = get_model_info(model=ROUTED_MODEL, custom_llm_provider="azure_ai") + return ( + ROUTED_USAGE.prompt_tokens * (routed_info["input_cost_per_token"] or 0.0), + ROUTED_USAGE.completion_tokens * (routed_info["output_cost_per_token"] or 0.0), + ) + + +@pytest.mark.usefixtures("local_model_cost_map") class TestAzureModelRouterFlatCost: - """Test Azure AI Foundry Model Router flat cost calculation.""" + """cost_per_token charges the router fee once, for whichever router name the caller gives it.""" - def test_model_router_flat_cost_basic(self): - """Test that flat cost is added for Model Router requests.""" - model = "azure-model-router" - usage = Usage( - prompt_tokens=1000, - completion_tokens=500, - total_tokens=1500, - ) + def test_unmapped_router_deployment_name_prices_the_fee(self) -> None: + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + prompt_cost, completion_cost_usd = cost_per_token(model="azure-model-router", usage=usage) + assert prompt_cost == pytest.approx(1000 * ROUTER_FEE_PER_TOKEN, rel=1e-9) + assert completion_cost_usd == 0.0 - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # Calculate expected flat cost - expected_flat_cost = ( - usage.prompt_tokens - * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - / 1_000_000 - ) - - # Flat cost should be $0.00014 (1000 tokens × $0.14 / 1M tokens) - assert expected_flat_cost == pytest.approx(0.00014, rel=1e-9) - - # Prompt cost should include the flat cost - # (plus any base cost from the actual model used, which might be 0 if not in model_cost) - assert prompt_cost >= expected_flat_cost - print( - f"Model Router flat cost for {usage.prompt_tokens} tokens: ${expected_flat_cost:.6f}" - ) - print(f"Total prompt cost: ${prompt_cost:.6f}") - - def test_model_router_flat_cost_large_request(self): - """Test flat cost calculation for larger requests.""" - model = "model-router" - usage = Usage( - prompt_tokens=100_000, - completion_tokens=50_000, - total_tokens=150_000, - ) - - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # Calculate expected flat cost - expected_flat_cost = ( - usage.prompt_tokens - * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - / 1_000_000 - ) - - # Flat cost should be $0.014 (100k tokens × $0.14 / 1M tokens) - assert expected_flat_cost == pytest.approx(0.014, rel=1e-9) - # Use approx for floating-point comparison - assert prompt_cost >= expected_flat_cost or prompt_cost == pytest.approx( - expected_flat_cost, rel=1e-9 - ) - print( - f"Model Router flat cost for {usage.prompt_tokens} tokens: ${expected_flat_cost:.6f}" - ) - print(f"Total prompt cost: ${prompt_cost:.6f}") - - def test_model_router_flat_cost_1m_tokens(self): - """Test flat cost for exactly 1 million input tokens.""" - model = "azure-model-router" - usage = Usage( - prompt_tokens=1_000_000, - completion_tokens=100_000, - total_tokens=1_100_000, - ) - - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # Calculate expected flat cost - expected_flat_cost = AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - - # Flat cost should be exactly $0.14 for 1M tokens - assert expected_flat_cost == pytest.approx(0.14, rel=1e-9) - assert prompt_cost >= expected_flat_cost - print(f"Model Router flat cost for 1M tokens: ${expected_flat_cost:.6f}") - print(f"Total prompt cost: ${prompt_cost:.6f}") - - def test_non_model_router_no_flat_cost(self): - """Test that non-Model Router models don't get the flat cost.""" - model = "gpt-4o" - usage = Usage( - prompt_tokens=1000, - completion_tokens=500, - total_tokens=1500, - ) - - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # No flat cost should be added for non-Model Router models - # The cost might be 0 or based on the model's pricing - print(f"Non-Model Router prompt cost: ${prompt_cost:.6f}") - # We just ensure it doesn't crash and returns valid values - assert prompt_cost >= 0 - assert completion_cost >= 0 - - def test_model_router_with_cached_tokens(self): - """Test Model Router flat cost with cached tokens.""" - model = "azure-model-router" + def test_unmapped_router_deployment_name_charges_the_fee_over_cached_prompt_tokens_too(self) -> None: usage = Usage( prompt_tokens=2000, completion_tokens=800, @@ -189,268 +146,165 @@ class TestAzureModelRouterFlatCost: cache_read_input_tokens=500, cache_creation_input_tokens=200, ) + prompt_cost, completion_cost_usd = cost_per_token(model="azure-model-router", usage=usage) + assert prompt_cost == pytest.approx(2000 * ROUTER_FEE_PER_TOKEN, rel=1e-9) + assert completion_cost_usd == 0.0 - prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) - - # Flat cost is based on ALL prompt tokens (including cached) - expected_flat_cost = ( - usage.prompt_tokens - * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - / 1_000_000 + def test_router_deployment_name_as_both_names_charges_the_fee_once(self) -> None: + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + prompt_cost, completion_cost_usd = cost_per_token( + model="model_router/my-deployment", usage=usage, request_model="azure_ai/model_router/my-deployment" ) + assert prompt_cost == pytest.approx(1000 * ROUTER_FEE_PER_TOKEN, rel=1e-9) + assert completion_cost_usd == 0.0 - assert expected_flat_cost == pytest.approx(0.00028, rel=1e-9) - assert prompt_cost >= expected_flat_cost - print( - f"Model Router flat cost with caching for {usage.prompt_tokens} tokens: ${expected_flat_cost:.6f}" + @pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"]) + def test_router_entry_prices_its_own_fee(self, router_entry_name: str) -> None: + usage = Usage(prompt_tokens=1_000_000, completion_tokens=0, total_tokens=1_000_000) + prompt_cost, completion_cost_usd = cost_per_token(model=router_entry_name, usage=usage) + assert prompt_cost == pytest.approx(0.14, rel=1e-9) + assert completion_cost_usd == 0.0 + + def test_routed_model_is_priced_as_itself(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + prompt_cost, completion_cost_usd = cost_per_token(model=ROUTED_MODEL, usage=ROUTED_USAGE) + assert routed_prompt_cost > 0 + assert prompt_cost == pytest.approx(routed_prompt_cost, rel=1e-9) + assert completion_cost_usd == pytest.approx(routed_completion_cost, rel=1e-9) + + def test_unmapped_model_that_is_not_a_router_name_raises(self) -> None: + usage = Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20) + with pytest.raises(Exception, match="no-such-azure-ai-model"): + cost_per_token(model="no-such-azure-ai-model", usage=usage) + + def test_request_model_through_the_router_adds_the_fee_once(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + prompt_cost, completion_cost_usd = cost_per_token( + model=ROUTED_MODEL, usage=ROUTED_USAGE, request_model="azure_ai/model-router" ) - print(f"Total prompt cost: ${prompt_cost:.6f}") + assert prompt_cost == pytest.approx(routed_prompt_cost + ROUTED_FEE, rel=1e-9) + assert completion_cost_usd == pytest.approx(routed_completion_cost, rel=1e-9) - def test_router_flat_cost_when_response_has_actual_model(self): - """ - Test that router flat cost is added when request was via router but response - contains the actual model (e.g., gpt-5-nano). + def test_request_model_that_is_not_the_router_adds_nothing(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + assert cost_per_token( + model=ROUTED_MODEL, usage=ROUTED_USAGE, request_model=f"azure_ai/{ROUTED_MODEL}" + ) == pytest.approx((routed_prompt_cost, routed_completion_cost), rel=1e-9) - This is the key fix: Azure returns the actual model in the response, but we - must still add the router flat cost because the request was made via model router. - """ - usage = Usage( - prompt_tokens=10000, - completion_tokens=5000, - total_tokens=15000, + @pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"]) + def test_request_model_does_not_double_the_router_entry(self, router_entry_name: str) -> None: + prompt_cost, completion_cost_usd = cost_per_token( + model=router_entry_name, usage=ROUTED_USAGE, request_model=f"azure_ai/{router_entry_name}" ) + assert prompt_cost == pytest.approx(ROUTED_FEE, rel=1e-9) + assert completion_cost_usd == 0.0 - # Response model is the actual model Azure used (not a router name) - response_model = "gpt-5-nano-2025-08-07" - # Request model is the router - user called azure_ai/model_router/model-router - request_model = "azure_ai/model_router/model-router" - - prompt_cost, completion_cost = cost_per_token( - model=response_model, - usage=usage, - request_model=request_model, + def test_public_cost_per_token_keeps_the_request_model_keyword(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + prompt_cost, completion_cost_usd = litellm.cost_per_token( + model=ROUTED_MODEL, + custom_llm_provider="azure_ai", + usage_object=ROUTED_USAGE, + request_model="azure_ai/model-router", ) + assert prompt_cost == pytest.approx(routed_prompt_cost + ROUTED_FEE, rel=1e-9) + assert completion_cost_usd == pytest.approx(routed_completion_cost, rel=1e-9) - # Expected: model cost (from gpt-5-nano) + router flat cost - expected_flat_cost = ( - usage.prompt_tokens - * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS - / 1_000_000 + def test_flat_cost_helper(self) -> None: + assert calculate_azure_model_router_flat_cost( + model="azure-model-router", prompt_tokens=10_000 + ) == pytest.approx(0.0014, rel=1e-9) + assert calculate_azure_model_router_flat_cost(model="gpt-5-nano", prompt_tokens=10_000) == 0.0 + + def test_flat_cost_reads_the_fee_from_the_deployment_named_entry(self) -> None: + litellm.register_model( + {"azure_ai/model-router": {"input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "mode": "chat"}} ) - assert expected_flat_cost == pytest.approx(0.0014, rel=1e-9) - - # Total cost should be model cost + flat cost - total_cost = prompt_cost + completion_cost - assert total_cost >= expected_flat_cost - - # Prompt cost should include both model prompt cost and router flat cost - assert prompt_cost >= expected_flat_cost + litellm.get_model_info.cache_clear() + assert calculate_azure_model_router_flat_cost(model="model-router", prompt_tokens=1_000_000) == pytest.approx( + 0.2, rel=1e-9 + ) + assert calculate_azure_model_router_flat_cost( + model="azure-model-router", prompt_tokens=1_000_000 + ) == pytest.approx(0.14, rel=1e-9) +@pytest.mark.usefixtures("local_model_cost_map") class TestAzureModelRouterCostBreakdown: - """Test that Azure Model Router flat cost is tracked in cost breakdown.""" + """completion_cost charges the router fee exactly once: as the breakdown's additional cost line when a routed + model is priced as itself, inside the input cost when the priced name is the router.""" - def test_flat_cost_calculation_helper(self): - """Test that flat cost can be calculated using the helper function.""" - from litellm.llms.azure_ai.cost_calculator import ( - calculate_azure_model_router_flat_cost, - ) - - model = "azure-model-router" - prompt_tokens = 10000 - - # Calculate flat cost using helper function - flat_cost = calculate_azure_model_router_flat_cost( - model=model, prompt_tokens=prompt_tokens - ) - - # Expected flat cost - expected_flat_cost = ( - prompt_tokens * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 - ) - - assert flat_cost > 0 - assert flat_cost == pytest.approx(expected_flat_cost, rel=1e-9) - print(f"Flat cost calculated: ${flat_cost:.6f}") - - def test_flat_cost_integration_with_completion_cost(self): - """Test that flat cost is properly integrated into completion_cost calculation.""" - import litellm - from litellm.cost_calculator import completion_cost - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - # Create a mock response for azure_ai model router - response = ModelResponse( - id="test-123", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - role="assistant", - content="Test response", - ), - ) - ], - created=1234567890, - model="azure-model-router", - object="chat.completion", - usage=Usage( - prompt_tokens=5000, - completion_tokens=2000, - total_tokens=7000, - ), - ) - - # Set hidden params for provider - response._hidden_params = {"custom_llm_provider": "azure_ai"} - - # Calculate cost + def test_unmapped_router_deployment_name_costs_only_the_fee(self) -> None: cost = completion_cost( - completion_response=response, + completion_response=_azure_ai_response("azure-model-router"), model="azure-model-router", custom_llm_provider="azure_ai", ) + assert cost == pytest.approx(ROUTED_FEE, rel=1e-9) - # Expected flat cost - expected_flat_cost = ( - 5000 * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 - ) - - # Cost should include the flat cost (use approx for floating-point comparison) - assert cost >= expected_flat_cost or cost == pytest.approx( - expected_flat_cost, rel=1e-9 - ) - print(f"Total cost with flat fee: ${cost:.6f}") - print(f"Expected minimum flat cost: ${expected_flat_cost:.6f}") - - def test_additional_costs_in_cost_breakdown(self): - """Test that Azure Model Router flat cost appears in additional_costs dict.""" - from datetime import datetime - - from litellm.cost_calculator import completion_cost - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - # Create logging object with required parameters - logging_obj = Logging( - model="azure-model-router", - messages=[{"role": "user", "content": "Hello"}], - stream=False, - call_type="completion", - start_time=datetime.now(), - litellm_call_id="test-123", - function_id="test-function", - ) - - # Create a mock response for azure_ai model router - response = ModelResponse( - id="test-123", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - role="assistant", - content="Test response", - ), - ) - ], - created=1234567890, - model="azure-model-router", - object="chat.completion", - usage=Usage( - prompt_tokens=5000, - completion_tokens=2000, - total_tokens=7000, - ), - ) - - # Set hidden params for provider - response._hidden_params = {"custom_llm_provider": "azure_ai"} - - # Calculate cost with logging object + def test_unmapped_router_name_carries_the_fee_as_its_input_cost(self) -> None: + logging_obj = _router_logging("azure-model-router") cost = completion_cost( - completion_response=response, + completion_response=_azure_ai_response("azure-model-router"), model="azure-model-router", custom_llm_provider="azure_ai", litellm_logging_obj=logging_obj, ) + breakdown = logging_obj.cost_breakdown + assert breakdown is not None + assert breakdown["input_cost"] == pytest.approx(ROUTED_FEE, rel=1e-9) + assert "additional_costs" not in breakdown + assert cost == pytest.approx(ROUTED_FEE, rel=1e-9) - # Check that cost breakdown contains additional_costs - assert hasattr(logging_obj, "cost_breakdown") - assert logging_obj.cost_breakdown is not None - assert "additional_costs" in logging_obj.cost_breakdown - assert isinstance(logging_obj.cost_breakdown["additional_costs"], dict) - - # Check that the Azure Model Router flat cost is in additional_costs - additional_costs = logging_obj.cost_breakdown["additional_costs"] - assert "Azure Model Router Flat Cost" in additional_costs - - # Verify the flat cost value - expected_flat_cost = ( - 5000 * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 - ) - actual_flat_cost = additional_costs["Azure Model Router Flat Cost"] - assert actual_flat_cost == pytest.approx(expected_flat_cost, rel=1e-9) - - print(f"Additional costs in breakdown: {additional_costs}") - print(f"Azure Model Router Flat Cost: ${actual_flat_cost:.6f}") - - def test_additional_costs_when_response_has_actual_model_via_hidden_params(self): - """additional_costs populated when response has actual model but request was via model router (hidden_params).""" - from datetime import datetime - - from litellm.cost_calculator import completion_cost - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - logging_obj = Logging( - model="gpt-4.1-nano-2025-04-14", - messages=[{"role": "user", "content": "Hello"}], - stream=False, - call_type="completion", - start_time=datetime.now(), - litellm_call_id="test-123", - function_id="test-function", - ) - response = ModelResponse( - id="test-123", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message(role="assistant", content="Hello"), - ) - ], - created=1234567890, - model="gpt-4.1-nano-2025-04-14", - object="chat.completion", - usage=Usage(prompt_tokens=5000, completion_tokens=2000, total_tokens=7000), - ) - response._hidden_params = { - "custom_llm_provider": "azure_ai", - "litellm_model_name": "azure_ai/model-router", - } + def test_router_request_with_routed_response_charges_the_fee_once(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + logging_obj = _router_logging("model-router") cost = completion_cost( - completion_response=response, - model="gpt-4.1-nano-2025-04-14", + completion_response=_azure_ai_response(ROUTED_MODEL), + model=ROUTED_MODEL, custom_llm_provider="azure_ai", litellm_logging_obj=logging_obj, ) - expected_flat_cost = ( - 5000 * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + breakdown = logging_obj.cost_breakdown + assert breakdown is not None + assert breakdown["input_cost"] == pytest.approx(routed_prompt_cost, rel=1e-9) + assert breakdown["output_cost"] == pytest.approx(routed_completion_cost, rel=1e-9) + assert breakdown.get("additional_costs") == pytest.approx( + {"Azure Model Router Flat Cost": ROUTED_FEE}, rel=1e-9 ) - assert cost >= expected_flat_cost - assert logging_obj.cost_breakdown is not None - assert "additional_costs" in logging_obj.cost_breakdown - assert ( - "Azure Model Router Flat Cost" - in logging_obj.cost_breakdown["additional_costs"] + assert cost == pytest.approx(routed_prompt_cost + routed_completion_cost + ROUTED_FEE, rel=1e-9) + + def test_routed_response_named_by_hidden_params_charges_the_fee_once(self) -> None: + routed_prompt_cost, routed_completion_cost = _routed_model_cost() + logging_obj = _router_logging(ROUTED_MODEL) + cost = completion_cost( + completion_response=_azure_ai_response(ROUTED_MODEL, litellm_model_name="azure_ai/model-router"), + model=ROUTED_MODEL, + custom_llm_provider="azure_ai", + litellm_logging_obj=logging_obj, ) - assert logging_obj.cost_breakdown["additional_costs"][ - "Azure Model Router Flat Cost" - ] == pytest.approx(expected_flat_cost, rel=1e-9) + breakdown = logging_obj.cost_breakdown + assert breakdown is not None + assert breakdown["input_cost"] == pytest.approx(routed_prompt_cost, rel=1e-9) + assert breakdown.get("additional_costs") == pytest.approx( + {"Azure Model Router Flat Cost": ROUTED_FEE}, rel=1e-9 + ) + assert cost == pytest.approx(routed_prompt_cost + routed_completion_cost + ROUTED_FEE, rel=1e-9) + + @pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"]) + def test_response_priced_as_the_router_entry_charges_the_fee_once(self, router_entry_name: str) -> None: + logging_obj = _router_logging(router_entry_name) + cost = completion_cost( + completion_response=_azure_ai_response(router_entry_name), + model=router_entry_name, + custom_llm_provider="azure_ai", + litellm_logging_obj=logging_obj, + ) + breakdown = logging_obj.cost_breakdown + assert breakdown is not None + assert "additional_costs" not in breakdown + assert breakdown["input_cost"] == pytest.approx(ROUTED_FEE, rel=1e-9) + assert cost == pytest.approx(ROUTED_FEE, rel=1e-9) class TestAzureAIServiceTierCostCalculation: @@ -459,26 +313,27 @@ class TestAzureAIServiceTierCostCalculation: @pytest.fixture(autouse=True) def register_test_model(self): import litellm - litellm.register_model(model_cost={ - "test-azure-ai-model": { - "input_cost_per_token": 0.001, - "output_cost_per_token": 0.002, - "input_cost_per_token_priority": 0.01, - "output_cost_per_token_priority": 0.02, - "input_cost_per_token_flex": 0.0005, - "output_cost_per_token_flex": 0.001, - "litellm_provider": "azure_ai", - "max_tokens": 8192, + + litellm.register_model( + model_cost={ + "test-azure-ai-model": { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "input_cost_per_token_priority": 0.01, + "output_cost_per_token_priority": 0.02, + "input_cost_per_token_flex": 0.0005, + "output_cost_per_token_flex": 0.001, + "litellm_provider": "azure_ai", + "max_tokens": 8192, + } } - }) + ) def test_service_tier_priority_higher_cost(self): """Priority tier should cost more than standard for azure_ai.""" usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - standard_prompt, standard_completion = cost_per_token( - model="test-azure-ai-model", usage=usage - ) + standard_prompt, standard_completion = cost_per_token(model="test-azure-ai-model", usage=usage) priority_prompt, priority_completion = cost_per_token( model="test-azure-ai-model", usage=usage, service_tier="priority" ) @@ -490,12 +345,8 @@ class TestAzureAIServiceTierCostCalculation: """Flex tier should cost less than standard for azure_ai.""" usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - standard_prompt, standard_completion = cost_per_token( - model="test-azure-ai-model", usage=usage - ) - flex_prompt, flex_completion = cost_per_token( - model="test-azure-ai-model", usage=usage, service_tier="flex" - ) + standard_prompt, standard_completion = cost_per_token(model="test-azure-ai-model", usage=usage) + flex_prompt, flex_completion = cost_per_token(model="test-azure-ai-model", usage=usage, service_tier="flex") assert flex_prompt < standard_prompt assert flex_completion < standard_completion diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py new file mode 100644 index 00000000000..84d5cd2a7d4 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py @@ -0,0 +1,113 @@ +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm import completion_cost, cost_per_token, get_model_info +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.types.utils import TranscriptionResponse + +REPO_ROOT: Final = Path(__file__).parents[4] +MAIN_COST_MAP: Final = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_COST_MAP: Final = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" +COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) +AZURE_PRICING_PREFIX: Final = "https://azure.microsoft.com/en-us/pricing/details/" +A_MILLION: Final = 1_000_000 +AN_HOUR_IN_SECONDS: Final = 3600 + +TOKEN_PRICED_NAMES: Final = ( + "gpt-chat-latest", + "codex-mini", + "model-router", + "cohere-command-a", + "grok-4-20-reasoning", + "grok-4-20-non-reasoning", +) +GROK_4_20_NAMES: Final = ("grok-4-20-reasoning", "grok-4-20-non-reasoning") +CATALOG_NAMES: Final = TOKEN_PRICED_NAMES + ("whisper",) + + +def _cost_map_entry(path: Path, catalog_name: str) -> dict[str, object]: + return COST_MAP_ADAPTER.validate_json(path.read_bytes())[f"azure_ai/{catalog_name}"] + + +def _whisper_transcription_cost(duration_seconds: int) -> float: + transcription: Final = TranscriptionResponse(text="hello") + transcription._hidden_params = { # pyright: ignore[reportPrivateUsage] # TranscriptionResponse exposes no public hidden-params setter + "custom_llm_provider": "azure_ai", + "model": "azure_ai/whisper", + "audio_transcription_duration": duration_seconds, + } + return completion_cost( + completion_response=transcription, + model="azure_ai/whisper", + custom_llm_provider="azure_ai", + call_type="atranscription", + ) + + +@pytest.mark.parametrize("catalog_name", CATALOG_NAMES) +def test_azure_ai_catalog_name_routes_to_azure_ai(catalog_name: str) -> None: + routed_model, provider, _, _ = get_llm_provider(model=f"azure_ai/{catalog_name}") + assert (routed_model, provider) == (catalog_name, "azure_ai") + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("catalog_name", TOKEN_PRICED_NAMES) +def test_azure_ai_catalog_name_charges_its_own_entry_per_token(catalog_name: str) -> None: + entry: Final = get_model_info(f"azure_ai/{catalog_name}") + prompt_cost, completion_cost_usd = cost_per_token( + model=f"azure_ai/{catalog_name}", prompt_tokens=A_MILLION, completion_tokens=A_MILLION + ) + assert prompt_cost > 0 + assert prompt_cost == pytest.approx(A_MILLION * entry["input_cost_per_token"]) + assert completion_cost_usd == pytest.approx(A_MILLION * entry["output_cost_per_token"]) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("catalog_name", TOKEN_PRICED_NAMES) +def test_azure_ai_catalog_name_prices_the_same_in_any_casing(catalog_name: str) -> None: + lowercase_cost = cost_per_token(model=f"azure_ai/{catalog_name}", prompt_tokens=A_MILLION, completion_tokens=0) + upper_cost = cost_per_token(model=f"azure_ai/{catalog_name.upper()}", prompt_tokens=A_MILLION, completion_tokens=0) + assert upper_cost == lowercase_cost + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("catalog_name", GROK_4_20_NAMES) +def test_azure_ai_grok_4_20_bills_cached_prompt_tokens_at_the_input_price(catalog_name: str) -> None: + uncached_prompt_cost, _ = cost_per_token(model=f"azure_ai/{catalog_name}", prompt_tokens=A_MILLION, completion_tokens=0) + cached_prompt_cost, _ = cost_per_token( + model=f"azure_ai/{catalog_name}", + prompt_tokens=A_MILLION, + completion_tokens=0, + cache_read_input_tokens=A_MILLION, + ) + assert uncached_prompt_cost > 0 + assert cached_prompt_cost == pytest.approx(uncached_prompt_cost) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_azure_ai_whisper_catalog_name_is_priced_per_second() -> None: + one_second_cost: Final = _whisper_transcription_cost(1) + one_hour_cost: Final = _whisper_transcription_cost(AN_HOUR_IN_SECONDS) + assert one_second_cost > 0 + assert one_hour_cost == pytest.approx(AN_HOUR_IN_SECONDS * one_second_cost) + + +@pytest.mark.parametrize("catalog_name", CATALOG_NAMES) +def test_azure_ai_catalog_entry_source_and_backup_match(catalog_name: str) -> None: + main_entry = _cost_map_entry(MAIN_COST_MAP, catalog_name) + backup_entry = _cost_map_entry(BACKUP_COST_MAP, catalog_name) + + assert str(main_entry["source"]).startswith(AZURE_PRICING_PREFIX) + assert backup_entry == main_entry + + +def test_azure_ai_model_router_spellings_share_one_entry() -> None: + underscore_entry = _cost_map_entry(MAIN_COST_MAP, "model_router") + hyphen_entry = _cost_map_entry(MAIN_COST_MAP, "model-router") + + assert {k: v for k, v in underscore_entry.items() if k != "comment"} == { + k: v for k, v in hyphen_entry.items() if k != "comment" + } diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py index f3618572622..1b2ca298694 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py @@ -12,112 +12,6 @@ from importlib.resources import files import pytest -FW_MODELS = { - "azure_ai/FW-Kimi-K2.5": { - "input_cost_per_token": 6.6e-07, - "output_cost_per_token": 3.3e-06, - "cache_read_input_token_cost": 1.1e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_vision": True, - }, - "azure_ai/FW-Kimi-K2.6": { - "input_cost_per_token": 1.045e-06, - "output_cost_per_token": 4.4e-06, - "cache_read_input_token_cost": 1.76e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_vision": True, - }, - "azure_ai/FW-Kimi-K2.7-Code": { - "input_cost_per_token": 1.05e-06, - "output_cost_per_token": 4.4e-06, - "cache_read_input_token_cost": 2.1e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_vision": True, - }, - "azure_ai/FW-Kimi-K3": { - "input_cost_per_token": 3.3e-06, - "output_cost_per_token": 1.65e-05, - "cache_read_input_token_cost": 3.3e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 131072, - "supports_vision": True, - }, - "azure_ai/FW-Inkling": { - "input_cost_per_token": 1e-06, - "output_cost_per_token": 4.05e-06, - "cache_read_input_token_cost": 1.7e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 1048576, - }, - "azure_ai/FW-DeepSeek-V3.2": { - "input_cost_per_token": 6.2e-07, - "output_cost_per_token": 1.85e-06, - "cache_read_input_token_cost": 3.1e-07, - "max_input_tokens": 163840, - "max_output_tokens": 163840, - }, - "azure_ai/FW-DeepSeek-V4-Pro": { - "input_cost_per_token": 1.925e-06, - "output_cost_per_token": 3.828e-06, - "cache_read_input_token_cost": 1.65e-07, - "max_input_tokens": 1000000, - "max_output_tokens": 384000, - }, - "azure_ai/FW-MiniMax-M3": { - "input_cost_per_token": 3.3e-07, - "output_cost_per_token": 1.32e-06, - "cache_read_input_token_cost": 6.6e-08, - "max_input_tokens": 512000, - "max_output_tokens": 512000, - "supports_vision": True, - }, - "azure_ai/FW-MiniMax-M2.5": { - "input_cost_per_token": 3.3e-07, - "output_cost_per_token": 1.32e-06, - "cache_read_input_token_cost": 3.3e-08, - "max_input_tokens": 1000000, - "max_output_tokens": 1000000, - }, - "azure_ai/FW-Nemotron-3-Ultra-NVFP4": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2.4e-06, - "cache_read_input_token_cost": 1.19e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - }, - "azure_ai/FW-GLM-5.2-Fast": { - "input_cost_per_token": 2.1e-06, - "output_cost_per_token": 6.6e-06, - "cache_read_input_token_cost": 2.1e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 131072, - }, - "azure_ai/FW-GLM-5.2": { - "input_cost_per_token": 1.54e-06, - "output_cost_per_token": 4.84e-06, - "cache_read_input_token_cost": 1.5e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 131072, - }, - "azure_ai/FW-GLM-5.1": { - "input_cost_per_token": 1.54e-06, - "output_cost_per_token": 4.84e-06, - "cache_read_input_token_cost": 2.86e-07, - "max_input_tokens": 202800, - "max_output_tokens": 131072, - }, - "azure_ai/FW-GLM-5": { - "input_cost_per_token": 1.1e-06, - "output_cost_per_token": 3.52e-06, - "cache_read_input_token_cost": 2.2e-07, - "max_input_tokens": 200000, - "max_output_tokens": 128000, - }, -} - @pytest.fixture(scope="module") def use_local_model_cost_map(): @@ -144,28 +38,6 @@ def use_local_model_cost_map(): monkeypatch.undo() -@pytest.mark.parametrize("model_key,expected", list(FW_MODELS.items())) -def test_azure_ai_fw_model_info(use_local_model_cost_map, model_key, expected): - model_info = use_local_model_cost_map.get_model_info(model=model_key) - - assert model_info["litellm_provider"] == "azure_ai" - assert model_info["mode"] == "chat" - assert model_info["input_cost_per_token"] == pytest.approx(expected["input_cost_per_token"]) - assert model_info["output_cost_per_token"] == pytest.approx(expected["output_cost_per_token"]) - assert model_info["cache_read_input_token_cost"] == pytest.approx( - expected["cache_read_input_token_cost"] - ) - assert model_info["max_input_tokens"] == expected["max_input_tokens"] - assert model_info["max_output_tokens"] == expected["max_output_tokens"] - assert model_info["max_tokens"] == expected["max_output_tokens"] - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_prompt_caching"] is True - if expected.get("supports_vision"): - assert model_info["supports_vision"] is True - - @pytest.mark.parametrize( "model_name,expected_prompt,expected_completion", [ @@ -197,22 +69,6 @@ def test_azure_ai_fw_cost_per_token( assert completion_cost == pytest.approx(expected_completion) -def test_azure_ai_fw_nemotron_lightning_model_info(use_local_model_cost_map): - model_info = use_local_model_cost_map.get_model_info(model="azure_ai/FW-Nemotron-Lightning-3.5-30B-A3B") - - assert model_info["litellm_provider"] == "azure_ai" - assert model_info["mode"] == "chat" - assert model_info["input_cost_per_token"] == pytest.approx(6e-08) - assert model_info["output_cost_per_token"] == pytest.approx(2.2e-07) - assert model_info["cache_read_input_token_cost"] == pytest.approx(1e-08) - assert model_info["max_input_tokens"] == 262144 - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_prompt_caching"] is True - assert model_info["supports_vision"] is False - - def test_azure_ai_fw_nemotron_lightning_supports_tool_choice(use_local_model_cost_map): from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py index 812b9288ca8..cbcc2a94043 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py @@ -33,33 +33,6 @@ def use_local_model_cost_map(): monkeypatch.undo() -def test_azure_ai_kimi_k26_model_info(use_local_model_cost_map): - model_info = use_local_model_cost_map.get_model_info(model="azure_ai/kimi-k2.6") - - assert model_info["litellm_provider"] == "azure_ai" - assert model_info["mode"] == "chat" - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - assert model_info["max_tokens"] == 262144 - assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07) - assert model_info["output_cost_per_token"] == pytest.approx(4e-06) - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_vision"] is True - - -def test_azure_ai_kimi_k26_raw_model_cost_entry(use_local_model_cost_map): - model_info = use_local_model_cost_map.model_cost["azure_ai/kimi-k2.6"] - - assert model_info["supported_modalities"] == ["text", "image"] - assert model_info["supported_output_modalities"] == ["text"] - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_vision"] is True - - def test_azure_ai_kimi_k26_cost_per_token(use_local_model_cost_map): from litellm.llms.azure_ai.cost_calculator import cost_per_token from litellm.types.utils import Usage diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 2a9b7a6d138..03daafcad72 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -8,7 +8,7 @@ the tests don't hit AWS. from __future__ import annotations -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch import pytest @@ -480,3 +480,93 @@ def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3): fake_client.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN) assert batch.status == "cancelled" + + +class _TagGatedSTSClient: + """Stands in for STS behind a trust policy that only admits sessions carrying ``tags``.""" + + def __init__(self, tags: list[dict[str, str]], access_key_id: str) -> None: + self._tags = tags + self._access_key_id = access_key_id + + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + from botocore.exceptions import ClientError + + if list(params.get("Tags", ())) != self._tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": self._access_key_id, + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.now(timezone.utc) + timedelta(minutes=30), + } + } + + +def test_handle_model_invocation_job_status_builds_the_client_from_the_tagged_session(monkeypatch): + """Status polling must assume the role with the deployment's session tags, like every other call.""" + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + bedrock_client_kwargs: list[dict] = [] + fake_bedrock = MagicMock() + fake_bedrock.get_model_invocation_job.return_value = _fake_boto3_response() + + def boto3_client(service_name, **kwargs): + if service_name == "sts": + return _TagGatedSTSClient(tags, "ASIABATCHSTATUSTAGGED") + bedrock_client_kwargs.append(kwargs) + return fake_bedrock + + with patch("boto3.client", side_effect=boto3_client): + batch = BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=JOB_ARN, + aws_access_key_id="AKIABATCHSTATUSCALLER", + aws_secret_access_key="pod-caller-secret", + aws_role_name="arn:aws:iam::999999999999:role/litellm-batch-role", + aws_session_name="litellm-batch-session", + aws_session_tags=tags, + ) + + assert batch.status == "completed" + assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHSTATUSTAGGED"] + + +def test_cancel_batch_stops_and_polls_the_job_with_the_tagged_session(monkeypatch): + """Cancelling on a tag-gated role must forward the deployment's session tags to both the stop and status calls.""" + import litellm + + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + bedrock_client_kwargs: list[dict] = [] + fake_bedrock = MagicMock() + fake_bedrock.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped") + + def boto3_client(service_name, **kwargs): + if service_name == "sts": + return _TagGatedSTSClient(tags, "ASIABATCHCANCELTAGGED") + bedrock_client_kwargs.append(kwargs) + return fake_bedrock + + with patch("boto3.client", side_effect=boto3_client): + batch = litellm.cancel_batch( + batch_id=JOB_ARN, + custom_llm_provider="bedrock", + aws_access_key_id="AKIABATCHCANCELCALLER", + aws_secret_access_key="pod-caller-secret", + aws_role_name="arn:aws:iam::999999999999:role/litellm-batch-role", + aws_session_name="litellm-batch-session", + aws_session_tags=tags, + ) + + fake_bedrock.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN) + assert batch.status == "cancelled" + assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHCANCELTAGGED"] * 2 diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py index 531f334e460..bcd1df26020 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py @@ -15,6 +15,7 @@ AWS_AUTH_PARAMS = { "aws_sts_endpoint": "https://sts.us-west-2.amazonaws.com", "aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-west-2.amazonaws.com", "aws_external_id": "external", + "aws_session_tags": [{"Key": "team", "Value": "genai"}], } diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index aba51689094..96a2fa6ec67 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -1,5 +1,7 @@ import json +from unittest.mock import MagicMock +import httpx import pytest @@ -177,3 +179,58 @@ def test_guardrail_config_flows_to_headers_not_request_body(model): assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "ff6ujrregl1q" assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "DRAFT" assert headers["X-Amzn-Bedrock-Trace"] == "DISABLED" + + +def test_get_error_class_preserves_provider_headers(): + """The invoke handler path hands real provider headers to get_error_class (LIT-5428).""" + error = AmazonInvokeConfig().get_error_class( + error_message="Amazon Bedrock is unable to process your request.", + status_code=500, + headers={"x-amzn-RequestId": "req-invoke-500"}, + ) + + assert isinstance(error, BedrockError) + assert error.headers == {"x-amzn-RequestId": "req-invoke-500"} + assert error.response.headers["x-amzn-requestid"] == "req-invoke-500" + + +def test_transform_response_hands_json_mode_to_nova(): + """The invoke dispatcher forwards its json_mode argument to Nova instead of dropping it.""" + from litellm.types.utils import ModelResponse + + response_json = { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_nova_json", + "name": "json_tool_call", + "input": {"city": "Paris", "temperature": 21}, + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9}, + } + raw_response = httpx.Response(200, json=response_json, request=httpx.Request("POST", "https://bedrock")) + + result = AmazonInvokeConfig().transform_response( + model="invoke/amazon.nova-lite-v1:0", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "weather"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + json_mode=True, + ) + + assert result.choices[0].message.tool_calls is None + assert json.loads(result.choices[0].message.content) == {"city": "Paris", "temperature": 21} diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 0d7573a2536..bcba4bf7711 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1,15 +1,19 @@ import asyncio import json +import uuid from unittest.mock import patch +import httpx import pytest # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. +import litellm from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler def test_get_supported_params_thinking(): @@ -53,6 +57,7 @@ def test_aws_params_filtered_from_request_body(): "aws_sts_endpoint": "https://sts.amazonaws.com", "aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-west-2.amazonaws.com", "aws_external_id": "external-id-123", + "aws_session_tags": [{"Key": "team", "Value": "genai"}], } # Transform the request @@ -101,6 +106,9 @@ def test_aws_params_filtered_from_request_body(): assert ( "aws_external_id" not in result_json ), "AWS external ID should not be in request body" + assert ( + "aws_session_tags" not in result_json + ), "AWS session tags should not be in request body" # Also check that the sensitive values themselves are not in the response assert ( @@ -714,3 +722,95 @@ def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking assert result["thinking"] == {"type": "adaptive"} assert result["output_config"] == {"effort": "high"} + + +async def test_bedrock_invoke_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] + + +async def test_bedrock_invoke_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch): + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "A lease"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this document?"}, + {"type": "document", "source": {"type": "url", "url": pdf_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "A lease" + assert async_only_image_fetch.fetched == [pdf_url] + assert { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png}, + } in captured["body"]["messages"][0]["content"] diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py new file mode 100644 index 00000000000..a8448f5fa7a --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -0,0 +1,99 @@ +import json +import uuid + +import httpx + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + +async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/mantle/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] + + +async def test_bedrock_mantle_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch): + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "A lease"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/mantle/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this document?"}, + {"type": "document", "source": {"type": "url", "url": pdf_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "A lease" + assert async_only_image_fetch.fetched == [pdf_url] + assert { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png}, + } in captured["body"]["messages"][0]["content"] diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index c4d6896b17b..4c2aa4ec4cf 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -6,16 +6,21 @@ extension, and AWS credential resolution is stubbed so nothing reaches STS. from __future__ import annotations +import asyncio +from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch +import boto3 import httpx import pytest from botocore.credentials import Credentials +from botocore.exceptions import ClientError from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import chat_completions as bridge from litellm.types.utils import ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe RUST_RESPONSE = { "created": 1_700_000_000, @@ -48,7 +53,7 @@ RESOLVED_CREDENTIALS = Credentials( @pytest.fixture(autouse=True) def reset_bridge(monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") bridge.set_rust_chat_completions( chat_completions=None, achat_completions=None, decline=None ) @@ -87,7 +92,7 @@ def _completion_kwargs(**overrides): "optional_params": {"maxTokens": 16}, "acompletion": False, "timeout": 30.0, - "litellm_params": {"rust": True}, + "litellm_params": {}, "extra_headers": None, "client": None, "api_key": None, @@ -157,7 +162,8 @@ def test_the_core_receives_the_untranslated_openai_messages(): ] -def test_without_the_opt_in_the_core_is_never_consulted(): +def test_without_the_opt_in_the_core_is_never_consulted(monkeypatch): + monkeypatch.setenv("LITELLM_RUST", "0") seen = _inject() try: _run(litellm_params={}) @@ -307,7 +313,9 @@ CONVERSE_RESPONSE = { } -async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj): +async def _drive_async_completion( + *, skip_pre_call_logging: bool, logging_obj, credentials: Credentials = RESOLVED_CREDENTIALS +): """Run the real `async_completion` with a stubbed transport.""" import httpx as _httpx @@ -334,7 +342,7 @@ async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj): stream=None, optional_params={"maxTokens": 16}, litellm_params={"aws_region_name": "us-west-2"}, - credentials=RESOLVED_CREDENTIALS, + credentials=credentials, headers={}, client=client, skip_pre_call_logging=skip_pre_call_logging, @@ -356,6 +364,23 @@ async def test_async_completion_logs_pre_call_by_default(): assert logging_obj.pre_call.call_count == 1 +@pytest.mark.asyncio +async def test_async_completion_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: botocore refreshes expiring credentials inside SigV4 signing with a + blocking HTTP call, so `async_completion` must sign on a worker thread to keep the loop serving.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + response = await _drive_async_completion( + skip_pre_call_logging=False, logging_obj=MagicMock(), credentials=probe.credentials() + ) + await release + + assert response.choices[0].message.content == "hi" + assert probe.served_during_refresh is True + + def _sync_client_returning_converse_response(): client = MagicMock() client.post.side_effect = lambda **_kwargs: httpx.Response( @@ -401,9 +426,10 @@ def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines(): assert logging_obj.pre_call.call_count == 1 -def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(): +def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(monkeypatch): """The suppression must not swallow the log on a request the gate declined, so a deployment with no `rust` flag keeps exactly the log it always had.""" + monkeypatch.setenv("LITELLM_RUST", "0") logging_obj = MagicMock() response = _run( logging_obj=logging_obj, @@ -491,6 +517,7 @@ def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monke """With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no credentials at all. Preparing the Rust handoff must not dereference that None: the bearer token signs the request on its own.""" + monkeypatch.setenv("LITELLM_RUST", "0") monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") client = _sync_client_returning_converse_response() @@ -520,6 +547,7 @@ def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, co """The deployment's AWS profile does not exist, so resolving SigV4 credentials raises; a bearer-token deployment must still serve the request, since the bearer token alone signs it.""" + monkeypatch.setenv("LITELLM_RUST", "0") if configured_through == "env_var": monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") else: @@ -537,3 +565,54 @@ def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, co assert response.choices[0].message.content == "hi" assert client.post.call_args.kwargs["headers"]["Authorization"] == "Bearer bedrock-bearer-token" + + +def test_session_tags_sign_the_request_and_stay_out_of_the_body(monkeypatch): + """The tagged STS session signs the Converse call and the tags never reach the request body (#34069).""" + monkeypatch.setenv("LITELLM_RUST", "0") + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIACONVERSETAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.now(timezone.utc) + timedelta(minutes=30), + } + } + + client = _sync_client_returning_converse_response() + with patch.object(boto3, "client", return_value=FakeSTSClient()): + response = BedrockConverseLLM().completion( + **_completion_kwargs( + optional_params={ + "maxTokens": 16, + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIACONVERSECALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-converse-role", + "aws_session_name": "litellm-converse-session", + "aws_session_tags": tags, + }, + litellm_params={}, + client=client, + ) + ) + + assert response.choices[0].message.content == "hi" + sent = client.post.call_args.kwargs + assert "Credential=ASIACONVERSETAGGED/" in sent["headers"]["Authorization"] + assert "aws_session_tags" not in sent["data"] diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index cb05cdb9451..2e9ea90f3b8 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2,6 +2,7 @@ import asyncio import json import os +import httpx import pytest from fastapi.testclient import TestClient @@ -381,6 +382,8 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): "us.openai.gpt-5.6-sol", "global.openai.gpt-5.6-terra", "bedrock/converse/us.openai.gpt-5.6-luna", + "us.openai.gpt-6-astra", + "bedrock/converse/global.openai.gpt-6-astra", ], ) def test_reasoning_effort_maps_to_reasoning_effort_for_openai_gpt5_converse(model, local_model_cost_map): @@ -411,6 +414,7 @@ def test_reasoning_effort_maps_to_reasoning_effort_for_openai_gpt5_converse(mode [ "us.openai.gpt-5.6-sol", "bedrock/converse/global.openai.gpt-5.6-luna", + "us.openai.gpt-6-astra", ], ) def test_openai_gpt5_converse_never_forwards_thinking(model, local_model_cost_map): @@ -862,6 +866,191 @@ def test_get_supported_openai_params(): assert "reasoning_effort" in supported_params +@pytest.mark.parametrize( + "model", + [ + "bedrock/us.deepseek.r1-v1:0", + "bedrock/converse/us.deepseek.r1-v1:0", + "bedrock/deepseek.v3-v1:0", + "bedrock/deepseek.v3.2", + ], +) +def test_bedrock_deepseek_does_not_advertise_thinking(model): + """DeepSeek reasons natively on Bedrock and does not take the Anthropic-shaped `thinking` + field (R1 400s on it, V3 ignores it), so it must not be advertised as supported.""" + config = AmazonConverseConfig() + supported_params = config.get_supported_openai_params(model=model) + assert "thinking" not in supported_params + assert "output_config" not in supported_params + + +@pytest.mark.parametrize("model", ["bedrock/us.deepseek.r1-v1:0", "bedrock/converse/us.deepseek.r1-v1:0"]) +def test_bedrock_deepseek_r1_does_not_advertise_reasoning_effort(model): + """DeepSeek R1 always reasons and returns a 400 for any reasoning_effort shape.""" + config = AmazonConverseConfig() + assert "reasoning_effort" not in config.get_supported_openai_params(model=model) + + +@pytest.mark.parametrize("model", ["bedrock/deepseek.v3-v1:0", "bedrock/deepseek.v3.2", "bedrock/us.deepseek.v3.2"]) +def test_bedrock_deepseek_v3_advertises_reasoning_effort(model): + """DeepSeek V3 on Bedrock accepts a raw reasoning_effort in additionalModelRequestFields.""" + config = AmazonConverseConfig() + assert "reasoning_effort" in config.get_supported_openai_params(model=model) + + +@pytest.mark.parametrize("model", ["us.deepseek.r1-v1:0", "deepseek.v3.2"]) +def test_bedrock_deepseek_thinking_raises_without_drop_params(model): + """Passing `thinking` to Bedrock DeepSeek must fail client-side with a clear + UnsupportedParamsError instead of leaking through to Bedrock.""" + with pytest.raises(litellm.UnsupportedParamsError): + litellm.utils.get_optional_params( + model=model, + custom_llm_provider="bedrock", + thinking={"type": "enabled", "budget_tokens": 1024}, + ) + + +def test_bedrock_deepseek_r1_reasoning_effort_raises_without_drop_params(): + with pytest.raises(litellm.UnsupportedParamsError): + litellm.utils.get_optional_params( + model="us.deepseek.r1-v1:0", + custom_llm_provider="bedrock", + reasoning_effort="high", + ) + + +@pytest.mark.parametrize("model", ["us.deepseek.r1-v1:0", "deepseek.v3.2"]) +def test_bedrock_deepseek_thinking_dropped_does_not_leak_into_request(model): + """With drop_params, `thinking` is dropped rather than forwarded into + additionalModelRequestFields for Bedrock DeepSeek.""" + optional_params = litellm.utils.get_optional_params( + model=model, + custom_llm_provider="bedrock", + thinking={"type": "enabled", "budget_tokens": 1024}, + drop_params=True, + ) + assert "thinking" not in optional_params + + config = AmazonConverseConfig() + request = config._transform_request( + model=f"bedrock/converse/{model}", + messages=[{"role": "user", "content": "Say hi in one word."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + assert "thinking" not in (request.get("additionalModelRequestFields") or {}) + + +@pytest.mark.parametrize("param", ["thinking", "reasoning_effort"]) +def test_bedrock_deepseek_r1_reasoning_params_not_forwarded_by_map(param): + """Even when map_openai_params is called directly (bypassing the supported-params + gate), DeepSeek R1 must not forward thinking/reasoning_effort into + additionalModelRequestFields, since Bedrock rejects both with a 400.""" + config = AmazonConverseConfig() + model = "bedrock/converse/us.deepseek.r1-v1:0" + value = {"type": "enabled", "budget_tokens": 1024} if param == "thinking" else "high" + + optional_params = config.map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model=model, + drop_params=False, + ) + assert "thinking" not in optional_params + assert "reasoning_effort" not in optional_params + + request = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Say hi in one word."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + assert request.get("additionalModelRequestFields") is None + + +def test_bedrock_deepseek_v3_reasoning_effort_forwarded_raw(): + """DeepSeek V3 takes reasoning_effort verbatim in additionalModelRequestFields, never + converted into the Anthropic `thinking` block that Claude models get.""" + config = AmazonConverseConfig() + model = "bedrock/deepseek.v3.2" + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "high", "max_tokens": 100}, + optional_params={}, + model=model, + drop_params=False, + ) + assert "thinking" not in optional_params + + request = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Say hi in one word."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + assert request["additionalModelRequestFields"] == {"reasoning_effort": "high"} + + +def test_bedrock_deepseek_v3_thinking_dropped_by_map(): + config = AmazonConverseConfig() + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 1024}, "max_tokens": 100}, + optional_params={}, + model="bedrock/deepseek.v3.2", + drop_params=False, + ) + assert "thinking" not in optional_params + assert "reasoning_effort" not in optional_params + + +@pytest.mark.parametrize( + "model, param, value, kept_key", + [ + ( + "bedrock/us.anthropic.claude-opus-4-20250514-v1:0", + "thinking", + {"type": "enabled", "budget_tokens": 1024}, + "thinking", + ), + ( + "bedrock/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + "thinking", + {"type": "enabled", "budget_tokens": 1024}, + "thinking", + ), + ( + "bedrock/openai.gpt-oss-safeguard-20b-1:0", + "reasoning_effort", + "high", + "reasoning_effort", + ), + ( + "bedrock/us.amazon.nova-2-lite-v1:0", + "reasoning_effort", + "high", + "reasoningConfig", + ), + ], +) +def test_bedrock_non_deepseek_reasoning_params_preserved(model, param, value, kept_key): + """The DeepSeek leak fix must only drop reasoning request params for DeepSeek. + + Claude behind an application-inference-profile ARN, gpt-oss-safeguard (absent from the + cost map so `supports_reasoning` is False), and Nova 2 all reason via a request param and + must keep it. Regression guard against gating the drop on a positive allowlist, which + silently degraded reasoning for anything the allowlist/ARN introspection missed.""" + config = AmazonConverseConfig() + optional_params = config.map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model=model, + drop_params=False, + ) + assert kept_key in optional_params + + def test_get_supported_openai_params_bedrock_converse(): """ Test that all documented bedrock converse models have the same set of supported openai params when using @@ -6039,6 +6228,8 @@ def test_transform_response_does_not_leak_body_on_parse_failure(): leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}} class MockResponse: + headers = httpx.Headers({"x-amzn-RequestId": "req-parse-failure"}) + def json(self): return leaky_body @@ -6067,6 +6258,7 @@ def test_transform_response_does_not_leak_body_on_parse_failure(): msg = str(exc_info.value) assert "secret content" not in msg assert "Error converting to valid response block" in msg + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-parse-failure" def test_converse_drops_sampling_params_for_models_that_removed_them(): @@ -6723,3 +6915,41 @@ def test_forced_tool_choice_forwarded_on_converse_models_that_support_it( ) assert result == {"any": {}} + + +def test_transform_response_honors_json_mode_kwarg_when_optional_params_lack_it(): + response_json = { + "metrics": {"latencyMs": 900}, + "output": { + "message": { + "content": [ + { + "toolUse": { + "input": {"city": "Paris", "population": 2100000}, + "name": "json_tool_call", + "toolUseId": "tooluse_invoke_nova_json", + } + } + ], + "role": "assistant", + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 40, "outputTokens": 20, "totalTokens": 60}, + } + raw_response = httpx.Response(200, json=response_json, request=httpx.Request("POST", "https://bedrock.test")) + logging_obj = MagicMock() + result = AmazonConverseConfig().transform_response( + model="bedrock/invoke/us.amazon.nova-micro-v1:0", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[], + optional_params={"tools": [{"type": "function", "function": {"name": "json_tool_call", "parameters": {}}}]}, + litellm_params={}, + encoding=None, + json_mode=True, + ) + assert result.choices[0].message.tool_calls is None + assert json.loads(result.choices[0].message.content) == {"city": "Paris", "population": 2100000} diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index 4bef59842f1..d0adabe7b4e 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -496,3 +496,171 @@ async def test_async_invoke_streaming_forwards_bedrock_response_headers(): assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987" + +def _bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={ + "x-amzn-RequestId": request_id, + "x-amzn-ErrorType": "InternalServerException", + }, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + +def test_invoke_streaming_error_forwards_bedrock_response_headers(): + error_response = _bedrock_stream_error_response(500, "req-stream-err-1") + client = HTTPHandler() + client.post = MagicMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-1" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_error_forwards_bedrock_response_headers(): + error_response = _bedrock_stream_error_response(500, "req-stream-err-2") + client = AsyncHTTPHandler() + client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-2" + + +def _unread_bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={ + "x-amzn-RequestId": request_id, + "x-amzn-ErrorType": "InternalServerException", + }, + stream=httpx.ByteStream(b'{"message":"Amazon Bedrock is unable to process your request."}'), + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + +def test_invoke_streaming_error_forwards_headers_when_body_was_never_read(): + """A retried streamed request raises HTTPStatusError over a body nobody read, so + reading it for the error message throws and loses the request id (LIT-5428).""" + error_response = _unread_bedrock_stream_error_response(500, "req-unread-sync") + client = HTTPHandler() + client.post = MagicMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-sync" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_error_forwards_headers_when_body_was_never_read(): + error_response = _unread_bedrock_stream_error_response(500, "req-unread-async") + client = AsyncHTTPHandler() + client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-async" + + +def test_invoke_streaming_non_200_forwards_bedrock_response_headers(): + """A caller-supplied client that returns a failure instead of raising still reaches the + provider's headers, and reading the streamed body for the message must not throw (LIT-5428).""" + error_response = _unread_bedrock_stream_error_response(500, "req-non200-sync") + client = HTTPHandler() + client.post = MagicMock(return_value=error_response) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-sync" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers(): + error_response = _unread_bedrock_stream_error_response(500, "req-non200-async") + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=error_response) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-async" diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py new file mode 100644 index 00000000000..3622ce7f212 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -0,0 +1,54 @@ +import asyncio +from unittest.mock import AsyncMock + +import httpx +import pytest +from botocore.credentials import RefreshableCredentials + +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe + + +class _ProbedCountTokensHandler(BedrockCountTokensHandler): + def __init__(self, probe: EventLoopProbe) -> None: + super().__init__() + self._probe = probe + + def get_credentials( + self, + **kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores + ) -> RefreshableCredentials: + return self._probe.credentials() + + +@pytest.mark.asyncio +async def test_handle_count_tokens_request_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the count_tokens handler signed on the loop, so botocore's blocking + credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + client = AsyncMock(spec=AsyncHTTPHandler) + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json={"inputTokens": 7}, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"), + ) + ) + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + result = await _ProbedCountTokensHandler(probe).handle_count_tokens_request( + request_data={ + "model": "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "messages": [{"role": "user", "content": "hi"}], + }, + litellm_params={"aws_region_name": "us-west-2"}, + resolved_model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + client=client, + ) + await release + + assert result == {"input_tokens": 7} + assert client.post.call_args.kwargs["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index 74a55cc1ef2..18f4b0f6ced 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -1,11 +1,15 @@ import json +import asyncio from unittest.mock import Mock, patch +import httpx import pytest +import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.base import HiddenParams +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Mock async invoke responses async_invoke_response = { @@ -184,6 +188,45 @@ class TestBedrockAsyncInvokeEmbedding: request_url = mock_post.call_args.kwargs.get("url", "") assert "/async-invoke" in request_url + def test_async_invoke_marengo_3_wraps_the_nested_payload_with_the_base_model_id(self): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(async_invoke_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/async_invoke/twelvelabs.marengo-embed-3-0-v1:0", + input="s3://test-bucket/clip.mp4", + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key="test-bearer-token-12345", + input_type="video", + embeddingOption=["visual", "audio"], + segmentation={"method": "fixed", "fixed": {"durationSec": 6}}, + bucketOwner="123456789012", + output_s3_uri="s3://test-bucket/async-invoke-output/", + ) + + assert response._hidden_params._invocation_arn == async_invoke_response["invocationArn"] + assert mock_post.call_args.kwargs["url"].endswith("/async-invoke") + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "modelId": "twelvelabs.marengo-embed-3-0-v1:0", + "modelInput": { + "inputType": "video", + "video": { + "mediaSource": {"s3Location": {"uri": "s3://test-bucket/clip.mp4", "bucketOwner": "123456789012"}}, + "segmentation": {"method": "fixed", "fixed": {"durationSec": 6}}, + "embeddingOption": ["visual", "audio"], + }, + }, + "outputDataConfig": {"s3OutputDataConfig": {"s3Uri": "s3://test-bucket/async-invoke-output/"}}, + } + @pytest.mark.asyncio async def test_async_invoke_twelvelabs_embedding_async_with_mock(self): """Test async invoke embedding with async calls.""" @@ -383,3 +426,34 @@ class TestBedrockAsyncInvokeEmbedding: async_endpoint == "https://bedrock-runtime.us-east-1.amazonaws.com/async-invoke" ) + + +@pytest.mark.asyncio +async def test_async_invoke_status_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the GetAsyncInvoke poll is a signed GET, and botocore refreshes + expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker + thread to keep the loop serving other requests.""" + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + embedder = BedrockEmbedding() + probe = EventLoopProbe() + + with ( + patch.object(embedder, "_load_credentials", return_value=(probe.credentials(), "us-east-1")), + respx.mock, + ): + route = respx.get(url__regex=r"https://bedrock-runtime\.us-east-1\.amazonaws\.com/async-invoke/.*").mock( + return_value=httpx.Response(200, json=async_invoke_status_response) + ) + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + status = await embedder._get_async_invoke_status( + invocation_arn=async_invoke_status_response["invocationArn"], aws_region_name="us-east-1" + ) + await release + + assert status["status"] == "InProgress" + assert "Authorization" in route.calls.last.request.headers + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 50f8bbcf584..e5a460e2f1a 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -1,11 +1,17 @@ import json +import asyncio import os from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, MagicMock import pytest +import httpx import litellm +from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.bedrock.embed.embedding import BedrockEmbedding +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Mock responses for different embedding models titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} @@ -1035,6 +1041,63 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch): assert "aws_external_id" not in optional_params +def test_embedding_session_tags_sign_the_request_and_stay_out_of_the_body(monkeypatch): + """The tagged STS session signs the InvokeModel call and the tags never reach the body (#34069).""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAEMBEDTAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + client = HTTPHandler() + with patch.object(boto3, "client", return_value=FakeSTSClient()), patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(titan_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/amazon.titan-embed-text-v1", + input=test_input, + client=client, + aws_region_name="us-east-1", + aws_access_key_id="AKIAEMBEDCALLERKEY", + aws_secret_access_key="pod-caller-secret", + aws_role_name="arn:aws:iam::999999999999:role/litellm-embed-role", + aws_session_name="litellm-embed-session", + aws_session_tags=tags, + ) + + assert response.data[0]["embedding"] == titan_embedding_response["embedding"] + sent = mock_post.call_args.kwargs + assert "Credential=ASIAEMBEDTAGGED/" in sent["headers"]["Authorization"] + assert "aws_session_tags" not in sent["data"] + + def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): """The deployment's AWS profile does not exist, so resolving SigV4 credentials raises; a bearer-token deployment must still serve the request, since the @@ -1059,3 +1122,217 @@ def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(mo assert response.data[0]["embedding"] == titan_embedding_response["embedding"] assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_async_single_func_embeddings_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: Titan, Nova, and TwelveLabs embeddings sign one SigV4 request per + input, and botocore refreshes expiring credentials inside that signing with a blocking HTTP call, + so each signing must run on a worker thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + client = MagicMock() + client.__class__ = AsyncHTTPHandler + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json=titan_embedding_response, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), + ) + ) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await BedrockEmbedding()._async_single_func_embeddings( + client=client, + timeout=None, + batch_data=[{"inputText": test_input}], + credentials=probe.credentials(), + extra_headers=None, + endpoint_url="https://bedrock-runtime.us-west-2.amazonaws.com/model/amazon.titan-embed-text-v1/invoke", + aws_region_name="us-west-2", + model="amazon.titan-embed-text-v1", + logging_obj=MagicMock(), + provider="amazon", + ) + await release + + assert response.data[0]["embedding"] == titan_embedding_response["embedding"] + assert "Authorization" in client.post.call_args.kwargs["headers"] + assert probe.served_during_refresh is True +marengo_3_embedding_response = {"data": [{"embedding": [0.01 * i for i in range(512)]}]} +MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw==" + + +@pytest.mark.parametrize( + "model,kwargs,expected_body,expected_usage_details", + [ + ( + "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "text"}, + {"inputType": "text", "text": {"inputText": "a duck on water"}}, + {"query_count": 1}, + ), + ( + "bedrock/twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "text"}, + {"inputType": "text", "text": {"inputText": "a duck on water"}}, + {"query_count": 1}, + ), + ( + "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "text_image", "media_source": MARENGO_3_DUCK}, + { + "inputType": "text_image", + "text_image": {"inputText": "a duck on water", "mediaSource": {"base64String": "ZHVjaw=="}}, + }, + {"query_count": 1, "image_count": 1}, + ), + ( + "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "multi_input", "media_sources": {"bird": MARENGO_3_DUCK}}, + { + "inputType": "multi_input", + "multi_input": { + "inputText": "a duck on water", + "mediaSources": [{"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="}], + }, + }, + {"query_count": 1, "image_count": 1}, + ), + ], +) +def test_marengo_3_embedding_sends_the_nested_payload_and_parses_512_dims( + model, kwargs, expected_body, expected_usage_details +): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(marengo_3_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model=model, + input="a duck on water", + client=client, + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + **kwargs, + ) + + assert json.loads(mock_post.call_args.kwargs["data"]) == expected_body + assert mock_post.call_args.kwargs["url"].endswith(f"/model/{model.removeprefix('bedrock/').replace(':', '%3A')}/invoke") + assert len(response.data[0]["embedding"]) == 512 + assert response.data[0]["embedding"][:2] == [0.0, 0.01] + assert response.usage.prompt_tokens == 0 + assert response.usage.total_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == expected_usage_details + + +def test_marengo_3_image_embedding_sends_the_media_under_the_image_key(): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(marengo_3_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + input=MARENGO_3_DUCK, + client=client, + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + input_type="image", + ) + + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "inputType": "image", + "image": {"mediaSource": {"base64String": "ZHVjaw=="}}, + } + assert len(response.data[0]["embedding"]) == 512 + assert response.data[0]["embedding"][:2] == [0.0, 0.01] + assert response.usage.prompt_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"image_count": 1} + + +def test_marengo_2_7_embedding_keeps_the_flat_payload(): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(twelvelabs_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0", + input="a duck on water", + client=client, + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + input_type="text", + ) + + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "inputType": "text", + "inputText": "a duck on water", + "textTruncate": "end", + } + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert response.usage.prompt_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1} + + +def test_marengo_usage_counts_text_requests_and_images_across_a_batch(): + duck = {"mediaType": "image", "base64String": "ZHVjaw=="} + response = TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=[marengo_3_embedding_response, marengo_3_embedding_response, marengo_3_embedding_response], + model="us.twelvelabs.marengo-embed-3-0-v1:0", + batch_data=[ + {"inputType": "text", "text": {"inputText": "a duck"}}, + {"inputType": "image", "image": {"mediaSource": {"base64String": "ZHVjaw=="}}}, + {"inputType": "multi_input", "multi_input": {"mediaSources": [{"name": "a", **duck}, {"name": "b", **duck}]}}, + ], + ) + + assert [item["index"] for item in response.data] == [0, 1, 2] + assert response.usage.prompt_tokens == 0 + assert response.usage.total_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1, "image_count": 3} + + +def test_marengo_usage_without_request_data_bills_nothing(): + response = TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=[marengo_3_embedding_response], model="us.twelvelabs.marengo-embed-3-0-v1:0" + ) + + assert len(response.data[0]["embedding"]) == 512 + assert response.usage.prompt_tokens == 0 + assert response.usage.prompt_tokens_details is None + + +def test_marengo_response_items_without_an_embedding_are_skipped(): + response = TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=[{"data": [{"embeddingOption": "visual-text", "startSec": 0.0}, {"embedding": [0.1, 0.2, 0.3]}]}], + model="us.twelvelabs.marengo-embed-3-0-v1:0", + ) + + assert [item["embedding"] for item in response.data] == [[0.1, 0.2, 0.3]] + assert response.data[0]["index"] == 0 + + +def test_marengo_3_text_image_without_media_source_is_a_bad_request(): + with pytest.raises(litellm.BadRequestError, match=r"text_image.*media_source"): + litellm.embedding( + model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + input="a duck on water", + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + input_type="text_image", + ) diff --git a/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py new file mode 100644 index 00000000000..f149953b6f1 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py @@ -0,0 +1,416 @@ +import json +from unittest.mock import Mock, patch + +import pytest + +import litellm +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( + MARENGO_2_7_ONLY_PARAMS, + build_marengo_3_request, + is_marengo_3_model, +) +from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import ( + TwelveLabsMarengoEmbeddingConfig, + drop_params_enabled, +) + +MARENGO_3_BASE = "twelvelabs.marengo-embed-3-0-v1:0" +MARENGO_3_US = "us.twelvelabs.marengo-embed-3-0-v1:0" +MARENGO_27_US = "us.twelvelabs.marengo-embed-2-7-v1:0" +DUCK_DATA_URL = "data:image/png;base64,ZHVjaw==" +OUTPUT_S3_URI = "s3://out-bucket/marengo/" + + +@pytest.mark.parametrize( + "model,expected", + [ + (MARENGO_3_BASE, True), + (MARENGO_3_US, True), + ("eu.twelvelabs.marengo-embed-3-0-v1:0", True), + ("async_invoke/twelvelabs.marengo-embed-3-0-v1:0", True), + (MARENGO_27_US, False), + ("twelvelabs.marengo-embed-2-7-v1:0", False), + ("twelvelabs.marengo-embed-30-v1:0", False), + (None, False), + ], +) +def test_is_marengo_3_model(model, expected): + assert is_marengo_3_model(model) is expected + + +def wire(request: object) -> object: + return json.loads(json.dumps(request)) + + +def test_text_request_nests_input_text_under_text(): + assert build_marengo_3_request("a dog on the beach", {"input_type": "text"}) == { + "inputType": "text", + "text": {"inputText": "a dog on the beach"}, + } + + +def test_missing_input_type_defaults_to_text(): + assert build_marengo_3_request("hello", {})["inputType"] == "text" + + +def test_camel_case_input_type_wins_over_snake_case(): + request = build_marengo_3_request(DUCK_DATA_URL, {"inputType": "image", "input_type": "text"}) + assert request["inputType"] == "image" + + +def test_image_request_strips_data_url_prefix(): + assert build_marengo_3_request(DUCK_DATA_URL, {"input_type": "image"}) == { + "inputType": "image", + "image": {"mediaSource": {"base64String": "ZHVjaw=="}}, + } + + +def test_image_request_from_s3_carries_bucket_owner(): + request = build_marengo_3_request("s3://media/duck.png", {"input_type": "image", "bucketOwner": "123456789012"}) + assert request == { + "inputType": "image", + "image": {"mediaSource": {"s3Location": {"uri": "s3://media/duck.png", "bucketOwner": "123456789012"}}}, + } + + +@pytest.mark.parametrize( + "input_media,params", + [ + ("s3://media/duck.png", {"input_type": "image"}), + ("s3://media/clip.mp4", {"input_type": "video"}), + ("a duck", {"input_type": "text_image", "media_source": "s3://media/duck.png"}), + ("a duck", {"input_type": "multi_input", "media_sources": {"img1": "s3://media/duck.png"}}), + ], +) +def test_s3_media_without_bucket_owner_is_rejected_naming_it(input_media, params): + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request(input_media, params) + assert excinfo.value.status_code == 400 + assert excinfo.value.message == ( + "s3:// media requires the 'bucketOwner' parameter, the account id that owns the bucket" + ) + + +def test_text_image_request_pairs_text_with_media_source(): + request = build_marengo_3_request( + "a duck", {"input_type": "text_image", "media_source": DUCK_DATA_URL, "output_s3_uri": OUTPUT_S3_URI} + ) + assert request == { + "inputType": "text_image", + "text_image": {"inputText": "a duck", "mediaSource": {"base64String": "ZHVjaw=="}}, + } + + +def test_text_image_request_requires_media_source(): + with pytest.raises(BedrockError, match=r"text_image.*media_source") as excinfo: + build_marengo_3_request("a duck", {"input_type": "text_image"}) + assert excinfo.value.status_code == 400 + + +def test_multi_input_request_names_each_media_source(): + request = build_marengo_3_request( + "a photo of <@bird> next to <@dog>", + { + "input_type": "multi_input", + "media_sources": {"bird": DUCK_DATA_URL, "dog": "s3://media/dog.png"}, + "bucketOwner": "123456789012", + }, + ) + assert wire(request) == { + "inputType": "multi_input", + "multi_input": { + "inputText": "a photo of <@bird> next to <@dog>", + "mediaSources": [ + {"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="}, + { + "name": "dog", + "mediaType": "image", + "s3Location": {"uri": "s3://media/dog.png", "bucketOwner": "123456789012"}, + }, + ], + }, + } + + +def test_multi_input_without_text_omits_input_text(): + request = build_marengo_3_request("", {"input_type": "multi_input", "media_sources": {"bird": DUCK_DATA_URL}}) + assert "inputText" not in request["multi_input"] + assert request["multi_input"]["mediaSources"][0]["name"] == "bird" + + +@pytest.mark.parametrize("params", [{"input_type": "multi_input"}, {"input_type": "multi_input", "media_sources": {}}]) +def test_multi_input_request_requires_media_sources(params): + with pytest.raises(BedrockError, match=r"multi_input.*media_sources") as excinfo: + build_marengo_3_request("<@bird>", params) + assert excinfo.value.status_code == 400 + + +@pytest.mark.parametrize("input_type", ["video", "audio"]) +def test_timed_media_request_nests_every_option_under_the_media_key(input_type): + request = build_marengo_3_request( + "s3://media/clip.mp4", + { + "input_type": input_type, + "startSec": 2, + "endSec": 12.5, + "segmentation": {"method": "dynamic", "dynamic": {"minDurationSec": 4}}, + "embeddingOption": ["visual", "audio"], + "embeddingType": ["fused_embedding"], + "embeddingScope": ["clip", "asset"], + "inferenceId": "req-42", + "bucketOwner": "123456789012", + }, + ) + assert wire(request) == { + "inputType": input_type, + input_type: { + "mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}}, + "startSec": 2.0, + "endSec": 12.5, + "segmentation": {"method": "dynamic", "dynamic": {"minDurationSec": 4}}, + "embeddingOption": ["visual", "audio"], + "embeddingType": ["fused_embedding"], + "embeddingScope": ["clip", "asset"], + }, + "inferenceId": "req-42", + } + + +def test_timed_media_request_without_options_carries_only_the_media_source(): + request = build_marengo_3_request("s3://media/clip.mp4", {"input_type": "video", "bucketOwner": "123456789012"}) + assert request["video"] == { + "mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}} + } + + +@pytest.mark.parametrize( + "params", + [ + {"input_type": "clip"}, + {"input_type": "video", "embeddingOption": ["visual-text"]}, + {"input_type": "video", "segmentation": {"method": "fixed", "dynamic": {"minDurationSec": 4}}}, + {"input_type": "multi_input", "media_sources": ["not", "a", "mapping"]}, + ], +) +def test_invalid_marengo_3_params_are_rejected_before_the_request_is_sent(params): + with pytest.raises(BedrockError, match=r"Invalid Marengo 3\.0 parameters") as excinfo: + build_marengo_3_request("s3://media/clip.mp4", params) + assert excinfo.value.status_code == 400 + + +def test_config_sends_the_nested_payload_for_marengo_3_and_the_flat_one_for_2_7(): + nested = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US)._transform_request( + input="hello", inference_params={"input_type": "text"} + ) + flat = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US)._transform_request( + input="hello", inference_params={"input_type": "text"} + ) + assert nested == {"inputType": "text", "text": {"inputText": "hello"}} + assert flat == {"inputType": "text", "inputText": "hello", "textTruncate": "end"} + + +def test_config_without_a_model_keeps_the_2_7_payload(): + request = TwelveLabsMarengoEmbeddingConfig()._transform_request(input="hello", inference_params={}) + assert request == {"inputType": "text", "inputText": "hello", "textTruncate": "end"} + + +@pytest.mark.parametrize("input_type", ["video", "audio"]) +def test_marengo_3_video_and_audio_still_require_the_async_route(input_type): + with pytest.raises(ValueError, match=f"Input type '{input_type}' requires async_invoke route"): + TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request( + input="s3://media/clip.mp4", inference_params={"input_type": input_type} + ) + + +def test_marengo_3_async_invoke_wraps_the_nested_payload_with_the_base_model_id(): + request = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request( + input="s3://media/clip.mp4", + inference_params={ + "input_type": "video", + "embeddingOption": ["visual"], + "bucketOwner": "123456789012", + "output_s3_uri": OUTPUT_S3_URI, + }, + async_invoke_route=True, + model_id="async_invoke%2Ftwelvelabs.marengo-embed-3-0-v1%3A0", + output_s3_uri=OUTPUT_S3_URI, + ) + assert wire(request) == { + "modelId": MARENGO_3_BASE, + "modelInput": { + "inputType": "video", + "video": { + "mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}}, + "embeddingOption": ["visual"], + }, + }, + "outputDataConfig": {"s3OutputDataConfig": {"s3Uri": OUTPUT_S3_URI}}, + } + + +def test_marengo_3_async_invoke_requires_an_output_s3_uri(): + with pytest.raises(ValueError, match="output_s3_uri cannot be empty"): + TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request( + input="hello", + inference_params={"input_type": "text"}, + async_invoke_route=True, + model_id=MARENGO_3_BASE, + output_s3_uri="", + ) + + +def test_encoding_format_float_no_longer_injects_2_7_embedding_options_for_marengo_3(): + marengo_3 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).map_openai_params( + non_default_params={"encoding_format": "float"}, optional_params={} + ) + marengo_27 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US).map_openai_params( + non_default_params={"encoding_format": "float"}, optional_params={} + ) + assert marengo_3 == {} + assert marengo_27 == {"embeddingOption": ["visual-text", "visual-image"]} + + +def test_marengo_3_only_params_are_forwarded_by_map_openai_params(): + mapped = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).map_openai_params( + non_default_params={ + "input_type": "text_image", + "media_source": DUCK_DATA_URL, + "media_sources": {"bird": DUCK_DATA_URL}, + "endSec": 5, + "segmentation": {"method": "fixed", "fixed": {"durationSec": 6}}, + "embeddingType": ["separate_embedding"], + "embeddingScope": ["clip"], + "inferenceId": "req-1", + }, + optional_params={}, + ) + assert mapped == { + "inputType": "text_image", + "media_source": DUCK_DATA_URL, + "media_sources": {"bird": DUCK_DATA_URL}, + "endSec": 5, + "segmentation": {"method": "fixed", "fixed": {"durationSec": 6}}, + "embeddingType": ["separate_embedding"], + "embeddingScope": ["clip"], + "inferenceId": "req-1", + } + + +@pytest.mark.parametrize( + "params,problem", + [ + ( + {"input_type": "clip"}, + "input_type: Input should be 'text', 'image', 'video', 'audio', 'text_image' or 'multi_input'", + ), + ({"input_type": "video", "embeddingOption": "visual"}, "embeddingOption: Input should be a valid tuple"), + ( + {"input_type": "multi_input", "media_sources": ["not", "a", "mapping"]}, + "media_sources: Input should be a valid dictionary", + ), + ], +) +def test_invalid_marengo_3_params_name_the_field_and_the_reason(params, problem): + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request("s3://media/clip.mp4", params) + assert excinfo.value.message == f"Invalid Marengo 3.0 parameters: {problem}" + + +MARENGO_2_7_ONLY_VALUES = {"textTruncate": "end", "lengthSec": 5, "useFixedLengthSec": True, "minClipSec": 2} + + +@pytest.mark.parametrize("name", MARENGO_2_7_ONLY_PARAMS) +def test_marengo_2_7_only_params_are_rejected_on_3_0_unless_dropped(name): + params = {"input_type": "text", name: MARENGO_2_7_ONLY_VALUES[name]} + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request("hello", params) + assert excinfo.value.status_code == 400 + assert excinfo.value.message == ( + f"Marengo 3.0 does not accept the Marengo 2.7 parameters {name}; set drop_params to drop them" + ) + assert build_marengo_3_request("hello", params, drop_params=True) == { + "inputType": "text", + "text": {"inputText": "hello"}, + } + + +def test_marengo_2_7_only_params_are_advertised_only_for_2_7(): + marengo_3 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).get_supported_openai_params() + marengo_27 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US).get_supported_openai_params() + assert set(MARENGO_2_7_ONLY_PARAMS).isdisjoint(marengo_3) + assert set(MARENGO_2_7_ONLY_PARAMS) <= set(marengo_27) + assert set(marengo_3) <= set(marengo_27) + + +def test_drop_params_comes_from_the_call_or_the_global(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + assert drop_params_enabled({}) is False + assert drop_params_enabled({"drop_params": True}) is True + monkeypatch.setattr(litellm, "drop_params", True) + assert drop_params_enabled({}) is True + + +def test_config_drops_marengo_2_7_only_params_only_when_asked(): + config = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US) + with pytest.raises(BedrockError, match=r"Marengo 2\.7 parameters textTruncate"): + config._transform_request("hello", {"textTruncate": "end"}) + assert config._transform_request("hello", {"textTruncate": "end"}, drop_params=True) == { + "inputType": "text", + "text": {"inputText": "hello"}, + } + + +@pytest.mark.parametrize( + "params", + [ + {"input_type": "text"}, + {"input_type": "image"}, + {"input_type": "text_image", "media_source": DUCK_DATA_URL}, + {"input_type": "multi_input", "media_sources": {"bird": DUCK_DATA_URL}}, + ], +) +def test_timed_media_options_are_rejected_on_untimed_input_types_unless_dropped(params): + timed = {**params, "startSec": 0, "embeddingOption": ["visual"]} + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request(DUCK_DATA_URL, timed) + assert excinfo.value.status_code == 400 + assert excinfo.value.message == ( + f"Input type '{params['input_type']}' does not accept startSec, embeddingOption; set drop_params to drop them" + ) + assert build_marengo_3_request(DUCK_DATA_URL, timed, drop_params=True) == build_marengo_3_request( + DUCK_DATA_URL, params + ) + + +def _embed_marengo_3_us(client: HTTPHandler, **params: object): + return litellm.embedding( + model=f"bedrock/{MARENGO_3_US}", + input="hello", + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key="test-bearer-token", + **params, + ) + + +def test_per_request_drop_params_reaches_the_marengo_3_builder(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + client = HTTPHandler() + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps({"data": [{"embedding": [0.1, 0.2]}]}) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + with pytest.raises(litellm.BadRequestError, match=r"Marengo 2\.7 parameters textTruncate"): + _embed_marengo_3_us(client, textTruncate="end") + assert mock_post.call_count == 0 + + response = _embed_marengo_3_us(client, textTruncate="end", drop_params=True) + + assert response.data[0]["embedding"] == [0.1, 0.2] + assert json.loads(mock_post.call_args.kwargs["data"]) == {"inputType": "text", "text": {"inputText": "hello"}} diff --git a/tests/test_litellm/llms/bedrock/event_loop_probe.py b/tests/test_litellm/llms/bedrock/event_loop_probe.py new file mode 100644 index 00000000000..c347247ec32 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/event_loop_probe.py @@ -0,0 +1,57 @@ +"""Refreshable credentials whose refresh only completes while the event loop keeps serving.""" + +from __future__ import annotations + +import asyncio +import threading +import time +from datetime import datetime, timedelta, timezone +from typing import Final + +from botocore.credentials import RefreshableCredentials + +REFRESH_RELEASE_TIMEOUT_SECONDS: Final = 2.0 +REFRESH_START_TIMEOUT_SECONDS: Final = 10.0 + + +class EventLoopProbe: + """Blocks inside botocore's credential refresh until a coroutine on the loop releases it. + + Signing on the event loop thread can never be released, so `served_during_refresh` reads False there + and True only when the refresh ran on another thread while the loop stayed responsive. + """ + + def __init__(self) -> None: + self.refresh_started: Final = threading.Event() + self.loop_served: Final = threading.Event() + self.served_during_refresh: bool | None = None + + def refresh(self) -> dict[str, str | None]: + self.refresh_started.set() + served: Final = self.loop_served.wait(timeout=REFRESH_RELEASE_TIMEOUT_SECONDS) + if self.served_during_refresh is None: + self.served_during_refresh = served + return { + "access_key": "AKIAREFRESHED", + "secret_key": "refreshed-secret", + "token": None, + "expiry_time": (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat(), + } + + def credentials(self) -> RefreshableCredentials: + return RefreshableCredentials( + access_key="AKIASTALE", + secret_key="stale-secret", + token=None, + expiry_time=datetime.now(timezone.utc) + timedelta(seconds=60), + refresh_using=self.refresh, + method="event-loop-probe", + ) + + async def release_refresh_from_the_loop(self) -> None: + deadline: Final = time.monotonic() + REFRESH_START_TIMEOUT_SECONDS + while not self.refresh_started.is_set(): + if time.monotonic() > deadline: + raise TimeoutError("signing finished without ever starting a credential refresh") + await asyncio.sleep(0.005) + self.loop_served.set() diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 541c0db15d8..3b01a4f2054 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -5,6 +5,8 @@ Test bedrock files transformation functionality import json import os from collections.abc import Mapping +from contextlib import AsyncExitStack, closing +from typing import Final from unittest.mock import MagicMock from urllib.parse import unquote, urlparse @@ -1855,6 +1857,104 @@ class TestBedrockBatchNonChatEndpointRecords: ] +class TestBedrockFileDeletion: + S3_URI: Final = "s3://my-bucket/litellm-bedrock-files-model-abc.jsonl" + URL: Final = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files-model-abc.jsonl" + + def test_interleaved_deletions_keep_their_own_file_ids(self, monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + config: Final = BedrockFilesConfig() + params: Final = { + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + } + file_ids: Final = (self.S3_URI, "s3://my-bucket/litellm-bedrock-files-model-second.jsonl") + for file_id in file_ids: + config.transform_delete_file_request(file_id=file_id, optional_params={}, litellm_params=params) + + deleted: Final = tuple( + config.transform_delete_file_response( + raw_response=httpx.Response(204), + logging_obj=MagicMock(model_call_details={"additional_args": {"file_id": file_id}}), + litellm_params=params, + ).id + for file_id in file_ids + ) + + assert deleted == file_ids + + def test_delete_file_sends_signed_delete_and_returns_matching_id(self, monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + import respx + + import litellm + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + with respx.mock, closing(HTTPHandler()) as client: + route: Final = respx.delete(self.URL).mock(return_value=httpx.Response(204)) + deleted: Final = litellm.file_delete( + file_id=self.S3_URI, custom_llm_provider="bedrock", client=client, + aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key="test-secret", aws_region_name="us-west-2", + ) + assert route.call_count == 1 + request: Final = route.calls[0].request + assert request.content == b"" + signed: Final = AWSRequest(method="DELETE", url=self.URL, headers={ + "X-Amz-Date": request.headers["X-Amz-Date"], + "X-Amz-Content-SHA256": request.headers["X-Amz-Content-SHA256"], + }) + signed.context["timestamp"] = request.headers["X-Amz-Date"] + auth: Final = S3SigV4Auth(Credentials("AKIAEXAMPLE", "test-secret"), "s3", "us-west-2") + signature: Final = auth.signature(auth.string_to_sign(signed, auth.canonical_request(signed)), signed) + assert request.headers["Authorization"].endswith(f"Signature={signature}") + assert deleted.id == self.S3_URI and deleted.deleted is True + + @pytest.mark.asyncio + async def test_adelete_file_propagates_s3_errors(self, monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + import respx + + import litellm + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + async with AsyncExitStack() as stack: + client: Final = AsyncHTTPHandler() + stack.push_async_callback(client.close) + with respx.mock: + route: Final = respx.delete(self.URL).mock( + return_value=httpx.Response(403, content=b"AccessDenied") + ) + from litellm.llms.bedrock.common_utils import BedrockError + + with pytest.raises(BedrockError, match="AccessDenied"): + await litellm.afile_delete( + file_id=self.S3_URI, custom_llm_provider="bedrock", client=client, + aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key="test-secret", aws_region_name="us-west-2", + ) + assert route.call_count == 1 + + @pytest.mark.parametrize("file_id, message", [ + ("s3://other-bucket/litellm-bedrock-files-model-abc.jsonl", "configured storage bucket"), + ("s3://my-bucket/private/data.jsonl", "LiteLLM-managed"), + ]) + def test_delete_rejects_untrusted_objects_before_signing( + self, file_id: str, message: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + with pytest.raises(ValueError, match=message): + BedrockFilesConfig().transform_delete_file_request(file_id=file_id, optional_params={}, litellm_params={}) + + class TestBedrockFileContentTransformation: """SigV4-signed S3 GetObject retrieval of Bedrock batch output files.""" @@ -1873,7 +1973,7 @@ class TestBedrockFileContentTransformation: import hashlib from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) @@ -1889,7 +1989,7 @@ class TestBedrockFileContentTransformation: assert url == self.EXPECTED_URL assert params == {} - signed_headers = litellm_params[S3_SIGNED_GET_HEADERS_PARAM] + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] content_hashes = { value for name, value in signed_headers.items() @@ -2139,7 +2239,7 @@ class TestBedrockFileContentTransformation: def test_s3_region_name_wins_for_content_signing(self, monkeypatch): """s3_region_name must override aws_region_name for both the URL and the signature.""" from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) @@ -2154,17 +2254,17 @@ class TestBedrockFileContentTransformation: ) assert url.startswith("https://s3.eu-west-1.amazonaws.com/") - authorization = litellm_params[S3_SIGNED_GET_HEADERS_PARAM]["Authorization"] + authorization = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]["Authorization"] assert "/eu-west-1/s3/aws4_request" in authorization def test_validate_environment_merges_and_pops_signed_get_headers(self): from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) litellm_params = { - S3_SIGNED_GET_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"} + S3_SIGNED_REQUEST_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"} } headers = BedrockFilesConfig().validate_environment( @@ -2179,7 +2279,7 @@ class TestBedrockFileContentTransformation: "x-custom": "kept", "Authorization": "AWS4-HMAC-SHA256 test", } - assert S3_SIGNED_GET_HEADERS_PARAM not in litellm_params + assert S3_SIGNED_REQUEST_HEADERS_PARAM not in litellm_params def test_transform_file_content_response_wraps_binary_content(self): import httpx @@ -2379,7 +2479,7 @@ class TestBedrockFilesS3SignatureEncoding: self, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) @@ -2402,7 +2502,7 @@ class TestBedrockFilesS3SignatureEncoding: method="GET", url=url, body=None, - headers=litellm_params[S3_SIGNED_GET_HEADERS_PARAM], + headers=litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM], ) @@ -2457,7 +2557,7 @@ def test_sign_s3_request_assumes_role_with_external_id(monkeypatch): assert "ASIAFILESPUTROLE" in authorization -def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch): +def test_sign_s3_request_without_body_assumes_role_with_external_id(monkeypatch): """A trust policy requiring sts:ExternalId must be satisfied when signing the S3 download request.""" import datetime from unittest.mock import patch @@ -2504,7 +2604,7 @@ def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch): assert request_params.aws_external_id == "external-id-files-get" with patch.object(boto3, "client", return_value=FakeSTSClient()): - signed_headers = BedrockFilesConfig()._sign_s3_get_request( + signed_headers = BedrockFilesConfig()._sign_s3_request_without_body( api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", aws_region_name="us-east-1", request_params=request_params, diff --git a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index b005d77ac8b..f2a9af11af7 100644 --- a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -1,7 +1,19 @@ +import base64 +import json +import struct +import tracemalloc +from binascii import crc32 +from datetime import datetime from unittest.mock import patch - +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig +from litellm.types.utils import ModelResponse + +CONVERSE_MODEL = "anthropic.claude-sonnet-4-5-20250929-v1:0" +CONVERSE_STREAM_ENDPOINT = f"/model/{CONVERSE_MODEL}/converse-stream" +INVOKE_STREAM_ENDPOINT = f"/model/{CONVERSE_MODEL}/invoke-with-response-stream" def test_bedrock_passthrough_get_complete_url_default_endpoint(): @@ -500,3 +512,186 @@ def test_bedrock_passthrough_model_id_without_arn(): f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model_id}/converse" ) assert url_str == expected_url + + +def _event_frame(event_type: str, payload: dict) -> bytes: + def header(name: str, value: str) -> bytes: + name_b, value_b = name.encode(), value.encode() + return struct.pack("!B", len(name_b)) + name_b + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b + + payload_b = json.dumps(payload, separators=(",", ":")).encode() + headers_b = ( + header(":event-type", event_type) + + header(":content-type", "application/json") + + header(":message-type", "event") + ) + prelude = struct.pack("!II", 12 + len(headers_b) + len(payload_b) + 4, len(headers_b)) + prelude_crc = crc32(prelude) & 0xFFFFFFFF + message = struct.pack("!I", prelude_crc) + headers_b + payload_b + return prelude + message + struct.pack("!I", crc32(message, prelude_crc) & 0xFFFFFFFF) + + +def _text_block(index: int, texts: list[str]) -> bytes: + return ( + _event_frame("contentBlockStart", {"contentBlockIndex": index, "start": {}}) + + b"".join( + _event_frame("contentBlockDelta", {"contentBlockIndex": index, "delta": {"text": text}}) for text in texts + ) + + _event_frame("contentBlockStop", {"contentBlockIndex": index}) + ) + + +def _stream_tail(stop_reason: str, output_tokens: int) -> bytes: + return _event_frame("messageStop", {"stopReason": stop_reason}) + _event_frame( + "metadata", + { + "metrics": {"latencyMs": 1234}, + "usage": {"inputTokens": 25, "outputTokens": output_tokens, "totalTokens": 25 + output_tokens}, + }, + ) + + +def _invoke_chunk(payload: dict) -> bytes: + return _event_frame("chunk", {"bytes": base64.b64encode(json.dumps(payload).encode()).decode()}) + + +def _stream_logging_obj(endpoint: str) -> Logging: + logging_obj = Logging( + model=CONVERSE_MODEL, + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "bedrock" + logging_obj.model_call_details["endpoint"] = endpoint + return logging_obj + + +def _converse_stream_logging_obj() -> Logging: + return _stream_logging_obj(CONVERSE_STREAM_ENDPOINT) + + +def _stream_collector(endpoint: str) -> PassthroughStreamCollector: + return BedrockPassthroughConfig().create_stream_collector( + model=CONVERSE_MODEL, custom_llm_provider="bedrock", endpoint=endpoint + ) + + +def _converse_stream_collector() -> PassthroughStreamCollector: + return _stream_collector(CONVERSE_STREAM_ENDPOINT) + + +def _feed(collector: PassthroughStreamCollector, stream: bytes, chunk_size: int = 16384) -> None: + for offset in range(0, len(stream), chunk_size): + collector.add(stream[offset : offset + chunk_size]) + + +def test_converse_stream_collector_keeps_usage_without_retaining_the_stream(): + texts = [f"tok{i} " for i in range(4000)] + stream = _event_frame("messageStart", {"role": "assistant"}) + _text_block(0, texts) + _stream_tail("end_turn", 4000) + _feed(_converse_stream_collector(), stream) + + tracemalloc.start() + try: + base = tracemalloc.get_traced_memory()[0] + collector = _converse_stream_collector() + _feed(collector, stream) + retained = tracemalloc.get_traced_memory()[0] - base + finally: + tracemalloc.stop() + + assert retained < len(stream) // 4 + + response = collector.build_logged_response(_converse_stream_logging_obj()) + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "".join(texts) + assert response.choices[0].finish_reason == "stop" + assert (response.usage.prompt_tokens, response.usage.completion_tokens) == (25, 4000) + + +def test_converse_stream_collector_keeps_tool_calls_between_text_runs(): + stream = ( + _event_frame("messageStart", {"role": "assistant"}) + + _text_block(0, ["Let me ", "check."]) + + _event_frame( + "contentBlockStart", + {"contentBlockIndex": 1, "start": {"toolUse": {"toolUseId": "tool-1", "name": "get_weather"}}}, + ) + + _event_frame("contentBlockDelta", {"contentBlockIndex": 1, "delta": {"toolUse": {"input": '{"city": '}}}) + + _event_frame("contentBlockDelta", {"contentBlockIndex": 1, "delta": {"toolUse": {"input": '"Paris"}'}}}) + + _event_frame("contentBlockStop", {"contentBlockIndex": 1}) + + _text_block(2, ["Done", "."]) + + _stream_tail("tool_use", 12) + ) + collector = _converse_stream_collector() + _feed(collector, stream, chunk_size=7) + + response = collector.build_logged_response(_converse_stream_logging_obj()) + assert isinstance(response, ModelResponse) + message = response.choices[0].message + assert message.content == "Let me check.Done." + assert [(call.function.name, call.function.arguments) for call in message.tool_calls] == [ + ("get_weather", '{"city": "Paris"}') + ] + assert response.choices[0].finish_reason == "tool_calls" + assert (response.usage.prompt_tokens, response.usage.completion_tokens) == (25, 12) + + +def test_invoke_stream_collector_keeps_usage_without_retaining_the_stream(): + texts = [f"tok{i} " for i in range(4000)] + stream = ( + _invoke_chunk( + { + "type": "message_start", + "message": { + "id": "msg-1", + "type": "message", + "role": "assistant", + "model": CONVERSE_MODEL, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 25, "output_tokens": 1}, + }, + } + ) + + _invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}) + + b"".join( + _invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}) + for text in texts + ) + + _invoke_chunk({"type": "content_block_stop", "index": 0}) + + _invoke_chunk( + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4000}} + ) + + _invoke_chunk({"type": "message_stop"}) + ) + _feed(_stream_collector(INVOKE_STREAM_ENDPOINT), stream) + + tracemalloc.start() + try: + base = tracemalloc.get_traced_memory()[0] + collector = _stream_collector(INVOKE_STREAM_ENDPOINT) + _feed(collector, stream) + retained = tracemalloc.get_traced_memory()[0] - base + finally: + tracemalloc.stop() + + assert retained < len(stream) // 4 + + response = collector.build_logged_response(_stream_logging_obj(INVOKE_STREAM_ENDPOINT)) + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "".join(texts) + assert response.choices[0].finish_reason == "stop" + assert (response.usage.prompt_tokens, response.usage.completion_tokens) == (25, 4000) + + +def test_stream_collector_logs_nothing_for_an_unrecognized_endpoint(): + collector = BedrockPassthroughConfig().create_stream_collector( + model=CONVERSE_MODEL, custom_llm_provider="bedrock", endpoint=f"/model/{CONVERSE_MODEL}/rerank" + ) + collector.add(_event_frame("messageStart", {"role": "assistant"})) + + assert collector.build_logged_response(_converse_stream_logging_obj()) is None diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index f854d806bdc..c5b8e7ecc9d 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1,4 +1,6 @@ +import asyncio import json +from concurrent.futures import ThreadPoolExecutor import os import threading import time @@ -15,14 +17,17 @@ from unittest.mock import MagicMock, patch from botocore.awsrequest import AWSPreparedRequest, AWSRequest from botocore.auth import SigV4Auth from botocore.credentials import Credentials -from botocore.exceptions import NoCredentialsError +from botocore.exceptions import ClientError, NoCredentialsError import litellm from litellm.llms.bedrock.base_aws_llm import ( AwsAuthError, BaseAWSLLM, Boto3CredentialsInfo, + run_aws_signing, + sign_request_off_loop_if_aws, ) +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Global variable for the base_aws_llm.py file path @@ -38,6 +43,14 @@ def flush_shared_bedrock_iam_cache(): yield +@pytest.fixture(autouse=True) +def _clean_ssl_env(monkeypatch): + """get_ssl_verify reads these, so the sts client's verify= would otherwise depend on + the ambient environment. The published images set SSL_CERT_FILE.""" + for env_var in ("SSL_CERT_FILE", "SSL_VERIFY"): + monkeypatch.delenv(env_var, raising=False) + + def test_base_aws_llm_instances_share_process_wide_iam_cache(): """Regression LIT-2662: new instances must reuse iam_cache (Bedrock passthrough is per-request).""" first = BaseAWSLLM() @@ -2395,6 +2408,283 @@ def test_assume_role_without_external_id(): ) +_SESSION_TAGS = ({"Key": "team", "Value": "genai"}, {"Key": "env", "Value": "prod"}) +_SORTED_SESSION_TAGS = ({"Key": "env", "Value": "prod"}, {"Key": "team", "Value": "genai"}) +_TAGGED_ROLE_ARN = "arn:aws:iam::123456789012:role/TaggedRole" + + +class _TagAwareSTSClient: + """STS stand-in for a trust policy that only admits sessions carrying exactly the expected tags.""" + + def __init__(self, expected_tags: tuple = (), access_key: str = "ASIATAGGEDSESSION") -> None: + self.expected_tags = expected_tags + self.access_key = access_key + self.assume_role_calls: list = [] + + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role_with_web_identity(self, **params): + return { + "Credentials": { + "AccessKeyId": "ASIAIRSATEMP", + "SecretAccessKey": "irsa-temp-secret-key", + "SessionToken": "irsa-temp-session-token", + "Expiration": datetime.now(timezone.utc) + timedelta(hours=1), + } + } + + def assume_role(self, **params): + self.assume_role_calls.append(params) + if tuple(params.get("Tags", ())) != self.expected_tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": self.access_key, + "SecretAccessKey": "assumed-secret-key", + "SessionToken": "assumed-session-token", + "Expiration": datetime.now(timezone.utc) + timedelta(hours=1), + } + } + + +def _irsa_env(tmp_path, irsa_role_arn: str) -> dict: + token_file = tmp_path / "web-identity-token" + token_file.write_text("test-web-identity-token") + return { + "AWS_WEB_IDENTITY_TOKEN_FILE": str(token_file), + "AWS_ROLE_ARN": irsa_role_arn, + "AWS_REGION": "us-east-1", + } + + +def test_assume_role_sends_session_tags(): + """The STS session carries the configured tags, so a trust policy gated on sts:TagSession admits it.""" + sts = _TagAwareSTSClient(expected_tags=_SESSION_TAGS) + + with patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_session_tags=list(_SESSION_TAGS), + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + {"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "test-session", "Tags": _SESSION_TAGS} + ] + + +def test_assume_role_sends_session_tags_alongside_external_id(): + sts = _TagAwareSTSClient(expected_tags=_SESSION_TAGS) + + with patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_external_id="UniqueExternalID123", + aws_session_tags=_SESSION_TAGS, + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + { + "RoleArn": _TAGGED_ROLE_ARN, + "RoleSessionName": "test-session", + "ExternalId": "UniqueExternalID123", + "Tags": _SESSION_TAGS, + } + ] + + +@pytest.mark.parametrize("aws_session_tags", [None, [], ()], ids=["none", "empty-list", "empty-tuple"]) +def test_assume_role_omits_the_tags_key_without_session_tags(aws_session_tags): + """Nothing configured means the AssumeRole request looks exactly as it did before tags existed.""" + sts = _TagAwareSTSClient(expected_tags=()) + + with patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_session_tags=aws_session_tags, + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [{"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "test-session"}] + + +def test_irsa_cross_account_assume_role_sends_session_tags(tmp_path): + irsa_role_arn = "arn:aws:iam::111111111111:role/eks-service-account-role" + sts = _TagAwareSTSClient(expected_tags=_SESSION_TAGS) + + with patch.dict(os.environ, _irsa_env(tmp_path, irsa_role_arn)), patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_session_tags=_SESSION_TAGS, + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + {"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "test-session", "Tags": _SESSION_TAGS} + ] + + +def test_irsa_same_account_assume_role_sends_session_tags(tmp_path): + sts = _TagAwareSTSClient(expected_tags=_SESSION_TAGS) + + with patch.dict(os.environ, _irsa_env(tmp_path, _TAGGED_ROLE_ARN)), patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_session_tags=_SESSION_TAGS, + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + {"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "test-session", "Tags": _SESSION_TAGS} + ] + + +def test_get_credentials_canonicalizes_session_tag_order_for_the_cache(): + """Two deployments listing the same tags in a different order share one STS session.""" + base_aws_llm = BaseAWSLLM() + sts = _TagAwareSTSClient(expected_tags=_SORTED_SESSION_TAGS) + + with patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True), patch("boto3.client", return_value=sts): + first = base_aws_llm.get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=list(_SESSION_TAGS), + ) + second = base_aws_llm.get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=list(reversed(_SESSION_TAGS)), + ) + + assert first.access_key == second.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + {"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "team-session", "Tags": _SORTED_SESSION_TAGS} + ] + + +def test_get_credentials_scopes_the_cache_per_session_tag_set(): + """Different tag sets are different principals to AWS, so each gets its own STS session.""" + base_aws_llm = BaseAWSLLM() + mock_sts_client = _assume_role_sts_mock() + mock_sts_client.assume_role.side_effect = [ + { + "Credentials": { + "AccessKeyId": f"assumed-access-key-{team}", + "SecretAccessKey": "assumed-secret-key", + "SessionToken": f"assumed-session-token-{team}", + "Expiration": datetime.now(timezone.utc) + timedelta(hours=1), + } + } + for team in ("genai", "platform") + ] + + with ( + patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True), + patch("boto3.client", return_value=mock_sts_client), + ): + genai = base_aws_llm.get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=[{"Key": "team", "Value": "genai"}], + ) + platform = base_aws_llm.get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=[{"Key": "team", "Value": "platform"}], + ) + + assert genai.access_key == "assumed-access-key-genai" + assert platform.access_key == "assumed-access-key-platform" + assert [call.kwargs["Tags"] for call in mock_sts_client.assume_role.call_args_list] == [ + ({"Key": "team", "Value": "genai"},), + ({"Key": "team", "Value": "platform"},), + ] + + +@pytest.mark.parametrize( + "aws_session_tags", + [ + "team=genai", + {"team": "genai"}, + [["team", "genai"]], + [{"key": "team", "value": "genai"}], + [{"Key": "team"}], + [{"Key": 1, "Value": "genai"}], + ], + ids=["string", "flat-dict", "pair-list", "lowercase-keys", "missing-value", "non-string-key"], +) +def test_get_credentials_rejects_malformed_session_tags(aws_session_tags): + with pytest.raises(ValueError, match="Invalid 'aws_session_tags' value"): + BaseAWSLLM().get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=aws_session_tags, + ) + + +def test_get_boto_credentials_from_optional_params_consumes_session_tags(): + """Tags feed the STS call and must not linger in optional_params to be serialized into the body.""" + sts = _TagAwareSTSClient(expected_tags=_SORTED_SESSION_TAGS) + optional_params = { + "aws_region_name": "us-east-1", + "aws_role_name": _TAGGED_ROLE_ARN, + "aws_session_name": "team-session", + "aws_session_tags": list(_SESSION_TAGS), + } + + with patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True), patch("boto3.client", return_value=sts): + target = BaseAWSLLM()._get_boto_credentials_from_optional_params(optional_params) + + assert target.credentials.access_key == "ASIATAGGEDSESSION" + assert "aws_session_tags" not in optional_params + + +def test_sign_request_signs_with_the_tagged_sts_session(): + sts = _TagAwareSTSClient(expected_tags=_SORTED_SESSION_TAGS) + optional_params = { + "aws_region_name": "us-east-1", + "aws_role_name": _TAGGED_ROLE_ARN, + "aws_session_name": "team-session", + "aws_session_tags": list(_SESSION_TAGS), + } + + with patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True), patch("boto3.client", return_value=sts): + headers, _body = BaseAWSLLM()._sign_request( + service_name="bedrock", + headers={}, + optional_params=optional_params, + request_data={"prompt": "hi"}, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-opus-5/invoke", + ) + + assert "Credential=ASIATAGGEDSESSION/" in headers["Authorization"] + + def test_converse_handler_external_id_extraction(): """Test that BedrockConverseLLM properly extracts and passes aws_external_id parameter""" from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM @@ -3215,3 +3505,53 @@ class TestGetRequestHeadersResign: extra_headers={"Authorization": "Bearer foo"}, ) assert prepped.headers["Authorization"] == "Bearer foo" + + +@pytest.mark.asyncio +async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credentials_refresh(): + """Regression for issue #40165: an AWS provider's signing (and the botocore credential refresh + inside it) must run off the event loop, so other requests keep being served meanwhile.""" + probe = EventLoopProbe() + + def sign(headers: dict[str, str]) -> dict[str, str]: + request = AWSRequest( + method="POST", url="https://bedrock-runtime.us-west-2.amazonaws.com/", data="{}", headers=headers + ) + SigV4Auth(probe.credentials(), "bedrock", "us-west-2").add_auth(request) + return dict(request.headers) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + signed = await sign_request_off_loop_if_aws(BaseAWSLLM(), sign, headers={"Content-Type": "application/json"}) + await release + + assert "Authorization" in signed + assert probe.served_during_refresh is True + + +def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers(): + """A signing parked on botocore's refresh lock must not hold a default-executor thread, since every + other provider's async entry point hops through that same executor. The scenario runs on its own loop + so the one-thread default executor it pins never leaks into the session loop.""" + + async def scenario() -> tuple[str, str]: + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + signing_parked = asyncio.Event() + refresh_done = threading.Event() + + def sign() -> str: + loop.call_soon_threadsafe(signing_parked.set) + refresh_done.wait() + return threading.current_thread().name + + signing = asyncio.create_task(run_aws_signing(sign)) + try: + await asyncio.wait_for(signing_parked.wait(), timeout=5) + other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5) + finally: + refresh_done.set() + return other_provider.name, await signing + + other_provider, signing_thread = asyncio.run(scenario()) + assert other_provider != signing_thread + assert signing_thread.startswith("aws-signing") diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 9302dc01abe..a8a21e2cd37 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -486,6 +486,7 @@ def test_merge_bedrock_aws_request_params_strips_caller_identity_when_deployment "aws_role_name": "arn:aws:iam::123456789012:role/caller", "aws_session_token": "caller-token", "aws_web_identity_token": "caller-web-identity", + "aws_session_tags": [{"Key": "team", "Value": "caller-chosen"}], "timeout": 600, }, ) @@ -500,6 +501,7 @@ def test_merge_bedrock_aws_request_params_strips_caller_identity_when_deployment "aws_role_name", "aws_session_token", "aws_web_identity_token", + "aws_session_tags", ): assert stripped not in merged @@ -614,3 +616,314 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] assert "ASIABATCHSIGNROLE" in authorization assert signed_data == b'{"jobName": "litellm-batch-job"}' + + +def test_sign_aws_request_assumes_role_with_session_tags(monkeypatch): + """Batch and file signing must carry the deployment's session tags into the AssumeRole call too.""" + import datetime + from unittest.mock import patch + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils + + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIABATCHSIGNTAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIABATCHSIGNCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-batch-sign-role", + "aws_session_name": "litellm-batch-sign-session", + "aws_session_tags": tags, + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + signed_headers, _signed_data = CommonBatchFilesUtils().sign_aws_request( + service_name="bedrock", + data={"jobName": "litellm-batch-job"}, + endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job", + optional_params=optional_params, + ) + + authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] + assert "Credential=ASIABATCHSIGNTAGGED/" in authorization + + +# --------------------------------------------------------------------------- # +# Provider error headers (LIT-5428) # +# --------------------------------------------------------------------------- # + + +def _bedrock_chat_error_configs(): + from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.llms.bedrock.chat.invoke_agent.transformation import AmazonInvokeAgentConfig + from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( + AmazonMoonshotConfig, + ) + from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, + ) + from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig, + ) + + return [ + AmazonInvokeConfig, + AmazonConverseConfig, + AmazonMoonshotConfig, + AmazonBedrockOpenAIConfig, + AmazonAgentCoreConfig, + AmazonInvokeAgentConfig, + ] + + +@pytest.mark.parametrize("config", _bedrock_chat_error_configs()) +def test_bedrock_chat_get_error_class_keeps_provider_headers(config): + """Every Bedrock chat route must carry x-amzn-RequestId out to the caller (LIT-5428). + + A config that drops the headers it is handed shadows the fix for its own models. + """ + error = config().get_error_class( + error_message="Amazon Bedrock is unable to process your request.", + status_code=500, + headers={"x-amzn-RequestId": "req-chat-500"}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-chat-500" + + +def test_error_response_text_reads_a_read_response(): + import httpx + + from litellm.llms.bedrock.common_utils import error_response_text + + response = httpx.Response(status_code=500, text="Amazon Bedrock is unable to process your request.") + + assert error_response_text(response) == "Amazon Bedrock is unable to process your request." + + +def test_error_response_text_falls_back_when_a_streamed_response_was_never_read(): + """A retried streamed request raises HTTPStatusError over an unread body; reading it + throws ResponseNotRead and would lose the status and headers this fix preserves.""" + import httpx + + from litellm.llms.bedrock.common_utils import error_response_text + + request = httpx.Request(method="POST", url="https://bedrock-runtime.amazonaws.com") + response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-unread-500"}, + stream=httpx.ByteStream(b"never read"), + request=request, + ) + + with pytest.raises(httpx.ResponseNotRead): + _ = response.text + + assert error_response_text(response) == "Internal Server Error" + + +def test_bedrock_error_skips_header_values_httpx_cannot_carry(): + """The shared HTTP handler copies an arbitrary exception's header values in verbatim, + so a non-str value must not take down the whole error (LIT-5428).""" + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError( + status_code=500, + message="boom", + headers={"x-amzn-RequestId": "req-mixed-500", "x-retry-count": 3, "x-nothing": None}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-mixed-500" + assert "x-retry-count" not in error.response.headers + assert isinstance(error.response, httpx.Response) + + +def test_bedrock_error_keeps_duplicate_httpx_header_values(): + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError( + status_code=500, + message="boom", + headers=httpx.Headers([("x-amzn-RequestId", "req-dup-500"), ("set-cookie", "a=1"), ("set-cookie", "b=2")]), + ) + + assert error.response.headers.get_list("set-cookie") == ["a=1", "b=2"] + + +def _bedrock_httpx_status_error_sites(): + """Every `except httpx.HTTPStatusError as err` that raises a BedrockError, across bedrock.""" + import ast + import pathlib + + sites = [] + for path in sorted(pathlib.Path("litellm/llms/bedrock").rglob("*.py")): + tree = ast.parse(path.read_text()) + for handler in (n for n in ast.walk(tree) if isinstance(n, ast.ExceptHandler)): + caught = ast.unparse(handler.type) if handler.type is not None else "" + if "HTTPStatusError" not in caught or handler.name is None: + continue + for call in ( + n + for n in ast.walk(handler) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "BedrockError" + ): + sites.append((str(path), call.lineno, handler.name, {k.arg for k in call.keywords})) + return sites + + +def test_every_bedrock_httpx_status_error_site_keeps_provider_headers(): + """A raise site holding the provider's failed response must hand its headers on (LIT-5428). + + These sites are the only place x-amzn-RequestId still exists; a site that drops it + silently shadows the fix for that whole surface. + """ + sites = _bedrock_httpx_status_error_sites() + + assert len(sites) >= 12 + dropped = [f"{path}:{lineno}" for path, lineno, _, kwargs in sites if "headers" not in kwargs] + assert dropped == [] + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_bedrock_embedding_call_keeps_provider_headers(is_async): + """The embeddings surface raises from the same shape as chat and lost the same header.""" + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + failure = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-embed-500"}, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + class _SyncUpstream(HTTPHandler): + def post(self, *args, **kwargs): + return failure + + class _AsyncUpstream(AsyncHTTPHandler): + async def post(self, *args, **kwargs): + return failure + + async def _drive(): + embedding = BedrockEmbedding() + kwargs = dict( + timeout=None, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/", + headers={}, + data={}, + ) + if is_async: + return await embedding._make_async_call(client=_AsyncUpstream(), **kwargs) + return embedding._make_sync_call(client=_SyncUpstream(), **kwargs) + + with pytest.raises(BedrockError) as exc_info: + await _drive() + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-embed-500" + + +def _bedrock_mantle_error_configs(): + from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig + from litellm.llms.bedrock_mantle.responses.transformation import BedrockMantleResponsesAPIConfig + + return [BedrockMantleChatConfig, BedrockMantleResponsesAPIConfig] + + +@pytest.mark.parametrize("config", _bedrock_mantle_error_configs()) +def test_bedrock_mantle_get_error_class_keeps_provider_headers(config): + """bedrock_mantle rides the OpenAI-compatible surfaces, whose errors drop the headers. + + A chat request for a responses-API model is bridged onto the responses config, so + fixing only the chat one leaves the model the customer actually calls uncovered. + """ + error = config().get_error_class( + error_message="prompt tokens exceed model maximum", + status_code=400, + headers={"x-amzn-RequestId": "req-mantle-400"}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-mantle-400" + + +def _bedrock_configs_with_get_error_class(): + import importlib + import inspect + import pathlib + + import litellm + + llms_root = pathlib.Path(inspect.getfile(litellm)).parent / "llms" + configs = [] + for package in ("bedrock", "bedrock_mantle"): + for path in sorted((llms_root / package).rglob("*.py")): + module_name = "litellm.llms." + ".".join(path.relative_to(llms_root).with_suffix("").parts) + module = importlib.import_module(module_name) + for name, obj in vars(module).items(): + if not inspect.isclass(obj) or obj.__module__ != module_name: + continue + if getattr(obj, "get_error_class", None) is None: + continue + configs.append(pytest.param(obj, id=f"{module_name}.{name}")) + return configs + + +@pytest.mark.parametrize("config", _bedrock_configs_with_get_error_class()) +def test_every_bedrock_config_get_error_class_keeps_provider_headers(config): + """Every bedrock surface must classify errors through BedrockError, not a header-dropping base. + + A config that inherits get_error_class from a provider-agnostic base builds a blank + response, so the request id is gone before the proxy ever reads it. + """ + try: + instance = config() + except Exception: + instance = config.__new__(config) + + try: + error = instance.get_error_class( + error_message="boom", + status_code=500, + headers={"x-amzn-RequestId": "req-audit-500"}, + ) + except Exception as raised: # some bases raise the exception instead of returning it + error = raised + + assert error.response.headers["x-amzn-requestid"] == "req-audit-500" + + +def test_bedrock_get_error_class_audit_covers_every_surface(): + assert len(_bedrock_configs_with_get_error_class()) >= 30 diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index 5f12ae8566c..fda3c8ceb8f 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -1,8 +1,5 @@ """Test Bedrock cross-region inference profile model mapping""" -import json -from functools import lru_cache -from pathlib import Path from typing import NamedTuple import pytest @@ -102,13 +99,6 @@ GPT_5_6_PROFILES = [ ] -@lru_cache(maxsize=1) -def _packaged_cost_map(): - """The map litellm actually resolves against, for fields ModelInfoBase drops.""" - path = Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json" - return json.loads(path.read_text()) - - def _bedrock_response(model, usage): return ModelResponse( id="test", @@ -126,17 +116,6 @@ def _bedrock_response(model, usage): ) -def test_bedrock_cross_region_inference_profile_mapping(): - """Test that bedrock cross-region inference profile model is mapped""" - model = "bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0" - - model_info = _get_model_info_helper(model=model, custom_llm_provider="bedrock") - - assert model_info is not None - assert model_info["litellm_provider"] == "bedrock" - assert model_info["input_cost_per_token"] == 8e-07 - - def test_proxy_cost_calculation_scenario(): """Test exact GitHub issue scenario: proxy cost calculation""" model = "litellm_proxy/bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0" @@ -176,38 +155,6 @@ def test_bedrock_gpt_5_6_profiles_route_to_converse(profile, local_model_cost_ma assert BedrockModelInfo.get_bedrock_route(f"bedrock/{profile.model_id}") == "converse" -@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) -def test_bedrock_gpt_5_6_published_rates(profile, local_model_cost_map): - """Geo and Global profiles carry their own published rates, per context tier.""" - model_info = _get_model_info_helper( - model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" - ) - - assert model_info["litellm_provider"] == "bedrock_converse" - assert model_info["mode"] == "chat" - assert model_info["max_input_tokens"] == 1000000 - assert model_info["input_cost_per_token"] == profile.input_cost - assert ( - model_info["input_cost_per_token_above_272k_tokens"] - == profile.input_cost_above_272k - ) - assert model_info["output_cost_per_token"] == profile.output_cost - assert ( - model_info["output_cost_per_token_above_272k_tokens"] - == profile.output_cost_above_272k - ) - assert model_info["cache_creation_input_token_cost"] == profile.cache_write - assert ( - model_info["cache_creation_input_token_cost_above_272k_tokens"] - == profile.cache_write_above_272k - ) - assert model_info["cache_read_input_token_cost"] == profile.cache_read - assert ( - model_info["cache_read_input_token_cost_above_272k_tokens"] - == profile.cache_read_above_272k - ) - - def test_bedrock_gpt_5_6_above_272k_tier_applies_to_cost(local_model_cost_map): """A prompt over 272K tokens is billed at the long-context rate, not the base rate.""" response = _bedrock_response( @@ -267,31 +214,6 @@ def test_bedrock_gpt_5_6_bills_cache_write_tokens(local_model_cost_map): assert cost == pytest.approx(expected, rel=1e-9) -@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) -def test_bedrock_gpt_5_6_advertises_only_converse_supported_features( - profile, local_model_cost_map -): - model_info = _get_model_info_helper( - model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" - ) - - assert model_info["supports_function_calling"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_vision"] is True - - # Bedrock rejects an explicit cachePoint block for these models, so the flag that - # offers caller-driven caching stays off even though the cache rates are declared. - assert not model_info.get("supports_prompt_caching") - - # ModelInfoBase drops these two, so they are read from the map litellm resolves. - raw = _packaged_cost_map()[profile.model_id] - assert raw["supported_modalities"] == ["text", "image"] - assert raw["supported_output_modalities"] == ["text"] - # No bedrock_converse entry declares supported_endpoints; these models are reachable - # on chat completions and on the Responses API without it. - assert "supported_endpoints" not in raw - - @pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(profile, local_model_cost_map): """GPT-5.x on Converse maps reasoning_effort to reasoning.effort, so reasoning_effort diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 1f23d39c631..9457d5faaff 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -8,9 +8,7 @@ gate, the URL construction for both paths, and the shared Bearer auth. """ import copy -import json import logging -from pathlib import Path import pytest from botocore.exceptions import ( @@ -159,7 +157,6 @@ class TestBedrockMantleResponsesURL: assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses" assert url.count("/responses") == 1 - def test_url_aws_region_name_overrides_stale_api_base(self, monkeypatch): monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) @@ -1777,53 +1774,6 @@ class TestBedrockMantleResponsesSigV4: class TestBedrockMantleResponsesPricing: - def test_gpt_5_5_pricing_and_mode(self, local_cost_map): - info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(5.5e-06) - assert info["output_cost_per_token"] == pytest.approx(3.3e-05) - assert info["cache_read_input_token_cost"] == pytest.approx(5.5e-07) - assert info["max_input_tokens"] == 1050000 - - def test_gpt_5_4_pricing_and_mode(self, local_cost_map): - info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.4") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(2.75e-06) - assert info["output_cost_per_token"] == pytest.approx(1.65e-05) - assert info["cache_read_input_token_cost"] == pytest.approx(2.75e-07) - assert info["max_input_tokens"] == 1050000 - - def test_gpt_5_6_cyber_pricing_and_mode(self, local_cost_map): - info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.6-cyber") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(1.375e-05) - assert info["cache_creation_input_token_cost"] == pytest.approx(1.71875e-05) - assert info["cache_read_input_token_cost"] == pytest.approx(1.375e-06) - assert info["output_cost_per_token"] == pytest.approx(8.25e-05) - assert info["max_input_tokens"] == 272000 - - @pytest.mark.parametrize( - "model, input_cost, cache_creation_cost, cache_read_cost, output_cost", - [ - ("openai.gpt-5.6-sol", 5.5e-06, 6.875e-06, 5.5e-07, 3.3e-05), - ("openai.gpt-5.6-terra", 2.2e-06, 2.75e-06, 2.2e-07, 1.32e-05), - ("openai.gpt-5.6-luna", 2.2e-07, 2.75e-07, 2.2e-08, 1.32e-06), - ], - ) - def test_gpt_5_6_pricing_and_mode( - self, local_cost_map, model, input_cost, cache_creation_cost, cache_read_cost, output_cost - ): - info = litellm.get_model_info(f"bedrock_mantle/{model}") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(input_cost) - assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost) - assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost) - assert info["output_cost_per_token"] == pytest.approx(output_cost) - assert info["max_input_tokens"] == 1050000 - assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2) - assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2) - assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2) - assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.5) @pytest.mark.parametrize( "model, input_cost, output_cost", @@ -1861,58 +1811,3 @@ class TestBedrockMantleResponsesPricing: def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models - - -def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]: - repo_root = Path(__file__).resolve().parents[4] - paths = { - "root": repo_root / "model_prices_and_context_window.json", - "bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json", - } - return json.loads(paths[map_name].read_text()) - - -class TestMantleGptRegistryEntries: - """Locks the OpenAI GPT entries to Bedrock Mantle's live behavior. - - Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna - and for gpt-5.5 and gpt-5.4 (oversize requests 400 with "prompt tokens (N) - exceed model maximum (1050000)", and a 1,030,590-token request completes - on every one of them), while the AWS model cards still quote 272K for - gpt-5.5 and gpt-5.4. mode must stay "responses": Mantle's native - /v1/chat/completions rejects function tools unless reasoning_effort is - "none", so chat traffic has to keep bridging to the Responses API - (see the responses_api_bridge tests above). - """ - - @pytest.mark.parametrize("map_name", ("root", "bundled_backup")) - @pytest.mark.parametrize( - "key", - ( - "bedrock_mantle/openai.gpt-5.6-sol", - "bedrock_mantle/openai.gpt-5.6-terra", - "bedrock_mantle/openai.gpt-5.6-luna", - ), - ) - def test_entry_matches_mantle_enforced_limits(self, map_name, key): - entry = _repo_cost_map(map_name)[key] - assert entry["max_input_tokens"] == 1050000 - assert entry["max_output_tokens"] == 128000 - assert entry["mode"] == "responses" - assert entry["use_openai_responses_path"] is True - assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] - - @pytest.mark.parametrize("map_name", ("root", "bundled_backup")) - @pytest.mark.parametrize( - "key", - ( - "bedrock_mantle/openai.gpt-5.5", - "bedrock_mantle/openai.gpt-5.4", - ), - ) - def test_gpt_55_and_54_entries_match_mantle_enforced_limits(self, map_name, key): - entry = _repo_cost_map(map_name)[key] - assert entry["max_input_tokens"] == 1050000 - assert entry["max_output_tokens"] == 128000 - assert entry["mode"] == "responses" - assert entry["use_openai_responses_path"] is True diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index b88c27e64b9..e83a844c87e 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -6,15 +6,20 @@ API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.ht """ import json +import asyncio from unittest.mock import patch import httpx import pytest +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest import litellm from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig +from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws from litellm.types.utils import LlmProviders +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.fixture @@ -684,40 +689,6 @@ class TestBedrockMantleProviderResolution: class TestBedrockMantlePricing: """Tests that verify Bedrock Mantle uses correct AWS Bedrock pricing, not OpenAI pricing.""" - def test_gpt_oss_120b_pricing(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - # Bedrock pricing: $0.15/M input, $0.60/M output - assert info["input_cost_per_token"] == pytest.approx(1.5e-7) - assert info["output_cost_per_token"] == pytest.approx(6e-7) - - def test_gpt_oss_20b_pricing(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-20b") - # Bedrock pricing: $0.075/M input, $0.30/M output - assert info["input_cost_per_token"] == pytest.approx(7.5e-8) - assert info["output_cost_per_token"] == pytest.approx(3e-7) - - def test_pricing_significantly_cheaper_than_openai_native(self, monkeypatch): - """ - Verify Bedrock Mantle pricing is cheaper than OpenAI's direct API pricing. - This is the core issue the provider addition fixes — previously users were being - billed at OpenAI rates instead of the cheaper Bedrock rates. - """ - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - bedrock_info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - # OpenAI direct pricing for gpt-oss-120b is ~$0.039/M input, $0.190/M output - # Bedrock should be cheaper at $0.15/M input and $0.60/M output... wait - # Actually, Bedrock ADDS value not reduces cost vs OpenAI direct for these models. - # The key fix is that we now use Bedrock-specific prices instead of mapping to - # some unrelated OpenAI model (like gpt-4) pricing. - # Just validate the pricing is as expected from AWS docs. - assert bedrock_info["input_cost_per_token"] == pytest.approx(1.5e-7) - assert bedrock_info["output_cost_per_token"] == pytest.approx(6e-7) - def test_safeguard_models_have_larger_output_tokens(self, monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") litellm.add_known_models() @@ -727,49 +698,6 @@ class TestBedrockMantlePricing: ) assert info_safeguard["max_output_tokens"] > info_120b["max_output_tokens"] - def test_reasoning_support(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - assert info.get("supports_reasoning") is True - - def test_context_window(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - assert info["max_input_tokens"] == 131072 - - -@pytest.mark.parametrize( - "model_id,input_cost,output_cost,max_tokens", - [ - ("google.gemma-4-31b", 1.4e-07, 4e-07, 256000), - ("google.gemma-4-26b-a4b", 1.3e-07, 4e-07, 256000), - ("google.gemma-4-e2b", 4e-08, 8e-08, 128000), - ], -) -def test_gemma_4_bedrock_mantle_model_metadata( - local_cost_map, model_id, input_cost, output_cost, max_tokens -): - full_model_name = f"bedrock_mantle/{model_id}" - info = litellm.get_model_info(full_model_name) - - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == pytest.approx(input_cost) - assert info["output_cost_per_token"] == pytest.approx(output_cost) - assert info["max_input_tokens"] == max_tokens - assert info["max_output_tokens"] == max_tokens - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert ( - litellm.supports_parallel_function_calling( - model=full_model_name, custom_llm_provider="bedrock_mantle" - ) - is False - ) - @pytest.mark.parametrize( "model_id", @@ -787,3 +715,26 @@ def test_gemma_4_models_register_under_bedrock_mantle(local_cost_map, model_id): resolved_model, provider, _, _ = litellm.get_llm_provider(full_model_name) assert provider == "bedrock_mantle" assert resolved_model == model_id + + +@pytest.mark.asyncio +async def test_mantle_signing_runs_off_the_event_loop(): + """Regression for issue #40165: Mantle signs with SigV4 through a composed BaseAWSLLM, so the + off-loop gate must recognise it too, or its credential refresh blocks the loop like Bedrock's did.""" + probe = EventLoopProbe() + + def sign(headers: dict[str, str]) -> dict[str, str]: + request = AWSRequest( + method="POST", url="https://bedrock-mantle.us-east-1.api.aws/v1/responses", data="{}", headers=headers + ) + SigV4Auth(probe.credentials(), "bedrock", "us-east-1").add_auth(request) + return dict(request.headers) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + signed = await sign_request_off_loop_if_aws( + BedrockMantleChatConfig(), sign, headers={"Content-Type": "application/json"} + ) + await release + + assert "Authorization" in signed + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py index ec243b7058d..d9d6e813d86 100644 --- a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py +++ b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py @@ -16,6 +16,9 @@ import httpx import pytest +from litellm.llms.black_forest_labs.image_edit import ( + transformation as bfl_transformation, +) from litellm.llms.black_forest_labs.image_edit.transformation import ( BlackForestLabsImageEditConfig, ) @@ -186,7 +189,7 @@ class TestBlackForestLabsImageEditTransformation: assert data["output_format"] == "jpeg" # BFL uses JSON, not multipart - files should be empty - assert files == [] + assert files == () def test_transform_image_edit_request_with_mask(self): """Test request transformation with mask for inpainting.""" @@ -299,3 +302,76 @@ class TestBlackForestLabsImageEditTransformation: def test_use_multipart_form_data_returns_false(self): """Test that use_multipart_form_data returns False for BFL.""" assert self.config.use_multipart_form_data() is False + + +async def test_async_transform_image_edit_request_downloads_url_images_with_the_async_fetcher(monkeypatch): + served = b"png-bytes-from-cdn" + fetched = [] + + def forbid_sync_fetch(client, url, **kwargs): + raise AssertionError(f"sync image fetch ran on the event loop: {url}") + + async def serve(client, url, **kwargs): + fetched.append((url, kwargs.get("timeout"))) + return httpx.Response(200, content=served, request=httpx.Request("GET", url)) + + monkeypatch.setattr(bfl_transformation, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(bfl_transformation, "async_safe_get", serve) + + data, files = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image="https://cdn.example/photo.png", + image_edit_optional_request_params={"mask": "https://cdn.example/mask.png", "seed": 7}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert base64.b64decode(data["input_image"]) == served + assert base64.b64decode(data["mask"]) == served + assert data["seed"] == 7 + assert files == () + assert fetched == [("https://cdn.example/photo.png", 60.0), ("https://cdn.example/mask.png", 60.0)] + + +async def test_async_transform_image_edit_request_never_fetches_for_local_images(monkeypatch): + def refuse(*args, **kwargs): + raise AssertionError("no network fetch expected for local image bytes") + + monkeypatch.setattr(bfl_transformation, "safe_get", refuse) + monkeypatch.setattr(bfl_transformation, "async_safe_get", refuse) + + data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image=[BytesIO(b"first"), BytesIO(b"other")], + image_edit_optional_request_params={"mask": b"mask-bytes"}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert base64.b64decode(data["input_image"]) == b"first" + assert base64.b64decode(data["mask"]) == b"mask-bytes" + + +async def test_async_transform_image_edit_request_downloads_only_the_first_url_of_a_list(monkeypatch): + fetched = [] + + async def serve(client, url, **kwargs): + fetched.append(url) + return httpx.Response(200, content=b"first-bytes", request=httpx.Request("GET", url)) + + monkeypatch.setattr(bfl_transformation, "safe_get", lambda *args, **kwargs: pytest.fail("sync fetch ran")) + monkeypatch.setattr(bfl_transformation, "async_safe_get", serve) + + data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image=["https://cdn.example/a.png", "https://cdn.example/b.png"], + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert fetched == ["https://cdn.example/a.png"] + assert base64.b64decode(data["input_image"]) == b"first-bytes" diff --git a/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py b/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py index 09718b1e6e0..a47180e9511 100644 --- a/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py +++ b/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py @@ -1,3 +1,6 @@ +import pytest + +import litellm from litellm.llms.cerebras.chat import CerebrasConfig @@ -59,3 +62,23 @@ def test_map_openai_params_preserves_max_retries_zero_falsy() -> None: assert "max_retries" in result and result["max_retries"] == 0, ( f"max_retries=0 (falsy) must not be silently omitted; got: {result!r}" ) + + +def test_qwen_3_8_27b_cost_and_tokens(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + model = "cerebras/qwen-3.8-27b" + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, + prompt_tokens=1000, + completion_tokens=1000, + ) + assert abs(prompt_cost - 0.00099) < 1e-9 + assert abs(completion_cost - 0.00149) < 1e-9 + + model_info = litellm.get_model_info(model) + assert model_info["max_input_tokens"] == 65536 + assert model_info["max_output_tokens"] == 32768 + assert model_info["supports_vision"] is True + assert model_info["supports_reasoning"] is True + assert model_info["supports_parallel_function_calling"] is True diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py index ca79c8d7025..12b5f03aedc 100644 --- a/tests/test_litellm/llms/chat/test_converse_handler.py +++ b/tests/test_litellm/llms/chat/test_converse_handler.py @@ -308,3 +308,64 @@ def test_completion_plumbs_stream_chunk_size_through_converse(): stream_chunk_size=2048, ) iter_bytes_spy.assert_called_once_with(chunk_size=2048) + + +def _bedrock_error_response(status_code: int, request_id: str) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={ + "x-amzn-RequestId": request_id, + "x-amzn-ErrorType": "InternalServerException", + }, + text=json.dumps({"message": "Amazon Bedrock is unable to process your request."}), + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + +def test_converse_completion_error_forwards_bedrock_response_headers(): + error_response = _bedrock_error_response(500, "req-err-123") + client = HTTPHandler() + client.post = MagicMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-123" + + +@pytest.mark.asyncio +async def test_async_converse_completion_error_forwards_bedrock_response_headers(): + error_response = _bedrock_error_response(500, "req-err-456") + client = AsyncHTTPHandler() + client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-456" diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 8e0415d50de..a7520bd5955 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -10,18 +10,23 @@ from unittest.mock import MagicMock, patch import httpx import pytest - +import litellm +from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig from litellm.llms.openai.common_utils import OpenAIError +from litellm.main import responses_api_bridge_check from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager -from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig class TestChatGPTResponsesAPITransformation: @pytest.mark.parametrize( "model_name", [ + "chatgpt/gpt-5.5", + "chatgpt/gpt-5.6-luna", + "chatgpt/gpt-5.6-sol", + "chatgpt/gpt-5.6-terra", "chatgpt/gpt-5.4", "chatgpt/gpt-5.4-pro", "chatgpt/gpt-5.3-chat-latest", @@ -40,6 +45,52 @@ class TestChatGPTResponsesAPITransformation: assert isinstance(config, ChatGPTResponsesAPIConfig) assert config.custom_llm_provider == LlmProviders.CHATGPT + @pytest.mark.parametrize( + "model_name", + [ + "chatgpt/gpt-5.5", + "chatgpt/gpt-5.6-luna", + "chatgpt/gpt-5.6-sol", + "chatgpt/gpt-5.6-terra", + ], + ) + def test_chatgpt_responses_model_metadata(self, model_name: str, local_model_cost_map: None) -> None: + model_info = litellm.get_model_info(model_name) + + assert model_info["litellm_provider"] == "chatgpt" + assert model_info["mode"] == "responses" + assert model_info["supported_endpoints"] == [ + "/v1/chat/completions", + "/v1/responses", + ] + assert model_info["max_input_tokens"] == 1050000 + assert model_info["max_output_tokens"] == 128000 + + @pytest.mark.parametrize( + "model_name", + [ + "gpt-5.5", + "gpt-5.6-luna", + "gpt-5.6-sol", + "gpt-5.6-terra", + ], + ) + def test_chatgpt_models_bridge_chat_completions_to_responses( + self, model_name: str, local_model_cost_map: None + ) -> None: + """A chat completions request for these models must take the Responses bridge. + + `gpt-5.6-*` also exists as an openai chat model, so an unregistered + chatgpt model resolves to mode "chat" here and never reaches the bridge. + """ + model_info, resolved_model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="chatgpt", + ) + + assert model_info["mode"] == "responses" + assert resolved_model == model_name + @patch("litellm.llms.chatgpt.responses.transformation.Authenticator") def test_chatgpt_responses_endpoint_url(self, mock_authenticator_class): mock_auth_instance = MagicMock() diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 9e64bfafa54..f8868cfaf83 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1547,3 +1547,131 @@ def test_sync_force_ipv4_https_proxy_mount_uses_handler_ca_bundle( handler.close() assert response.text == "ok-tls" + + +@pytest.mark.asyncio +async def test_put_can_refuse_to_follow_a_redirect(): + """The client follows redirects by default; a caller uploading to a URL it did not choose must be able to opt out.""" + hops: list[str] = [] # mutable-ok: the fake transport records the paths it was asked for + + async def mock_handler(request: httpx.Request) -> httpx.Response: + hops.append(request.url.path) + if request.url.path == "/first": + return httpx.Response(302, request=request, headers={"location": "/second"}) + return httpx.Response(200, request=request) + + handler = AsyncHTTPHandler() + await handler.client.aclose() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler), follow_redirects=True) + try: + followed = await handler.put("https://uploads.example/first", data=b"x") + assert followed.status_code == 200 + assert hops == ["/first", "/second"] + + hops.clear() + with pytest.raises(MaskedHTTPStatusError) as refused: + await handler.put("https://uploads.example/first", data=b"x", follow_redirects=False) + assert refused.value.status_code == 302 + assert hops == ["/first"] + finally: + await handler.close() + + +@pytest.mark.asyncio +async def test_a_retried_put_stays_a_put_and_still_refuses_redirects(): + """ + The connection-error retry used to resend as POST through a client that follows redirects. + + Storage answers a POST to a presigned PUT url with 403 or 405, so the batch looked + permanently rejected, and the redirect refusal the caller asked for was silently lost. + """ + attempts: list[tuple[str, str]] = [] # mutable-ok: the fake transports record what they were asked for + + async def refusing_transport(request: httpx.Request) -> httpx.Response: + attempts.append((request.method, request.url.path)) + raise httpx.ConnectError("connection reset", request=request) + + async def retry_transport(request: httpx.Request) -> httpx.Response: + attempts.append((request.method, request.url.path)) + if request.url.path == "/first": + return httpx.Response(302, request=request, headers={"location": "/second"}) + return httpx.Response(200, request=request) + + class HandlerWithFakeRetryClient(AsyncHTTPHandler): + def create_client(self, *args, **kwargs) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(retry_transport), follow_redirects=True) + + handler = HandlerWithFakeRetryClient() + await handler.client.aclose() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(refusing_transport)) + try: + with pytest.raises(MaskedHTTPStatusError) as refused: + await handler.put("https://uploads.example/first", data=b"x", follow_redirects=False) + + assert refused.value.status_code == 302 + assert attempts == [("PUT", "/first"), ("PUT", "/first")] + finally: + await handler.client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target", ["https://example.com/final.json?next=1", "https://other.example/final.json?next=1"]) +async def test_bounded_get_preserves_sdk_redirect_auth_and_query_handling(respx_mock, monkeypatch, target): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + respx_mock.get("https://example.com/spec.json?original=1").respond(302, headers={"location": target}) + destination = respx_mock.get(target).respond(200, json={"paths": {}}) + handler = AsyncHTTPHandler() + try: + response = await handler.get( + "https://example.com/spec.json?original=1", max_response_bytes=100, follow_redirects=True, + headers={"Authorization": "Bearer sentinel", "Accept-Encoding": "gzip"}, timeout=2.0, + ) + finally: + await handler.close() + assert response.json() == {"paths": {}} + request = destination.calls[0].request + assert request.headers.get("authorization") == (None if "other.example" in target else "Bearer sentinel") + assert request.headers["accept-encoding"] == "identity" + assert str(request.url) == target + assert request.extensions["timeout"]["read"] == 2.0 + + +@pytest.mark.asyncio +async def test_bounded_get_stops_redirect_loops(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + route = respx_mock.get("https://example.com/spec.json").respond(302, headers={"location": "/spec.json"}) + handler = AsyncHTTPHandler() + try: + with pytest.raises(ValueError, match="Too many redirects"): + await handler.get("https://example.com/spec.json", max_response_bytes=100, follow_redirects=True) + finally: + await handler.close() + assert route.call_count == 11 + + +@pytest.mark.asyncio +async def test_bounded_get_closes_stream_on_cancellation(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + started = asyncio.Event() + closed = asyncio.Event() + + class SlowStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b"x" + started.set() + await asyncio.Event().wait() + + async def aclose(self): + closed.set() + + respx_mock.get("https://example.com/slow.json").respond(200, stream=SlowStream()) + handler = AsyncHTTPHandler() + try: + task = asyncio.create_task(handler.get("https://example.com/slow.json", max_response_bytes=100)) + await asyncio.wait_for(started.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + await handler.close() + assert closed.is_set() diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 023e2d8843f..c39779972c0 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,11 +1,14 @@ import asyncio import json import logging +import threading import time +from typing import Final from unittest.mock import AsyncMock, Mock, patch import httpx import pytest +from botocore.credentials import RefreshableCredentials import litellm from litellm._logging import verbose_logger @@ -17,7 +20,11 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, ) -from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS +from litellm.llms.brave.search.transformation import BraveSearchConfig +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -27,15 +34,173 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _rust_responses_websocket_enabled, ) from litellm.llms.azure.videos.transformation import AzureVideoConfig +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, +) +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig +from litellm.llms.tinyfish.search.transformation import TinyfishSearchConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import TranscriptionResponse +from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +async def _get_search_with_client( + client: HTTPHandler | AsyncHTTPHandler, provider_config: BaseSearchConfig | None = None +) -> SearchResponse: + result: Final = BaseLLMHTTPHandler().search( + query="test", + optional_params={}, + timeout=5, + logging_obj=Mock(), + api_key="test-key", + api_base="https://search.example.test/", + custom_llm_provider="tinyfish" if isinstance(provider_config, TinyfishSearchConfig) else "brave", + client=client, + asearch=isinstance(client, AsyncHTTPHandler), + provider_config=provider_config or BraveSearchConfig(), + ) + return await result if asyncio.iscoroutine(result) else result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.parametrize("status_code", (400, 401, 403, 422, 429, 500)) +async def test_get_search_raises_provider_http_errors(is_async: bool, status_code: int) -> None: + upstream_response: Final = httpx.Response( + status_code, json={"error": "rejected request"}, headers={"retry-after": "7"} + ) + transport: Final = httpx.MockTransport(lambda request: upstream_response) + async with httpx.AsyncClient(transport=transport) as async_client: + with httpx.Client(transport=transport) as sync_client: + client: Final = AsyncHTTPHandler() if is_async else HTTPHandler(client=sync_client) + if isinstance(client, AsyncHTTPHandler): + await client.close() + client.client = async_client + with pytest.raises(BaseLLMException) as error: + await _get_search_with_client(client) + assert error.value.status_code == status_code + assert "rejected request" in error.value.message + assert error.value.headers is not None + assert error.value.headers["retry-after"] == "7" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.parametrize("has_results", (False, True)) +async def test_get_search_preserves_successful_results(is_async: bool, has_results: bool) -> None: + results: Final = ( + [{"title": "Example", "url": "https://example.com", "description": "Example snippet"}] if has_results else [] + ) + transport: Final = httpx.MockTransport(lambda request: httpx.Response(200, json={"web": {"results": results}})) + async with httpx.AsyncClient(transport=transport) as async_client: + with httpx.Client(transport=transport) as sync_client: + client: Final = AsyncHTTPHandler() if is_async else HTTPHandler(client=sync_client) + if isinstance(client, AsyncHTTPHandler): + await client.close() + client.client = async_client + response: Final = await _get_search_with_client(client) + assert response.object == "search" + assert len(response.results) == int(has_results) + if has_results: + assert response.results[0].title == "Example" + assert response.results[0].url == "https://example.com" + assert response.results[0].snippet == "Example snippet" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +async def test_get_search_preserves_tinyfish_http_error_formatting(is_async: bool) -> None: + upstream_response: Final = httpx.Response( + 429, + json={"error": {"code": "RATE_LIMIT_EXCEEDED", "message": "rate limit exceeded"}}, + headers={"retry-after": "7"}, + ) + transport: Final = httpx.MockTransport(lambda request: upstream_response) + async with httpx.AsyncClient(transport=transport) as async_client: + with httpx.Client(transport=transport) as sync_client: + client: Final = AsyncHTTPHandler() if is_async else HTTPHandler(client=sync_client) + if isinstance(client, AsyncHTTPHandler): + await client.close() + client.client = async_client + with pytest.raises(BaseLLMException) as error: + await _get_search_with_client(client, TinyfishSearchConfig()) + assert error.value.status_code == 429 + assert error.value.message == ( + "TinyFish Search: rate limit exceeded. See https://docs.tinyfish.ai/search-api for details." + ) + assert error.value.headers is not None + assert error.value.headers["retry-after"] == "7" + + +OCR_RESPONSE = { + "pages": [{"index": 0, "markdown": "OCR output", "images": []}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, +} + + +def _ocr_sync_client() -> HTTPHandler: + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE))) + return client + + +def _ocr_async_client() -> AsyncHTTPHandler: + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient( + transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE)) + ) + return client + + +def test_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = BaseLLMHTTPHandler().ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_sync_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + + +@pytest.mark.asyncio +async def test_async_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = await BaseLLMHTTPHandler().async_ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_async_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + + def test_prepare_fake_stream_request(): # Initialize the BaseLLMHTTPHandler handler = BaseLLMHTTPHandler() @@ -747,6 +912,65 @@ async def test_anthropic_messages_streaming_response_aclose_closes_agentic_upstr assert tracker.closed is True +class _ProbedBedrockMessagesConfig(AmazonAnthropicClaudeMessagesConfig): + def __init__(self, probe: EventLoopProbe) -> None: + super().__init__() + self._probe = probe + + def get_credentials( + self, + **kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores + ) -> RefreshableCredentials: + return self._probe.credentials() + + +@pytest.mark.asyncio +async def test_async_anthropic_messages_handler_signs_bedrock_off_the_event_loop(monkeypatch): + """Regression for issue #40165: /v1/messages on Bedrock signed on the loop, so botocore's blocking + credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + handler = BaseLLMHTTPHandler() + upstream_response = httpx.Response( + 200, + json={ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "model": "claude-haiku-4-5-20251001", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"), + ) + mock_client = AsyncMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=upstream_response) + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.dynamic_success_callbacks = None + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + await handler.async_anthropic_messages_handler( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_provider_config=_ProbedBedrockMessagesConfig(probe), + anthropic_messages_optional_request_params={"max_tokens": 16}, + custom_llm_provider="bedrock", + litellm_params=GenericLiteLLMParams(aws_region_name="us-west-2"), + logging_obj=mock_logging_obj, + client=mock_client, + stream=False, + kwargs={}, + ) + await release + + sent_headers = mock_client.post.call_args.kwargs["headers"] + assert sent_headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True + + @pytest.mark.asyncio async def test_async_anthropic_messages_handler_passes_litellm_metadata(): """Ensure litellm_metadata from kwargs is forwarded via update_from_kwargs. @@ -2689,20 +2913,18 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h @pytest.mark.parametrize( - "custom_llm_provider, litellm_params, expected", - [ - ("openai", GenericLiteLLMParams(rust=True), True), - ("openai", GenericLiteLLMParams(), False), - ("openai", GenericLiteLLMParams(rust=False), False), - ("azure", GenericLiteLLMParams(rust=True), False), - ("hosted_vllm", GenericLiteLLMParams(rust=True), False), - (None, GenericLiteLLMParams(rust=True), False), - ], + "custom_llm_provider, enabled, expected", + [("openai", True, True), ("openai", False, False), ("azure", True, False), + ("hosted_vllm", True, False), (None, True, False)], ) -def test_the_rust_responses_websocket_needs_both_openai_and_the_rust_flag( - custom_llm_provider, litellm_params, expected +def test_the_rust_responses_websocket_needs_openai_and_process_enablement( + custom_llm_provider, enabled, expected, monkeypatch ): - assert _rust_responses_websocket_enabled(custom_llm_provider, litellm_params) is expected + from litellm.rust_bridge import configuration + + configuration.reset_rust_configuration() + monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0") + assert _rust_responses_websocket_enabled(custom_llm_provider) is expected def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch): @@ -3186,3 +3408,308 @@ async def test_async_container_list_handler_transforms_success_response(): assert [container.id for container in response.data] == ["cntr_a"] assert response.has_more is True + + +class _TransformRecordingConfig(BaseConfig): + def __init__(self, transform_async: bool): + self.transform_async = transform_async + self.transform_calls = [] + self.sign_threads = [] + + @property + def uses_async_transform_request(self) -> bool: + return self.transform_async + + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, non_default_params, optional_params, model, drop_params): + return optional_params + + def validate_environment( + self, headers, model, messages, optional_params, litellm_params, api_key=None, api_base=None + ): + return {} + + def transform_request(self, model, messages, optional_params, litellm_params, headers): + self.transform_calls.append("sync") + return {"transformed_by": "sync"} + + async def async_transform_request(self, model, messages, optional_params, litellm_params, headers): + self.transform_calls.append("async") + return {"transformed_by": "async"} + + def sign_request( + self, headers, optional_params, request_data, api_base, api_key=None, model=None, stream=None, fake_stream=None + ): + self.sign_threads.append(threading.current_thread()) + return headers, None + + def transform_response( + self, + model, + raw_response, + model_response, + logging_obj, + request_data, + messages, + optional_params, + litellm_params, + encoding, + api_key=None, + json_mode=None, + ): + model_response.choices[0].message.content = raw_response.json()["transformed_by"] + return model_response + + def get_error_class(self, error_message, status_code, headers): + return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + + def get_model_response_iterator(self, streaming_response, sync_stream, json_mode=False): + return litellm.OpenAIGPTConfig().get_model_response_iterator( + streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode + ) + + +def _start_async_completion(config, logging_obj=None): + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=captured["body"]) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + pending = BaseLLMHTTPHandler().completion( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + api_base="https://llm.example/v1/chat", + custom_llm_provider="openai", + model_response=ModelResponse(), + encoding=None, + logging_obj=logging_obj if logging_obj is not None else Mock(dynamic_success_callbacks=None, model_call_details={}), + optional_params={}, + timeout=10.0, + litellm_params={}, + acompletion=True, + client=client, + provider_config=config, + ) + return pending, captured + + +async def test_completion_awaits_async_transform_request_when_config_opts_in(): + config = _TransformRecordingConfig(transform_async=True) + + pending, captured = _start_async_completion(config) + assert config.transform_calls == [] + + response = await pending + + assert config.transform_calls == ["async"] + assert captured["body"] == {"transformed_by": "async"} + assert response.choices[0].message.content == "async" + + +async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_transform(): + config = _TransformRecordingConfig(transform_async=True) + loop_thread = threading.current_thread() + pre_call_threads = [] + logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={}) + logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread()) + + pending, captured = _start_async_completion(config, logging_obj) + response = await pending + + assert response.choices[0].message.content == "async" + assert captured["body"] == {"transformed_by": "async"} + assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) + assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads) + assert not any(thread.name.startswith("aws-signing") for thread in config.sign_threads + pre_call_threads) + + +class _AWSTransformRecordingConfig(SignsRequestsWithAWS, _TransformRecordingConfig): + pass + + +async def test_completion_signs_aws_configs_on_the_aws_signing_pool_after_the_async_transform(): + config = _AWSTransformRecordingConfig(transform_async=True) + pre_call_threads = [] + logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={}) + logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread()) + + pending, captured = _start_async_completion(config, logging_obj) + response = await pending + + assert response.choices[0].message.content == "async" + assert captured["body"] == {"transformed_by": "async"} + assert config.sign_threads and all(thread.name.startswith("aws-signing") for thread in config.sign_threads) + assert pre_call_threads and all(thread.name.startswith("aws-signing") for thread in pre_call_threads) + + +async def test_completion_keeps_sync_transform_request_before_returning_by_default(): + config = _TransformRecordingConfig(transform_async=False) + + pending, captured = _start_async_completion(config) + assert config.transform_calls == ["sync"] + + response = await pending + + assert config.transform_calls == ["sync"] + assert captured["body"] == {"transformed_by": "sync"} + assert response.choices[0].message.content == "sync" + + +def _sse_echoing_transformed_by(request): + transformed_by = json.loads(request.content)["transformed_by"] + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "stub-model", + "choices": [{"index": 0, "delta": {"content": transformed_by}, "finish_reason": None}], + } + return httpx.Response( + 200, + content=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode(), + headers={"content-type": "text/event-stream"}, + request=request, + ) + + +def _streaming_logging_obj(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="async-transform-stream", + function_id="f", + ) + logging_obj.update_environment_variables( + model="stub-model", user="", optional_params={}, litellm_params={}, custom_llm_provider="openai" + ) + return logging_obj + + +async def test_completion_streams_after_the_async_transform_request(): + config = _TransformRecordingConfig(transform_async=True) + loop_thread = threading.current_thread() + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_sse_echoing_transformed_by)) + + stream = await BaseLLMHTTPHandler().completion( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + api_base="https://llm.example/v1/chat", + custom_llm_provider="openai", + model_response=ModelResponse(), + encoding=None, + logging_obj=_streaming_logging_obj(), + optional_params={}, + timeout=10.0, + litellm_params={}, + acompletion=True, + stream=True, + client=client, + provider_config=config, + ) + collected = [chunk async for chunk in stream] + + assert config.transform_calls == ["async"] + assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) + assert "".join(chunk.choices[0].delta.content or "" for chunk in collected) == "async" + + +class _ImageEditRecordingConfig(BaseImageEditConfig): + def __init__(self): + self.transform_calls = [] + + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, image_edit_optional_params, model, drop_params): + return dict(image_edit_optional_params) + + def validate_environment(self, headers, model, api_key=None, litellm_params=None, api_base=None): + return {} + + def get_complete_url(self, model, api_base, litellm_params): + return "https://images.example/v1/edits" + + def use_multipart_form_data(self): + return False + + def transform_image_edit_request( + self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers + ): + self.transform_calls.append("sync") + return {"transformed_by": "sync"}, [] + + async def async_transform_image_edit_request( + self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers + ): + self.transform_calls.append("async") + return {"transformed_by": "async"}, [] + + def transform_image_edit_response(self, model, raw_response, logging_obj): + return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["transformed_by"])]) + + +def _echo_json_transport(captured): + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=captured["body"]) + + return httpx.MockTransport(handle) + + +async def test_async_image_edit_handler_awaits_the_async_transform(): + config = _ImageEditRecordingConfig() + captured = {} + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_echo_json_transport(captured)) + + response = await BaseLLMHTTPHandler().async_image_edit_handler( + model="edit-model", + image=b"raw-image", + prompt="add a hat", + image_edit_provider_config=config, + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert config.transform_calls == ["async"] + assert captured["body"] == {"transformed_by": "async"} + assert response.data[0].b64_json == "async" + + +def test_image_edit_handler_keeps_the_sync_transform(): + config = _ImageEditRecordingConfig() + captured = {} + client = HTTPHandler() + client.client = httpx.Client(transport=_echo_json_transport(captured)) + + response = BaseLLMHTTPHandler().image_edit_handler( + model="edit-model", + image=b"raw-image", + prompt="add a hat", + image_edit_provider_config=config, + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert config.transform_calls == ["sync"] + assert captured["body"] == {"transformed_by": "sync"} + assert response.data[0].b64_json == "sync" diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py b/tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py index 936de812bc6..4466e5b8767 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py @@ -25,19 +25,43 @@ class TestDashScopeRerankURL: url = self.config.get_complete_url(api_base=None, model="qwen3-rerank") assert url == DEFAULT_RERANK_URL - def test_explicit_v1_base_appends_reranks(self): + def test_chat_shaped_base_remaps_to_rerank_route(self): url = self.config.get_complete_url( api_base="https://dashscope.aliyuncs.com/compatible-mode/v1", model="qwen3-rerank", ) - assert url == "https://dashscope.aliyuncs.com/compatible-mode/v1/reranks" + assert url == "https://dashscope.aliyuncs.com/compatible-api/v1/reranks" - def test_intl_v1_base_appends_reranks(self): + def test_intl_chat_shaped_base_remaps_to_intl_rerank_route(self): url = self.config.get_complete_url( api_base="https://dashscope-intl.aliyuncs.com/compatible-mode/v1", model="qwen3-rerank", ) - assert url == "https://dashscope-intl.aliyuncs.com/compatible-mode/v1/reranks" + assert url == "https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks" + + def test_chat_shaped_base_with_trailing_slash_remaps(self): + url = self.config.get_complete_url( + api_base="https://dashscope-intl.aliyuncs.com/compatible-mode/v1/", + model="qwen3-rerank", + ) + assert url == "https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks" + + def test_rerank_env_var_wins_over_chat_shaped_base(self, monkeypatch): + monkeypatch.setenv( + "DASHSCOPE_API_BASE_RERANK", "https://rerank.example.com/v1/reranks" + ) + url = self.config.get_complete_url( + api_base="https://dashscope-intl.aliyuncs.com/compatible-mode/v1", + model="qwen3-rerank", + ) + assert url == "https://rerank.example.com/v1/reranks" + + def test_non_aliyun_chat_path_base_not_remapped(self): + url = self.config.get_complete_url( + api_base="https://gateway.example.com/compatible-mode/v1", + model="qwen3-rerank", + ) + assert url == "https://gateway.example.com/compatible-mode/v1/reranks" def test_already_complete_url_passthrough(self): full = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks" diff --git a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py index 064d9d58f0c..7862297bcd5 100644 --- a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py +++ b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py @@ -209,6 +209,11 @@ class TestQwenBrandDefaultUrls: url = brand["rerank_config"]().get_complete_url(api_base=None, model="gte-rerank-v2") assert url == "https://rerank.example.com/v1/reranks" + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_rerank_remaps_chat_shaped_default_base(self, brand): + url = brand["rerank_config"]().get_complete_url(api_base=brand["default_base"], model="gte-rerank-v2") + assert url == brand["default_rerank_base"] + @pytest.mark.parametrize("brand", BRAND_CASES) def test_image_generation_complete_url(self, brand): url = brand["image_config"]().get_complete_url( diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 02655fb7f77..caf7bed7385 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -255,6 +255,19 @@ def test_transform_messages_sanitizes_empty_content(): assert result[1]["content"] == "Hi" +def test_transform_request_preserves_unity_model_service_name(): + config = DatabricksConfig() + result = config.transform_request( + model="system.ai.kimi-k3", + messages=[{"role": "user", "content": "hello"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["model"] == "system.ai.kimi-k3" + + def test_transform_request_strips_thinking_blocks_and_reasoning_content(): """Regression for LIT-6762: replaying an assistant turn that litellm decorated with `thinking_blocks` / `reasoning_content` made Databricks 400 with @@ -590,3 +603,87 @@ def test_chunk_parser_without_usage_still_parses_content(): assert result.id == "chatcmpl-test" assert result.model == "databricks-claude-sonnet-5" assert result.choices[0]["delta"]["content"] == "hi" + + +@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"]) +def test_transform_choices_surfaces_top_level_reasoning_content(reasoning_key: str) -> None: + config = DatabricksConfig() + databricks_choices = [ + { + "message": { + "role": "assistant", + "content": "391", + reasoning_key: "We need answer just number. 17*23=391.", + }, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.content == "391" + assert choices[0].message.reasoning_content == "We need answer just number. 17*23=391." + assert getattr(choices[0].message, "thinking_blocks", None) is None + + +def test_transform_choices_parses_think_tags_in_string_content(): + config = DatabricksConfig() + databricks_choices = [ + { + "message": {"role": "assistant", "content": "17 times 23391"}, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.content == "391" + assert choices[0].message.reasoning_content == "17 times 23" + + +def test_transform_choices_prefers_reasoning_blocks_over_top_level_field(): + config = DatabricksConfig() + databricks_choices = [ + { + "message": { + "role": "assistant", + "content": [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "from block"}]}, + {"type": "text", "text": "391"}, + ], + "reasoning_content": "from field", + }, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.reasoning_content == "from block" + assert choices[0].message.content == "391" + + +@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"]) +def test_chunk_parser_surfaces_top_level_reasoning_delta(reasoning_key: str) -> None: + iterator = DatabricksChatResponseIterator(None, sync_stream=True) + chunk = { + "id": "1", + "object": "chat.completion.chunk", + "created": 0, + "model": "lit-qa-deepseek-v4-flash", + "choices": [ + { + "delta": {"role": "assistant", "content": None, reasoning_key: "We need answer"}, + "index": 0, + "finish_reason": None, + } + ], + } + + parsed = iterator.chunk_parser(chunk) + + assert parsed.choices[0].delta.reasoning_content == "We need answer" + assert parsed.choices[0].delta.content is None diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py index 39198bb20f3..c6ad78366f7 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py +++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py @@ -657,6 +657,77 @@ class TestEndpointURLConstruction: assert api_base.endswith("/chat/completions") + def test_chat_gateway_endpoint_for_unity_model_on_legacy_base(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + url = DatabricksConfig().get_complete_url( + api_base="https://test.net/serving-endpoints", + api_key="test-key", + model="system.ai.kimi-k3", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions" + + def test_chat_gateway_endpoint_preserves_explicit_gateway_base(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + url = DatabricksConfig().get_complete_url( + api_base="https://test.net/ai-gateway/mlflow/v1/", + api_key="test-key", + model="system.ai.kimi-k3", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions" + + def test_chat_gateway_preserves_unity_model_service_name_with_explicit_base(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + config = DatabricksConfig() + request = config.transform_request( + model="catalog.schema.kimi-k3", + messages=[{"role": "user", "content": "hello"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert config.get_complete_url( + api_base="https://test.net/ai-gateway/mlflow/v1", + api_key="test-key", + model="catalog.schema.kimi-k3", + optional_params={}, + litellm_params={}, + ) == "https://test.net/ai-gateway/mlflow/v1/chat/completions" + assert request["model"] == "catalog.schema.kimi-k3" + + def test_chat_legacy_endpoint_remains_default(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + url = DatabricksConfig().get_complete_url( + api_base="https://test.net/serving-endpoints", + api_key="test-key", + model="databricks-kimi-k3", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://test.net/serving-endpoints/chat/completions" + def test_embeddings_endpoint(self, monkeypatch): """Embeddings endpoint is correctly appended.""" monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index e6fe01be4ba..d4ef4282b27 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch import pytest import litellm -from litellm import get_model_info, supports_reasoning, supports_vision +from litellm import supports_reasoning, supports_vision from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id @@ -16,17 +16,6 @@ from litellm.types.utils import ( ) -@pytest.fixture(autouse=True) -def force_local_model_cost(monkeypatch): - """Force local model cost map usage for all tests in this file.""" - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - # Refresh model_cost from local map - import litellm - from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map - - litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url) - - def test_validate_environment_sets_session_affinity_from_litellm_session_id(): config = FireworksAIConfig() @@ -404,15 +393,6 @@ def test_get_supported_openai_params_parallel_tool_calls_without_tool_choice( assert "tool_choice" not in supported_params -def test_get_model_info_respects_explicit_fireworks_capabilities(): - """Test that get_model_info preserves explicit capability flags from the model map.""" - model_info = get_model_info("fireworks_ai/accounts/fireworks/models/glm-5p1") - - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - - def test_get_provider_info_omits_false_supports_reasoning(monkeypatch): """Test that Fireworks only overrides supports_reasoning for supported models.""" config = FireworksAIConfig() diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py new file mode 100644 index 00000000000..d0697ca9b0e --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -0,0 +1,625 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, TypedDict, cast +from unittest.mock import MagicMock, patch +from urllib.parse import quote + +import httpx +import pytest +from openai.types.responses import ( + EasyInputMessage, + ResponseFunctionToolCall, + ResponseOutputMessage, + ResponseOutputText, + ResponseReasoningItem, +) +from openai.types.responses.response_input_param import FunctionCallOutput +from openai.types.responses.tool_param import Mcp +from typing_extensions import ReadOnly + +import litellm +from litellm.llms.fireworks_ai.responses.transformation import FireworksAIResponsesAPIConfig +from litellm.responses.file_search.emulated_handler import should_use_emulated_file_search +from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponseInputParam, ResponsesAPIResponse +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +FIREWORKS_RESPONSES_URL: Final = "https://api.fireworks.ai/inference/v1/responses" +HTTPX_CLIENT_FACTORY: Final = "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client" +NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) +NO_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _GeneratedSessionMetadata(TypedDict): + litellm_session_id_generated: ReadOnly[bool] + + +def _fireworks_response(model: str) -> Mapping[str, object]: + return ResponsesAPIResponse( + id="resp_0e946f2d46bf4b49bf8b29ff78083583", + object="response", + created_at=1788550000, + model=model, + status="completed", + output=( + ResponseReasoningItem(id="rs_1", summary=(), type="reasoning"), + ResponseOutputMessage( + id="msg_1", + status="completed", + role="assistant", + type="message", + content=(ResponseOutputText(type="output_text", text="Paris is clear and 21C.", annotations=()),), + ), + ResponseFunctionToolCall( + id="fc_1", + call_id="call_abc123", + name="get_weather", + arguments='{"city": "Paris"}', + status="completed", + type="function_call", + ), + ), + usage=ResponseAPIUsage( + input_tokens=179, + output_tokens=100, + total_tokens=279, + input_tokens_details=InputTokensDetails(cached_tokens=0), + ), + ).model_dump(mode="json", exclude_none=True) + + +def _mock_http_client(response_body: Mapping[str, object]) -> MagicMock: + client: Final = MagicMock() + response: Final = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers((("content-type", "application/json"),)) + response.json.return_value = response_body + response.text = json.dumps(response_body) + client.post.return_value = response + return client + + +def _sent_request(client: MagicMock) -> tuple[str, Mapping[str, str], Mapping[str, object]]: + kwargs: Final = client.post.call_args.kwargs + body: Final = kwargs["json"] if "json" in kwargs else json.loads(kwargs["data"]) + return kwargs["url"], kwargs["headers"], body + + +@pytest.fixture(autouse=True) +def fireworks_env(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "FIREWORKS_API_KEY", + "FIREWORKS_AI_API_KEY", + "FIREWORKSAI_API_KEY", + "FIREWORKS_AI_TOKEN", + "FIREWORKS_API_BASE", + ): + monkeypatch.delenv(name, raising=False) + + +def test_fireworks_ai_provider_config_registration() -> None: + config: Final = ProviderConfigManager.get_provider_responses_api_config( + model="accounts/fireworks/models/kimi-k3", provider=LlmProviders.FIREWORKS_AI + ) + assert isinstance(config, FireworksAIResponsesAPIConfig) + assert config.custom_llm_provider == LlmProviders.FIREWORKS_AI + + +def test_responses_call_hits_native_endpoint_with_mcp_tool_untouched() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + mcp_tool: Final[Mcp] = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", + } + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + response: Final = litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input="What is litellm?", + tools=[mcp_tool], # mutable-ok: the Responses API takes tools as a JSON list + api_key="fw-test-key", + ) + url, headers, body = _sent_request(client) + assert url == FIREWORKS_RESPONSES_URL + assert headers["Authorization"] == "Bearer fw-test-key" + assert body["model"] == "accounts/fireworks/models/kimi-k3" + assert tuple(body["tools"]) == (mcp_tool,) + assert "messages" not in body + assert isinstance(response, ResponsesAPIResponse) + function_calls: Final = tuple(item for item in response.output if getattr(item, "type", None) == "function_call") + assert getattr(function_calls[0], "call_id", None) == "call_abc123" + + +def test_responses_call_expands_bare_model_name_to_fireworks_resource() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/glm-5p3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses(model="fireworks_ai/glm-5p3", input="hi", api_key="fw-test-key") + _, _, body = _sent_request(client) + assert body["model"] == "accounts/fireworks/models/glm-5p3" + + +def test_responses_call_forwards_previous_response_id_and_store() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + tool_output: Final[FunctionCallOutput] = { + "type": "function_call_output", + "call_id": "call_abc123", + "output": "{}", + } + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/kimi-k3", + input=[tool_output], # mutable-ok: the Responses API takes input items as a JSON list + previous_response_id="resp_0e946f2d46bf4b49bf8b29ff78083583", + store=True, + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert body["previous_response_id"] == "resp_0e946f2d46bf4b49bf8b29ff78083583" + assert body["store"] is True + assert body["input"][0]["call_id"] == "call_abc123" + + +def test_responses_call_folds_developer_items_into_instructions() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input=[ # mutable-ok: the Responses API takes input as a JSON list + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]}, + ], + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert body["instructions"] == "Answer with exactly one word." + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]}, + ) + + +def test_responses_call_folds_instructions_and_developer_item_into_instructions_with_reasoning_replayed() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", + instructions="You are a coding agent running in the Codex CLI.", + input=[ # mutable-ok: the Responses API takes input as a JSON list + { + "role": "developer", + "content": [{"type": "input_text", "text": "read-only"}], + }, + {"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]}, + {"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "A trivial question."}]}, + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Paris is the capital of France.", "annotations": []}], + }, + {"role": "user", "content": [{"type": "input_text", "text": "And of Spain?"}]}, + ], + store=False, + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert body["instructions"] == ( + "You are a coding agent running in the Codex CLI.\n\nread-only" + ) + assert tuple(body["input"]) == ( + {"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]}, + {"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "A trivial question."}]}, + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Paris is the capital of France.", "annotations": []}], + }, + {"role": "user", "content": [{"type": "input_text", "text": "And of Spain?"}]}, + ) + + +def test_responses_call_folds_instructions_and_developer_item_with_previous_response_id() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", + instructions="You are a terse assistant.", + input=[ # mutable-ok: the Responses API takes input as a JSON list + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "And of Spain?"}, + ], + previous_response_id="resp_0e946f2d46bf4b49bf8b29ff78083583", + store=True, + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert body["instructions"] == "You are a terse assistant.\n\nAnswer with exactly one word." + assert body["previous_response_id"] == "resp_0e946f2d46bf4b49bf8b29ff78083583" + assert tuple(body["input"]) == ({"role": "user", "content": "And of Spain?"},) + + +def test_responses_call_keeps_a_closing_developer_item_after_an_assistant_turn_in_place() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) + assistant_turn: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Paris.", "annotations": []}], + } + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", + instructions="Be terse.", + input=[ # mutable-ok: the Responses API takes input as a JSON list + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + assistant_turn, + {"role": "developer", "content": "Now restate it in French."}, + ], + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert body["instructions"] == "Be terse.\n\nAnswer with exactly one word." + assert tuple(body["input"]) == ( + {"role": "user", "content": "What is the capital of France?"}, + assistant_turn, + {"role": "system", "content": "Now restate it in French.", "type": "message"}, + ) + + +def test_responses_call_keeps_a_mid_conversation_system_item_in_place() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input=[ # mutable-ok: the Responses API takes input as a JSON list + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Switch to French."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert "instructions" not in body + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Switch to French."}, + {"role": "user", "content": "What is the capital of France?"}, + ) + + +def test_responses_call_keeps_a_developer_item_with_non_text_parts_in_place_as_a_system_item() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) + developer_item: Final = { + "role": "developer", + "content": [ + {"type": "input_text", "text": "Match the style of this reference image."}, + {"type": "input_image", "image_url": "data:image/png;base64,iVBORw0KGgo=", "detail": "auto"}, + ], + } + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", + instructions="Answer with one word.", + input=[developer_item, {"role": "user", "content": "What is the capital of France?"}], # mutable-ok: JSON list + store=False, + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert body["instructions"] == "Answer with one word." + assert tuple(body["input"]) == ( + {"role": "system", "content": developer_item["content"], "type": "message"}, + {"role": "user", "content": "What is the capital of France?"}, + ) + + +def test_responses_call_forwards_string_input_and_instructions_unchanged() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + instructions="Answer with exactly one word.", + input="What is the capital of France?", + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert body["instructions"] == "Answer with exactly one word." + assert body["input"] == "What is the capital of France?" + + +def test_transform_request_forwards_non_string_instructions_and_input_untouched() -> None: + developer_item: Final = {"role": "developer", "content": "Answer with exactly one word."} + user_item: Final = {"role": "user", "content": "What is the capital of France?"} + request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request( + model="accounts/fireworks/models/kimi-k3", + input=cast(ResponseInputParam, [developer_item, user_item]), # mutable-ok: JSON list + response_api_optional_request_params={"instructions": ["not", "a", "string"]}, # mutable-ok: base takes a dict + litellm_params=GenericLiteLLMParams(), + headers={}, # mutable-ok: base takes a dict + ) + assert request["instructions"] == ["not", "a", "string"] + assert tuple(request["input"]) == ( + {"role": "system", "content": "Answer with exactly one word.", "type": "message"}, + user_item, + ) + + +def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + pydantic_input: Final = cast( + ResponseInputParam, + [ # mutable-ok: the Responses API takes input as a JSON list + EasyInputMessage(role="developer", content="Answer with exactly one word.", type="message"), + ResponseReasoningItem(id="rs_1", summary=(), type="reasoning"), + ResponseFunctionToolCall( + id="fc_1", + call_id="call_abc123", + name="get_weather", + arguments='{"city": "Paris"}', + status="completed", + type="function_call", + ), + FunctionCallOutput(type="function_call_output", call_id="call_abc123", output="21C"), + ], + ) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", input=pydantic_input, api_key="fw-test-key" + ) + _, _, body = _sent_request(client) + assert body["instructions"] == "Answer with exactly one word." + assert tuple(body["input"]) == ( + {"id": "rs_1", "summary": [], "type": "reasoning"}, + { + "id": "fc_1", + "call_id": "call_abc123", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + "type": "function_call", + }, + {"type": "function_call_output", "call_id": "call_abc123", "output": "21C"}, + ) + + +def test_file_search_tools_take_litellm_emulated_search_not_fireworks() -> None: + config: Final = FireworksAIResponsesAPIConfig() + file_search: Final = ({"type": "file_search", "vector_store_ids": ("vs_kb",)},) + function_tool: Final = ({"type": "function", "name": "get_weather", "parameters": {"type": "object"}},) + assert should_use_emulated_file_search(tools=file_search, provider_config=config) + assert not should_use_emulated_file_search(tools=function_tool, provider_config=config) + + +def test_responses_call_sends_session_affinity_for_caller_session_id() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-test-key", litellm_session_id="sess-42") + _, headers, _ = _sent_request(client) + assert headers["x-session-affinity"] == "sess-42" + + +def test_responses_call_keeps_caller_supplied_session_affinity_header() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + pinned: Final[Mapping[str, str]] = MappingProxyType({"x-session-affinity": "explicit-node"}) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/kimi-k3", + input="hi", + api_key="fw-test-key", + litellm_session_id="sess-42", + extra_headers=pinned, + ) + _, headers, _ = _sent_request(client) + assert headers["x-session-affinity"] == "explicit-node" + + +def test_responses_call_maps_provider_errors_to_fireworks_ai() -> None: + client: Final = MagicMock() + request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL) + client.post.side_effect = httpx.HTTPStatusError( + "unauthorized", + request=request, + response=httpx.Response(401, text='{"error": {"message": "invalid api key"}}', request=request), + ) + with patch(HTTPX_CLIENT_FACTORY, return_value=client), pytest.raises(litellm.AuthenticationError) as raised: + litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-bad-key") + assert raised.value.llm_provider == "fireworks_ai" + assert raised.value.status_code == 401 + assert "invalid api key" in str(raised.value) + + +def test_responses_call_skips_session_affinity_for_proxy_generated_session_id() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + generated: Final[_GeneratedSessionMetadata] = {"litellm_session_id_generated": True} + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/kimi-k3", + input="hi", + api_key="fw-test-key", + litellm_session_id="generated-1", + litellm_metadata=generated, + ) + _, headers, _ = _sent_request(client) + assert "x-session-affinity" not in headers + + +@pytest.mark.parametrize( + "api_base, expected", + ( + (None, FIREWORKS_RESPONSES_URL), + ("https://api.fireworks.ai/inference/v1", FIREWORKS_RESPONSES_URL), + ("https://api.fireworks.ai/inference/v1/", FIREWORKS_RESPONSES_URL), + ("https://gateway.example.com/fireworks", "https://gateway.example.com/fireworks/responses"), + ), +) +def test_get_complete_url(api_base: str | None, expected: str) -> None: + assert FireworksAIResponsesAPIConfig().get_complete_url(api_base=api_base, litellm_params=NO_PARAMS) == expected + + +def test_responses_call_reads_fireworks_api_base_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FIREWORKS_API_BASE", "https://self-hosted.example.com/v1") + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-test-key") + url, _, _ = _sent_request(client) + assert url == "https://self-hosted.example.com/v1/responses" + + +@pytest.mark.parametrize( + "env_name", ("FIREWORKS_API_KEY", "FIREWORKS_AI_API_KEY", "FIREWORKSAI_API_KEY", "FIREWORKS_AI_TOKEN") +) +def test_validate_environment_reads_every_fireworks_key_name(monkeypatch: pytest.MonkeyPatch, env_name: str) -> None: + monkeypatch.setenv(env_name, "env-key") + headers: Final = FireworksAIResponsesAPIConfig().validate_environment( + headers=NO_HEADERS, model="accounts/fireworks/models/kimi-k3", litellm_params=GenericLiteLLMParams() + ) + assert headers["Authorization"] == "Bearer env-key" + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_prefers_explicit_api_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FIREWORKS_API_KEY", "env-key") + headers: Final = FireworksAIResponsesAPIConfig().validate_environment( + headers=NO_HEADERS, + model="accounts/fireworks/models/kimi-k3", + litellm_params=GenericLiteLLMParams(api_key="explicit"), + ) + assert headers["Authorization"] == "Bearer explicit" + + +def test_validate_environment_without_any_key_raises() -> None: + with pytest.raises(ValueError, match="FIREWORKS_API_KEY"): + FireworksAIResponsesAPIConfig().validate_environment( + headers=NO_HEADERS, model="accounts/fireworks/models/kimi-k3", litellm_params=None + ) + + +def test_delete_responses_maps_fireworks_message_only_body_to_deleted_result() -> None: + response_id: Final = ( + "resp_xFaIJR9Nc_OXmqKRqL78UuAGj2Te5GY5BT_knpZiMrYoNOVmu5oc2mQW1HI7hCtEYB4mcx2lEYS0DYP1U5yEQskHunuB4==" + ) + request: Final = httpx.Request("DELETE", f"{FIREWORKS_RESPONSES_URL}/{quote(response_id, safe='')}") + client: Final = MagicMock() + client.delete.return_value = httpx.Response(200, json={"message": "Response deleted successfully"}, request=request) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + result: Final = litellm.delete_responses( + response_id=response_id, custom_llm_provider="fireworks_ai", api_key="fw-test-key" + ) + assert client.delete.call_args.kwargs["url"] == str(request.url) + assert (result.id, result.object, result.deleted) == (response_id, "response", True) + + +def _fireworks_stream_response(status: str, output: tuple[Mapping[str, object], ...]) -> Mapping[str, object]: + return { + "id": "resp_htnkJ8piNKeOHkn9LfAusC38O2OgcDQs4S8trSOJ6anLeqjUDGqu2PkWmg5N", + "object": "response", + "created_at": 1788567245, + "model": "accounts/fireworks/models/kimi-k3", + "status": status, + "output": output, + "usage": None + if status == "in_progress" + else { + "input_tokens": 95, + "output_tokens": 89, + "total_tokens": 184, + "input_tokens_details": {"cached_tokens": 94}, + }, + } + + +FIREWORKS_SSE_EVENTS: Final[tuple[Mapping[str, object], ...]] = ( + {"type": "response.created", "sequence_number": 0, "response": _fireworks_stream_response("in_progress", ())}, + { + "type": "response.output_item.added", + "sequence_number": 1, + "output_index": 0, + "item": {"id": "rs_1", "type": "reasoning", "summary": []}, + }, + { + "type": "response.reasoning_summary_text.delta", + "sequence_number": 2, + "item_id": "rs_1", + "output_index": 0, + "summary_index": 0, + "delta": "pong", + }, + { + "type": "response.output_item.added", + "sequence_number": 3, + "output_index": 1, + "item": {"id": "msg_1", "type": "message", "role": "assistant", "status": "in_progress", "content": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 4, + "item_id": "msg_1", + "output_index": 1, + "content_index": 0, + "delta": "po", + }, + { + "type": "response.output_text.delta", + "sequence_number": 5, + "item_id": "msg_1", + "output_index": 1, + "content_index": 0, + "delta": "ng", + }, + { + "type": "response.completed", + "sequence_number": 6, + "response": _fireworks_stream_response( + "completed", + ( + {"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "pong"}]}, + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], + }, + ), + ), + }, +) + + +def _sse_body(events: tuple[Mapping[str, object], ...]) -> bytes: + return b"".join(f"data: {json.dumps(dict(event))}\n\n".encode() for event in events) + b"data: [DONE]\n\n" + + +def test_streaming_responses_call_hits_native_endpoint_and_yields_every_fireworks_event() -> None: + request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL) + client: Final = MagicMock() + client.post.return_value = httpx.Response( + 200, content=_sse_body(FIREWORKS_SSE_EVENTS), headers={"content-type": "text/event-stream"}, request=request + ) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + received: Final = tuple( + litellm.responses( + model="fireworks_ai/kimi-k3", + input="Reply with the single word pong.", + stream=True, + api_key="fw-test-key", + ) + ) + url, _, body = _sent_request(client) + assert (url, body["model"], body["stream"], client.post.call_args.kwargs["stream"]) == ( + FIREWORKS_RESPONSES_URL, + "accounts/fireworks/models/kimi-k3", + True, + True, + ) + assert tuple(event.type for event in received) == tuple(event["type"] for event in FIREWORKS_SSE_EVENTS) + assert "".join(event.delta for event in received if event.type == "response.output_text.delta") == "pong" + assert received[-1].response.usage.output_tokens == 89 diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py index 5641439aa54..41f6ad9d99d 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py @@ -56,17 +56,6 @@ def use_local_model_cost_map(): monkeypatch.undo() -@pytest.mark.parametrize("alias", KIMI_ALIASES) -def test_fireworks_kimi_raw_cost_entry_limits(use_local_model_cost_map, alias): - entry = use_local_model_cost_map.model_cost[alias] - - assert entry["litellm_provider"] == "fireworks_ai" - assert entry["max_input_tokens"] == CONTEXT_WINDOW - assert entry["max_output_tokens"] == OUTPUT_LIMIT - assert entry["max_tokens"] == OUTPUT_LIMIT - assert entry["max_output_tokens"] < entry["max_input_tokens"] - - @pytest.mark.parametrize("alias", KIMI_ALIASES) def test_fireworks_kimi_get_model_info_limits(use_local_model_cost_map, alias): model_info = use_local_model_cost_map.get_model_info(model=alias) diff --git a/tests/test_litellm/llms/gemini/google_genai/__init__.py b/tests/test_litellm/llms/gemini/google_genai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py new file mode 100644 index 00000000000..4119ce99423 --- /dev/null +++ b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py @@ -0,0 +1,234 @@ +""" +Tests for the Google GenAI generateContent guardrail translation handler. +""" + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.llms.gemini.google_genai.guardrail_translation.handler import ( + GoogleGenAIGenerateContentHandler, +) +from litellm.types.utils import CallTypes + + +class GuardrailBlockedError(Exception): + pass + + +def _mock_guardrail(returned_texts): + guardrail = MagicMock() + guardrail.apply_guardrail = AsyncMock(return_value={"texts": returned_texts}) + return guardrail + + +@pytest.mark.asyncio +async def test_input_contents_text_is_guardrailed_and_written_back(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["masked question"]) + data = { + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "raw question"}]}], + } + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + call_kwargs = guardrail.apply_guardrail.call_args.kwargs + assert call_kwargs["inputs"]["texts"] == ["raw question"] + assert call_kwargs["inputs"]["model"] == "gemini-2.5-flash" + assert call_kwargs["input_type"] == "request" + assert result["contents"][0]["parts"][0]["text"] == "masked question" + + +@pytest.mark.asyncio +async def test_input_system_instruction_text_is_scanned_and_written_back(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["masked instruction", "masked question"]) + data = { + "model": "gemini-2.5-flash", + "systemInstruction": {"role": "system", "parts": [{"text": "prohibited instruction"}]}, + "contents": [{"role": "user", "parts": [{"text": "benign question"}]}], + } + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == [ + "prohibited instruction", + "benign question", + ] + assert result["systemInstruction"]["parts"][0]["text"] == "masked instruction" + assert result["contents"][0]["parts"][0]["text"] == "masked question" + + +@pytest.mark.asyncio +async def test_input_config_nested_snake_case_system_instruction_is_scanned(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["clean"]) + instruction_part = SimpleNamespace(text="prohibited instruction") + data = { + "contents": [], + "config": SimpleNamespace(system_instruction=SimpleNamespace(parts=[instruction_part])), + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["prohibited instruction"] + assert instruction_part.text == "clean" + + +@pytest.mark.asyncio +async def test_input_without_text_skips_guardrail(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail([]) + data = {"model": "gemini-2.5-flash", "contents": [{"role": "user", "parts": [{"inlineData": {}}]}]} + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + guardrail.apply_guardrail.assert_not_called() + assert result is data + + +@pytest.mark.asyncio +async def test_output_dict_response_text_is_guardrailed_and_written_back(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["masked answer"]) + response = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "harmful answer"}]}, + "finishReason": "STOP", + } + ] + } + + result = await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + request_data={"model": "gemini-2.5-flash"}, + ) + + call_kwargs = guardrail.apply_guardrail.call_args.kwargs + assert call_kwargs["inputs"]["texts"] == ["harmful answer"] + assert call_kwargs["input_type"] == "response" + assert call_kwargs["request_data"]["response"] is response + assert result["candidates"][0]["content"]["parts"][0]["text"] == "masked answer" + + +@pytest.mark.asyncio +async def test_output_object_response_text_is_guardrailed_and_written_back(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["masked answer"]) + part = SimpleNamespace(text="harmful answer") + response = SimpleNamespace( + candidates=[SimpleNamespace(content=SimpleNamespace(parts=[part]), finish_reason="STOP")] + ) + + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"] + assert part.text == "masked answer" + + +@pytest.mark.asyncio +async def test_output_without_text_skips_guardrail(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail([]) + + result = await handler.process_output_response(response={"candidates": []}, guardrail_to_apply=guardrail) + + guardrail.apply_guardrail.assert_not_called() + assert result == {"candidates": []} + + +@pytest.mark.asyncio +async def test_output_blocking_guardrail_exception_propagates(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = MagicMock() + guardrail.apply_guardrail = AsyncMock(side_effect=GuardrailBlockedError("blocked")) + response = {"candidates": [{"content": {"parts": [{"text": "harmful answer"}]}}]} + + with pytest.raises(GuardrailBlockedError): + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + +@pytest.mark.asyncio +async def test_streaming_dict_chunks_accumulate_text(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["clean"]) + chunks = [ + {"candidates": [{"content": {"parts": [{"text": "harmful "}]}}]}, + {"candidates": [{"content": {"parts": [{"text": "answer"}]}}]}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + ) + + call_kwargs = guardrail.apply_guardrail.call_args.kwargs + assert call_kwargs["inputs"]["texts"] == ["harmful answer"] + assert call_kwargs["input_type"] == "response" + assert result is chunks + + +@pytest.mark.asyncio +async def test_streaming_raw_sse_chunks_accumulate_text_across_split_frames(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["clean"]) + frame_one = "data: " + json.dumps({"candidates": [{"content": {"parts": [{"text": "harmful "}]}}]}) + "\r\n\r\n" + frame_two = "data: " + json.dumps({"candidates": [{"content": {"parts": [{"text": "answer"}]}}]}) + "\r\n\r\n" + split_at = len(frame_one) // 2 + chunks = [frame_one[:split_at], frame_one[split_at:] + frame_two[:5], frame_two[5:].encode("utf-8")] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + ) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"] + + +@pytest.mark.asyncio +async def test_streaming_blocking_guardrail_exception_propagates(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = MagicMock() + guardrail.apply_guardrail = AsyncMock(side_effect=GuardrailBlockedError("blocked")) + chunks = [{"candidates": [{"content": {"parts": [{"text": "harmful"}]}}]}] + + with pytest.raises(GuardrailBlockedError): + await handler.process_output_streaming_response(responses_so_far=chunks, guardrail_to_apply=guardrail) + + +@pytest.mark.asyncio +async def test_streaming_without_text_skips_guardrail(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail([]) + + result = await handler.process_output_streaming_response(responses_so_far=[], guardrail_to_apply=guardrail) + + guardrail.apply_guardrail.assert_not_called() + assert result == [] + + +def test_generate_content_call_types_are_registered(): + from litellm.llms.gemini.google_genai.guardrail_translation import ( + guardrail_translation_mappings, + ) + + for call_type in ( + CallTypes.generate_content, + CallTypes.agenerate_content, + CallTypes.generate_content_stream, + CallTypes.agenerate_content_stream, + ): + assert guardrail_translation_mappings[call_type] is GoogleGenAIGenerateContentHandler + + +def test_discovery_finds_generate_content_handler(): + from litellm.llms import load_guardrail_translation_mappings + + mappings = load_guardrail_translation_mappings() + assert mappings[CallTypes.agenerate_content] is GoogleGenAIGenerateContentHandler + assert mappings[CallTypes.agenerate_content_stream] is GoogleGenAIGenerateContentHandler diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index 3d8200bc474..2b3b6343fad 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,13 +1,11 @@ import json -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock -import httpx import pytest import litellm from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig -from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents def test_gemini_realtime_transformation_session_created(): @@ -308,20 +306,6 @@ def test_gemini_realtime_transformation_generation_complete(): assert contains_audio_done_event, "Expected audio done event" -def test_gemini_3_1_flash_live_preview_model_cost_map_entry(): - for key in ( - "gemini-3.1-flash-live-preview", - "gemini/gemini-3.1-flash-live-preview", - ): - assert key in litellm.model_cost - info = litellm.model_cost[key] - assert "/v1/realtime" in info.get("supported_endpoints", []) - assert info.get("max_input_tokens") == 131072 - assert info.get("max_output_tokens") == 65536 - assert "video" in info.get("supported_modalities", []) - assert info.get("supports_function_calling") is True - - def test_gemini_realtime_tool_call_transformation(): """Test transformation of Gemini toolCall to OpenAI function_call_arguments.done format.""" config = GeminiRealtimeConfig() @@ -1845,19 +1829,6 @@ def test_is_audio_only_live_model_uses_cost_map(model, expected, patch_gemini_au assert GeminiRealtimeConfig._is_audio_only_live_model(model) == expected -def test_gemini_live_native_audio_entry_is_vertex_only(): - import json - from pathlib import Path - from typing import Final - - catalog_path: Final = Path(__file__).parents[5] / "model_prices_and_context_window.json" - catalog: Final = json.loads(catalog_path.read_text()) - vertex_key: Final = "gemini-live-2.5-flash-native-audio" - assert catalog[vertex_key]["litellm_provider"] == "vertex_ai-language-models" - assert catalog[vertex_key].get("gemini_native_audio") is True - assert "gemini/gemini-live-2.5-flash-native-audio" not in catalog, "the Gemini API does not serve this model" - - def test_is_setup_message_and_is_content_message(): config = GeminiRealtimeConfig() assert config.is_setup_message({"setup": {}}) is True diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py new file mode 100644 index 00000000000..bc0ea23e249 --- /dev/null +++ b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py @@ -0,0 +1,155 @@ +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config +from litellm.llms.hosted_vllm.image_edit.transformation import HostedVLLMImageEditConfig +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +PNG_BYTES = b"\x89PNG\r\n\x1a\nfakepng" +MODEL = "Qwen/Qwen-Image-Edit-2511" + + +@pytest.fixture(autouse=True) +def _clear_hosted_vllm_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + + +def test_provider_config_registration(): + config = ProviderConfigManager.get_provider_image_edit_config( + model=f"hosted_vllm/{MODEL}", + provider=LlmProviders.HOSTED_VLLM, + ) + + assert isinstance(config, HostedVLLMImageEditConfig) + assert isinstance(get_hosted_vllm_image_edit_config(MODEL), HostedVLLMImageEditConfig) + + +@pytest.mark.parametrize( + "api_base", + ["http://localhost:8091", "http://localhost:8091/", "http://localhost:8091/v1", "http://localhost:8091/v1/"], +) +def test_get_complete_url_appends_images_edits(api_base: str): + config = HostedVLLMImageEditConfig() + + assert ( + config.get_complete_url(model=MODEL, api_base=api_base, litellm_params={}) + == "http://localhost:8091/v1/images/edits" + ) + + +def test_get_complete_url_falls_back_to_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HOSTED_VLLM_API_BASE", "http://vllm-omni:8000/v1") + config = HostedVLLMImageEditConfig() + + assert ( + config.get_complete_url(model=MODEL, api_base=None, litellm_params={}) + == "http://vllm-omni:8000/v1/images/edits" + ) + + +def test_get_complete_url_requires_api_base(): + config = HostedVLLMImageEditConfig() + + with pytest.raises(ValueError, match="api_base not set"): + config.get_complete_url(model=MODEL, api_base=None, litellm_params={}) + + +def test_validate_environment_defaults_to_fake_api_key(): + headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL) + + assert headers == {"Authorization": "Bearer fake-api-key"} + + +def test_validate_environment_uses_provided_api_key_and_keeps_headers(): + headers = HostedVLLMImageEditConfig().validate_environment( + headers={"X-Test": "1"}, + model=MODEL, + api_key="my-custom-key", + ) + + assert headers == {"X-Test": "1", "Authorization": "Bearer my-custom-key"} + + +def test_validate_environment_falls_back_to_env_api_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HOSTED_VLLM_API_KEY", "env-key") + + headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL) + + assert headers["Authorization"] == "Bearer env-key" + + +def test_image_edit_posts_multipart_to_vllm_omni(): + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + response = litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + api_key="test-key", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + seed=42, + ) + + assert response.data + assert len(captured) == 1 + request = captured[0] + assert str(request.url) == "http://localhost:8091/v1/images/edits" + assert request.headers["authorization"] == "Bearer test-key" + assert request.headers["content-type"].startswith("multipart/form-data") + assert b'name="image[]"' in request.content + assert PNG_BYTES in request.content + assert f'name="model"\r\n\r\n{MODEL}'.encode() in request.content + assert b'name="prompt"\r\n\r\nadd a hat' in request.content + assert b'name="seed"\r\n\r\n42' in request.content + + +@pytest.mark.parametrize("param", ["mask", "quality", "input_fidelity"]) +def test_params_vllm_omni_ignores_are_not_advertised(param: str): + supported = HostedVLLMImageEditConfig().get_supported_openai_params(MODEL) + + assert param not in supported + assert {"image", "prompt", "n", "size", "response_format", "background", "user"} <= set(supported) + + +def test_image_edit_rejects_quality_unless_dropped(): + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + + with pytest.raises(litellm.UnsupportedParamsError, match="quality"): + litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + client=client, + quality="low", + ) + assert captured == [] + + litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + client=client, + quality="low", + drop_params=True, + ) + + assert len(captured) == 1 + assert b'name="quality"' not in captured[0].content + assert b'name="prompt"\r\n\r\nadd a hat' in captured[0].content diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py index cff3c6be940..04813143fae 100644 --- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py @@ -7,6 +7,7 @@ import os from unittest import mock import httpx +import pytest import litellm from litellm.llms.inception.chat.transformation import InceptionChatConfig @@ -231,26 +232,6 @@ def test_inception_in_provider_lists(): assert "https://api.inceptionlabs.ai/v1" in litellm.openai_compatible_endpoints -def test_inception_model_configuration(monkeypatch): - from litellm import get_model_info - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.inception_models = set() - litellm.add_known_models() - - info = get_model_info("inception/mercury-2") - assert info.get("litellm_provider") == "inception" - assert info.get("mode") == "chat" - assert info.get("max_input_tokens") == 128000 - assert info.get("input_cost_per_token") == 2.5e-07 - assert info.get("output_cost_per_token") == 7.5e-07 - assert info.get("cache_read_input_token_cost") == 2.5e-08 - assert info.get("supports_function_calling") is True - assert info.get("supports_tool_choice") is True - assert info.get("supports_response_schema") is True - - def test_inception_model_list_populated(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") @@ -258,6 +239,7 @@ def test_inception_model_list_populated(monkeypatch): litellm.add_known_models() assert "inception/mercury-2" in litellm.inception_models + assert "inception/mercury-2.5" in litellm.inception_models for model in litellm.inception_models: assert model.startswith("inception/") @@ -324,3 +306,24 @@ def test_inception_completion_targets_inception_endpoint(): assert captured["body"]["model"] == "mercury-2" assert captured["body"]["tool_choice"] == "auto" assert response.choices[0].message.content == "hi" + + +def test_inception_mercury_2_5_cost_and_tokens(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + model = "inception/mercury-2.5" + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, + prompt_tokens=1000, + completion_tokens=500, + ) + assert abs(prompt_cost - 0.0002) < 1e-9 + assert abs(completion_cost - 0.000375) < 1e-9 + + model_info = litellm.get_model_info(model) + assert model_info["max_input_tokens"] == 260000 + assert model_info["max_output_tokens"] == 65536 + assert model_info["litellm_provider"] == "inception" + assert model_info["mode"] == "chat" + assert model_info["supports_function_calling"] is True + assert model_info["supports_response_schema"] is True diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py index 62688a13c35..ed3f34fc744 100644 --- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py @@ -143,24 +143,6 @@ async def test_inception_fim_async(): assert r.choices[0].text == "a + b" -def test_inception_fim_model_configuration(monkeypatch): - from litellm import get_model_info - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.text_completion_inception_models = set() - litellm.add_known_models() - - assert ( - "text-completion-inception/mercury-edit-2" - in litellm.text_completion_inception_models - ) - info = get_model_info("text-completion-inception/mercury-edit-2") - assert info.get("litellm_provider") == "text-completion-inception" - assert info.get("mode") == "completion" - assert info.get("max_input_tokens") == 32000 - - def test_inception_fim_targets_fim_endpoint(): """ End-to-end: a FIM request must hit `/v1/fim/completions` (NOT diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py new file mode 100644 index 00000000000..a0f22a59f0c --- /dev/null +++ b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py @@ -0,0 +1,436 @@ +import asyncio +import json +from collections.abc import Sequence +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from openai import APIConnectionError + +import litellm +from litellm.llms.litellm_proxy.skills.skill_search import ( + MAX_SKILL_SEARCH_TEXT_CHARS, + SkillSearchEmbeddingFailed, + SkillSearchHits, + SkillSearchIndex, + SkillSearchNotConfigured, + search_skills, + skill_search_text, +) +from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.common_utils.semantic_text_index import Vector, cosine_similarity + +CALLER: Final = UserAPIKeyAuth(api_key="hashed-caller-key", team_id="team-1", user_id="user-1") + +TRANSLATOR: Final = LiteLLM_SkillsTable( + skill_id="translate-file", + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", +) +SQL_ANALYST: Final = LiteLLM_SkillsTable( + skill_id="warehouse-sql-analyst", + display_title="Warehouse SQL Analyst", + description="Runs SQL against the inventory database", +) +TRIP_PLANNER: Final = LiteLLM_SkillsTable( + skill_id="trip-planner", + display_title="Trip Planner", + description="Books flights and hotels", +) +SKILLS: Final = (TRANSLATOR, SQL_ANALYST, TRIP_PLANNER) + +VECTORS: Final = MappingProxyType( + { + "language translation": (1.0, 0.0, 0.0), + skill_search_text(TRANSLATOR): (0.9, 0.1, 0.0), + skill_search_text(SQL_ANALYST): (0.0, 1.0, 0.0), + skill_search_text(TRIP_PLANNER): (0.3, 0.0, 1.0), + } +) + + +def _pass_through_key_limits() -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + return limits + + +def _embedding_router() -> MagicMock: + router = MagicMock() + router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], + ) + ) + return router + + +class FakeEmbedder: + def __init__(self) -> None: + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + return tuple(VECTORS[text] for text in texts) + + +class FixedDimensionEmbedder: + def __init__(self, dimensions: int) -> None: + self.dimensions: Final = dimensions + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + await asyncio.sleep(0) + return tuple((1.0,) * self.dimensions for _ in texts) + + +class TestSkillSearchText: + def test_joins_title_description_and_instructions(self) -> None: + assert skill_search_text(TRANSLATOR) == ( + "Document Translator\n" + "Converts files from one language into another\n" + "Take an uploaded document and produce it in the target language" + ) + + def test_missing_fields_fall_back_to_whatever_is_present(self) -> None: + assert skill_search_text(LiteLLM_SkillsTable(skill_id="bare", display_title="bare")) == "bare" + + def test_all_fields_absent_is_an_empty_string(self) -> None: + assert skill_search_text(LiteLLM_SkillsTable(skill_id="empty")) == "" + + def test_oversized_instructions_are_cut_so_one_skill_cannot_blow_up_the_embedding_batch(self) -> None: + bloated = LiteLLM_SkillsTable( + skill_id="bloated", display_title="Bloated", instructions="x" * (MAX_SKILL_SEARCH_TEXT_CHARS * 3) + ) + text = skill_search_text(bloated) + assert len(text) == MAX_SKILL_SEARCH_TEXT_CHARS + assert text.startswith("Bloated\n") + + +class TestCosineSimilarity: + def test_identical_direction_scores_one(self) -> None: + assert cosine_similarity((2.0, 0.0), (1.0, 0.0)) == pytest.approx(1.0) + + def test_orthogonal_scores_zero(self) -> None: + assert cosine_similarity((1.0, 0.0), (0.0, 1.0)) == pytest.approx(0.0) + + def test_zero_vector_scores_zero_instead_of_dividing(self) -> None: + assert cosine_similarity((0.0, 0.0), (1.0, 0.0)) == 0.0 + + +class TestSkillSearchIndex: + @pytest.mark.asyncio + async def test_ranks_by_similarity_and_truncates_to_top_k(self) -> None: + outcome = await SkillSearchIndex().search( + "language translation", SKILLS, top_k=2, embed=FakeEmbedder(), embedding_model="m" + ) + assert isinstance(outcome, SkillSearchHits) + assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file", "trip-planner"] + assert outcome.hits[0].score > outcome.hits[1].score + + @pytest.mark.asyncio + async def test_second_search_only_embeds_the_query(self) -> None: + index = SkillSearchIndex() + embedder = FakeEmbedder() + await index.search("language translation", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("language translation", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert len(embedder.calls[0]) == 1 + len(SKILLS) + assert embedder.calls[1] == ("language translation",) + + @pytest.mark.asyncio + async def test_switching_embedding_models_does_not_reuse_cached_vectors(self) -> None: + index = SkillSearchIndex() + await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="small") + wide = FixedDimensionEmbedder(2) + outcome = await index.search("language translation", SKILLS, top_k=5, embed=wide, embedding_model="wide") + assert isinstance(outcome, SkillSearchHits) + assert len(wide.calls[0]) == 1 + len(SKILLS) + + @pytest.mark.asyncio + async def test_cached_vectors_of_another_dimension_are_re_embedded(self) -> None: + index = SkillSearchIndex() + await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="m") + fallback = FixedDimensionEmbedder(2) + outcome = await index.search("language translation", SKILLS, top_k=5, embed=fallback, embedding_model="m") + assert isinstance(outcome, SkillSearchHits) + assert fallback.calls == [ + ("language translation",), + ("language translation", *(skill_search_text(skill) for skill in SKILLS)), + ] + + @pytest.mark.asyncio + async def test_re_embedding_a_subset_drops_the_other_skills_old_vectors(self) -> None: + index = SkillSearchIndex() + await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="m") + wide = FixedDimensionEmbedder(2) + await index.search("language translation", SKILLS[:1], top_k=5, embed=wide, embedding_model="m") + await index.search("language translation", SKILLS, top_k=5, embed=wide, embedding_model="m") + assert wide.calls[-1] == ("language translation", *(skill_search_text(skill) for skill in SKILLS[1:])) + + @pytest.mark.asyncio + async def test_concurrent_searches_keep_each_others_vectors(self) -> None: + index = SkillSearchIndex() + embedder = FixedDimensionEmbedder(3) + await asyncio.gather( + index.search("q", SKILLS[:1], top_k=5, embed=embedder, embedding_model="m"), + index.search("q", SKILLS[1:], top_k=5, embed=embedder, embedding_model="m"), + ) + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q",) + + @pytest.mark.asyncio + async def test_least_recently_searched_skills_are_evicted_once_the_index_is_full(self) -> None: + index = SkillSearchIndex(max_entries=len(SKILLS)) + embedder = FixedDimensionEmbedder(3) + newcomer = LiteLLM_SkillsTable(skill_id="newcomer", display_title="Newcomer") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS[:1], top_k=5, embed=embedder, embedding_model="m") + await index.search("q", (newcomer,), top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[1])) + + @pytest.mark.asyncio + async def test_deleted_skills_stop_occupying_the_index_after_enough_new_ones(self) -> None: + index = SkillSearchIndex(max_entries=2) + embedder = FixedDimensionEmbedder(3) + for generation in range(50): + skill = LiteLLM_SkillsTable(skill_id=f"gen-{generation}", display_title=f"Generation {generation}") + await index.search("q", (skill,), top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[0])) + + @pytest.mark.asyncio + async def test_mixed_dimensions_in_one_batch_become_embedding_failed(self) -> None: + async def mixed(texts: Sequence[str]) -> Sequence[Vector]: + return ((1.0, 0.0), *((1.0, 0.0, 0.0) for _ in texts[1:])) + + outcome = await SkillSearchIndex().search("q", SKILLS, top_k=5, embed=mixed, embedding_model="m") + assert isinstance(outcome, SkillSearchEmbeddingFailed) + assert "mixed dimensions" in outcome.reason + + @pytest.mark.asyncio + async def test_no_accessible_skills_returns_no_hits_without_embedding(self) -> None: + embedder = FakeEmbedder() + outcome = await SkillSearchIndex().search("anything", (), top_k=5, embed=embedder, embedding_model="m") + assert outcome == SkillSearchHits(hits=()) + assert embedder.calls == [] + + @pytest.mark.asyncio + async def test_provider_error_becomes_embedding_failed(self) -> None: + async def failing(texts: Sequence[str]) -> Sequence[Vector]: + raise APIConnectionError(request=MagicMock()) + + outcome = await SkillSearchIndex().search("q", SKILLS, top_k=5, embed=failing, embedding_model="m") + assert isinstance(outcome, SkillSearchEmbeddingFailed) + assert "embedding the search query failed" in outcome.reason + + @pytest.mark.asyncio + async def test_wrong_vector_count_becomes_embedding_failed(self) -> None: + async def short(texts: Sequence[str]) -> Sequence[Vector]: + return ((1.0, 0.0, 0.0),) + + outcome = await SkillSearchIndex().search("q", SKILLS, top_k=5, embed=short, embedding_model="m") + assert isinstance(outcome, SkillSearchEmbeddingFailed) + + +class TestSearchSkills: + @pytest.mark.asyncio + async def test_no_embedding_model_is_not_configured(self) -> None: + outcome = await search_skills( + "q", + SKILLS, + 5, + router=MagicMock(), + embedding_model=None, + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + assert isinstance(outcome, SkillSearchNotConfigured) + assert "skill_search_embedding_model" in outcome.reason + + @pytest.mark.asyncio + async def test_no_router_is_not_configured(self) -> None: + outcome = await search_skills( + "q", + SKILLS, + 5, + router=None, + embedding_model="m", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + assert isinstance(outcome, SkillSearchNotConfigured) + + @pytest.mark.asyncio + async def test_router_embeddings_are_read_from_the_response(self) -> None: + router = _embedding_router() + outcome = await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + assert isinstance(outcome, SkillSearchHits) + assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file"] + assert router.aembedding.await_args.kwargs["model"] == "text-embedding-3-small" + + @pytest.mark.asyncio + async def test_embedding_spend_is_attributed_to_the_calling_key(self) -> None: + router = _embedding_router() + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + metadata = router.aembedding.await_args.kwargs["metadata"] + assert metadata["user_api_key"] == "hashed-caller-key" + assert metadata["user_api_key_team_id"] == "team-1" + assert metadata["user_api_key_user_id"] == "user-1" + + @pytest.mark.asyncio + async def test_key_limits_are_checked_against_the_real_embedding_call_before_it_runs(self) -> None: + router = _embedding_router() + key_limits = _pass_through_key_limits() + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + checked = key_limits.pre_call_hook.await_args.kwargs + assert checked["user_api_key_dict"] is CALLER + assert checked["call_type"] == "aembedding" + assert checked["data"]["model"] == "text-embedding-3-small" + assert checked["data"]["input"] == router.aembedding.await_args.kwargs["input"] + assert checked["data"]["metadata"]["user_api_key"] == "hashed-caller-key" + + @pytest.mark.asyncio + async def test_the_embedding_model_sees_the_request_as_the_guardrails_rewrote_it(self) -> None: + router = MagicMock() + router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": [1.0, 0.0, 0.0]} for i in range(len(input))], + ) + ) + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type: { + **data, + "input": ["[MASKED]" for _ in data["input"]], + "metadata": {**data["metadata"], "guardrail": "masked"}, + } + ) + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + sent = tuple(call.kwargs for call in router.aembedding.await_args_list) + assert sent + assert all(set(call["input"]) == {"[MASKED]"} for call in sent) + assert all(call["metadata"]["guardrail"] == "masked" for call in sent) + + @pytest.mark.asyncio + async def test_a_key_over_its_limit_never_reaches_the_embedding_model(self) -> None: + router = _embedding_router() + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock(side_effect=ProxyRateLimitError(detail="rpm exceeded")) + with pytest.raises(ProxyRateLimitError) as raised: + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + assert raised.value.status_code == 429 + router.aembedding.assert_not_awaited() + + +@pytest.fixture +def accessible_skills(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + list_for_search = AsyncMock(return_value=list(SKILLS)) + monkeypatch.setattr( + "litellm.llms.litellm_proxy.skills.handler.LiteLLMSkillsHandler.list_skills_for_search", list_for_search + ) + return list_for_search + + +@pytest.fixture +def key_limits(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + limits = _pass_through_key_limits() + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", limits) + return limits + + +@pytest.fixture +def embedding_router(monkeypatch: pytest.MonkeyPatch, key_limits: MagicMock) -> MagicMock: + router = _embedding_router() + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "skill_search_embedding_model", "text-embedding-3-small") + return router + + +class TestHandleSkillSearchMCP: + @pytest.mark.asyncio + async def test_top_k_is_clamped_to_the_same_ceiling_as_rest( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + from litellm.llms.litellm_proxy.skills.skill_search import MAX_SKILL_SEARCH_TOP_K + from litellm.proxy._experimental.mcp_server.tool_search import handle_skill_search + + many_skills: Final = tuple( + LiteLLM_SkillsTable( + skill_id=f"skill-{i}", display_title=TRIP_PLANNER.display_title, description=TRIP_PLANNER.description + ) + for i in range(MAX_SKILL_SEARCH_TOP_K + 50) + ) + accessible_skills.return_value = list(many_skills) + + result = await handle_skill_search( + query="language translation", top_k=10_000, user_api_key_dict=UserAPIKeyAuth(user_id="u") + ) + assert result.isError is False + assert len(json.loads(result.content[0].text)) == MAX_SKILL_SEARCH_TOP_K + + @pytest.mark.asyncio + async def test_top_k_below_one_is_raised_to_one( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + from litellm.proxy._experimental.mcp_server.tool_search import handle_skill_search + + result = await handle_skill_search( + query="language translation", top_k=0, user_api_key_dict=UserAPIKeyAuth(user_id="u") + ) + assert result.isError is False + assert len(json.loads(result.content[0].text)) == 1 diff --git a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py new file mode 100644 index 00000000000..d819a79cef1 --- /dev/null +++ b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py @@ -0,0 +1,197 @@ +import base64 +from typing import Final +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig +from litellm.llms.mistral.audio_speech.transformation import ( + MistralTextToSpeechConfig, + MistralTextToSpeechException, +) +from litellm.utils import ProviderConfigManager + +SPEECH_URL: Final = "https://api.mistral.ai/v1/audio/speech" + + +def test_mistral_text_to_speech_config_installed(): + config: Final = ProviderConfigManager.get_provider_text_to_speech_config( + model="voxtral-mini-tts-2603", + provider=litellm.LlmProviders.MISTRAL, + ) + assert isinstance(config, BaseTextToSpeechConfig) + assert isinstance(config, MistralTextToSpeechConfig) + + +def test_map_openai_params_drops_speed_and_instructions(): + config: Final = MistralTextToSpeechConfig() + voice, params = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={"response_format": "wav", "speed": 1.5, "instructions": "sound cheerful"}, + voice="en_paul_neutral", + ) + assert voice == "en_paul_neutral" + assert params == {"response_format": "wav"} + + +def test_map_openai_params_accepts_voice_dict_and_ref_audio(): + config: Final = MistralTextToSpeechConfig() + voice, params = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice={"voice_id": "1f3a8b0c-voice-uuid"}, + kwargs={"ref_audio": "bXktdm9pY2Utc2FtcGxl"}, + ) + assert voice == "1f3a8b0c-voice-uuid" + assert params == {"ref_audio": "bXktdm9pY2Utc2FtcGxl"} + + +def test_transform_request_builds_mistral_body(): + config: Final = MistralTextToSpeechConfig() + data: Final = config.transform_text_to_speech_request( + model="voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + optional_params={"response_format": "wav"}, + litellm_params={}, + headers={}, + ) + assert data["dict_body"] == { + "model": "voxtral-mini-tts-2603", + "input": "hello from litellm", + "voice_id": "en_paul_neutral", + "response_format": "wav", + } + assert data["headers"] == {"Content-Type": "application/json"} + + +def test_transform_request_omits_voice_for_ref_audio_cloning(): + config: Final = MistralTextToSpeechConfig() + data: Final = config.transform_text_to_speech_request( + model="voxtral-mini-tts-2603", + input="clone me", + voice=None, + optional_params={"ref_audio": "bXktdm9pY2Utc2FtcGxl"}, + litellm_params={}, + headers={}, + ) + assert data["dict_body"] == { + "model": "voxtral-mini-tts-2603", + "input": "clone me", + "ref_audio": "bXktdm9pY2Utc2FtcGxl", + } + + +def test_get_complete_url_default_base(): + config: Final = MistralTextToSpeechConfig() + url: Final = config.get_complete_url(model="voxtral-mini-tts-2603", api_base=None, litellm_params={}) + assert url == SPEECH_URL + + +@pytest.mark.parametrize( + "api_base", + ["https://custom.api.example.com/v1/", "https://custom.api.example.com/v1", "https://custom.api.example.com"], +) +def test_get_complete_url_custom_base_always_versioned(api_base: str): + config: Final = MistralTextToSpeechConfig() + url: Final = config.get_complete_url(model="voxtral-mini-tts-2603", api_base=api_base, litellm_params={}) + assert url == "https://custom.api.example.com/v1/audio/speech" + + +def test_validate_environment_sets_bearer_header(): + config: Final = MistralTextToSpeechConfig() + headers: Final = config.validate_environment( + headers={"x-custom": "1"}, + model="voxtral-mini-tts-2603", + api_key="sk-mistral-test", + ) + assert headers == { + "x-custom": "1", + "Authorization": "Bearer sk-mistral-test", + "Content-Type": "application/json", + } + + +def test_validate_environment_requires_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("MISTRAL_API_KEY", raising=False) + config: Final = MistralTextToSpeechConfig() + with pytest.raises(MistralTextToSpeechException, match="MISTRAL_API_KEY"): + config.validate_environment(headers={}, model="voxtral-mini-tts-2603") + + +def test_transform_response_decodes_base64_audio(): + config: Final = MistralTextToSpeechConfig() + audio_bytes: Final = b"RIFF-fake-wav-bytes" + raw_response: Final = httpx.Response( + 200, + json={"audio_data": base64.b64encode(audio_bytes).decode()}, + headers={"x-request-id": "req-123"}, + request=httpx.Request( + "POST", + SPEECH_URL, + json={"model": "voxtral-mini-tts-2603", "input": "hi", "response_format": "wav"}, + ), + ) + result: Final = config.transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + assert result.content == audio_bytes + assert result.response.headers["content-type"] == "audio/wav" + assert result.response.headers["content-length"] == str(len(audio_bytes)) + assert result.response.headers["x-request-id"] == "req-123" + + +def test_transform_response_missing_audio_data_raises(): + config: Final = MistralTextToSpeechConfig() + raw_response: Final = httpx.Response( + 200, + json={"detail": "unexpected"}, + request=httpx.Request("POST", SPEECH_URL, json={"model": "voxtral-mini-tts-2603", "input": "hi"}), + ) + with pytest.raises(MistralTextToSpeechException, match="audio_data"): + config.transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + +def test_map_openai_params_maps_openai_voice_aliases(): + config: Final = MistralTextToSpeechConfig() + alloy_voice, _ = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice="alloy", + ) + nova_voice, _ = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice="Nova", + ) + passthrough_voice, _ = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice="en_paul_happy", + ) + assert alloy_voice == "en_paul_neutral" + assert nova_voice == "gb_jane_sarcasm" + assert passthrough_voice == "en_paul_happy" + + +def test_transform_response_invalid_base64_raises(): + config: Final = MistralTextToSpeechConfig() + raw_response: Final = httpx.Response( + status_code=200, + json={"audio_data": "QUJD!QUJD"}, + request=httpx.Request("POST", SPEECH_URL), + ) + with pytest.raises(MistralTextToSpeechException, match="base64"): + config.transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) diff --git a/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py b/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py index f5d31c0da54..9de473fb1f0 100644 --- a/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py +++ b/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py @@ -1,1537 +1,222 @@ -import asyncio -import gc -import sys -import threading -import weakref -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +import json +from collections.abc import Mapping +from typing import Final +from unittest.mock import MagicMock import httpx import pytest import litellm -from litellm.exceptions import BadRequestError, ServiceUnavailableError, Timeout -from litellm.llms.mongodb.common_utils import ( - _MAX_CACHED_CLIENTS, - _async_clients, - _sync_clients, - MongoClientKey, - index_not_ready_error, - missing_index_error, - get_async_client, - get_sync_client, - reset_client_cache, - translate_mongo_error, -) -from litellm.llms.mongodb.vector_stores.transformation import ( - MongoDBVectorStoreConfig, - _MongoDBSearchParams, -) -from litellm.types.utils import LlmProviders -from litellm.utils import ProviderConfigManager +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.mongodb.vector_stores.transformation import MongoDBVectorStoreConfig +from litellm.types.utils import EmbeddingResponse +from litellm.types.vector_stores import VectorStoreSearchOptionalRequestParams, VectorStoreSearchResponse -CONNECTION_STRING = "mongodb+srv://user:pw@cluster.example.mongodb.net" -INDEX = "movies_vector_index" - -BASE_PARAMS = { - "litellm_embedding_model": "openai/text-embedding-ada-002", - "mongodb_connection_string": CONNECTION_STRING, - "mongodb_database": "sample_mflix", - "mongodb_collection": "embedded_movies", +BASE_PARAMS: Final = { + "api_base": "https://sidecar.example/prefix", + "api_key": "test-sidecar-key", + "litellm_embedding_model": "embedding-alias", + "mongodb_database": "policies", + "mongodb_collection": "documents", +} +RESULT: Final = { + "object": "vector_store.search_results.page", + "search_query": "travel policy", + "data": [ + {"score": 0.9, "file_id": "123", "filename": "123", "content": [{"type": "text", "text": "Use code BLUE-42"}]} + ], } -READY_INDEX = [{"name": INDEX, "status": "READY", "queryable": True}] +class RecordingEmbeddingExecutor: + def __init__(self) -> None: + self.call: Final = MagicMock(return_value=EmbeddingResponse(data=[{"embedding": [0.1, 0.2, 0.3]}])) + + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + return self.call(model, query, configuration) + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + return self.call(model, query, configuration) -class RecordingClient: - """Stands in for pymongo's client class so the cache tests inject a fake rather than - patching the importer, and so they can assert what the client was actually built with.""" - - def __init__(self, connection_string, **kwargs): - self.connection_string = connection_string - self.kwargs = kwargs - - -class FakeCollection: - def __init__(self, documents, error=None, search_indexes=None): - self.documents = documents - self.error = error - self.search_indexes = READY_INDEX if search_indexes is None else search_indexes - self.pipeline = None - self.listed_indexes = [] - - def aggregate(self, pipeline): - self.pipeline = pipeline - if self.error is not None: - raise self.error - return iter(self.documents) - - def list_search_indexes(self, name): - self.listed_indexes.append(name) - return iter(self.search_indexes) - - -class FakeAsyncCollection(FakeCollection): - async def aggregate(self, pipeline): - self.pipeline = pipeline - if self.error is not None: - raise self.error - - async def cursor(): - for document in self.documents: - yield document - - return cursor() - - async def list_search_indexes(self, name): - self.listed_indexes.append(name) - - async def cursor(): - for entry in self.search_indexes: - yield entry - - return cursor() - - -class FakeDatabase: - def __init__(self, collection): - self.collection = collection - self.requested_collection = None - - def __getitem__(self, name): - self.requested_collection = name - return self.collection - - -class FakeClient: - def __init__(self, collection): - self.database = FakeDatabase(collection) - self.requested_database = None - - def __getitem__(self, name): - self.requested_database = name - return self.database - - -class FakeEmbeddingExecutor: - def __init__(self, embedding): - self.embedding = embedding - self.captured = None - - def _respond(self, model, query, configuration): - self.captured = SimpleNamespace(model=model, query=query, configuration=configuration) - return SimpleNamespace(data=[{"embedding": self.embedding}] if self.embedding is not None else []) - - def embed(self, model, query, configuration): - return self._respond(model, query, configuration) - - async def aembed(self, model, query, configuration): - return self._respond(model, query, configuration) - - -def _config(documents=(), embedding=(0.1, 0.2, 0.3), error=None, search_indexes=None): - collection = FakeCollection(list(documents), error, search_indexes) - client = FakeClient(collection) - config = MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor(list(embedding) if embedding is not None else None), - sync_client_factory=lambda key: client, - ) - return config, client, collection - - -def _async_config(documents=(), embedding=(0.1, 0.2, 0.3), error=None, search_indexes=None): - collection = FakeAsyncCollection(list(documents), error, search_indexes) - client = FakeClient(collection) - config = MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor(list(embedding) if embedding is not None else None), - async_client_factory=lambda key: client, - ) - return config, client, collection - - -def _search(config, query="a lone astronaut", optional_params=None, litellm_params=None, timeout=None): - return config.execute_search_vector_store_request( - vector_store_id=INDEX, - query=query, - vector_store_search_optional_params=optional_params or {}, - litellm_logging_obj=MagicMock(), - litellm_params={**BASE_PARAMS, **(litellm_params or {})}, - timeout=timeout, - ) - - -async def _asearch(config, query="a lone astronaut", optional_params=None, litellm_params=None): - return await config.aexecute_search_vector_store_request( - vector_store_id=INDEX, - query=query, - vector_store_search_optional_params=optional_params or {}, - litellm_logging_obj=MagicMock(), - litellm_params={**BASE_PARAMS, **(litellm_params or {})}, - ) - - -def _stage(collection, name): - return next(stage[name] for stage in collection.pipeline if name in stage) - - -def test_search_builds_vector_search_stage_against_the_named_index(): - config, client, collection = _config() - - _search(config, optional_params={"max_num_results": 5}) - - assert client.requested_database == "sample_mflix" - assert client.database.requested_collection == "embedded_movies" - assert _stage(collection, "$vectorSearch") == { - "index": INDEX, - "path": "embedding", - "queryVector": (0.1, 0.2, 0.3), - "numCandidates": 100, - "limit": 5, +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("limit,candidates", [(None, 100), (1, 100), (50, 500)]) +@pytest.mark.asyncio +async def test_search_preserves_embedding_and_http_contract( + asynchronous: bool, limit: int | None, candidates: int +) -> None: + executor: Final = RecordingEmbeddingExecutor() + config: Final = MongoDBVectorStoreConfig(executor) + params: Final = { + **BASE_PARAMS, + "mongodb_text_field": "metadata.body", + "mongodb_embedding_field": "stored_vector", + "litellm_embedding_config": {"dimensions": 3}, + "timeout": 0.75, } - - -def test_the_pipeline_reaches_pymongo_as_a_list(): - """pymongo's common.validate_list rejects any other sequence with - 'pipeline must be a list, not ', so the outer container is part of the contract.""" - config, _, collection = _config() - - _search(config) - - assert isinstance(collection.pipeline, list) - - -def test_search_projects_the_text_field_and_the_similarity_score(): - config, _, collection = _config() - - _search(config) - - assert _stage(collection, "$project") == {"text": 1, "score": {"$meta": "vectorSearchScore"}} - - -def test_search_defaults_to_ten_results(): - config, _, collection = _config() - - _search(config) - - assert _stage(collection, "$vectorSearch")["limit"] == 10 - - -def test_search_honors_custom_field_names(): - config, _, collection = _config() - - _search( - config, - litellm_params={"mongodb_embedding_field": "plot_embedding", "mongodb_text_field": "plot"}, - ) - - assert _stage(collection, "$vectorSearch")["path"] == "plot_embedding" - assert _stage(collection, "$project") == {"plot": 1, "score": {"$meta": "vectorSearchScore"}} - - -def test_num_candidates_scales_with_the_requested_limit(): - config, _, collection = _config() - - _search(config, optional_params={"max_num_results": 40}) - - assert _stage(collection, "$vectorSearch")["numCandidates"] == 400 - - -def test_num_candidates_can_be_overridden(): - config, _, collection = _config() - - _search(config, optional_params={"max_num_results": 5}, litellm_params={"mongodb_num_candidates": 250}) - - assert _stage(collection, "$vectorSearch")["numCandidates"] == 250 - - -@pytest.mark.parametrize("configured", [4, 10_001]) -def test_num_candidates_below_the_limit_or_above_the_ceiling_is_rejected(configured): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_num_candidates"): - _search(config, optional_params={"max_num_results": 5}, litellm_params={"mongodb_num_candidates": configured}) - - -def test_list_query_is_joined_into_one_embedding_input(): - config, _, _ = _config() - - _search(config, query=["deep", "space", "rescue"]) - - assert config.embedding_executor.captured.query == "deep space rescue" - - -def test_embedding_config_is_expanded_into_the_embedding_call(): - config, _, _ = _config() - - _search(config, litellm_params={"litellm_embedding_config": {"api_base": "https://example.test", "timeout": 7}}) - - captured = config.embedding_executor.captured - assert captured.configuration == {"api_base": "https://example.test", "timeout": 7} - assert captured.model == "openai/text-embedding-ada-002" - - -def test_response_maps_documents_to_openai_shaped_results(): - documents = [ - {"_id": "abc123", "text": "an astronaut adrift", "score": 0.94}, - {"_id": "def456", "text": "a robot dog", "score": 0.81}, - ] - config, _, _ = _config(documents=documents) - - response = _search(config) - - assert response["object"] == "vector_store.search_results.page" - assert response["search_query"] == "a lone astronaut" - assert [result["score"] for result in response["data"]] == [0.94, 0.81] - assert [result["content"][0]["text"] for result in response["data"]] == ["an astronaut adrift", "a robot dog"] - assert [result["file_id"] for result in response["data"]] == ["abc123", "def456"] - assert [result["filename"] for result in response["data"]] == ["abc123", "def456"] - assert response["data"][0]["content"][0]["type"] == "text" - - -def test_response_reads_a_dotted_text_field_path(): - config, _, _ = _config(documents=[{"_id": 1, "metadata": {"body": "nested text"}, "score": 0.5}]) - - response = _search(config, litellm_params={"mongodb_text_field": "metadata.body"}) - - assert response["data"][0]["content"][0]["text"] == "nested text" - - -def test_a_dotted_path_resolves_three_levels_deep(): - config, _, _ = _config(documents=[{"_id": 1, "a": {"b": {"c": "deep text"}}, "score": 0.5}]) - - response = _search(config, litellm_params={"mongodb_text_field": "a.b.c"}) - - assert response["data"][0]["content"][0]["text"] == "deep text" - - -def test_a_dotted_path_that_runs_through_a_scalar_counts_as_absent(): - """Walking 'plot.nope' when plot is a string must report the misconfiguration, not - stringify the scalar and hand the model text from the wrong field.""" - config, _, _ = _config(documents=[{"_id": 1, "plot": "a plain string", "score": 0.5}]) - - with pytest.raises(BadRequestError, match=r"has a 'plot\.nope' field"): - _search(config, litellm_params={"mongodb_text_field": "plot.nope"}) - - -def test_a_non_string_text_field_is_stringified(): - config, _, _ = _config(documents=[{"_id": 1, "year": 1979, "score": 0.5}]) - - response = _search(config, litellm_params={"mongodb_text_field": "year"}) - - assert response["data"][0]["content"][0]["text"] == "1979" - - -def test_a_null_text_field_counts_as_absent(): - config, _, _ = _config(documents=[{"_id": 1, "text": None, "score": 0.5}]) - - with pytest.raises(BadRequestError, match="has a 'text' field"): - _search(config) - - -def test_response_tolerates_a_sparse_document_missing_the_text_field(): - config, _, _ = _config(documents=[{"_id": 1, "score": 0.5}, {"_id": 2, "text": "has text", "score": 0.4}]) - - response = _search(config) - - assert response["data"][0]["content"][0]["text"] == "" - assert response["data"][1]["content"][0]["text"] == "has text" - - -def test_a_present_but_empty_text_field_is_not_treated_as_a_misconfiguration(): - config, _, _ = _config(documents=[{"_id": 1, "text": "", "score": 0.5}]) - - response = _search(config) - - assert response["data"][0]["content"][0]["text"] == "" - - -def test_matches_that_all_lack_the_text_field_name_the_setting_to_fix(): - """Atlas matches on the vector, so a mistyped mongodb_text_field returns confidently - scored results whose content is empty and hands the model an empty context.""" - config, _, _ = _config(documents=[{"_id": 1, "score": 0.9}, {"_id": 2, "score": 0.8}]) - - with pytest.raises(BadRequestError, match="mongodb_text_field"): - _search(config) - - -def test_response_tolerates_a_document_missing_a_score(): - config, _, _ = _config(documents=[{"_id": 1, "text": "no score"}]) - - response = _search(config) - - assert response["data"][0]["score"] is None - - -def test_response_stringifies_a_non_string_document_id(): - config, _, _ = _config(documents=[{"_id": 12345, "text": "numeric id", "score": 0.5}]) - - response = _search(config) - - assert response["data"][0]["file_id"] == "12345" - - -def test_search_requires_an_embedding_model(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="litellm_embedding_model is required"): - config.execute_search_vector_store_request( - vector_store_id=INDEX, - query="q", - vector_store_search_optional_params={}, + kwargs: Final = { + "vector_store_id": "exact index", + "query": ["travel", "policy"], + "vector_store_search_optional_params": {"max_num_results": limit}, + "api_base": BASE_PARAMS["api_base"], + "litellm_logging_obj": MagicMock(), + "litellm_params": params, + } + if asynchronous: + url, body = await config.atransform_search_vector_store_request(**kwargs) + else: + url, body = config.transform_search_vector_store_request(**kwargs) + assert url == "https://sidecar.example/prefix/v1/vector_stores/exact%20index/search" + assert body == { + "query": "travel policy", + "query_vector": (0.1, 0.2, 0.3), + "mongodb_database": "policies", + "mongodb_collection": "documents", + "mongodb_text_field": "metadata.body", + "mongodb_embedding_field": "stored_vector", + "mongodb_num_candidates": candidates, + "max_num_results": limit or 10, + "timeout_ms": 750, + } + executor.call.assert_called_once_with("embedding-alias", "travel policy", {"dimensions": 3}) + assert config.transform_search_vector_store_response(httpx.Response(200, json=RESULT), MagicMock()) == RESULT + + +@pytest.mark.parametrize( + "query,overrides,options", + [ + ("", {}, {}), + (" ", {}, {}), + ("x" * 32_001, {}, {}), + ("travel", {"litellm_embedding_model": None}, {}), + ("travel", {"mongodb_database": None}, {}), + ("travel", {"mongodb_collection": None}, {}), + ("travel", {"mongodb_connection_string": "mongodb://obsolete-secret"}, {}), + ("travel", {"mongodb_filter": {"private": True}}, {}), + ("travel", {"mongodb_num_candidates": 9}, {}), + ("travel", {"mongodb_num_candidates": 10_001}, {}), + ("travel", {}, {"max_num_results": 0}), + ("travel", {}, {"max_num_results": 51}), + ("travel", {}, {"filters": {}}), + ("travel", {}, {"ranking_options": {}}), + ("travel", {}, {"rewrite_query": False}), + ], +) +def test_invalid_search_is_rejected_before_embedding( + query: str, overrides: Mapping[str, object], options: VectorStoreSearchOptionalRequestParams +) -> None: + executor: Final = RecordingEmbeddingExecutor() + config: Final = MongoDBVectorStoreConfig(executor) + with pytest.raises(litellm.BadRequestError) as error: + config.transform_search_vector_store_request( + vector_store_id="policy_index", + query=query, + vector_store_search_optional_params=options, + api_base=BASE_PARAMS["api_base"], litellm_logging_obj=MagicMock(), - litellm_params={k: v for k, v in BASE_PARAMS.items() if k != "litellm_embedding_model"}, + litellm_params={**BASE_PARAMS, **overrides}, ) + assert "obsolete-secret" not in str(error.value) + executor.call.assert_not_called() -def test_missing_embedding_model_message_names_the_field_being_searched(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match=r"embedded_movies\.embedding"): - config.execute_search_vector_store_request( - vector_store_id=INDEX, - query="q", - vector_store_search_optional_params={}, - litellm_logging_obj=MagicMock(), - litellm_params={k: v for k, v in BASE_PARAMS.items() if k != "litellm_embedding_model"}, - ) - - -def test_search_requires_a_connection_string(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_connection_string is required"): - _search(config, litellm_params={"mongodb_connection_string": None}) - - -@pytest.mark.parametrize("connection_string", ["postgres://host/db", "https://cluster.mongodb.net", "redis://host"]) -def test_search_rejects_a_non_mongodb_connection_scheme(connection_string): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="must start with 'mongodb://' or 'mongodb\\+srv://'"): - _search(config, litellm_params={"mongodb_connection_string": connection_string}) - - -def test_search_accepts_the_plain_mongodb_scheme(): - config, _, collection = _config() - - _search(config, litellm_params={"mongodb_connection_string": "mongodb://localhost:27017"}) - - assert collection.pipeline is not None - - -def test_search_requires_a_database(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_database is required"): - _search(config, litellm_params={"mongodb_database": None}) - - -def test_search_requires_a_collection(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_collection is required"): - _search(config, litellm_params={"mongodb_collection": None}) - - -def test_search_rejects_filters_rather_than_silently_ignoring_them(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="does not support the filters parameter"): - _search(config, optional_params={"filters": {"genre": "sci-fi"}}) - - +@pytest.mark.parametrize( + "status,body,error_type", + [ + (400, {"error": {"message": "Index is not queryable"}}, litellm.BadRequestError), + (401, {}, litellm.AuthenticationError), + (408, {}, litellm.Timeout), + (503, {}, litellm.ServiceUnavailableError), + (200, {}, litellm.ServiceUnavailableError), + (200, {**RESULT, "data": [{"score": "wrong"}]}, litellm.ServiceUnavailableError), + (0, {}, litellm.Timeout), + (-1, {}, litellm.BadRequestError), + (-2, {"api_base": "http://sidecar.example"}, litellm.BadRequestError), + (-2, {"api_base": "http://10.0.0.10:8080"}, litellm.BadRequestError), + (-2, {"api_base": "http://localhost:8080"}, litellm.BadRequestError), + (200, RESULT, None), + ], +) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("timeout", [0.75, 120.0]) +@pytest.mark.parametrize("api_base", ["https://sidecar.example/prefix", "http://127.0.0.1:8080", "http://[::1]:8080"]) @pytest.mark.asyncio -async def test_async_search_rejects_filters_rather_than_silently_ignoring_them(): - config, _, _ = _async_config() - - with pytest.raises(BadRequestError, match="does not support the filters parameter"): - await _asearch(config, optional_params={"filters": {"genre": "sci-fi"}}) - - -def test_search_rejects_ranking_options_rather_than_silently_ignoring_them(): - """A score_threshold that is quietly dropped is worse than an error: the caller asked for - results above 0.9, gets results scoring 0.5, and nothing says the threshold never ran.""" - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="does not support the ranking_options parameter"): - _search(config, optional_params={"ranking_options": {"score_threshold": 0.9}}) - - -def test_search_rejects_rewrite_query_rather_than_silently_ignoring_it(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="does not support the rewrite_query parameter"): - _search(config, optional_params={"rewrite_query": True}) - - -@pytest.mark.asyncio -async def test_async_search_rejects_ranking_options_rather_than_silently_ignoring_them(): - config, _, _ = _async_config() - - with pytest.raises(BadRequestError, match="does not support the ranking_options parameter"): - await _asearch(config, optional_params={"ranking_options": {"score_threshold": 0.9}}) - - -@pytest.mark.parametrize("query", ["", " ", "\n\t", []]) -def test_search_rejects_an_empty_query(query): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="query must not be empty"): - _search(config, query=query) - - -def test_search_rejects_an_oversized_query(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="at most 32000 characters"): - _search(config, query="x" * 32_001) - - -def test_search_accepts_a_query_at_the_size_ceiling(): - config, _, collection = _config() - - _search(config, query="x" * 32_000) - - assert collection.pipeline is not None - - -@pytest.mark.parametrize("max_num_results", [0, -1, 51, 1000]) -def test_search_rejects_out_of_range_max_num_results(max_num_results): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="max_num_results must be between 1 and 50"): - _search(config, optional_params={"max_num_results": max_num_results}) - - -@pytest.mark.parametrize("max_num_results", [1, 50]) -def test_search_allows_max_num_results_at_the_bounds(max_num_results): - config, _, collection = _config() - - _search(config, optional_params={"max_num_results": max_num_results}) - - assert _stage(collection, "$vectorSearch")["limit"] == max_num_results - - -def test_search_treats_an_explicit_null_max_num_results_as_the_default(): - config, _, collection = _config() - - _search(config, optional_params={"max_num_results": None}) - - assert _stage(collection, "$vectorSearch")["limit"] == 10 - - -def test_search_fails_when_the_embedding_model_returns_nothing(): - config, _, _ = _config(embedding=None) - - with pytest.raises(BadRequestError, match="returned no embedding"): - _search(config) - - -def test_validation_runs_before_any_connection_is_opened(): - opened = [] - config = MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1]), - sync_client_factory=lambda key: opened.append(key) or FakeClient(FakeCollection([])), - ) - - with pytest.raises(BadRequestError, match="query must not be empty"): - _search(config, query="") - - assert opened == [] - - -def test_create_vector_store_is_not_supported_and_says_why(): - """litellm.exception_type only passes its own exception types through untouched, so a - NotImplementedError here reaches the caller as APIConnectionError, which the proxy serves - as a 500 with a traceback. Refusing an unsupported operation is a client error.""" - config = MongoDBVectorStoreConfig() - - with pytest.raises(BadRequestError, match="search-only"): - config.transform_create_vector_store_request({}, "https://example.test") - - with pytest.raises(BadRequestError, match="search-only"): - config.transform_create_vector_store_response(httpx.Response(200)) - - -def test_the_create_refusal_survives_the_public_sdk_error_wrapper(): - import litellm - - with pytest.raises(BadRequestError) as raised: - litellm.vector_stores.create(custom_llm_provider="mongodb", name="anything") - - assert "search-only" in str(raised.value) - - -def test_provider_config_manager_returns_the_mongodb_config(): - config = ProviderConfigManager.get_provider_vector_stores_config(LlmProviders.MONGODB) - - assert isinstance(config, MongoDBVectorStoreConfig) - - -@pytest.mark.asyncio -async def test_async_search_builds_the_same_pipeline_and_maps_the_response(): - documents = [{"_id": "abc123", "text": "an astronaut adrift", "score": 0.94}] - config, client, collection = _async_config(documents=documents) - - response = await _asearch(config, optional_params={"max_num_results": 3}) - - assert client.requested_database == "sample_mflix" - assert client.database.requested_collection == "embedded_movies" - assert _stage(collection, "$vectorSearch")["limit"] == 3 - assert _stage(collection, "$vectorSearch")["queryVector"] == (0.1, 0.2, 0.3) - assert response["data"][0]["content"][0]["text"] == "an astronaut adrift" - assert response["data"][0]["score"] == 0.94 - - -@pytest.mark.asyncio -async def test_async_search_requires_an_embedding_model(): - config, _, _ = _async_config() - - with pytest.raises(BadRequestError, match="litellm_embedding_model is required"): - await config.aexecute_search_vector_store_request( - vector_store_id=INDEX, - query="q", - vector_store_search_optional_params={}, - litellm_logging_obj=MagicMock(), - litellm_params={k: v for k, v in BASE_PARAMS.items() if k != "litellm_embedding_model"}, - ) - - -class TestClientCache: - def setup_method(self): - reset_client_cache() - - def teardown_method(self): - reset_client_cache() - - def _key(self, connection_string=CONNECTION_STRING, socket_timeout_ms=30_000): - return MongoClientKey( - connection_string=connection_string, - connect_timeout_ms=10_000, - socket_timeout_ms=socket_timeout_ms, - server_selection_timeout_ms=10_000, - ) - - def test_the_same_connection_reuses_one_client(self): - first = get_sync_client(self._key(), RecordingClient) - second = get_sync_client(self._key(), RecordingClient) - - assert first is second - assert first.connection_string == CONNECTION_STRING - assert first.kwargs["socketTimeoutMS"] == 30_000 - assert first.kwargs["connectTimeoutMS"] == 10_000 - assert first.kwargs["appname"] == "litellm" - - def test_a_different_connection_gets_its_own_client(self): - first = get_sync_client(self._key(), RecordingClient) - second = get_sync_client(self._key(connection_string="mongodb://other.example.test"), RecordingClient) - - assert first is not second - assert second.connection_string == "mongodb://other.example.test" - - def test_a_different_timeout_gets_its_own_client(self): - first = get_sync_client(self._key(), RecordingClient) - second = get_sync_client(self._key(socket_timeout_ms=5_000), RecordingClient) - - assert first is not second - assert second.kwargs["socketTimeoutMS"] == 5_000 - - @pytest.mark.asyncio - async def test_async_clients_are_cached_per_event_loop(self): - first = get_async_client(self._key(), RecordingClient) - second = get_async_client(self._key(), RecordingClient) - - assert first is second - assert first.connection_string == CONNECTION_STRING - - - def _fill_cache(self): - for slot in range(_MAX_CACHED_CLIENTS): - get_sync_client(self._key(f"mongodb://cold-{slot}:27017"), RecordingClient) - - def test_a_store_added_after_the_cache_filled_is_still_cached(self): - """Rebuilding a client costs an SRV lookup, a TLS handshake and topology discovery, so a - store that misses the cache on every single search pays that on every search.""" - self._fill_cache() - latecomer = self._key("mongodb://latecomer:27017") - - first = get_sync_client(latecomer, RecordingClient) - - assert get_sync_client(latecomer, RecordingClient) is first - - def test_the_cache_evicts_the_least_recently_used_client(self): - self._fill_cache() - oldest = self._key("mongodb://cold-0:27017") - newest = self._key(f"mongodb://cold-{_MAX_CACHED_CLIENTS - 1}:27017") - kept = get_sync_client(newest, RecordingClient) - - get_sync_client(self._key("mongodb://latecomer:27017"), RecordingClient) - - assert get_sync_client(newest, RecordingClient) is kept - assert oldest not in _sync_clients - - def test_concurrent_searches_never_trip_over_an_eviction(self): - """Async searches run the sync client through executor threads, so a key can be evicted - between the lookup and the reordering that follows it.""" - errors = [] - churn = _MAX_CACHED_CLIENTS + 2 - - def hammer(offset): - try: - for step in range(3_000): - get_sync_client(self._key(f"mongodb://h-{(step + offset) % churn}:27017"), RecordingClient) - except Exception as e: - errors.append(repr(e)) - - previous = sys.getswitchinterval() - sys.setswitchinterval(1e-9) - try: - threads = [threading.Thread(target=hammer, args=(offset,)) for offset in range(16)] - for thread in threads: - thread.start() - for thread in threads: - thread.join() - finally: - sys.setswitchinterval(previous) - - assert errors == [] - - def test_the_cache_never_grows_past_its_cap(self): - for slot in range(_MAX_CACHED_CLIENTS * 3): - get_sync_client(self._key(f"mongodb://host-{slot}:27017"), RecordingClient) - - assert len(_sync_clients) == _MAX_CACHED_CLIENTS - - def test_a_new_loop_never_inherits_a_closed_loop_client(self): - """CPython recycles id() so aggressively that a fresh event loop almost always lands on - the id of one already collected: measured at 37 of 40 rounds. Keying the cache on the id - alone therefore hands the new loop an AsyncMongoClient bound to a closed loop, and every - operation on it raises "Event loop is closed".""" - - class LoopAgnosticClient: - """Holds no reference to the loop, unlike pymongo's, whose own reference happens to - keep ids from being recycled and hides the bug until the cache fills.""" - - def __init__(self, *args, **kwargs): - self.built_on = None - - key = self._key() - clients_handed_out = [] - - async def fetch(): - return get_async_client(key, LoopAgnosticClient) - - for _ in range(20): - loop = asyncio.new_event_loop() - client = loop.run_until_complete(fetch()) - clients_handed_out.append((client, client.built_on, loop.is_closed())) - client.built_on = weakref.ref(loop) - loop.close() - del loop - gc.collect() - - stale = [ - handed_out - for client, built_on, _ in clients_handed_out - if built_on is not None and (built_on() is None or built_on().is_closed()) - for handed_out in (client,) - ] - assert stale == [], f"{len(stale)} of 20 loops were handed a client built on a closed loop" - - def test_the_cache_releases_clients_built_on_closed_loops(self): - """pymongo's AsyncMongoClient keeps a reference to the loop it was built on, so an entry - for a closed loop holds that client, and its sockets, for the life of the process. A - script calling asyncio.run per search fills the cache to its cap that way: measured live - against Atlas at 32 pinned clients and 212 open descriptors after 40 loops.""" - - class LoopHoldingClient: - def __init__(self, *args, **kwargs): - self.loop = asyncio.get_running_loop() - - key = self._key() - - async def fetch(): - return get_async_client(key, LoopHoldingClient) - - for _ in range(_MAX_CACHED_CLIENTS + 8): - loop = asyncio.new_event_loop() - loop.run_until_complete(fetch()) - loop.close() - - assert len(_async_clients) == 1, f"{len(_async_clients)} closed-loop clients are still cached" - - -class TestClientKeyDerivation: - def test_no_timeout_uses_the_bounded_defaults(self): - key = MongoDBVectorStoreConfig._client_key(_MongoDBSearchParams.model_validate(BASE_PARAMS), None) - - assert key.connect_timeout_ms == 10_000 - assert key.socket_timeout_ms == 30_000 - assert key.server_selection_timeout_ms == 10_000 - - def test_a_numeric_timeout_bounds_the_connect_phase(self): - key = MongoDBVectorStoreConfig._client_key(_MongoDBSearchParams.model_validate(BASE_PARAMS), 3.0) - - assert key.socket_timeout_ms == 3_000 - assert key.connect_timeout_ms == 3_000 - - def test_a_short_timeout_also_shortens_server_selection(self): - """Server selection runs before the connect attempt, so leaving it at the 10s default - would let a caller asking for a 3s budget block for 10s before anything is tried.""" - key = MongoDBVectorStoreConfig._client_key(_MongoDBSearchParams.model_validate(BASE_PARAMS), 3.0) - - assert key.server_selection_timeout_ms == 3_000 - - def test_a_generous_timeout_does_not_raise_server_selection_above_the_default(self): - key = MongoDBVectorStoreConfig._client_key(_MongoDBSearchParams.model_validate(BASE_PARAMS), 120.0) - - assert key.socket_timeout_ms == 120_000 - assert key.server_selection_timeout_ms == 10_000 - - def test_an_httpx_timeout_maps_connect_and_read_separately(self): - key = MongoDBVectorStoreConfig._client_key( - _MongoDBSearchParams.model_validate(BASE_PARAMS), httpx.Timeout(connect=2.0, read=45.0, write=5.0, pool=5.0) - ) - - assert key.connect_timeout_ms == 2_000 - assert key.socket_timeout_ms == 45_000 - - -class TestErrorTranslation: - def _translate(self, error): - return translate_mongo_error(error, index_name=INDEX, database="sample_mflix", collection="embedded_movies") - - def test_server_selection_timeout_points_at_the_atlas_access_list(self): - from pymongo.errors import ServerSelectionTimeoutError - - translated = self._translate(ServerSelectionTimeoutError("no servers")) - - assert "IP access list" in str(translated) - assert "paused cluster" in str(translated) - - def test_authentication_failure_points_at_the_connection_string_credentials(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("auth failed", code=18)) - - assert "rejected the credentials" in str(translated) - - def test_a_dropped_connection_stays_retryable(self): - """A replica set failover reaches the driver as AutoReconnect. litellm only retries 408, - 409, 429 and 5xx, so classifying it as a client error would turn one failover into a - permanently failed search.""" - from pymongo.errors import AutoReconnect - - translated = self._translate(AutoReconnect("connection closed")) - - assert litellm._should_retry(translated.status_code) - assert "dropped or refused" in str(translated) - - def test_a_dropped_connection_still_names_the_misconfigurations_behind_it(self): - """Atlas answers a URI with no credentials by closing the connection rather than failing - auth, so the retryable message still has to name that.""" - from pymongo.errors import AutoReconnect - - translated = self._translate(AutoReconnect("connection closed")) - - assert "no username and password" in str(translated) - assert "mongod is listening" in str(translated) - - def test_the_retryable_classification_survives_the_public_sdk_error_wrapper(self): - """litellm.exception_type only passes its own exception types through; anything else becomes - an APIConnectionError and a 500, which would drop the retryable classification.""" - from pymongo.errors import AutoReconnect - - translated = self._translate(AutoReconnect("connection closed")) - - wrapped = litellm.exception_type( - model=None, - original_exception=translated, - custom_llm_provider="mongodb", - completion_kwargs={}, - extra_kwargs={}, - ) - - assert isinstance(wrapped, ServiceUnavailableError) - assert litellm._should_retry(wrapped.status_code) - - def test_a_pool_wait_queue_timeout_stays_retryable(self): - from pymongo.errors import WaitQueueTimeoutError - - translated = self._translate(WaitQueueTimeoutError("timed out waiting for a connection")) - - assert litellm._should_retry(translated.status_code) - - def test_server_selection_timeout_still_wins_over_the_connection_branch(self): - from pymongo.errors import ServerSelectionTimeoutError - - translated = self._translate(ServerSelectionTimeoutError("no servers")) - - assert isinstance(translated, Timeout) - assert "dropped or refused" not in str(translated) - - def test_network_timeout_still_wins_over_the_connection_branch(self): - from pymongo.errors import NetworkTimeout - - translated = self._translate(NetworkTimeout("socket timed out")) - - assert isinstance(translated, Timeout) - assert "dropped or refused" not in str(translated) - - def test_an_unescaped_password_character_is_a_400_not_a_500(self): - """pymongo's URI parser raises a plain ValueError, not a PyMongoError, for an unusable port, - which is also what an unescaped ':' in a password produces. It must not be a 500.""" - translated = self._translate(ValueError("Port contains non-digit characters")) - - assert isinstance(translated, BadRequestError) - assert "percent-encoded" in str(translated) - - def test_unauthorized_points_at_the_database_user_permissions(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("not authorized", code=13)) - - assert "sample_mflix.embedded_movies" in str(translated) - - def test_code_13_alone_is_enough_without_a_recognisable_message(self): - """The other unauthorized case carries "not authorized", which the message markers also - match, so it cannot tell whether the code is still being checked at all.""" - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("user lacks privileges on this namespace", code=13)) - - assert "rejected the credentials" in str(translated) - assert "sample_mflix.embedded_movies" in str(translated) - - def test_a_missing_index_names_the_index_and_the_collection(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("Index not found for name movies_vector_index", code=27)) - - assert INDEX in str(translated) - assert "READY" in str(translated) - - def test_a_dimension_mismatch_points_at_the_embedding_model(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("queryVector has 1536 dimensions, index expects 2048")) - - assert "litellm_embedding_model must be the same model" in str(translated) - - def test_an_unrecognised_operation_failure_still_names_the_target(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("something else entirely")) - - assert "sample_mflix.embedded_movies" in str(translated) - assert INDEX in str(translated) - - def test_a_configuration_error_points_at_the_connection_string(self): - from pymongo.errors import ConfigurationError - - translated = self._translate(ConfigurationError("bad uri")) - - assert "not a usable MongoDB connection string" in str(translated) - - def test_a_non_driver_error_is_returned_unchanged(self): - original = RuntimeError("unrelated") - - assert self._translate(original) is original - - def test_search_surfaces_a_translated_driver_error(self): - from pymongo.errors import ServerSelectionTimeoutError - - config, _, _ = _config(error=ServerSelectionTimeoutError("no servers")) - - with pytest.raises(Timeout, match="IP access list"): - _search(config) - - @pytest.mark.asyncio - async def test_async_search_surfaces_a_translated_driver_error(self): - from pymongo.errors import OperationFailure - - config, _, _ = _async_config(error=OperationFailure("auth failed", code=18)) - - with pytest.raises(BadRequestError, match="rejected the credentials"): - await _asearch(config) - - -class TestMissingDriver: - def test_the_sync_import_names_the_extra_to_install(self): - from litellm.llms.mongodb.common_utils import import_sync_mongo_client - - with patch.dict(sys.modules, {"pymongo": None}): - with pytest.raises(BadRequestError, match=r"pip install litellm\[mongodb\]"): - import_sync_mongo_client() - - def test_the_async_import_names_the_extra_to_install(self): - from litellm.llms.mongodb.common_utils import import_async_mongo_client - - with patch.dict(sys.modules, {"pymongo": None}): - with pytest.raises(BadRequestError, match=r"pip install litellm\[mongodb\]"): - import_async_mongo_client() - - def test_error_translation_degrades_gracefully_without_the_driver(self): - original = RuntimeError("boom") - - with patch.dict(sys.modules, {"pymongo.errors": None}): - assert translate_mongo_error(original, INDEX, "db", "col") is original - - -class TestEmptyResultsAreDisambiguated: - """$vectorSearch returns zero documents for a missing database, collection or index just as it - does for a query that matched nothing, so an empty result set is checked against the index - catalogue before it is reported as 'no matches'.""" - - def test_a_missing_index_becomes_an_error_rather_than_an_empty_page(self): - config, _, collection = _config(documents=[], search_indexes=[]) - - with pytest.raises(BadRequestError, match="No queryable MongoDB Vector Search index"): - _search(config) - - assert collection.listed_indexes == [INDEX] - - def test_the_missing_index_error_explains_why_mongodb_reported_no_results(self): - config, _, _ = _config(documents=[], search_indexes=[]) - - with pytest.raises(BadRequestError, match="returns no results rather than an error"): - _search(config) - - def test_an_index_still_building_becomes_an_error_naming_its_status(self): - config, _, _ = _config( - documents=[], search_indexes=[{"name": INDEX, "status": "PENDING", "queryable": False}] - ) - - with pytest.raises(BadRequestError, match="not queryable yet; its status is PENDING"): - _search(config) - - def test_a_genuine_no_match_against_a_ready_index_returns_an_empty_page(self): - config, _, collection = _config(documents=[]) - - response = _search(config) - - assert response["data"] == [] - assert response["object"] == "vector_store.search_results.page" - assert collection.listed_indexes == [INDEX] - - def test_the_catalogue_is_not_consulted_when_the_search_returned_hits(self): - config, _, collection = _config(documents=[{"_id": 1, "text": "hit", "score": 0.9}]) - - _search(config) - - assert collection.listed_indexes == [] - - @pytest.mark.asyncio - async def test_async_missing_index_becomes_an_error_rather_than_an_empty_page(self): - config, _, collection = _async_config(documents=[], search_indexes=[]) - - with pytest.raises(BadRequestError, match="No queryable MongoDB Vector Search index"): - await _asearch(config) - - assert collection.listed_indexes == [INDEX] - - @pytest.mark.asyncio - async def test_async_index_still_building_becomes_an_error_naming_its_status(self): - config, _, _ = _async_config( - documents=[], search_indexes=[{"name": INDEX, "status": "PENDING", "queryable": False}] - ) - - with pytest.raises(BadRequestError, match="not queryable yet; its status is PENDING"): - await _asearch(config) - - @pytest.mark.asyncio - async def test_async_genuine_no_match_returns_an_empty_page(self): - config, _, _ = _async_config(documents=[]) - - response = await _asearch(config) - - assert response["data"] == [] - - @pytest.mark.asyncio - async def test_async_catalogue_is_not_consulted_when_the_search_returned_hits(self): - config, _, collection = _async_config(documents=[{"_id": 1, "text": "hit", "score": 0.9}]) - - await _asearch(config) - - assert collection.listed_indexes == [] - - def test_a_failure_while_checking_the_catalogue_is_translated_too(self): - from pymongo.errors import OperationFailure - - class ExplodingCollection(FakeCollection): - def list_search_indexes(self, name): - raise OperationFailure("not authorized", code=13) - - collection = ExplodingCollection([], None, []) - config = MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1]), - sync_client_factory=lambda key: FakeClient(collection), - ) - - with pytest.raises(BadRequestError, match="lacks read access"): - _search(config) - - -class TestAtlasPlanExecutorErrors: - """Atlas reports a wrong vector path and a dimension mismatch through the same error code, so - each one has to be told apart by its message or both come back as a generic index failure.""" - - def _translate(self, message): - from pymongo.errors import OperationFailure - - return translate_mongo_error( - OperationFailure(message, code=8), - index_name=INDEX, - database="sample_mflix", - collection="embedded_movies", - ) - - def test_a_wrong_vector_path_points_at_the_embedding_field_setting(self): - translated = self._translate( - "PlanExecutor error during aggregation :: caused by :: nope is not indexed as vector" - ) - - assert "mongodb_embedding_field names a field" in str(translated) - - def test_a_dimension_mismatch_is_not_reported_as_a_wrong_path(self): - translated = self._translate( - "PlanExecutor error during aggregation :: caused by :: vector field is indexed with " - "1536 dimensions but queried with 3072" - ) - - assert "does not match the vector dimensions" in str(translated) - assert "mongodb_embedding_field" not in str(translated) - - -class TestErrorsCarryTheRightHttpStatus: - """litellm.exception_type passes a litellm exception through untouched but wraps anything - else into APIConnectionError, which the proxy serves as a 500 with a Python traceback in the - body. A misconfigured connection string is the caller's to fix, so it has to arrive as a 400. - """ - - @pytest.mark.parametrize( - "invoke", - [ - pytest.param(lambda: _search(_config()[0], query=" "), id="empty-query"), - pytest.param( - lambda: _search(_config()[0], optional_params={"max_num_results": 999}), - id="max-num-results-out-of-range", - ), - pytest.param( - lambda: _search(_config()[0], optional_params={"filters": {"genre": "Action"}}), - id="unsupported-filters", - ), - pytest.param( - lambda: _search(_config()[0], litellm_params={"mongodb_connection_string": "postgres://host/db"}), - id="wrong-uri-scheme", - ), - pytest.param( - lambda: _search(_config()[0], litellm_params={"mongodb_database": None}), id="missing-database" - ), - pytest.param( - lambda: _search(_config()[0], litellm_params={"litellm_embedding_model": None}), - id="missing-embedding-model", - ), - ], - ) - def test_configuration_failures_are_400(self, invoke): - with pytest.raises(BadRequestError) as excinfo: - invoke() - assert excinfo.value.status_code == 400 - assert excinfo.value.llm_provider == "mongodb" - - def test_missing_index_is_400(self): - error = missing_index_error("idx", "db", "coll") - assert error.status_code == 400 - assert error.llm_provider == "mongodb" - - def test_index_still_building_is_400(self): - error = index_not_ready_error("idx", "db", "coll", "PENDING") - assert error.status_code == 400 - - def test_unreachable_deployment_is_a_timeout_not_a_bad_request(self): - from pymongo.errors import ServerSelectionTimeoutError - - translated = translate_mongo_error( - ServerSelectionTimeoutError("no servers"), index_name="idx", database="db", collection="coll" - ) - assert isinstance(translated, Timeout) - assert translated.status_code == 408 - - def test_query_execution_timeout_is_a_timeout(self): - from pymongo.errors import ExecutionTimeout - - translated = translate_mongo_error( - ExecutionTimeout("too slow"), index_name="idx", database="db", collection="coll" - ) - assert isinstance(translated, Timeout) - assert translated.status_code == 408 - - def test_unrecognised_errors_are_not_relabelled_as_bad_requests(self): - original = RuntimeError("something else entirely") - assert ( - translate_mongo_error(original, index_name="idx", database="db", collection="coll") - is original - ) - - -def test_atlas_rejected_credentials_are_named_even_though_the_code_is_8000(): - """Atlas answers a wrong password with code 8000 "AtlasError", not the 18 that a - self-hosted deployment returns, so a code-only check reports it as a generic - rejected search and never tells the caller to look at their connection string.""" - from pymongo.errors import OperationFailure - - error = OperationFailure( - "bad auth : authentication failed", - code=8000, - details={"ok": 0, "errmsg": "bad auth : authentication failed", "code": 8000, "codeName": "AtlasError"}, - ) - translated = translate_mongo_error(error, index_name="idx", database="sample_mflix", collection="embedded_movies") - - assert isinstance(translated, BadRequestError) - assert "mongodb_connection_string" in str(translated) - assert "sample_mflix.embedded_movies" in str(translated) - - -def test_a_rejected_search_that_is_not_an_auth_failure_keeps_the_generic_message(): - from pymongo.errors import OperationFailure - - error = OperationFailure("PlanExecutor error", code=8, details={"errmsg": "PlanExecutor error"}) - translated = translate_mongo_error(error, index_name="idx", database="db", collection="coll") - - assert "mongodb_connection_string" not in str(translated) - - -class TestUnrecognisedParameters: - """litellm_params carries plenty of keys this provider does not own, so the params model has - to ignore extras. That turns a mistyped mongodb_collection into 'mongodb_collection is - required', pointing the reader at a key they can see they have set.""" - - def test_a_mistyped_parameter_is_named(self): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_collectoin"): - _search(config, litellm_params={"mongodb_collectoin": "embedded_movies"}) - - def test_the_supported_names_are_listed(self): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_connection_string"): - _search(config, litellm_params={"mongodb_databse": "sample_mflix"}) - - def test_unrelated_litellm_params_are_still_ignored(self): - config, _, _ = _config(documents=[{"_id": 1, "text": "hit", "score": 0.9}]) - - response = _search( - config, - litellm_params={"use_litellm_proxy": False, "use_in_pass_through": False, "vector_store_id": "x"}, - ) - - assert len(response["data"]) == 1 - - @pytest.mark.asyncio - async def test_the_async_path_rejects_them_too(self): - config, _, _ = _async_config() - - with pytest.raises(BadRequestError, match="mongodb_collectoin"): - await _asearch(config, litellm_params={"mongodb_collectoin": "embedded_movies"}) - - -class TestClientConstructionFailures: - """Building the client parses the URI and, for mongodb+srv://, performs a DNS SRV lookup, so it - fails on exactly the inputs a user is most likely to get wrong. Constructing it outside the - translation boundary let those escape as raw pymongo errors, which litellm.exception_type then - wrapped into a 500 with a traceback in the body.""" - - def _config_that_fails_to_connect(self, error): - def factory(_key): - raise error - - return MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1, 0.2, 0.3]), sync_client_factory=factory - ) - - def _async_config_that_fails_to_connect(self, error): - def factory(_key): - raise error - - return MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1, 0.2, 0.3]), async_client_factory=factory - ) - - def test_a_malformed_uri_is_a_bad_request_not_a_500(self): - from pymongo.errors import InvalidURI - - config = self._config_that_fails_to_connect(InvalidURI("Invalid URI scheme")) - - with pytest.raises(BadRequestError, match="not a usable MongoDB connection string"): - _search(config) - - def test_an_unresolvable_cluster_name_says_so(self): - from pymongo.errors import ConfigurationError - - config = self._config_that_fails_to_connect(ConfigurationError("The DNS query name does not exist")) - - with pytest.raises(BadRequestError, match="does not exist in DNS"): - _search(config) - - def test_a_dns_lookup_that_ran_out_of_time_is_a_timeout(self): - from pymongo.errors import ConfigurationError - - config = self._config_that_fails_to_connect( - ConfigurationError("The resolution lifetime expired after 0.291 seconds") - ) - - with pytest.raises(Timeout, match="did not finish in time"): - _search(config) - - @pytest.mark.asyncio - async def test_the_async_path_translates_them_too(self): - from pymongo.errors import InvalidURI - - config = self._async_config_that_fails_to_connect(InvalidURI("Invalid URI scheme")) - - with pytest.raises(BadRequestError, match="not a usable MongoDB connection string"): - await _asearch(config) - - -class TestSelfManagedDeploymentsAreFirstClass: - """mongod serves $vectorSearch identically whether mongot runs under Atlas or beside a - self-managed deployment, so an operator without an Atlas account has to be able to act on - every message. Guidance that only names Atlas remedies sends them looking for an IP access - list and a paused cluster that do not exist in their deployment.""" - - def _config_that_fails_to_connect(self, error): - def factory(_key): - raise error - - return MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1, 0.2, 0.3]), sync_client_factory=factory - ) - - def test_a_plain_mongodb_uri_without_srv_or_credentials_is_accepted(self): - params = _MongoDBSearchParams.model_validate( - {**BASE_PARAMS, "mongodb_connection_string": "mongodb://mongod.internal:27017"} - ) - - assert params.require_connection_string() == "mongodb://mongod.internal:27017" - - def test_an_unreachable_deployment_names_a_self_managed_remedy(self): - from pymongo.errors import ServerSelectionTimeoutError - - config = self._config_that_fails_to_connect(ServerSelectionTimeoutError("connection refused")) - - with pytest.raises(Timeout) as excinfo: - _search(config) - - assert "self-managed" in str(excinfo.value) - assert "host or port" in str(excinfo.value) - - def test_a_refused_connection_names_a_self_managed_remedy(self): - from pymongo.errors import ConnectionFailure - - config = self._config_that_fails_to_connect(ConnectionFailure("connection closed")) - - with pytest.raises(ServiceUnavailableError) as excinfo: - _search(config) - - assert "self-managed" in str(excinfo.value) - assert "mongod is listening" in str(excinfo.value) - - def test_an_unresolvable_hostname_names_a_self_managed_remedy(self): - from pymongo.errors import ConfigurationError - - config = self._config_that_fails_to_connect(ConfigurationError("The DNS query name does not exist")) - - with pytest.raises(BadRequestError) as excinfo: - _search(config) - - assert "self-managed" in str(excinfo.value) - - def test_the_missing_index_message_does_not_claim_atlas(self): - message = str(missing_index_error(INDEX, "sample_mflix", "embedded_movies")) - - assert "MongoDB Vector Search index" in message - assert "Atlas" not in message - - def test_the_not_ready_message_does_not_claim_atlas(self): - message = str(index_not_ready_error(INDEX, "sample_mflix", "embedded_movies", "PENDING")) - - assert "MongoDB Vector Search index" in message - assert "Atlas" not in message - - def test_the_search_only_refusal_does_not_claim_atlas(self): - config = MongoDBVectorStoreConfig() - - with pytest.raises(BadRequestError) as excinfo: - config.transform_create_vector_store_request({}, api_base="") - - assert "Atlas" not in str(excinfo.value) - - def test_a_dimension_mismatch_does_not_claim_atlas(self): - from pymongo.errors import OperationFailure - - error = OperationFailure("vector field is indexed with 128 dimensions but queried with 256") - translated = translate_mongo_error(error, index_name=INDEX, database="db", collection="c") - - assert "Atlas" not in str(translated) - assert "dimensions the index was built for" in str(translated) - - def test_an_uncovered_embedding_field_does_not_claim_atlas(self): - from pymongo.errors import OperationFailure - - error = OperationFailure("embedding is not indexed as vector") - translated = translate_mongo_error(error, index_name=INDEX, database="db", collection="c") - - assert "MongoDB Vector Search index does not cover" in str(translated) - assert "Atlas" not in str(translated) - - def test_a_self_managed_auth_failure_is_still_recognised_by_code_18(self): - from pymongo.errors import OperationFailure - - error = OperationFailure("Authentication failed.", code=18, details={"code": 18}) - translated = translate_mongo_error(error, index_name=INDEX, database="db", collection="c") - - assert isinstance(translated, BadRequestError) - assert "rejected the credentials" in str(translated) - - -class TestUnescapedCredentialsAreDiagnosed: - """Self-managed deployments usually carry a generated password, so '@', '/', ':' and '%' in one - are routine. pymongo reports those as a port, a database name or an RFC 3986 complaint, none of - which points the operator at their password, so each has to be named for what it is. The errors - here come from pymongo's real parser rather than a synthetic stand-in.""" - - @staticmethod - def _real_parse_error(uri): - from pymongo import MongoClient - - try: - MongoClient(uri, serverSelectionTimeoutMS=1) - except Exception as e: - return e - raise AssertionError(f"expected {uri!r} to fail parsing") - - def _translated(self, uri): - return translate_mongo_error( - self._real_parse_error(uri), index_name=INDEX, database="db", collection="c" - ) - - @pytest.mark.parametrize( - "uri", - [ - "mongodb://user:pa@ss@host:27017/", - "mongodb://user:pa:ss@host:27017/", - "mongodb://user:pa%ss@host:27017/", - "mongodb://user@x:pw@host:27017/", - ], - ) - def test_rfc_3986_complaints_tell_the_operator_to_encode_the_password(self, uri): - translated = self._translated(uri) - - assert isinstance(translated, BadRequestError) - assert "percent-encoded per RFC 3986" in str(translated) - - @pytest.mark.parametrize( - "uri", - ["mongodb://user:pa/ss@host:27017/", "mongodb://user/x:pw@host:27017/"], - ) - def test_a_slash_in_the_credentials_is_not_reported_as_a_database_name(self, uri): - translated = self._translated(uri) - - assert isinstance(translated, BadRequestError) - assert "percent-encoded per RFC 3986" in str(translated) - - def test_an_unusable_port_names_the_host_and_port_not_the_database(self): - translated = self._translated("mongodb://host:99999/") - - assert isinstance(translated, BadRequestError) - assert "host and port" in str(translated) - - def test_a_genuinely_bad_database_name_still_mentions_the_uri_path(self): - translated = self._translated("mongodb://host:27017/has space") - - assert isinstance(translated, BadRequestError) - assert "database name in the URI path" in str(translated) - - -class TestUnreadableTlsFilesAreDiagnosed: - """A private CA is how self-managed deployments present TLS, so tlsCAFile and - tlsCertificateKeyFile are on-prem options in practice. pymongo opens those files itself and - lets OSError out, which is not a PyMongoError, so before this they reached the caller as a 500 - with a traceback. The errors here come from pymongo's real TLS setup.""" - - @staticmethod - def _real_tls_error(uri): - from pymongo import MongoClient - - try: - MongoClient(uri, serverSelectionTimeoutMS=1500).admin.command("ping") - except Exception as e: - return e - raise AssertionError(f"expected {uri!r} to fail") - - def _translated(self, uri): - return translate_mongo_error(self._real_tls_error(uri), index_name=INDEX, database="db", collection="c") - - @pytest.mark.parametrize( - "path", - ["/nonexistent-directory-for-tests/ca.pem", "/tmp"], - ) - def test_an_unreadable_ca_file_is_a_400_naming_the_path(self, path): - translated = self._translated(f"mongodb://localhost:27717/?tls=true&tlsCAFile={path}") - - assert isinstance(translated, BadRequestError) - assert path in str(translated) - assert "tlsCAFile" in str(translated) - - def test_an_unreadable_client_certificate_is_a_400_naming_the_path(self): - path = "/nonexistent-directory-for-tests/client.pem" - translated = self._translated(f"mongodb://localhost:27717/?tls=true&tlsCertificateKeyFile={path}") - - assert isinstance(translated, BadRequestError) - assert path in str(translated) - - def test_an_oserror_carrying_no_filename_is_left_for_the_other_branches(self): - translated = translate_mongo_error(OSError("socket hung up"), index_name=INDEX, database="db", collection="c") - - assert not isinstance(translated, BadRequestError) - - -class TestTheCallerSuppliedEmbeddingExecutorIsUsed: - """litellm.vector_stores.search always hands a direct provider an embedding_executor, so the - provider has to accept it and route the query through it rather than its own default.""" - - def test_the_supplied_executor_produces_the_query_vector(self): - config, _, collection = _config(embedding=(0.9, 0.9, 0.9), search_indexes=READY_INDEX) - caller = FakeEmbeddingExecutor([0.4, 0.5, 0.6]) - - config.execute_search_vector_store_request( - vector_store_id=INDEX, - query="a lone astronaut", - vector_store_search_optional_params={}, - litellm_logging_obj=MagicMock(), - litellm_params=BASE_PARAMS, - embedding_executor=caller, - ) - - assert caller.captured.query == "a lone astronaut" - assert _stage(collection, "$vectorSearch")["queryVector"] == (0.4, 0.5, 0.6) - - @pytest.mark.asyncio - async def test_the_supplied_executor_produces_the_query_vector_on_the_async_path(self): - config, _, collection = _async_config(embedding=(0.9, 0.9, 0.9), search_indexes=READY_INDEX) - caller = FakeEmbeddingExecutor([0.4, 0.5, 0.6]) - - await config.aexecute_search_vector_store_request( - vector_store_id=INDEX, - query="a lone astronaut", - vector_store_search_optional_params={}, - litellm_logging_obj=MagicMock(), - litellm_params=BASE_PARAMS, - embedding_executor=caller, - ) - - assert caller.captured.query == "a lone astronaut" - assert _stage(collection, "$vectorSearch")["queryVector"] == (0.4, 0.5, 0.6) +async def test_public_sdk_preserves_http_errors_response_and_timeout( + status: int, + body: Mapping[str, object], + error_type: type[Exception] | None, + asynchronous: bool, + timeout: float, + api_base: str, +) -> None: + executor: Final = RecordingEmbeddingExecutor() + if status == -1: + if asynchronous: + with pytest.raises(litellm.BadRequestError, match="search-only"): + await litellm.vector_stores.acreate(custom_llm_provider="mongodb") + else: + with pytest.raises(litellm.BadRequestError, match="search-only"): + litellm.vector_stores.create(custom_llm_provider="mongodb") + return + if status == -2: + rejected_params: Final = {**BASE_PARAMS, "api_base": str(body["api_base"])} + if asynchronous: + with pytest.raises(litellm.BadRequestError, match="requires HTTPS"): + await litellm.vector_stores.asearch( + vector_store_id="policy_index", + query="travel policy", + custom_llm_provider="mongodb", + _direct_vector_store_embedding_executor=executor, + **rejected_params, + ) + else: + with pytest.raises(litellm.BadRequestError, match="requires HTTPS"): + litellm.vector_stores.search( + vector_store_id="policy_index", + query="travel policy", + custom_llm_provider="mongodb", + _direct_vector_store_embedding_executor=executor, + **rejected_params, + ) + executor.call.assert_not_called() + return + + def respond(request: httpx.Request) -> httpx.Response: + assert request.url == f"{api_base}/v1/vector_stores/policy_index/search" + assert request.headers["authorization"] == "Bearer test-sidecar-key" + assert request.extensions["timeout"]["read"] == timeout + payload: Final = json.loads(request.content) + assert payload["timeout_ms"] == int(timeout * 1000) + assert payload["query_vector"] == [0.1, 0.2, 0.3] + if status == 0: + raise httpx.ReadTimeout("timed out", request=request) + return httpx.Response(status, json=body) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as async_transport: + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + client: Final = AsyncHTTPHandler() if asynchronous else HTTPHandler(client=transport) + if isinstance(client, AsyncHTTPHandler): + await client.client.aclose() + client.client = async_transport + + async def search() -> VectorStoreSearchResponse: + kwargs: Final = { + **BASE_PARAMS, + "api_base": api_base, + "vector_store_id": "policy_index", + "query": "travel policy", + "custom_llm_provider": "mongodb", + "_direct_vector_store_embedding_executor": executor, + "client": client, + "timeout": timeout, + } + if asynchronous: + return await litellm.vector_stores.asearch(**kwargs) + return litellm.vector_stores.search(**kwargs) + + if error_type is not None: + with pytest.raises(error_type): + await search() + else: + assert await search() == RESULT + executor.call.assert_called_once_with("embedding-alias", "travel policy", {}) diff --git a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py index 8c8bea00dea..d484fa437ae 100644 --- a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py +++ b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py @@ -708,38 +708,6 @@ class TestKimiK26ModelRegistry: """Load directly from the bundled backup so tests don't depend on remote fetch.""" return GetModelCostMap.load_local_model_cost_map() - def test_kimi_k26_in_model_cost_map(self, model_cost_map): - """kimi-k2.6 should be present in the model cost map.""" - assert "moonshot/kimi-k2.6" in model_cost_map, "moonshot/kimi-k2.6 not found in model_cost" - - def test_kimi_k26_pricing(self, model_cost_map): - """kimi-k2.6 pricing should match official Kimi API rates.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07) - assert model_info["output_cost_per_token"] == pytest.approx(4e-06) - assert model_info["cache_read_input_token_cost"] == pytest.approx(1.6e-07) - - def test_kimi_k26_context_window(self, model_cost_map): - """kimi-k2.6 should have a 256K (262144 token) context window.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - assert model_info["max_tokens"] == 262144 - - def test_kimi_k26_capabilities(self, model_cost_map): - """kimi-k2.6 should support function calling, vision, video input, tool choice, and reasoning.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info.get("supports_function_calling") is True - assert model_info.get("supports_tool_choice") is True - assert model_info.get("supports_vision") is True - assert model_info.get("supports_video_input") is True - assert model_info.get("supports_reasoning") is True - - def test_kimi_k26_provider(self, model_cost_map): - """kimi-k2.6 should be assigned to the moonshot provider.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info["litellm_provider"] == "moonshot" - class TestMoonshotResponseSchemaSupport: """Every model currently live on api.moonshot.ai supports json_schema @@ -762,10 +730,6 @@ class TestMoonshotResponseSchemaSupport: def model_cost_map(self): return GetModelCostMap.load_local_model_cost_map() - @pytest.mark.parametrize("model", LIVE_MODELS) - def test_live_model_supports_response_schema(self, model, model_cost_map): - assert model_cost_map[model].get("supports_response_schema") is True - def test_supports_response_schema_utility_reports_true(self, model_cost_map, monkeypatch): monkeypatch.setattr(litellm, "model_cost", model_cost_map) assert litellm.utils.supports_response_schema(model="moonshot/kimi-k2.5") is True diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py index 86c534c73c2..4c9bd29b337 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py @@ -1900,3 +1900,198 @@ class TestOCIImageUrlTransformation: adapt_messages_to_generic_oci_standard(messages) assert "image_url" in str(exc_info.value) + + +import itertools +from unittest.mock import patch + +from litellm.llms.oci.chat.transformation import OCIStreamWrapper, _iter_sse_events + +_STREAM_GENERIC_MODEL = "xai.grok-4" +_STREAM_COHERE_MODEL = "cohere.command-latest" + +_GENERIC_TEXT_EVENT = ( + 'data: {{"index":0,"message":{{"role":"ASSISTANT","content":[{{"type":"TEXT","text":"{text}"}}]}},"pad":"aaa"}}' +) +_GENERIC_TERMINAL_EVENT = ( + 'data: {"message":{"role":"ASSISTANT","content":[{"type":"TEXT","text":""}]},"finishReason":"stop","pad":"a"}' +) +_COHERE_TEXT_EVENT = 'data: {{"apiFormat":"COHERE","text":"{text}","pad":"aaaaaa"}}' +_COHERE_TERMINAL_EVENT = ( + 'data: {"apiFormat":"COHERE","text":"123","finishReason":"COMPLETE",' + '"chatHistory":[{"role":"USER","message":"count"},{"role":"CHATBOT","message":"123"}]}' +) + + +def _make_stream_wrapper(model: str) -> OCIStreamWrapper: + logging_obj = MagicMock() + logging_obj.model_call_details = {"custom_llm_provider": "oci", "litellm_params": {}} + return OCIStreamWrapper( + completion_stream=iter([]), + model=model, + custom_llm_provider="oci", + logging_obj=logging_obj, + ) + + +def _ticking_clock(): + """A ``time.time`` stand-in that advances a full second on every call. + + Without it the whole test runs inside one wall-clock second, so a per-chunk + ``created`` would coincidentally match and the drift would go unnoticed. + """ + return itertools.count(1_700_000_000.0) + + +class TestOCIStreamWrapperIdentityPinning: + """One OCI streaming completion must present one id, one created and the + wrapper's model on every chunk, the way every other provider does.""" + + def test_generic_stream_shares_one_id_created_and_model(self): + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + events = [ + _GENERIC_TEXT_EVENT.format(text="1"), + _GENERIC_TEXT_EVENT.format(text="2"), + _GENERIC_TEXT_EVENT.format(text="3"), + _GENERIC_TERMINAL_EVENT, + ] + + with patch("time.time", side_effect=_ticking_clock()): + chunks = [wrapper.chunk_creator(event) for event in events] + + assert len(chunks) == 4 + assert len({chunk.id for chunk in chunks}) == 1 + assert chunks[0].id.startswith("chatcmpl-") + assert len({chunk.created for chunk in chunks}) == 1 + assert {chunk.model for chunk in chunks} == {_STREAM_GENERIC_MODEL} + assert [chunk.choices[0].delta.content for chunk in chunks[:3]] == ["1", "2", "3"] + assert chunks[-1].choices[0].finish_reason == "stop" + assert all(chunk._hidden_params["custom_llm_provider"] == "oci" for chunk in chunks) + + def test_cohere_stream_shares_one_id_created_and_model(self): + """Rebuilding each chunk through the shared creator must not disturb the + Cohere bookkeeping that suppresses the terminal event's repeated text.""" + wrapper = _make_stream_wrapper(_STREAM_COHERE_MODEL) + events = [ + _COHERE_TEXT_EVENT.format(text="1"), + _COHERE_TEXT_EVENT.format(text="2"), + _COHERE_TEXT_EVENT.format(text="3"), + _COHERE_TERMINAL_EVENT, + ] + + with patch("time.time", side_effect=_ticking_clock()): + chunks = [wrapper.chunk_creator(event) for event in events] + + assert len(chunks) == 4 + assert len({chunk.id for chunk in chunks}) == 1 + assert len({chunk.created for chunk in chunks}) == 1 + assert {chunk.model for chunk in chunks} == {_STREAM_COHERE_MODEL} + assert [chunk.choices[0].delta.content for chunk in chunks[:3]] == ["1", "2", "3"] + assert chunks[-1].choices[0].finish_reason == "stop" + assert chunks[-1].choices[0].delta.content is None + assert wrapper._cohere_text_emitted is True + + def test_id_is_pinned_to_the_wrapper_response_id(self): + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + + first = wrapper.chunk_creator(_GENERIC_TEXT_EVENT.format(text="1")) + + assert wrapper.response_id == first.id + assert wrapper.created == first.created + + +class TestOCIStreamWrapperDoneSentinel: + """OCI's GENERIC apiFormat closes the stream with a literal `[DONE]` line; + parsing it as JSON turned every streaming completion into a 500.""" + + @pytest.mark.parametrize("done_event", ["data: [DONE]", "data:[DONE]", "data: [DONE] "]) + def test_done_sentinel_returns_none(self, done_event): + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + assert wrapper.chunk_creator(done_event) is None + + def test_done_sentinel_off_the_sse_splitter_is_skipped(self): + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + wire = ( + f"{_GENERIC_TEXT_EVENT.format(text='1')}\n\n" + f"{_GENERIC_TEXT_EVENT.format(text='2')}\n\n" + f"{_GENERIC_TERMINAL_EVENT}\n\n" + "data: [DONE]\n\n" + ) + + events = list(_iter_sse_events(iter([wire]))) + assert events[-1] == "data: [DONE]" + + chunks = [wrapper.chunk_creator(event) for event in events] + assert chunks[-1] is None + + emitted = [chunk for chunk in chunks if chunk is not None] + assert len(emitted) == 3 + assert len({chunk.id for chunk in emitted}) == 1 + + def test_unparseable_payload_still_raises_oci_error(self): + from litellm.llms.oci.common_utils import OCIError + + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"): + wrapper.chunk_creator("data: not-json-at-all") + + def test_done_lookalike_payload_still_raises_oci_error(self): + from litellm.llms.oci.common_utils import OCIError + + wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL) + with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"): + wrapper.chunk_creator("data: [DONE] trailing garbage") + + +_GENERIC_TOOL_CALL_EVENT = ( + 'data: {"index":0,"message":{"role":"ASSISTANT","content":[],' + '"toolCalls":[{"type":"FUNCTION","id":"call_1","name":"get_weather","arguments":"{}"}]}}' +) +_GENERIC_TOOL_TERMINAL_EVENT = ( + 'data: {"index":0,"message":{"role":"ASSISTANT","content":[]},"finishReason":"TOOL_CALLS"}' +) + + +def _drain_stream(model: str, events: list[str]) -> list: + logging_obj = MagicMock() + logging_obj.model_call_details = {"custom_llm_provider": "oci", "litellm_params": {}} + wrapper = OCIStreamWrapper( + completion_stream=iter(events), + model=model, + custom_llm_provider="oci", + logging_obj=logging_obj, + ) + return list(wrapper) + + +class TestOCIStreamWrapperTerminalChunk: + """OCI's ``chunk_creator`` override bypasses the shared handler's + finish-reason bookkeeping, so the shared end-of-stream finalizer used to + append a synthetic ``stop`` chunk after OCI's own terminal chunk, silently + downgrading a ``tool_calls`` completion for any client that reads the + finish reason off the last chunk.""" + + def test_generic_tool_call_stream_ends_on_tool_calls(self): + chunks = _drain_stream( + _STREAM_GENERIC_MODEL, + [_GENERIC_TOOL_CALL_EVENT, _GENERIC_TOOL_TERMINAL_EVENT, "data: [DONE]"], + ) + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "tool_calls"] + assert len({chunk.id for chunk in chunks}) == 1 + + def test_generic_text_stream_emits_exactly_one_finish_reason(self): + chunks = _drain_stream( + _STREAM_GENERIC_MODEL, + [_GENERIC_TEXT_EVENT.format(text="1"), _GENERIC_TERMINAL_EVENT, "data: [DONE]"], + ) + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"] + + def test_cohere_stream_emits_exactly_one_finish_reason(self): + chunks = _drain_stream( + _STREAM_COHERE_MODEL, + [_COHERE_TEXT_EVENT.format(text="123"), _COHERE_TERMINAL_EVENT], + ) + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"] diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py index 5dd44d72d68..729a2d25f41 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -445,6 +445,72 @@ class TestOCICohereToolCalls: assert result.choices[0].index == 0 assert result.choices[0].finish_reason == "stop" # COMPLETE is mapped to stop + _TOOL_TURN_TEXT = "I will use the tool to find out the weather in Paris." + _TOOL_TURN_DELTAS = [ + "I", " will", " use", " the", " tool", " to", " find", " out", " the", " weather", " in", " Paris", ".", + ] + _TOOL_TURN_CALLS = [{"name": "get_weather", "parameters": {"city": "Paris"}}] + _TOOL_TURN_HISTORY = [ + {"role": "USER", "message": "Briefly say what you will do, then find out the weather in Paris using the tool."}, + {"role": "CHATBOT", "message": _TOOL_TURN_TEXT, "toolCalls": _TOOL_TURN_CALLS}, + ] + _TOOL_TURN_TERMINAL_TOGETHER = [ + { + "apiFormat": "COHERE", + "text": _TOOL_TURN_TEXT, + "chatHistory": _TOOL_TURN_HISTORY, + "finishReason": "COMPLETE", + "toolCalls": _TOOL_TURN_CALLS, + }, + ] + _TOOL_TURN_TERMINAL_SPLIT = [ + { + "apiFormat": "COHERE", + "text": _TOOL_TURN_TEXT, + "chatHistory": _TOOL_TURN_HISTORY, + "toolCalls": _TOOL_TURN_CALLS, + }, + {"apiFormat": "COHERE", "finishReason": "COMPLETE"}, + ] + + @staticmethod + def _drain_cohere_stream(events): + wrapper = OCIStreamWrapper( + completion_stream=MagicMock(), model="cohere.command-a-03-2025", logging_obj=MagicMock() + ) + chunks = [wrapper.chunk_creator(f"data: {json.dumps(event)}") for event in events] + content = "".join(chunk.choices[0].delta.content or "" for chunk in chunks) + tool_calls = [call for chunk in chunks for call in (chunk.choices[0].delta.tool_calls or [])] + finish_reasons = [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] + return content, tool_calls, finish_reasons + + @pytest.mark.parametrize("terminal_events", [_TOOL_TURN_TERMINAL_TOGETHER, _TOOL_TURN_TERMINAL_SPLIT]) + def test_cohere_tool_turn_streams_the_answer_once(self, terminal_events): + """OCI restates the whole answer on the tool-calls chunk and again on the terminal chunk; + the client must read it exactly once, with one tool call and one finish reason.""" + deltas = [{"apiFormat": "COHERE", "text": token} for token in self._TOOL_TURN_DELTAS] + tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS} + + content, tool_calls, finish_reasons = self._drain_cohere_stream([*deltas, tool_calls_event, *terminal_events]) + + assert content == self._TOOL_TURN_TEXT + assert [(call["function"]["name"], call["function"]["arguments"]) for call in tool_calls] == [ + ("get_weather", '{"city": "Paris"}') + ] + assert finish_reasons == ["stop"] + + def test_cohere_tool_turn_without_preamble_deltas_keeps_the_only_text(self): + """When the tool-calls chunk carries the only copy of the text, dropping it would lose the answer.""" + tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS} + + content, tool_calls, finish_reasons = self._drain_cohere_stream( + [tool_calls_event, *self._TOOL_TURN_TERMINAL_TOGETHER] + ) + + assert content == self._TOOL_TURN_TEXT + assert len(tool_calls) == 1 + assert finish_reasons == ["stop"] + def test_cohere_parameter_mapping_excludes_tool_choice(self): """Test that tool_choice is excluded from Cohere parameter mapping""" config = OCIChatConfig() diff --git a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py index eadc2bc9541..b2071155f3f 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py @@ -2,9 +2,11 @@ import json from litellm._uuid import uuid from unittest.mock import MagicMock, patch +import httpx import pytest - +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.ollama.completion.transformation import ( OllamaConfig, OllamaTextCompletionResponseIterator, @@ -502,3 +504,43 @@ class TestOllamaTextCompletionResponseIterator: assert result["usage"]["prompt_tokens"] == 10 assert result["usage"]["completion_tokens"] == 5 assert result["usage"]["total_tokens"] == 15 + + +async def test_ollama_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"https://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "model": "llava", + "response": "Green", + "done": True, + "prompt_eval_count": 1, + "eval_count": 1, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="ollama/llava", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + api_base="http://ollama.example:11434", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert captured["body"]["images"] == [async_only_image_fetch.base64_png] diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 36e715d5804..5a29a96829f 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1074,6 +1074,306 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: # Should return the responses assert result == responses_so_far + @staticmethod + def _ended_stream_chunks() -> list: + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return [ + ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=Delta(content="Hello"), finish_reason=None)], + ), + ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop")], + ), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_text_back_into_chunks(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert chunks[0].choices[0].delta.content == "HELLO WORLD" + assert chunks[1].choices[0].delta.content in (None, "") + assert chunks[1].choices[0].finish_reason == "stop" + + @staticmethod + def _ended_tool_call_stream_chunks() -> list: + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk(tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None): + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(tool_calls=[tool_call] if tool_call else None), + finish_reason=finish_reason, + ) + ], + ) + + def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None): + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + return [ + chunk(fragment("", name="lookup_fruit", call_id="call_1")), + chunk(fragment('{"fruit":')), + chunk(fragment(' "persimmon"}')), + chunk(None, finish_reason="tool_calls"), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_call_arguments_back_into_chunks(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_tool_call_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]] + assert [fragment[0].function.arguments for fragment in fragments] == ['{"fruit": "PERSIMMON"}', "", ""] + assert fragments[0][0].function.name == "lookup_fruit" + assert fragments[0][0].id == "call_1" + assert chunks[3].choices[0].delta.tool_calls is None + assert chunks[3].choices[0].finish_reason == "tool_calls" + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_call_name_back_into_chunks(self): + class RenameTool(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call["function"]["name"] = "lookup_fruit_reviewed" + return inputs + + handler = OpenAIChatCompletionsHandler() + chunks = self._ended_tool_call_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=RenameTool(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + fragments = [chunk.choices[0].delta.tool_calls[0] for chunk in chunks[:3]] + assert [fragment.function.name for fragment in fragments] == ["lookup_fruit_reviewed", None, None] + assert json.loads("".join(fragment.function.arguments for fragment in fragments)) == {"fruit": "persimmon"} + assert fragments[0].id == "call_1" + + @pytest.mark.asyncio + async def test_ended_stream_tool_call_rewrite_leaves_chunks_untouched_by_default(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_tool_call_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]] + assert [fragment[0].function.arguments for fragment in fragments] == ["", '{"fruit":', ' "persimmon"}'] + + @pytest.mark.asyncio + async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert chunks[0].choices[0].delta.content == "Hello" + assert chunks[1].choices[0].delta.content == " world" + assert chunks[1].choices[0].finish_reason == "stop" + + @staticmethod + def _two_choice_stream_chunks() -> list: + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + def chunk(index: int, content: str, finish_reason: Optional[str] = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + return [ + chunk(0, "safe "), + chunk(1, "hello "), + chunk(0, "text", "stop"), + chunk(1, "world", "stop"), + ] + + @staticmethod + def _world_masking_guardrail() -> CustomGuardrail: + class MaskWorld(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts", []) + return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]} + + return MaskWorld(guardrail_name="test-mask") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrite_on_multi_choice_stream_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_stream_chunks() + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._world_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @staticmethod + def _two_choice_tool_call_stream_chunks() -> list: + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk( + choice_index: int, tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None + ) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=choice_index, + delta=Delta(tool_calls=[tool_call] if tool_call else None), + finish_reason=finish_reason, + ) + ], + ) + + def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None): + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + return [ + chunk(0, fragment("", name="lookup_fruit", call_id="call_1")), + chunk(1, fragment("", name="lookup_fruit", call_id="call_2")), + chunk(0, fragment('{"fruit": "persimmon"}')), + chunk(1, fragment('{"fruit": "durian"}')), + chunk(0, None, finish_reason="tool_calls"), + chunk(1, None, finish_reason="tool_calls"), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_tool_call_stream_chunks() + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert [c.choices[0].delta.content for c in chunks] == ["safe ", "hello ", "text", "world"] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrite_lands_on_nonzero_choice_index(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + + def chunk(content: str, finish_reason: Optional[str]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=1, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + chunks = [chunk("hello ", None), chunk("world", "stop")] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._world_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert chunks[0].choices[0].delta.content == "hello [MASKED]" + assert chunks[1].choices[0].delta.content in (None, "") + class TestUndecoratedGuardrailIsRecorded: """LIT-5983 regression: the handler calls apply_guardrail bare, so a custom guardrail diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 9737d63cc26..b110586ae5b 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -4,6 +4,7 @@ Tests for OpenAI GPT transformation (litellm/llms/openai/chat/gpt_transformation import pytest +from typing import Final import litellm @@ -1168,6 +1169,9 @@ class TestOpenAIPromptCacheBreakpointChatPath: assert "prompt_cache_options" not in request +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + + class TestToolSchemaCombinatorFlatteningForOpenAI: """ Regression tests for LIT-6488: OpenAI's chat completions validator rejects @@ -1281,3 +1285,50 @@ class TestToolSchemaCombinatorFlatteningForOpenAI: parameters = request["tools"][0]["function"]["parameters"] assert "anyOf" not in parameters assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + + @staticmethod + def _artifact_tool(): + return { + "type": "function", + "function": { + "name": "Artifact", + "parameters": { + "type": "object", + "properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}, + "required": ["field"], + }, + }, + } + + def test_drops_non_python_regex_pattern_for_hosted_openai(self): + tool = self._artifact_tool() + + request = self._transform(self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": None}, [tool]) + + assert request["tools"][0]["function"]["parameters"] == { + "type": "object", + "properties": {"field": {"type": "string"}}, + "required": ["field"], + } + assert tool == self._artifact_tool() + + def test_custom_api_base_drops_non_python_regex_pattern_but_keeps_union(self): + tool = self._anyof_tool() + tool["function"]["parameters"]["properties"]["id"]["pattern"] = _ARTIFACT_FIELD_PATTERN + + request = self._transform( + self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": "http://localhost:8000/v1"}, [tool] + ) + + parameters = request["tools"][0]["function"]["parameters"] + assert parameters["properties"]["id"] == {"type": "string"} + assert parameters["anyOf"] == self._anyof_tool()["function"]["parameters"]["anyOf"] + + def test_non_openai_provider_keeps_non_python_regex_pattern(self): + tool = self._artifact_tool() + + request = self._transform( + self.config, "some-oss-model", {"custom_llm_provider": "groq", "api_base": None}, [tool] + ) + + assert request["tools"][0] is tool diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index a453708040e..a4f0a77a9b6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -10,23 +10,33 @@ from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock +import logging + import pytest from fastapi import HTTPException -from openai.types.responses import ResponseFunctionToolCall +from pydantic import BaseModel +from openai.types.responses import ( + ResponseCustomToolCall, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, + ResponseFunctionToolCall, +) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.responses.main import GenericResponseOutputItem, OutputText +from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -56,6 +66,60 @@ class MockGuardrail(CustomGuardrail): return inputs +class PersimmonMaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + tool_calls = [ + { + **tool_call, + "function": { + **tool_call["function"], + "arguments": tool_call["function"]["arguments"].replace("persimmon", "[MASKED]"), + }, + } + for tool_call in inputs.get("tool_calls", []) + ] + return {**inputs, "tool_calls": tool_calls} + + +class FlatShapeGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + flat_tool_calls = [{"name": "exec", "input": "rm -rf /"} for _ in inputs.get("tool_calls", [])] + return {**inputs, "tool_calls": flat_tool_calls} + + +class DroppingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "tool_calls": []} + + +CUSTOM_TOOL_CALL_ITEM = { + "type": "custom_tool_call", + "id": "ctc_1", + "call_id": "call_exec_1", + "name": "exec", + "input": "echo persimmon", + "status": "completed", +} + + class TestOpenAIResponsesHandlerDiscovery: """Test that the handler is properly discovered by the guardrail system""" @@ -556,7 +620,7 @@ class TestOpenAIResponsesHandlerToolCallExtraction: texts_to_check: List[str] = [] images_to_check: List[str] = [] - tool_calls_to_check: List[Any] = [] + tool_calls_to_check: List[ChatCompletionToolCallChunk] = [] task_mappings: List[Tuple[int, int]] = [] # Extract tool calls @@ -627,6 +691,123 @@ class TestOpenAIResponsesHandlerToolCallExtraction: == '{"location":"Boston, MA","unit":"celsius"}' ) + @pytest.mark.parametrize( + "output_item", + [ + dict(CUSTOM_TOOL_CALL_ITEM), + CustomToolCallOutputItem(**CUSTOM_TOOL_CALL_ITEM), + ResponseCustomToolCall(**{key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "status"}), + ], + ids=["dict", "litellm_typed", "openai_typed"], + ) + def test_extract_custom_tool_call_input_as_arguments(self, output_item): + handler = OpenAIResponsesHandler() + texts_to_check: List[str] = [] + tool_calls_to_check: List[Any] = [] + + handler._extract_output_text_and_images( + output_item=output_item, + output_idx=2, + texts_to_check=texts_to_check, + images_to_check=[], + task_mappings=[], + tool_calls_to_check=tool_calls_to_check, + ) + + assert texts_to_check == [] + assert tool_calls_to_check == [ + { + "id": "call_exec_1", + "type": "function", + "function": {"name": "exec", "arguments": "echo persimmon"}, + "index": 2, + } + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("typed", [False, True], ids=["dict", "typed"]) + async def test_process_output_response_writes_tool_call_rewrites_back(self, typed): + handler = OpenAIResponsesHandler() + function_call = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_fn_1", + "name": "lookup_fruit", + "arguments": '{"fruit": "persimmon"}', + "status": "completed", + } + message = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "running persimmon", "annotations": []}], + } + payload = { + "id": "resp_1", + "created_at": 1, + "model": "gpt-5.6", + "object": "response", + "status": "completed", + "output": [message, function_call, dict(CUSTOM_TOOL_CALL_ITEM)], + } + response = ResponsesAPIResponse.model_validate(payload) if typed else payload + + result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask")) + + output = result.output if typed else result["output"] + function_item, custom_item = output[1], output[2] + assert (function_item.arguments if typed else function_item["arguments"]) == '{"fruit": "[MASKED]"}' + assert (custom_item.input if typed else custom_item["input"]) == "echo [MASKED]" + assert (custom_item.name if typed else custom_item["name"]) == "exec" + assert (output[0].content[0].text if typed else output[0]["content"][0]["text"]) == "running persimmon" + + @staticmethod + def _custom_tool_call_response(item: dict) -> dict: + return { + "id": "resp_1", + "created_at": 1, + "model": "gpt-5.6", + "object": "response", + "status": "completed", + "output": [item], + } + + @pytest.mark.asyncio + async def test_process_output_response_ignores_tool_call_rewrites_in_another_shape(self): + handler = OpenAIResponsesHandler() + response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM)) + + result = await handler.process_output_response(response, FlatShapeGuardrail(guardrail_name="flat")) + + assert result["output"][0]["input"] == "echo persimmon" + assert result["output"][0]["name"] == "exec" + + @pytest.mark.asyncio + async def test_process_output_response_warns_when_guardrail_drops_tool_calls(self, caplog): + handler = OpenAIResponsesHandler() + response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM)) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await handler.process_output_response(response, DroppingGuardrail(guardrail_name="dropper")) + + assert result["output"][0]["input"] == "echo persimmon" + assert any( + "dropper" in record.getMessage() and "0 tool calls for the 1 scanned" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.asyncio + async def test_process_output_response_keeps_a_nameless_custom_tool_call_nameless(self): + handler = OpenAIResponsesHandler() + nameless_item = {key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "name"} + response = self._custom_tool_call_response(nameless_item) + + result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask")) + + assert result["output"][0]["input"] == "echo [MASKED]" + assert "name" not in result["output"][0] + @pytest.mark.asyncio async def test_process_output_response_with_tool_calls(self): """Test processing output response containing function tool calls""" @@ -1128,6 +1309,555 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: output_text = result[-1]["response"]["output"][0]["content"][0]["text"] assert output_text == original_text + @staticmethod + def _ended_stream_events() -> List[dict]: + content = [{"type": "output_text", "text": "hello world"}] + item = { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": content, + } + return [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"}, + {"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"}, + { + "type": "response.content_part.done", + "output_index": 0, + "content_index": 0, + "part": {"type": "output_text", "text": "hello world"}, + }, + {"type": "response.output_item.done", "output_index": 0, "item": {**item, "content": [dict(c) for c in content]}}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "model": "gpt-4o", + "output": [{**item, "content": [dict(c) for c in content]}], + "status": "completed", + }, + }, + ] + + @staticmethod + def _masking_guardrail() -> CustomGuardrail: + class MaskWorld(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts", []) + return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]} + + return MaskWorld(guardrail_name="test-mask") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_all_stream_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["delta"] == "hello [MASKED]" + assert events[1]["delta"] == "" + assert events[2]["text"] == "hello [MASKED]" + assert events[3]["part"]["text"] == "hello [MASKED]" + assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" + assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + + @staticmethod + def _ended_function_call_stream_events() -> List[dict]: + def item(arguments: str, status: str) -> dict: + return { + "type": "function_call", + "id": "fc_123", + "call_id": "call_123", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.done", + "item_id": "fc_123", + "output_index": 0, + "arguments": '{"fruit": "persimmon"}', + }, + {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "created_at": 1, + "model": "gpt-4o", + "output": [item('{"fruit": "persimmon"}', "completed")], + "status": "completed", + }, + }, + ] + + @staticmethod + def _argument_masking_guardrail() -> CustomGuardrail: + class MaskArguments(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + tool_calls = [ + {**tool_call, "function": {**tool_call["function"], "arguments": '{"fruit": "[MASKED]"}'}} + for tool_call in inputs.get("tool_calls", []) + ] + return {**inputs, "tool_calls": tool_calls} + + return MaskArguments(guardrail_name="test-mask-arguments") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_function_call_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["item"]["arguments"] == "" + assert events[1]["delta"] == '{"fruit": "[MASKED]"}' + assert events[2]["delta"] == "" + assert events[3]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["response"]["output"][0]["name"] == "lookup_fruit" + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_typed_function_call_events(self): + from litellm.types.llms.openai import ( + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDoneEvent, + OutputItemAddedEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ) + + handler = OpenAIResponsesHandler() + typed_events: List[Any] = [ + model.model_validate(event) + for model, event in zip( + ( + OutputItemAddedEvent, + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDoneEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ), + self._ended_function_call_stream_events(), + ) + ] + completed_event = typed_events[5] + assert isinstance(completed_event, ResponseCompletedEvent) + assert isinstance(completed_event.response, ResponsesAPIResponse) + assert isinstance(completed_event.response.output[0], ResponseFunctionToolCall) + + await handler.process_output_streaming_response( + responses_so_far=typed_events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert typed_events[1].delta == '{"fruit": "[MASKED]"}' + assert typed_events[2].delta == "" + assert typed_events[3].arguments == '{"fruit": "[MASKED]"}' + assert typed_events[4].item.arguments == '{"fruit": "[MASKED]"}' + assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}' + assert completed_event.response.output[0].name == "lookup_fruit" + + @staticmethod + def _ended_custom_tool_call_stream_events() -> List[dict]: + def item(input_text: str, status: str) -> dict: + return {**CUSTOM_TOOL_CALL_ITEM, "input": input_text, "status": status} + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "echo "}, + {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "persimmon"}, + {"type": "response.custom_tool_call_input.done", "item_id": "ctc_1", "output_index": 0, "input": "echo persimmon"}, + {"type": "response.output_item.done", "output_index": 0, "item": item("echo persimmon", "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "created_at": 1, + "model": "gpt-5.6", + "output": [item("echo persimmon", "completed")], + "status": "completed", + }, + }, + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_custom_tool_call_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["item"]["input"] == "" + assert events[1]["delta"] == "echo [MASKED]" + assert events[2]["delta"] == "" + assert events[3]["input"] == "echo [MASKED]" + assert events[4]["item"]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["name"] == "exec" + assert "arguments" not in events[5]["response"]["output"][0] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_keep_a_nameless_custom_tool_call_nameless(self): + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + items = [events[0]["item"], events[4]["item"], events[5]["response"]["output"][0]] + for item in items: + del item["name"] + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert events[3]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]" + assert all("name" not in item for item in items) + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_typed_custom_tool_call_events(self): + from litellm.types.llms.openai import ( + OutputItemAddedEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ) + + handler = OpenAIResponsesHandler() + typed_events: List[BaseModel] = [ + model.model_validate({**event, "sequence_number": sequence_number}) + for sequence_number, (model, event) in enumerate( + zip( + ( + OutputItemAddedEvent, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ), + self._ended_custom_tool_call_stream_events(), + ) + ) + ] + completed_event = typed_events[5] + assert isinstance(completed_event.response.output[0], CustomToolCallOutputItem) + + await handler.process_output_streaming_response( + responses_so_far=typed_events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert typed_events[1].delta == "echo [MASKED]" + assert typed_events[2].delta == "" + assert typed_events[3].input == "echo [MASKED]" + assert typed_events[4].item.input == "echo [MASKED]" + assert completed_event.response.output[0].input == "echo [MASKED]" + assert completed_event.response.output[0].name == "exec" + + @pytest.mark.asyncio + async def test_deliver_ended_stream_custom_tool_call_rewrite_without_matching_events_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + events[5]["response"]["output"] = [{**events[5]["response"]["output"][0], "call_id": "call_999"}] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @staticmethod + def _bridged_function_call_stream_events() -> List[dict]: + reasoning = {"type": "reasoning", "id": "rs_1", "summary": []} + text = {"type": "output_text", "text": "Looking that up", "annotations": []} + message = {"type": "message", "id": "msg_1", "role": "assistant", "status": "completed", "content": [text]} + + def function_call(arguments: str, status: str) -> dict: + return { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": dict(reasoning)}, + {"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning)}, + {"type": "response.output_item.added", "output_index": 0, "item": {**message, "status": "in_progress", "content": []}}, + {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": "Looking that up"}, + {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": [dict(text)]}}, + {"type": "response.output_item.added", "output_index": 1, "item": function_call("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.done", + "item_id": "fc_1", + "output_index": 1, + "arguments": '{"fruit": "persimmon"}', + }, + {"type": "response.output_item.done", "output_index": 1, "item": function_call('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_1", + "model": "claude-haiku-4-5", + "output": [ + dict(reasoning), + {**message, "content": [dict(text)]}, + function_call('{"fruit": "persimmon"}', "completed"), + ], + }, + }, + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_keys_bridged_function_call_events_by_call_id(self): + handler = OpenAIResponsesHandler() + events = self._bridged_function_call_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert events[6]["delta"] == '{"fruit": "[MASKED]"}' + assert events[7]["delta"] == "" + assert events[8]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["item"]["name"] == "lookup_fruit" + assert events[9]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[10]["response"]["output"][2]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[3]["delta"] == "Looking that up" + assert events[4]["item"]["content"][0]["text"] == "Looking that up" + assert events[10]["response"]["output"][1]["content"][0]["text"] == "Looking that up" + assert events[1]["item"] == {"type": "reasoning", "id": "rs_1", "summary": []} + + @pytest.mark.asyncio + @pytest.mark.parametrize("mismatch", ["orphan_call_id", "duplicate_call_id"]) + async def test_deliver_ended_stream_function_call_rewrite_without_matching_events_fails_closed(self, mismatch): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + envelope_item = events[5]["response"]["output"][0] + if mismatch == "orphan_call_id": + events[5]["response"]["output"] = [{**envelope_item, "call_id": "call_999"}] + else: + events[5]["response"]["output"] = [dict(envelope_item), dict(envelope_item)] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self): + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + ) + + assert events[1]["delta"] == '{"fruit":' + assert events[3]["arguments"] == '{"fruit": "persimmon"}' + assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "persimmon"}' + + @pytest.mark.asyncio + @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"]) + async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events() + events[-1]["type"] = terminal_type + events[-1]["response"]["status"] = terminal_type.split(".")[-1] + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["delta"] == "hello [MASKED]" + assert events[1]["delta"] == "" + assert events[2]["text"] == "hello [MASKED]" + assert events[3]["part"]["text"] == "hello [MASKED]" + assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" + assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + + @pytest.mark.asyncio + async def test_fallback_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"}, + ] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_fallback_delta_only_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"}, + ] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_output_item_done_last_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_stream_events()[:-1] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_output_item_done_last_scans_text_with_delivery_expected(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events()[:-1] + guardrail = MockRecordingGuardrail(guardrail_name="test") + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hello world"]] + + @pytest.mark.asyncio + async def test_output_item_done_last_without_delivery_expected_skips_text(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events()[:-1] + guardrail = MockRecordingGuardrail(guardrail_name="test") + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert result is events + assert guardrail.seen_inputs == [] + + @pytest.mark.asyncio + async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self): + handler = OpenAIResponsesHandler() + events = [ + {"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + ) + + assert result is events + + @pytest.mark.asyncio + async def test_ended_stream_rewrite_leaves_delta_events_untouched_by_default(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + ) + + assert events[0]["delta"] == "hello " + assert events[1]["delta"] == "world" + assert events[2]["text"] == "hello world" + assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + @pytest.mark.asyncio async def test_failed_stream_scans_delta_text(self): """A stream ending in response.failed has text only in delta events; the @@ -2319,8 +3049,21 @@ class TestOpenAIResponsesHandlerStreamingScanKey: assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] assert ended_key != open_key + def test_completed_event_with_a_custom_tool_call_changes_the_key(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + ended_key = handler.get_streaming_scan_key( + [self._delta(0, "hi"), self._completed(1, [message, dict(CUSTOM_TOOL_CALL_ITEM)])] + ) + rewritten_key = handler.get_streaming_scan_key( + [self._delta(0, "hi"), self._completed(1, [message, {**CUSTOM_TOOL_CALL_ITEM, "input": "echo kumquat"}])] + ) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "echo persimmon" in ended_key.tool_calls[0] + assert rewritten_key != ended_key + def test_completed_event_reads_every_output_text_part(self): - from litellm.types.responses.main import GenericResponseOutputItem, OutputText + from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText item = GenericResponseOutputItem( type="message", diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index a5b748e391c..4cf8767764b 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -1,5 +1,6 @@ import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx @@ -20,6 +21,8 @@ from litellm.types.llms.openai import ( ) from litellm.types.router import GenericLiteLLMParams +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + class TestOpenAIResponsesAPIConfig: def setup_method(self): @@ -2020,3 +2023,188 @@ class TestFlattenToolSchemaCombinatorsWiring: assert result["tools"][0] is opaque_tool assert "anyOf" not in result["tools"][1]["parameters"] + + +class TestToolSchemaRegexPatternWiring: + """Claude Code's Artifact tool reaches /v1/responses (the /v1/messages bridge) with an + ECMA-262 ``pattern``; OpenAI compiles patterns with Python ``re`` and 400s + "'...' is not a 'regex'" for every model family, so the keyword is dropped. + """ + + def _artifact_tool(self): + return { + "type": "function", + "name": "Artifact", + "parameters": { + "type": "object", + "properties": { + "field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}, + "doc_id": {"type": "string", "pattern": r"^(?!\.\.?(?:/|$))[A-Za-z0-9_\-.~:@+]{1,200}$"}, + }, + "required": ["field"], + }, + } + + @pytest.mark.parametrize("model", ["gpt-5.6", "gpt-4o", "o3"]) + def test_openai_drops_only_the_pattern_python_re_rejects_for_every_family(self, model): + tool = self._artifact_tool() + + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model=model, + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + properties = result["tools"][0]["parameters"]["properties"] + assert properties["field"] == {"type": "string"} + assert properties["doc_id"] == tool["parameters"]["properties"]["doc_id"] + assert result["tools"][0]["parameters"]["required"] == ["field"] + assert tool["parameters"]["properties"]["field"]["pattern"] == _ARTIFACT_FIELD_PATTERN + assert json.loads(json.dumps(result["tools"])) == result["tools"] + + def test_openai_drops_patterns_inside_codex_namespace_tools(self): + namespace = {"type": "namespace", "name": "mcp__claude", "tools": [self._artifact_tool()]} + + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-5.6", + input="hi", + response_api_optional_request_params={"tools": [namespace]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0]["tools"][0]["parameters"]["properties"]["field"] == {"type": "string"} + + def test_openai_compact_request_drops_patterns(self): + _, data = OpenAIResponsesAPIConfig().transform_compact_response_api_request( + model="gpt-5.6", + input="hi", + response_api_optional_request_params={"tools": [self._artifact_tool()]}, + api_base="https://api.openai.com/v1/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert data["tools"][0]["parameters"]["properties"]["field"] == {"type": "string"} + + def test_non_openai_subclass_keeps_patterns(self): + from litellm.llms.hosted_vllm.responses.transformation import HostedVLLMResponsesAPIConfig + + tool = self._artifact_tool() + + result = HostedVLLMResponsesAPIConfig().transform_responses_api_request( + model="hosted_vllm/qwen", + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0] is tool + + +class TestReasoningFollowsModelSupport: + """Responses API clients like Codex send `reasoning` on every request, and OpenAI 400s it + on non-reasoning models like gpt-4o. drop_params must strip it there, the same way the + chat completions surface already strips reasoning_effort for those models. + """ + + @pytest.mark.parametrize( + "model, reasoning_survives", + [ + ("gpt-4o", False), + ("gpt-4.1", False), + ("gpt-4o-mini", False), + ("ft:gpt-4o-2024-08-06:my-org::abc123", False), + ("chat-latest", True), + ("gpt-5.6", True), + ("o3", True), + ("o3-deep-research", True), + ("o4-mini-deep-research", True), + ("codex-mini-latest", True), + ("ft:o4-mini-2025-04-16:my-org::abc123", True), + ("computer-use-preview", True), + ], + ) + def test_drop_params_strips_reasoning_by_model(self, local_model_cost_map, model, reasoning_survives): + mapped = OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": "auto"}}, + model=model, + drop_params=True, + ) + assert ("reasoning" in mapped) is reasoning_survives + + @pytest.mark.parametrize( + "model, reasoning_survives", + [ + ("gpt-4o", False), + ("chat-latest", True), + ("o3", True), + ("o3-deep-research", True), + ], + ) + def test_a_cost_map_older_than_this_release_never_strips_a_known_reasoning_model( + self, local_model_cost_map, monkeypatch, model, reasoning_survives + ): + lagging = { + name: {field: value for field, value in entry.items() if field != "supports_reasoning"} + for name, entry in litellm.model_cost.items() + } + monkeypatch.setattr(litellm, "model_cost", lagging) + mapped = OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium"}}, + model=model, + drop_params=True, + ) + assert ("reasoning" in mapped) is reasoning_survives + + def test_without_drop_params_the_error_is_litellms_400(self, local_model_cost_map, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as excinfo: + OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium"}}, + model="gpt-4o", + drop_params=False, + ) + assert excinfo.value.status_code == 400 + assert "reasoning.effort" in str(excinfo.value) + assert "cost map" in str(excinfo.value) + + @pytest.mark.parametrize("drop_params", [True, False]) + @pytest.mark.parametrize( + "reasoning", + [{"summary": "auto"}, {"effort": None, "summary": "auto"}, {}], + ) + def test_reasoning_without_an_effort_passes_through_on_non_reasoning_models( + self, local_model_cost_map, monkeypatch, drop_params, reasoning + ): + monkeypatch.setattr(litellm, "drop_params", drop_params) + mapped = OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": dict(reasoning)}, + model="gpt-4o", + drop_params=drop_params, + ) + assert mapped["reasoning"] == reasoning + + def test_an_explicit_supports_reasoning_false_beats_the_bundled_floor(self, local_model_cost_map, monkeypatch): + overridden = { + name: ({**entry, "supports_reasoning": False} if name == "o3" else entry) + for name, entry in litellm.model_cost.items() + } + monkeypatch.setattr(litellm, "model_cost", overridden) + mapped = OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium"}}, + model="o3", + drop_params=True, + ) + assert "reasoning" not in mapped + + def test_azure_deployments_keep_reasoning_even_on_a_non_reasoning_model_name(self, local_model_cost_map): + mapped = AzureOpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium"}}, + model="gpt-4o", + drop_params=True, + ) + assert mapped["reasoning"] == {"effort": "medium"} diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index c86ce4df2ac..b538fad71a2 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -276,6 +276,17 @@ def test_gpt5_drops_reasoning_effort_xhigh_when_requested(config: OpenAIConfig): assert "reasoning_effort" not in params +def test_gpt5_1_gpt5_2_gpt5_4_drop_minimal_reasoning_effort(config: OpenAIConfig): + for model in ("gpt-5.1", "gpt-5.2", "gpt-5.4", "gpt-5.4-pro"): + params = config.map_openai_params( + non_default_params={"reasoning_effort": "minimal"}, + optional_params={}, + model=model, + drop_params=True, + ) + assert "reasoning_effort" not in params + + # GPT-5.1 temperature handling tests def test_gpt5_1_model_detection(gpt5_config: OpenAIGPT5Config): """Test that models supporting reasoning_effort='none' are correctly detected via model map.""" @@ -388,26 +399,26 @@ def test_gpt5_4_mini_allows_reasoning_effort_none(config: OpenAIConfig): assert params["reasoning_effort"] == "none" -def test_gpt5_4_allows_reasoning_effort_minimal(config: OpenAIConfig): - """gpt-5.4 supports reasoning_effort='minimal'.""" - params = config.map_openai_params( - non_default_params={"reasoning_effort": "minimal"}, - optional_params={}, - model="gpt-5.4", - drop_params=False, - ) - assert params["reasoning_effort"] == "minimal" +def test_gpt5_4_rejects_reasoning_effort_minimal(config: OpenAIConfig): + """gpt-5.4 rejects reasoning_effort='minimal'.""" + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"reasoning_effort": "minimal"}, + optional_params={}, + model="gpt-5.4", + drop_params=False, + ) -def test_gpt5_4_pro_allows_reasoning_effort_minimal(config: OpenAIConfig): - """gpt-5.4-pro supports reasoning_effort='minimal'.""" - params = config.map_openai_params( - non_default_params={"reasoning_effort": "minimal"}, - optional_params={}, - model="gpt-5.4-pro", - drop_params=False, - ) - assert params["reasoning_effort"] == "minimal" +def test_gpt5_4_pro_rejects_reasoning_effort_minimal(config: OpenAIConfig): + """gpt-5.4-pro rejects reasoning_effort='minimal'.""" + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"reasoning_effort": "minimal"}, + optional_params={}, + model="gpt-5.4-pro", + drop_params=False, + ) def test_gpt5_4_mini_rejects_reasoning_effort_minimal(config: OpenAIConfig): @@ -468,13 +479,13 @@ def test_gpt5_minimal_dict_triggers_validation(config: OpenAIConfig): def test_gpt5_minimal_dict_accepted_for_supported_model(config: OpenAIConfig): - """Dict with effort='minimal' passes through for gpt-5.4+.""" + """Dict with effort='minimal' passes through for gpt-5.""" params = config.map_openai_params( non_default_params={ "reasoning_effort": {"effort": "minimal", "summary": "detailed"} }, optional_params={}, - model="gpt-5.4", + model="gpt-5", drop_params=False, ) assert params["reasoning_effort"] == "minimal" @@ -482,8 +493,8 @@ def test_gpt5_minimal_dict_accepted_for_supported_model(config: OpenAIConfig): def test_gpt5_supports_reasoning_effort_level_minimal(gpt5_config: OpenAIGPT5Config): """Test that _supports_reasoning_effort_level correctly identifies minimal support.""" - assert gpt5_config._supports_reasoning_effort_level("gpt-5.4", "minimal") - assert gpt5_config._supports_reasoning_effort_level("gpt-5.4-pro", "minimal") + assert not gpt5_config._supports_reasoning_effort_level("gpt-5.4", "minimal") + assert not gpt5_config._supports_reasoning_effort_level("gpt-5.4-pro", "minimal") assert not gpt5_config._supports_reasoning_effort_level("gpt-5.4-mini", "minimal") assert not gpt5_config._supports_reasoning_effort_level("gpt-5.4-nano", "minimal") @@ -504,10 +515,10 @@ def test_gpt5_minimal_explicitly_disabled_check(gpt5_config: OpenAIGPT5Config): assert gpt5_config._is_reasoning_effort_level_explicitly_disabled( "openai/gpt-5.4-mini", "minimal" ) - assert not gpt5_config._is_reasoning_effort_level_explicitly_disabled( + assert gpt5_config._is_reasoning_effort_level_explicitly_disabled( "gpt-5.4", "minimal" ) - assert not gpt5_config._is_reasoning_effort_level_explicitly_disabled( + assert gpt5_config._is_reasoning_effort_level_explicitly_disabled( "gpt-5.4-pro", "minimal" ) @@ -523,8 +534,8 @@ def test_is_explicitly_disabled_factory_minimal(): assert _is_explicitly_disabled_factory("gpt-5.4-mini", None, key) assert _is_explicitly_disabled_factory("gpt-5.4-nano", None, key) assert _is_explicitly_disabled_factory("openai/gpt-5.4-mini", None, key) - assert not _is_explicitly_disabled_factory("gpt-5.4", None, key) - assert not _is_explicitly_disabled_factory("gpt-5.4-pro", None, key) + assert _is_explicitly_disabled_factory("gpt-5.4", None, key) + assert _is_explicitly_disabled_factory("gpt-5.4-pro", None, key) assert not _is_explicitly_disabled_factory("gpt-5.4-turbo-preview", None, key) diff --git a/tests/test_litellm/llms/openai_like/test_cognition_provider.py b/tests/test_litellm/llms/openai_like/test_cognition_provider.py index 5c71b60e08a..d392abc6cc5 100644 --- a/tests/test_litellm/llms/openai_like/test_cognition_provider.py +++ b/tests/test_litellm/llms/openai_like/test_cognition_provider.py @@ -110,22 +110,6 @@ class TestCognitionProviderIdentity: class TestCognitionCostTracking: - @pytest.mark.parametrize( - "model, input_cost, output_cost, cache_read_cost", - [ - ("cognition/swe-1.6", 5e-07, 2.5e-06, 2e-07), - ("cognition/swe-1.7", 5e-07, 2.5e-06, 2e-07), - ("cognition/swe-1.7-lightning", 2.5e-06, 1.25e-05, 1e-06), - ], - ) - def test_cost_map_entries(self, model: str, input_cost: float, output_cost: float, cache_read_cost: float): - info = litellm.get_model_info(model=model) - - assert info["litellm_provider"] == "cognition" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cache_read_cost @pytest.mark.parametrize( "model, expected_prompt_cost, expected_completion_cost", diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/test_litellm/llms/openai_like/test_json_providers.py index c8743e1809d..d84cc8d3237 100644 --- a/tests/test_litellm/llms/openai_like/test_json_providers.py +++ b/tests/test_litellm/llms/openai_like/test_json_providers.py @@ -2,10 +2,9 @@ Tests for JSON-based provider configuration system. """ -import json import os import sys -from unittest.mock import MagicMock, patch +from unittest.mock import patch try: import pytest @@ -318,25 +317,6 @@ class TestDarkbloom: assert config is not None assert config.custom_llm_provider == "darkbloom" - def test_darkbloom_model_cost_map(self): - with open( - os.path.join(workspace_path, "model_prices_and_context_window.json") - ) as f: - model_cost = json.load(f) - - expected_models = { - "darkbloom/gemma-4-26b": (3e-08, 1.65e-07), - "darkbloom/gpt-oss-20b": (1.45e-08, 7e-08), - } - for model, (input_cost, output_cost) in expected_models.items(): - assert model in model_cost - assert model_cost[model]["litellm_provider"] == "darkbloom" - assert model_cost[model]["max_output_tokens"] == 32768 - assert model_cost[model]["supports_function_calling"] is True - assert model_cost[model]["supports_tool_choice"] is True - assert model_cost[model]["input_cost_per_token"] == input_cost - assert model_cost[model]["output_cost_per_token"] == output_cost - class TestPublicAIIntegration: """Integration tests for PublicAI provider""" diff --git a/tests/test_litellm/llms/openai_like/test_libertai_provider.py b/tests/test_litellm/llms/openai_like/test_libertai_provider.py index fdbe3046e9b..c17eaf7c87f 100644 --- a/tests/test_litellm/llms/openai_like/test_libertai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_libertai_provider.py @@ -59,23 +59,6 @@ class TestLibertAIProviderConfig: assert api_base == "https://custom.example.com/v1" assert api_key == "sk-test" - def test_libertai_model_cost_map(self): - """Test that libertai models are present in the model cost map""" - model_cost = litellm.model_cost - - assert "libertai/qwen3.6-27b" in model_cost - info = model_cost["libertai/qwen3.6-27b"] - assert info["litellm_provider"] == "libertai" - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 262144 - assert info["max_output_tokens"] == 262144 - - # thinking variants are marked as reasoning models - assert ( - model_cost["libertai/qwen3.6-27b-thinking"].get("supports_reasoning") - is True - ) - def test_libertai_router_config(self): """Test that libertai can be used in Router configuration""" from litellm import Router @@ -95,20 +78,6 @@ class TestLibertAIProviderConfig: assert len(router.model_list) == 1 assert router.model_list[0]["model_name"] == "libertai-chat" - def test_libertai_model_modes(self): - """Chat models carry mode 'chat'; the embedding model carries mode 'embedding'.""" - model_cost = litellm.model_cost - - # chat model - assert model_cost["libertai/qwen3.6-27b"]["mode"] == "chat" - - # embedding model (bge-m3) must be normalized to mode 'embedding' so - # /embeddings routing and the supported-endpoints matrix stay consistent - assert "libertai/bge-m3" in model_cost - bge = model_cost["libertai/bge-m3"] - assert bge["litellm_provider"] == "libertai" - assert bge["mode"] == "embedding" - def test_libertai_supported_endpoints_matrix(self): """The runtime-served backup matrix (GET /public/supported_endpoints) lists libertai.""" import json diff --git a/tests/test_litellm/llms/openai_like/test_meta_provider.py b/tests/test_litellm/llms/openai_like/test_meta_provider.py index 11b78828da6..c79e4b77cc5 100644 --- a/tests/test_litellm/llms/openai_like/test_meta_provider.py +++ b/tests/test_litellm/llms/openai_like/test_meta_provider.py @@ -193,19 +193,6 @@ class TestMetaAnthropicMessages: class TestMuseSparkModelInfo: - def test_muse_spark_pricing_and_capabilities(self): - info = litellm.get_model_info("meta/muse-spark-1.1") - - assert info["litellm_provider"] == "meta" - assert info["input_cost_per_token"] == 1.25e-06 - assert info["output_cost_per_token"] == 4.25e-06 - assert info["cache_read_input_token_cost"] == 1.5e-07 - assert info["max_input_tokens"] == 1048576 - assert info["supports_reasoning"] is True - assert info["supports_web_search"] is True - assert info["supports_vision"] is True - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True def test_muse_spark_cost_calculation(self): from litellm import completion_cost diff --git a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py index 6a6271e95e2..6ca7072e7ab 100644 --- a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py +++ b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py @@ -3,7 +3,6 @@ Unit tests for Perplexity embedding transformation logic. """ import base64 -import json import struct from unittest.mock import MagicMock @@ -298,25 +297,3 @@ class TestPerplexityEmbeddingProviderConfig: ) assert config is not None assert isinstance(config, PerplexityEmbeddingConfig) - - -class TestPerplexityEmbeddingModelInfo: - """Test that Perplexity embedding models are in model_prices_and_context_window.""" - - def test_model_info_available(self): - import litellm - - info = litellm.get_model_info("perplexity/pplx-embed-v1-0.6b") - assert info is not None - assert info["mode"] == "embedding" - assert info["max_input_tokens"] == 32768 - assert info["output_vector_size"] == 1024 - - def test_model_info_4b_available(self): - import litellm - - info = litellm.get_model_info("perplexity/pplx-embed-v1-4b") - assert info is not None - assert info["mode"] == "embedding" - assert info["max_input_tokens"] == 32768 - assert info["output_vector_size"] == 2560 diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py index 6630039e92e..7556b215e66 100644 --- a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py @@ -26,7 +26,6 @@ from litellm.types.utils import ( Usage, PromptTokensDetailsWrapper, ) -from litellm.utils import get_model_info class TestPerplexityCostCalculator: @@ -317,21 +316,6 @@ class TestPerplexityCostCalculator: assert math.isclose(total_cost, expected_total, rel_tol=1e-6) - def test_model_info_access(self): - """Test that model info correctly returns the new cost fields.""" - model_info = get_model_info( - model="sonar-deep-research", custom_llm_provider="perplexity" - ) - - # Check that the new fields are accessible - assert "citation_cost_per_token" in model_info - assert model_info["citation_cost_per_token"] == 2e-6 - assert model_info["search_context_cost_per_query"] == { - "search_context_size_low": 0.005, - "search_context_size_medium": 0.005, - "search_context_size_high": 0.005, - } - @pytest.mark.parametrize("citation_tokens", [0, 10, 25, 100]) @pytest.mark.parametrize("search_queries", [0, 1, 5, 10]) @pytest.mark.parametrize("reasoning_tokens", [0, 15, 30]) @@ -477,37 +461,6 @@ class TestPerplexityCostCalculator: assert math.isclose(prompt_cost, expected_prompt, rel_tol=1e-9) assert math.isclose(completion_cost, expected_completion, rel_tol=1e-9) - @pytest.mark.parametrize( - "model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read", - [ - ("deepseek-v4-flash-0731", 0.13, 0.26, 0.028), - ("glm-5.2", 1.4, 4.4, 0.14), - ("kimi-k3", 3.0, 15.0, 0.3), - ("kimi-k2.7-code", 0.95, 4.0, 0.19), - ], - ) - def test_agent_api_entries_carry_perplexity_published_rates( - self, model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read - ): - """The Agent API third-party models are priced from Perplexity's own catalog - (GET https://api.perplexity.ai/v1/models, `pricing` in usd_per_1m_tokens). - Perplexity's model id already starts with `perplexity/`, so the cost-map key - doubles the prefix. Regression: glm-5.2 shipped glm-5.3's 0.26 cache-read rate, - copied from the neighbouring catalog row, an 86% overcharge on cached input. - """ - info = get_model_info( - model=f"perplexity/{model_id}", custom_llm_provider="perplexity" - ) - - assert info["key"] == f"perplexity/perplexity/{model_id}" - assert info["litellm_provider"] == "perplexity" - assert info["mode"] == "responses" - assert math.isclose(info["input_cost_per_token"], usd_per_1m_input / 1e6, rel_tol=1e-9) - assert math.isclose(info["output_cost_per_token"], usd_per_1m_output / 1e6, rel_tol=1e-9) - assert math.isclose( - info["cache_read_input_token_cost"], usd_per_1m_cache_read / 1e6, rel_tol=1e-9 - ) - def test_agent_api_fallback_rates_price_a_response_without_metered_cost(self): """Perplexity meters cost on the response, but when `usage.cost` is absent the calculator falls back to the mapped per-token rates. Regression: that fallback diff --git a/tests/test_litellm/llms/reducto/test_parse_v3.py b/tests/test_litellm/llms/reducto/test_parse_v3.py index 140b9737dc0..1d0c826ef8b 100644 --- a/tests/test_litellm/llms/reducto/test_parse_v3.py +++ b/tests/test_litellm/llms/reducto/test_parse_v3.py @@ -1,8 +1,9 @@ import json -import litellm import pytest +import litellm + def _reducto_parse_response() -> dict: return { @@ -68,15 +69,11 @@ def disable_aiohttp_transport(): @pytest.mark.asyncio -async def test_parse_v3_file_upload_and_response_mapping( - disable_aiohttp_transport, respx_mock -): +async def test_parse_v3_file_upload_and_response_mapping(disable_aiohttp_transport, respx_mock): upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( json={"file_id": "reducto://uploaded.pdf"} ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( - json=_reducto_parse_response() - ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(json=_reducto_parse_response()) response = await litellm.aocr( model="reducto/parse-v3", @@ -123,15 +120,11 @@ async def test_parse_v3_file_upload_and_response_mapping( @pytest.mark.asyncio -async def test_parse_v3_reducto_id_passthrough_skips_upload( - disable_aiohttp_transport, respx_mock -): +async def test_parse_v3_reducto_id_passthrough_skips_upload(disable_aiohttp_transport, respx_mock): upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( json={"file_id": "reducto://should-not-upload.pdf"} ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( - json=_reducto_parse_response() - ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(json=_reducto_parse_response()) response = await litellm.aocr( model="reducto/parse-v3", @@ -150,3 +143,28 @@ async def test_parse_v3_reducto_id_passthrough_skips_upload( assert parse_request_body["input"] == "reducto://already-uploaded.pdf" assert parse_request_body["retrieval"]["chunk_mode"] == "section" assert response.pages[0].markdown.startswith("Page 1 block A") + + +@pytest.mark.asyncio +async def test_unknown_model_uses_current_protocol_without_local_rejection( + disable_aiohttp_transport, respx_mock +): + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( + json=_reducto_parse_response() + ) + + response = await litellm.aocr( + model="reducto/future-parse-model", + document={ + "type": "document_url", + "document_url": "reducto://already-uploaded.pdf", + }, + api_key="test-key", + api_base="https://platform.reducto.ai", + ) + + assert parse_route.called + assert json.loads(parse_route.calls[0].request.read()) == { + "input": "reducto://already-uploaded.pdf" + } + assert response.model == "future-parse-model" diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py index aa1a59d0e5c..bb891c06fa2 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py @@ -46,3 +46,45 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch): assert credentials.token == "assumed-session-token" assert aws_region_name == "us-east-1" assert "aws_external_id" not in optional_params + + +def test_load_credentials_assumes_role_with_session_tags(monkeypatch): + """A trust policy gated on sts:TagSession only admits the session when the deployment's tags are sent.""" + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIASMCHATTAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIASMCHATCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-sm-chat-role", + "aws_session_name": "litellm-sm-chat-session", + "aws_session_tags": tags, + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = SagemakerChatHandler()._load_credentials(optional_params) + + assert credentials.access_key == "ASIASMCHATTAGGED" + assert aws_region_name == "us-east-1" + assert "aws_session_tags" not in optional_params diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py index 881bac096b1..9a5f36e2081 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py @@ -219,3 +219,51 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch): assert credentials.token == "assumed-session-token" assert aws_region_name == "us-east-1" assert "aws_external_id" not in optional_params + + +def test_load_credentials_assumes_role_with_session_tags(monkeypatch): + """A trust policy gated on sts:TagSession only admits the session when the deployment's tags are sent.""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + from unittest.mock import patch + + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIASMCOMPTAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIASMCOMPCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-sm-completion-role", + "aws_session_name": "litellm-sm-completion-session", + "aws_session_tags": tags, + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = SagemakerLLM()._load_credentials(optional_params) + + assert credentials.access_key == "ASIASMCOMPTAGGED" + assert aws_region_name == "us-east-1" + assert "aws_session_tags" not in optional_params diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index 25a961c3413..5687a319f06 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -6,6 +6,7 @@ Tests tool calling request/response transformations and chat completions import asyncio import os import copy +import uuid import json from typing import Any, Dict, List @@ -945,3 +946,48 @@ class TestSnowflakeChatCompletion: assert len(chunks_received) > 0 content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content) + + +async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="snowflake/claude-sonnet-4-6", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + api_key="fake-jwt", + account_id="FAKE-ACCOUNT", + api_base=FAKE_API_BASE, + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py index eadd87d9c92..08e46b1ffac 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py @@ -322,8 +322,8 @@ class TestModelCostEntry: entry = json.load(f)["vertex_ai/gemini-3.5-transcribe-preview"] assert entry["mode"] == "audio_transcription" assert entry["litellm_provider"] == "vertex_ai" - assert entry["input_cost_per_audio_token"] == pytest.approx(2.5e-06) - assert entry["input_cost_per_token"] == pytest.approx(2.5e-06) + assert entry["input_cost_per_audio_token"] == pytest.approx(2e-06) + assert entry["input_cost_per_token"] == pytest.approx(2e-06) assert entry["output_cost_per_token"] == pytest.approx(1.2e-05) assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"] diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 38fde3caa63..24df3214da3 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -35,7 +35,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest - from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) @@ -178,6 +177,197 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client(): sync_client.post.assert_not_called() +def test_create_batch_sync_does_not_resolve_publisher_models(): + """Publisher-model jobs must not incur the endpoint-resolution GET, and the job model must + stay the publisher path untouched.""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get") as safe_get, + ): + out = h.create_batch( + _is_async=False, + create_batch_data=CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert isinstance(out, LiteLLMBatch) + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == "publishers/google/models/gemini-1.5-flash-001" + safe_get.assert_not_called() + + +ENDPOINT_ID = "7768560373388541952" +ENDPOINT_CREATE_DATA = { + "input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid" +} +TUNED_MODEL_RESOURCE = f"projects/{PROJECT}/locations/{LOCATION}/models/1234509876" + + +def _endpoint_get_response(deployed_models: list | None = None) -> MagicMock: + resp = MagicMock() + resp.status_code = 200 + resp.json.return_value = { + "name": f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + "deployedModels": ( + deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}] + ), + } + return resp + + +def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): + """A fine-tuned Gemini file id must produce a batch job against the endpoint's deployed + tuned model resource; the v1 batch API rejects endpoint resources in `model` (LIT-6899).""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response()) as safe_get, + ): + out = h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert isinstance(out, LiteLLMBatch) + get_args, get_kwargs = safe_get.call_args + assert get_args[1] == ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}" + ) + assert get_kwargs["headers"]["Authorization"] == f"Bearer {TOKEN}" + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == TUNED_MODEL_RESOURCE + + +@pytest.mark.parametrize( + "api_base, expected", + [ + ( + None, + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal", + f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal/v1", + f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal/vertex", + f"https://proxy.internal/vertex/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ], +) +def test_build_endpoint_resolution_url(api_base, expected): + """A custom api_base must replace the Google host for the endpoint-resolution GET without + producing a malformed url (no ':' grafting, no doubled /v1).""" + url = VertexAIBatchPrediction._build_endpoint_resolution_url( + api_base=api_base, + model=f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + vertex_location=LOCATION, + ) + assert url == expected + + +def test_create_batch_sync_endpoint_resolution_error_raises(): + h = _make_handler() + client = MagicMock() + resolve_response = MagicMock() + resolve_response.status_code = 404 + resolve_response.text = "endpoint not found" + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=resolve_response), + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 404 + client.post.assert_not_called() + + +def test_create_batch_custom_endpoint_raises_400_without_io(): + """custom_endpoint deployments have no Vertex batch surface; creating a job would target a + nonexistent publisher model, so the handler must 400 before any auth or HTTP work (LIT-6899).""" + h = _make_handler() + client = MagicMock() + + with patch(f"{HMOD}._get_httpx_client", return_value=client): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + assert exc_info.value.status_code == 400 + assert "custom_endpoint" in str(exc_info.value) + h._ensure_access_token.assert_not_called() + client.post.assert_not_called() + + +def test_create_batch_sync_endpoint_without_deployed_model_raises_400(): + h = _make_handler() + client = MagicMock() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response(deployed_models=[])), + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + assert "no deployed model" in str(exc_info.value) + client.post.assert_not_called() + + def test_create_batch_sync_httpstatuserror_propagates(): """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the sync create path must surface that error, not swallow it.""" diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index ccb2d7e310d..232c6413e78 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -34,6 +34,12 @@ INPUT_FILE = ( "models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8" ) +ENDPOINT_ID = "7768560373388541952" +ENDPOINT_INPUT_FILE = ( + f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/" + "e9412502-2c91-42a6-8e61-f5c294cc0fc8" +) + # =========================================================================== # # transform_openai_batch_request_to_vertex_ai_batch_request @@ -67,6 +73,41 @@ def test_transform_openai_request_missing_input_file_id_raises(): T.transform_openai_batch_request_to_vertex_ai_batch_request({}) +def test_transform_openai_request_fine_tuned_endpoint_builds_endpoint_resource(): + """A fine-tuned Gemini file id (endpoints/) must target the endpoint resource, + not a nonexistent publisher model (LIT-6899).""" + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": ENDPOINT_INPUT_FILE}, + vertex_project="my-project", + vertex_location="us-central1", + ) + assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}" + + +def test_transform_openai_request_fine_tuned_endpoint_defaults_location(): + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": ENDPOINT_INPUT_FILE}, + vertex_project="my-project", + ) + assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}" + + +def test_transform_openai_request_fine_tuned_endpoint_without_project_raises_400(): + with pytest.raises(VertexAIError) as exc_info: + T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": ENDPOINT_INPUT_FILE}) + assert exc_info.value.status_code == 400 + assert "vertex_project" in str(exc_info.value) + + +def test_transform_openai_request_publisher_model_ignores_project_and_location(): + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": INPUT_FILE}, + vertex_project="my-project", + vertex_location="europe-west4", + ) + assert job["model"] == "publishers/google/models/gemini-1.5-flash-001" + + @pytest.mark.parametrize( "input_file_id", [ @@ -321,6 +362,35 @@ def test_get_model_from_gcs_file_no_publishers_raises_400(): assert exc_info.value.status_code == 400 +def test_get_model_from_gcs_file_fine_tuned_endpoint(): + """The whole endpoint id must survive parsing; the old 3-segment publishers/ parse dropped it.""" + assert T._get_model_from_gcs_file(ENDPOINT_INPUT_FILE) == f"endpoints/{ENDPOINT_ID}" + + +def test_get_model_from_gcs_file_publisher_path_wins_over_endpoints_prefix(): + """A bucket prefix containing endpoints/ must not override the publisher model path + LiteLLM appended after it.""" + uri = "gs://bucket/team-endpoints/999/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/uuid" + assert T._get_model_from_gcs_file(uri) == "publishers/google/models/gemini-1.5-flash-001" + + +def test_get_model_from_gcs_file_last_endpoints_segment_wins(): + """With no publisher path, the endpoint id closest to the file (last occurrence) is the one + LiteLLM stored; an earlier prefix segment must not shadow it.""" + uri = f"gs://bucket/endpoints/999/litellm-vertex-files/endpoints/{ENDPOINT_ID}/uuid" + assert T._get_model_from_gcs_file(uri) == f"endpoints/{ENDPOINT_ID}" + + +def test_get_model_from_gcs_file_non_numeric_endpoints_segment_raises_400(): + with pytest.raises(VertexAIError) as exc_info: + T._get_model_from_gcs_file("gs://bucket/endpoints/not-a-number/file-uuid") + assert exc_info.value.status_code == 400 + + +def test_get_bare_model_name_from_gcs_file_fine_tuned_endpoint(): + assert T.get_bare_model_name_from_gcs_file(ENDPOINT_INPUT_FILE) == ENDPOINT_ID + + # =========================================================================== # # is_unmanaged_gcs_batch_input_file_id # =========================================================================== # @@ -334,6 +404,8 @@ def test_get_model_from_gcs_file_no_publishers_raises_400(): ("file-abc123", False), ("gs://bucket/no-model-here.jsonl", False), ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False), + (ENDPOINT_INPUT_FILE, True), + ("gs://bucket/endpoints/not-a-number/file-uuid", False), ], ) def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected): diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 3c2d56997b7..8a249820cbd 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -159,6 +159,91 @@ class TestCreateFileUrl: assert "?" not in object_name +class TestBatchObjectNaming: + def test_should_store_publisher_model_under_publishers_path(self, config): + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini-2.5-flash"}}]) + assert object_name.startswith("litellm-vertex-files/publishers/google/models/gemini-2.5-flash/") + + def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config): + """A numeric endpoint id must not be filed under publishers/google/models/gemini/, + which the batch transformation later mangles into a nonexistent publisher model (LIT-6899).""" + object_name = config._get_gcs_object_name_from_batch_jsonl( + [{"body": {"model": "gemini/7768560373388541952"}}] + ) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "publishers" not in object_name + + def test_should_store_bare_numeric_endpoint_under_endpoints_path(self, config): + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "7768560373388541952"}}]) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + + def test_deployment_model_overrides_jsonl_body_model(self, config): + """The stored path decides which Vertex model the batch later runs against with the + deployment's credentials, so a user-crafted JSONL body.model must not be able to redirect + an authorized deployment to a different endpoint.""" + object_name = config._get_gcs_object_name_from_batch_jsonl( + [{"body": {"model": "9999999999999999999"}}], + deployment_model="vertex_ai/gemini/7768560373388541952", + ) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "9999999999999999999" not in object_name + + def test_url_derives_object_path_from_configured_model(self, config): + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={ + "gcs_bucket_name": "my-bucket", + "model": "vertex_ai/gemini/7768560373388541952", + }, + data={ + "file": ("batch.jsonl", b'{"body": {"model": "9999999999999999999"}}', "application/jsonl"), + "purpose": "batch", + }, + ) + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "9999999999999999999" not in object_name + + +class TestCustomEndpointBatchUpload: + def test_should_reject_batch_upload_for_custom_endpoint_deployment(self, config): + """custom_endpoint deployments have no Vertex batch surface; the upload must 400 instead + of staging a file that can only produce a doomed batch job (LIT-6899).""" + from litellm.llms.vertex_ai.common_utils import VertexAIError + + with pytest.raises(VertexAIError) as exc_info: + config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True}, + data={ + "file": ("batch.jsonl", b'{"body": {"model": "openai/gemma-2-2b-it"}}', "application/jsonl"), + "purpose": "batch", + }, + ) + assert exc_info.value.status_code == 400 + assert "custom_endpoint" in str(exc_info.value) + + def test_should_allow_non_batch_upload_for_custom_endpoint_deployment(self, config): + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True}, + data={ + "file": ("notes.txt", b"hello", "text/plain"), + "purpose": "assistants", + }, + ) + assert "/b/my-bucket/" in url + + class TestTransformRetrieveFile: def test_should_build_correct_gcs_metadata_url(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py index fad310fc5c0..f135acd094f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py @@ -1,6 +1,13 @@ +import json +import uuid +from unittest.mock import Mock + +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.gemini import transformation from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -338,3 +345,133 @@ def test_map_function_enterprise_web_search_snake_case(): assert len(result) == 1 assert "enterpriseWebSearch" in result[0] + + +async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "Green"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + api_key="fake-gemini-key", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] + + +async def test_gemini_ai_studio_async_completion_passes_files_api_uris_through_unfetched(async_only_image_fetch): + files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + files_api_image = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "A report"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize these"}, + {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": files_api_image, "format": "image/png"}}, + ], + } + ], + api_key="fake-gemini-key", + client=client, + ) + + assert response.choices[0].message.content == "A report" + assert async_only_image_fetch.fetched == [] + file_parts = [part["file_data"] for part in captured["body"]["contents"][0]["parts"] if "file_data" in part] + assert file_parts == [ + {"mime_type": "application/pdf", "file_uri": files_api_pdf}, + {"mime_type": "image/png", "file_uri": files_api_image}, + ] + + +async def test_vertex_ai_async_transform_inlines_only_the_urls_gemini_cannot_fetch_itself(async_only_image_fetch): + plain_http_png = f"http://img.example/{uuid.uuid4()}.png" + extensionless_https = f"https://cdn.example/files/{uuid.uuid4().hex}" + https_png = f"https://img.example/{uuid.uuid4()}.png" + hinted_extensionless = f"https://cdn.example/files/{uuid.uuid4().hex}" + files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe these"}, + {"type": "image_url", "image_url": {"url": plain_http_png}}, + {"type": "image_url", "image_url": {"url": extensionless_https}}, + {"type": "image_url", "image_url": {"url": https_png}}, + {"type": "image_url", "image_url": {"url": hinted_extensionless, "mime_type": "image/webp"}}, + {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}}, + ], + } + ] + + body = await transformation.async_transform_request_body( + gemini_api_key=None, + messages=messages, + api_base=None, + model="gemini-3.8-flash", + client=None, + timeout=None, + extra_headers=None, + optional_params={}, + logging_obj=Mock(), + custom_llm_provider="vertex_ai", + litellm_params={}, + vertex_project="qa-project", + vertex_location="us-central1", + vertex_auth_header=None, + ) + + inlined = {"inline_data": {"mime_type": "image/png", "data": async_only_image_fetch.base64_png}} + assert body["contents"][0]["parts"] == [ + {"text": "Describe these"}, + inlined, + inlined, + {"file_data": {"mime_type": "image/png", "file_uri": https_png}}, + {"file_data": {"mime_type": "image/webp", "file_uri": hinted_extensionless}}, + {"file_data": {"mime_type": "application/pdf", "file_uri": files_api_pdf}}, + ] + assert sorted(async_only_image_fetch.fetched) == sorted([plain_http_png, extensionless_https]) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index d2788408e09..101f6e6fa5d 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5769,3 +5769,70 @@ def test_calculate_web_search_requests_counts_unique_queries(): assert VertexGeminiConfig._calculate_web_search_requests([]) is None assert VertexGeminiConfig._calculate_web_search_requests([{"webSearchQueries": ["", ""]}]) is None + + +@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai"]) +@pytest.mark.parametrize( + "model", + ["gemini-2.5-flash", "gemini-3-pro-preview"], + ids=["thinking_budget_mapper", "thinking_level_mapper"], +) +@pytest.mark.parametrize("reasoning_effort", ["banana", "xhigh"]) +def test_invalid_reasoning_effort_is_a_400_not_a_500(custom_llm_provider, model, reasoning_effort): + """Regression for #40474. + + Both reasoning_effort mappers used to end their if/elif chain in a bare `ValueError`, which + `exception_type()` has no branch for, so it fell through to `APIConnectionError` and the proxy + answered a malformed client request with a retryable HTTP 500. `xhigh` is covered alongside the + nonsense value because it is a member of litellm's own `REASONING_EFFORT` literal, so callers + bridging from OpenAI-shaped code reach it without typing anything wrong. + """ + from litellm.utils import get_optional_params + + with pytest.raises(litellm.BadRequestError) as exc_info: + get_optional_params( + model=model, + custom_llm_provider=custom_llm_provider, + reasoning_effort=reasoning_effort, + drop_params=True, + ) + + assert exc_info.value.status_code == 400 + message: Final = str(exc_info.value) + assert reasoning_effort in message + for supported in ("minimal", "low", "medium", "high", "none", "disable"): + assert supported in message + + +@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai"]) +def test_invalid_reasoning_effort_surfaces_as_400_through_completion(custom_llm_provider): + """The same request through `completion()` must not come back as a retryable 500. + + Needs no provider credentials: param mapping runs before any network call. + """ + with pytest.raises(litellm.BadRequestError) as exc_info: + completion( + model=f"{custom_llm_provider}/gemini-3-pro-preview", + messages=[{"role": "user", "content": "hi"}], + reasoning_effort="banana", + ) + + assert exc_info.value.status_code == 400 + assert not isinstance(exc_info.value, litellm.APIConnectionError) + + +@pytest.mark.parametrize("model", ["gemini-2.5-flash", "gemini-3-pro-preview"]) +def test_supported_reasoning_efforts_still_map(model): + """Guards the fix against over-rejecting: every advertised value must still produce a config.""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + SUPPORTED_REASONING_EFFORTS, + ) + + for effort in SUPPORTED_REASONING_EFFORTS: + result: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, + optional_params={}, + model=model, + drop_params=False, + ) + assert "thinkingConfig" in result diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index dddc95bf54a..c206fcec420 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -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") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py new file mode 100644 index 00000000000..98010021bca --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -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"] + ) diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index fba337b5f2c..b5eec42b569 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -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 diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py index fa286f6f609..98071594ebf 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py @@ -6,11 +6,15 @@ Vertex AI Anthropic models don't support URL sources for images. LiteLLM should convert image URLs to base64 when using Vertex AI Anthropic. """ +import json +import sys from unittest.mock import patch, MagicMock +import httpx import pytest - +import litellm +from litellm._uuid import uuid from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_messages_pt, convert_to_anthropic_tool_result, @@ -371,3 +375,59 @@ class TestToolMessageImageURLHandling: assert item["source"]["type"] == "url" return pytest.fail("Could not find image in tool result") + + +async def test_vertex_ai_anthropic_async_completion_inlines_https_images_off_the_event_loop(async_only_image_fetch): + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + image_url = f"https://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + vertexai = MagicMock() + vertexai.preview.language_models = MagicMock() + + with ( + patch.dict(sys.modules, {"vertexai": vertexai}), + patch.object( # test-quality-ok: litellm.acompletion has no seam for Vertex token minting + litellm.main.vertex_partner_models_chat_completion, + "_ensure_access_token", + return_value=("token", "test-project"), + ), + ): + response = await litellm.acompletion( + model="vertex_ai/claude-sonnet-4-6", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + vertex_project="test-project", + vertex_location="us-east5", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + sources = [part["source"] for part in captured["body"]["messages"][0]["content"] if part["type"] == "image"] + assert sources == [{"type": "base64", "media_type": "image/png", "data": async_only_image_fetch.base64_png}] diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 6ba8706b0d8..04e46eab1b7 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -136,19 +136,6 @@ class TestVertexAIVideoConfig: # Should NOT include endpoint assert not url.endswith(":predictLongRunning") - def test_veo_31_lite_model_cost_entries_match_pricing(self): - for path in (ROOT_MODEL_COST_PATH, BACKUP_MODEL_COST_PATH): - model_cost = _load_model_cost_map(path) - info = model_cost.get(VEO_31_LITE_VERTEX_MODEL) - - assert info is not None, f"{VEO_31_LITE_VERTEX_MODEL} missing from {path}" - assert info["litellm_provider"] == "vertex_ai-video-models" - assert info["mode"] == "video_generation" - assert info["max_input_tokens"] == 1024 - assert info["output_cost_per_second"] == 0.05 - assert info["output_cost_per_second_1080p"] == 0.08 - assert info["supported_modalities"] == ["text", "image"] - def test_veo_31_lite_provider_routing_from_local_model_map( self, monkeypatch: pytest.MonkeyPatch ): diff --git a/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py b/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py new file mode 100644 index 00000000000..d3e912ce6af --- /dev/null +++ b/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py @@ -0,0 +1,263 @@ +import json +from unittest.mock import MagicMock + +import pytest + + +class TestVoyageContextualEmbeddings: + def test_contextual_model_detection(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + assert VoyageContextualEmbeddingConfig.is_contextualized_embeddings("voyage-context-3") + assert VoyageContextualEmbeddingConfig.is_contextualized_embeddings("voyage-context-4") + assert not VoyageContextualEmbeddingConfig.is_contextualized_embeddings("voyage-3-lite") + + def test_url_generation(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + assert ( + config.get_complete_url(None, None, "voyage-context-4", {}, {}) + == "https://api.voyageai.com/v1/contextualizedembeddings" + ) + assert ( + config.get_complete_url("https://custom.api.com", None, "voyage-context-4", {}, {}) + == "https://custom.api.com/contextualizedembeddings" + ) + assert ( + config.get_complete_url( + "https://custom.api.com/contextualizedembeddings", + None, + "voyage-context-4", + {}, + {}, + ) + == "https://custom.api.com/contextualizedembeddings" + ) + + def test_get_supported_openai_params(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + assert config.get_supported_openai_params("voyage-context-4") == [ + "encoding_format", + "dimensions", + ] + + def test_map_openai_params(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + result = config.map_openai_params( + {"encoding_format": "float", "dimensions": 512}, {}, "voyage-context-4", False + ) + assert result["encoding_format"] == "float" + assert result["output_dimension"] == 512 + + def test_validate_environment_with_api_key(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + headers = config.validate_environment( + {}, "voyage-context-4", [], {}, {}, api_key="test-key" + ) + assert headers == {"Authorization": "Bearer test-key"} + + def test_validate_environment_secret_fallback(self, monkeypatch): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + monkeypatch.setenv("VOYAGE_API_KEY", "secret-key") + config = VoyageContextualEmbeddingConfig() + headers = config.validate_environment( + {}, "voyage-context-4", [], {}, {}, api_key=None + ) + assert headers == {"Authorization": "Bearer secret-key"} + + def test_nested_list_passthrough(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + nested = [["Hello", "world"], ["Test"]] + transformed = config.transform_embedding_request( + "voyage-context-4", nested, {}, {} + ) + assert transformed["inputs"] == nested + assert transformed["model"] == "voyage-context-4" + assert "enable_auto_chunking" not in transformed + + def test_flat_list_str_auto_chunked(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", ["Hello", "world"], {}, {} + ) + assert transformed["inputs"] == ["Hello", "world"] + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 32000 + assert transformed["input_type"] == "document" + + def test_flat_list_str_query_no_auto_chunk(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", ["Hello", "world"], {"input_type": "query"}, {} + ) + assert transformed["inputs"] == ["Hello", "world"] + assert transformed["input_type"] == "query" + assert "enable_auto_chunking" not in transformed + + def test_flat_list_str_document_preserves_input_type(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", ["Hello"], {"input_type": "document"}, {} + ) + assert transformed["input_type"] == "document" + assert transformed["enable_auto_chunking"] is True + + def test_flat_list_str_caller_chunk_params_win(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", + ["Hello", "world"], + {"input_type": "document", "chunk_size": 512, "chunk_overlap": 32}, + {}, + ) + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 512 + assert transformed["chunk_overlap"] == 32 + assert transformed["input_type"] == "document" + + def test_flat_list_str_caller_can_disable_auto_chunking(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", ["Hello"], {"enable_auto_chunking": False}, {} + ) + assert transformed["enable_auto_chunking"] is False + assert transformed["input_type"] == "document" + + def test_nested_list_keeps_caller_params(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", [["Hello", "world"]], {"input_type": "document", "output_dimension": 512}, {} + ) + assert transformed == { + "inputs": [["Hello", "world"]], + "model": "voyage-context-4", + "input_type": "document", + "output_dimension": 512, + } + + def test_single_string_auto_chunked(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", "Hello", {}, {} + ) + assert transformed["inputs"] == ["Hello"] + assert transformed["enable_auto_chunking"] is True + assert transformed["input_type"] == "document" + + def test_single_string_query_no_auto_chunk(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + transformed = config.transform_embedding_request( + "voyage-context-4", "Hello", {"input_type": "query"}, {} + ) + assert transformed["inputs"] == ["Hello"] + assert transformed["input_type"] == "query" + assert "enable_auto_chunking" not in transformed + + def test_response_transformation(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + from litellm.types.utils import EmbeddingResponse + + config = VoyageContextualEmbeddingConfig() + response_payload = { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "voyage-context-4", + "usage": {"total_tokens": 24}, + } + raw_response = MagicMock() + raw_response.json.return_value = response_payload + raw_response.status_code = 200 + raw_response.text = json.dumps(response_payload) + + model_response = EmbeddingResponse() + transformed = config.transform_embedding_response( + "voyage-context-4", raw_response, model_response, MagicMock() + ) + assert transformed.model == "voyage-context-4" + assert transformed.object == "list" + assert transformed.data == response_payload["data"] + assert transformed.usage.prompt_tokens == 24 + assert transformed.usage.total_tokens == 24 + + def test_error_response_and_error_class(self): + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + VoyageError, + ) + from litellm.types.utils import EmbeddingResponse + + config = VoyageContextualEmbeddingConfig() + raw_response = MagicMock() + raw_response.json.side_effect = ValueError("not json") + raw_response.status_code = 400 + raw_response.text = "bad request" + + with pytest.raises(VoyageError) as exc_info: + config.transform_embedding_response( + "voyage-context-4", raw_response, EmbeddingResponse(), MagicMock() + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.message == "bad request" + + error = config.get_error_class("rate limited", 429, {"x-test": "1"}) + assert isinstance(error, VoyageError) + assert error.status_code == 429 + assert error.message == "rate limited" diff --git a/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py b/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py index a5d1eccebe0..dd0d1bdbb9d 100644 --- a/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py +++ b/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py @@ -5,18 +5,92 @@ These tests validate the WandbInferenceConfig class which extends OpenAIGPTConfi Nebius AI Studio is an OpenAI-compatible provider with minor customizations. """ - +import json +from typing import Final import pytest +import respx import litellm from litellm import completion from litellm.llms.wandb.chat.transformation import WandbConfig +WANDB_REASONING_MODELS: Final = ( + "deepseek-ai/DeepSeek-V4-Flash", + "deepseek-ai/DeepSeek-V4-Flash-0731", + "deepseek-ai/DeepSeek-V4-Pro", + "deepseek-ai/DeepSeek-V4-Pro-0813", + "deepseek-ai/DeepSeek-V3.1", + "google/gemma-4-31B-it", + "ibm-granite/granite-4.2-8b", + "MiniMaxAI/MiniMax-M3", + "moonshotai/Kimi-K2.7-Code", + "moonshotai/Kimi-K2.6", + "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B", + "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B", + "openai/gpt-oss-120b", + "openai/gpt-oss-20b", + "Qwen/Qwen3.8-27B", + "Qwen/Qwen3.6-35B-A3B", + "Qwen/Qwen3.6-27B", + "Qwen/Qwen3.5-35B-A3B", + "zai-org/GLM-5.2", + "moonshotai/Kimi-K2.5", + "MiniMaxAI/MiniMax-M2.5", + "zai-org/GLM-4.5", + "Qwen/Qwen3-235B-A22B-Thinking-2507", + "deepseek-ai/DeepSeek-R1-0528", +) + + +@pytest.fixture +def wandb_test_config(local_model_cost_map, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "telemetry", False) + monkeypatch.setattr(litellm, "drop_params", False) + + +@pytest.fixture +def wandb_request_mock(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("https://api.inference.wandb.ai/v1/chat/completions").respond( + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "test-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Done"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + status_code=200, + ) + + class TestWandbConfig: """Test class for WandB Inference functionality""" + @pytest.mark.parametrize("model", WANDB_REASONING_MODELS) + def test_map_openai_params_preserves_reasoning_effort(self, wandb_test_config, model: str): + assert litellm.model_cost[f"wandb/{model}"].get("supports_reasoning") is True + supported_params = litellm.get_supported_openai_params(model=f"wandb/{model}") + assert supported_params is not None + assert "reasoning_effort" in supported_params + + result = WandbConfig().map_openai_params( + non_default_params={"reasoning_effort": "medium", "max_completion_tokens": 64}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert result == {"reasoning_effort": "medium", "max_tokens": 64} + def test_default_api_base(self): """Test that default API base is used when none is provided""" config = WandbConfig() @@ -139,3 +213,104 @@ class TestWandbConfig: # Check for specific content in the response assert "```python" in content assert "Hey from LiteLLM" in content + + @pytest.mark.respx() + @pytest.mark.parametrize( + "model,effort", + tuple((model, "medium") for model in WANDB_REASONING_MODELS) + + ( + ("Qwen/Qwen3.8-27B", "low"), + ("Qwen/Qwen3.8-27B", "xhigh"), + ), + ) + def test_wandb_completion_preserves_reasoning_effort_with_drop_params( + self, wandb_test_config, wandb_request_mock: respx.Route, model: str, effort: str + ): + completion( + model=f"wandb/{model}", + messages=[{"role": "user", "content": "Hello"}], + api_key="fake-wandb-key", + api_base="https://api.inference.wandb.ai/v1", + reasoning_effort=effort, + max_completion_tokens=64, + drop_params=True, + ) + + assert wandb_request_mock.call_count == 1 + request_body = json.loads(wandb_request_mock.calls[0].request.content) + assert request_body["model"] == model + assert request_body["reasoning_effort"] == effort + assert request_body["max_tokens"] == 64 + assert "max_completion_tokens" not in request_body + + @pytest.mark.respx(assert_all_called=False) + @pytest.mark.parametrize("drop_params", [True, False]) + @pytest.mark.parametrize( + "model,explicit_false", + [ + ("meta-llama/Llama-3.1-8B-Instruct", False), + ("openai/gpt-oss-20b", True), + ], + ) + def test_wandb_completion_without_reasoning_support( + self, + wandb_test_config, + wandb_request_mock: respx.Route, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + model: str, + explicit_false: bool, + drop_params: bool, + ): + with monkeypatch.context() as context: + if explicit_false: + context.setitem(litellm.model_cost[f"wandb/{model}"], "supports_reasoning", False) + + kwargs = { + "model": f"wandb/{model}", + "messages": [{"role": "user", "content": "Hello"}], + "api_key": "fake-wandb-key", + "api_base": "https://api.inference.wandb.ai/v1", + "reasoning_effort": "medium", + "drop_params": drop_params, + } + if not drop_params: + with pytest.raises(litellm.UnsupportedParamsError, match="reasoning_effort"): + completion(**kwargs) + assert len(respx_mock.calls) == 0 + return + + completion(**kwargs) + assert wandb_request_mock.call_count == 1 + request_body = json.loads(wandb_request_mock.calls[0].request.content) + assert request_body["model"] == model + assert "reasoning_effort" not in request_body + + supported_params = litellm.get_supported_openai_params(model=f"wandb/{model}") + assert supported_params is not None + assert "reasoning_effort" not in supported_params + + @pytest.mark.respx() + def test_wandb_completion_keeps_reasoning_effort_for_an_unregistered_model( + self, wandb_test_config, wandb_request_mock: respx.Route + ): + """A wandb id the registry has not named yet resolves through the + wandb-reasoning-baseline fallback generalization, so its reasoning_effort reaches + the provider instead of raising. W&B adds reasoning models faster than this + registry names them, and an exact entry still wins wherever one exists.""" + model: Final = "zai-org/GLM-6-Turbo" + assert f"wandb/{model}" not in litellm.model_cost + + completion( + model=f"wandb/{model}", + messages=[{"role": "user", "content": "Hello"}], + api_key="fake-wandb-key", + api_base="https://api.inference.wandb.ai/v1", + reasoning_effort="medium", + drop_params=False, + ) + + assert wandb_request_mock.call_count == 1 + request_body = json.loads(wandb_request_mock.calls[0].request.content) + assert request_body["model"] == model + assert request_body["reasoning_effort"] == "medium" diff --git a/tests/test_litellm/llms/watsonx/test_watsonx.py b/tests/test_litellm/llms/watsonx/test_watsonx.py index 8ac4472b22d..285afffefc0 100644 --- a/tests/test_litellm/llms/watsonx/test_watsonx.py +++ b/tests/test_litellm/llms/watsonx/test_watsonx.py @@ -356,6 +356,74 @@ async def test_watsonx_gpt_oss_uses_async_http_handler(): assert result["status"] == "success", "Should return success status" +@pytest.mark.parametrize("tokenizer_config_cached", [False, True], ids=["tokenizer_config", "cached_config_jinja"]) +async def test_watsonx_text_gpt_oss_async_completion_fetches_hf_template_off_the_event_loop( + monkeypatch, tokenizer_config_cached +): + import httpx + + from litellm._uuid import uuid + from litellm.litellm_core_utils.prompt_templates import huggingface_template_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + hf_model = f"openai/gpt-oss-{uuid.uuid4()}" + chat_template = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}" + if tokenizer_config_cached: + cached_config = {"status": "success", "tokenizer": {"bos_token": None, "eos_token": None}} + monkeypatch.setattr(litellm, "known_tokenizer_config", {hf_model: cached_config}) + expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/chat_template.jinja" + else: + monkeypatch.setattr(litellm, "known_tokenizer_config", {}) + expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/tokenizer_config.json" + hf_fetched = [] + captured = {} + + def forbid_sync_client(): + raise AssertionError("sync HuggingFace fetch ran on the request path") + + async def serve_hf_file(url, **kwargs): + hf_fetched.append(url) + if url.endswith(".jinja"): + return httpx.Response(200, content=chat_template.encode()) + return httpx.Response(200, json={"chat_template": chat_template, "bos_token": None, "eos_token": None}) + + monkeypatch.setattr(huggingface_template_handler, "_get_httpx_client", forbid_sync_client) + monkeypatch.setattr(huggingface_template_handler, "get_async_httpx_client", lambda **kwargs: Mock(get=serve_hf_file)) + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "model_id": hf_model, + "results": [ + { + "generated_text": "Hi", + "generated_token_count": 1, + "input_token_count": 1, + "stop_reason": "eos_token", + } + ], + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model=f"watsonx_text/{hf_model}", + messages=[{"role": "user", "content": "Hi there"}], + api_base="https://test-api.watsonx.ai", + project_id="test-project-id", + token="test-token", + client=client, + ) + + assert response.choices[0].message.content == "Hi" + assert hf_fetched == [expected_fetch] + assert captured["body"]["input"] == "<|user|>Hi there" + + def test_watsonx_chat_completion_with_reasoning_effort(monkeypatch): """ Test that 'reasoning_effort' is correctly passed through to the WatsonX API payload. diff --git a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py index 83c3bf1ecef..e591c1ae682 100644 --- a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py +++ b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py @@ -63,11 +63,6 @@ TIER_COST_FIELDS = ( "output_cost_per_token_above_200k_tokens", "cache_read_input_token_cost_above_200k_tokens", ) -STALE_TIER_FIELDS = ( - "input_cost_per_token_above_128k_tokens", - "output_cost_per_token_above_128k_tokens", - "cache_read_input_token_cost_above_128k_tokens", -) def expected_retirement_date(slug: str) -> str: @@ -102,11 +97,10 @@ def test_redirected_slug_keeps_its_retirement_date(cost_map: dict, slug: str): assert cost_map[slug]["deprecation_date"] == expected_retirement_date(slug) -@pytest.mark.parametrize("slug", REDIRECTED_SLUGS) -def test_no_slug_keeps_the_superseded_128k_tier(cost_map: dict, slug: str): - """The 128k tier belonged to the retired model; grok-4.3 tiers at 200k.""" - for field in STALE_TIER_FIELDS: - assert field not in cost_map[slug], field +def test_a_live_xai_model_is_untouched(cost_map: dict): + """Guard against the repricing leaking onto models xAI still serves directly.""" + assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"] + assert "deprecation_date" not in cost_map["xai/grok-4.6"] @pytest.mark.parametrize("slug", REDIRECTED_SLUGS) @@ -116,12 +110,7 @@ def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str entry = cost_map[slug] for field in TIER_COST_FIELDS: assert entry[field] == target[field], field - - -def test_a_live_xai_model_is_untouched(cost_map: dict): - """Guard against the repricing leaking onto models xAI still serves directly.""" - assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"] - assert "deprecation_date" not in cost_map["xai/grok-4.6"] + assert {k for k in entry if "_above_" in k} == {k for k in target if "_above_" in k} def test_both_cost_maps_agree_on_the_redirected_slugs(): diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index 38ddac8d510..069ac5727f6 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/tests/test_litellm/llms/zai/test_zai_provider.py @@ -2,11 +2,9 @@ Tests for Z.AI (Zhipu AI) provider - GLM models """ -import json import math import pytest -import respx import litellm from litellm import completion @@ -57,32 +55,9 @@ def test_zai_in_provider_lists(): assert "zai" in litellm.provider_list -def test_zai_models_in_model_cost(local_model_cost_map): - """Test that ZAI models are in the model cost map""" - - zai_models = [ - "zai/glm-4.7", - "zai/glm-4.6", - "zai/glm-4.5", - "zai/glm-4.5v", - "zai/glm-4.5-x", - "zai/glm-4.5-air", - "zai/glm-4.5-airx", - "zai/glm-4-32b-0414-128k", - "zai/glm-4.5-flash", - ] - - for model in zai_models: - assert model in litellm.model_cost, f"Model {model} not found in model_cost" - assert litellm.model_cost[model]["litellm_provider"] == "zai" - - def test_zai_glm46_cost_calculation(local_model_cost_map): """Test the cost calculation for glm-4.6""" - key = "zai/glm-4.6" - info = litellm.model_cost[key] - prompt_cost, completion_cost = cost_per_token( model="zai/glm-4.6", prompt_tokens=1000000, # 1M tokens @@ -94,26 +69,6 @@ def test_zai_glm46_cost_calculation(local_model_cost_map): assert math.isclose(completion_cost, 2.2, rel_tol=1e-6) -def test_zai_flash_model_is_free(local_model_cost_map): - """Test that glm-4.5-flash has zero cost""" - - key = "zai/glm-4.5-flash" - info = litellm.model_cost[key] - - assert info["input_cost_per_token"] == 0 - assert info["output_cost_per_token"] == 0 - - -def test_glm47_supports_reasoning(local_model_cost_map): - """Test that GLM-4.7 supports reasoning""" - - key = "zai/glm-4.7" - assert key in litellm.model_cost, f"Model {key} not found in model_cost" - - info = litellm.model_cost[key] - assert info["supports_reasoning"] is True - - def test_glm47_cost_calculation(local_model_cost_map): """Test cost calculation for GLM-4.7""" diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index 9ae9b732066..aa6449c98dd 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -8,6 +8,7 @@ import pytest from pydantic import BaseModel, TypeAdapter from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.models.autorouter_session import LiteLLM_AutoRouterSession from litellm.models.budget import ( LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, @@ -588,3 +589,35 @@ class TestManagedTables: ) assert table.vector_store_id == "vs1" assert table.custom_llm_provider == "openai" + + +class TestAutoRouterSession: + @staticmethod + def _row(baseline_models: dict) -> LiteLLM_AutoRouterSession: + return LiteLLM_AutoRouterSession( + api_key="k", + session_id="s", + router_name="auto", + router_type="complexity", + first_turn_at=datetime(2026, 9, 1, 12, 0, 0), + last_turn_at=datetime(2026, 9, 1, 12, 5, 0), + last_model="anthropic/claude-sonnet-5", + turns=3, + spend=0.14, + saved_spend=0.24, + classifier_cost=0.0, + tier_turns={}, + baseline_models=baseline_models, + ) + + def test_the_baseline_label_is_the_one_most_turns_were_priced_against(self): + assert self._row({"anthropic/claude-opus-5": 2, "anthropic/claude-sonnet-5": 1}).baseline_model == ( + "anthropic/claude-opus-5" + ) + + def test_a_tie_between_baselines_is_broken_deterministically(self): + assert self._row({"b-model": 1, "a-model": 1}).baseline_model == "b-model" + assert self._row({"a-model": 1, "b-model": 1}).baseline_model == "b-model" + + def test_a_row_whose_turns_recorded_no_baseline_has_no_label(self): + assert self._row({}).baseline_model is None diff --git a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py index 0c8b1cc2836..460aff3e8d1 100644 --- a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py +++ b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py @@ -1,20 +1,18 @@ """ -Regression tests for Azure Document Intelligence api_base resolution in OCR. +Regression tests for Azure Document Intelligence api_base ownership in OCR. `azure_ai` exposes two OCR services on one provider; the `doc-intelligence` -sub-route must resolve to `AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT`, not to the -generic `AZURE_AI_API_BASE` fallback that `get_llm_provider` injects. These tests -pin that routing and guard the backwards-compatibility contract that an explicitly -supplied api_base is always honoured. +sub-route must defer environment resolution to Rust, not accept the generic +`AZURE_AI_API_BASE` fallback that `get_llm_provider` injects. An explicitly +supplied api_base is still always honoured. """ from litellm.llms.azure_ai.ocr.common_utils import ( is_azure_document_intelligence_model, ) -from litellm.ocr.main import _prepare_ocr_request, _rust_bridge_api_base +from litellm.ocr.main import _prepare_ocr_request _DOC = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} -_DOC_INTELLIGENCE_ENDPOINT = "https://di.cognitiveservices.azure.com" _AZURE_AI_API_BASE = "https://generic-azure-ai.example.com" @@ -23,13 +21,6 @@ class _FakeLogging: return None -def _resolve_secret(name: str) -> str | None: - return { - "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT": _DOC_INTELLIGENCE_ENDPOINT, - "AZURE_AI_API_BASE": _AZURE_AI_API_BASE, - }.get(name) - - def _prepare(model: str, api_base: str | None): return _prepare_ocr_request( model=model, @@ -56,15 +47,13 @@ class TestIsAzureDocumentIntelligenceModel: class TestDocIntelligenceApiBaseResolution: def test_generic_azure_ai_base_does_not_hijack_doc_intelligence(self, monkeypatch): - """Without an explicit api_base, the AZURE_AI_API_BASE fallback must not - overwrite the endpoint, so it resolves to the Document Intelligence one.""" + """The generic Azure base must not overwrite Rust-owned DI resolution.""" monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", raising=False) prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", None) assert prepared.api_base is None - assert _rust_bridge_api_base(prepared, _resolve_secret) == _DOC_INTELLIGENCE_ENDPOINT def test_explicit_api_base_is_honoured_for_doc_intelligence(self, monkeypatch): """A caller-supplied api_base must always win, even for doc-intelligence.""" @@ -74,7 +63,6 @@ class TestDocIntelligenceApiBaseResolution: prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", custom) assert prepared.api_base == custom - assert _rust_bridge_api_base(prepared, _resolve_secret) == custom def test_generic_azure_ai_base_still_applies_to_mistral_ocr(self, monkeypatch): """Non doc-intelligence azure_ai models keep using AZURE_AI_API_BASE.""" diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/ocr/test_ocr_native_format.py index 249fbda713e..46e9a4d3729 100644 --- a/tests/test_litellm/ocr/test_ocr_native_format.py +++ b/tests/test_litellm/ocr/test_ocr_native_format.py @@ -1,52 +1,60 @@ """ -Tests for the OCR `req_format` option in the SDK request path: -providers that don't support a native response must reject it, and the Rust -bridge (which only returns the normalized shape) must not serve native requests. +Tests for the OCR `req_format` option in the SDK request path. """ -import dataclasses -from unittest.mock import MagicMock - import pytest import litellm -from litellm.llms.azure_ai.ocr.cohere_parse_transformation import AzureAICohereParseConfig -from litellm.llms.cohere.ocr.transformation import CohereParseConfig -from litellm.ocr.main import _PreparedOCRRequest, _rust_ocr_supported +from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.rust_bridge.ocr import LiteLLMOcrRequest DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} -def _prepared(optional_params: dict[str, object]) -> _PreparedOCRRequest: - return _PreparedOCRRequest( - model="doc-intelligence/prebuilt-layout", - document=dict(DOCUMENT), +def _request( + optional_params: dict[str, object], model: str = "azure_ai/doc-intelligence/prebuilt-layout" +) -> LiteLLMOcrRequest: + return LiteLLMOcrRequest( + model=model, + document=DOCUMENT, api_key="fake-key", - api_base="https://example.cognitiveservices.azure.com", - custom_llm_provider="azure_ai", + api_base=None, + custom_llm_provider=None, extra_headers=None, - provider_config=MagicMock(), - optional_params=optional_params, - litellm_params={}, - effective_timeout=60.0, - litellm_logging_obj=MagicMock(), + timeout=60.0, + kwargs=optional_params, ) @pytest.mark.parametrize("optional_params", [{}, {"req_format": "litellm"}]) def test_rust_ocr_serves_default_format(optional_params): - assert _rust_ocr_supported(_prepared(optional_params)) is True + assert rust_ocr_bridge.supported(_request(optional_params)) is True -def test_rust_ocr_skipped_for_native_format(): - assert _rust_ocr_supported(_prepared({"req_format": "native"})) is False +def test_rust_ocr_serves_native_format_for_document_intelligence(): + assert rust_ocr_bridge.supported(_request({"req_format": "native"})) is True -@pytest.mark.parametrize("provider_config", [CohereParseConfig(), AzureAICohereParseConfig()]) -def test_rust_ocr_skipped_for_configs_without_bridge_support(provider_config): - prepared = dataclasses.replace(_prepared({}), provider_config=provider_config) +def test_rust_ocr_response_retains_provider_native_response(): + provider_response = {"status": "succeeded", "analyzeResult": {"content": "native"}} + response = rust_ocr_bridge._response( + { + "pages": [], + "model": "prebuilt-layout", + "document_annotation": None, + "usage_info": {"pages_processed": 0}, + "object": "ocr", + "provider_native_response": provider_response, + } + ) - assert _rust_ocr_supported(prepared) is False + assert response.get_provider_native_response() == provider_response + assert response.model_dump().get("provider_native_response") is None + + +@pytest.mark.parametrize("model", ["cohere/cohere-parse", "azure_ai/cohere-parse"]) +def test_rust_ocr_skipped_for_unsupported_models(model): + assert rust_ocr_bridge.supported(_request({}, model)) is False @pytest.mark.asyncio diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 1c2e07e0d24..dbb4f822d0b 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -3,12 +3,12 @@ import builtins import importlib import types -from typing import Any import httpx import pytest import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge import configuration @@ -17,6 +17,7 @@ from litellm.rust_bridge import configuration # explicitly via importlib rather than attribute traversal. ocr_main = importlib.import_module("litellm.ocr.main") rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") +rust_bridge_bindings = importlib.import_module("litellm.rust_bridge.bindings") rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" @@ -38,6 +39,10 @@ class CapturedException(Exception): pass +class RustUpstreamError(Exception): + pass + + class RecordingBridge: """A fake ``RustOcr`` callable that records the args it was handed.""" @@ -53,6 +58,7 @@ class RecordingBridge: custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: self.calls.append( @@ -64,6 +70,7 @@ class RecordingBridge: "custom_llm_provider": custom_llm_provider, "extra_headers": extra_headers, "optional_params": optional_params, + "input_sources": input_sources, "timeout_seconds": timeout_seconds, } ) @@ -85,6 +92,7 @@ class RecordingAsyncBridge: custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: self.calls.append( @@ -96,6 +104,7 @@ class RecordingAsyncBridge: "custom_llm_provider": custom_llm_provider, "extra_headers": extra_headers, "optional_params": optional_params, + "input_sources": input_sources, "timeout_seconds": timeout_seconds, } ) @@ -112,6 +121,7 @@ class RaisingBridge: custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -127,6 +137,7 @@ class RaisingAsyncBridge: custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -138,6 +149,9 @@ class RecordingLogging: def __init__(self) -> None: self.pre_call_kwargs: dict[str, object] | None = None + def update_from_kwargs(self, **kwargs: object) -> None: + self.update_kwargs = kwargs + def pre_call( self, *, @@ -152,74 +166,45 @@ class RecordingLogging: } -class FakeOCRConfig: - """A stand-in ``BaseOCRConfig`` that echoes the request it would build.""" - - def __init__(self, api_key_env_var: str = "MISTRAL_API_KEY") -> None: - self.api_key_env_var = api_key_env_var - - def get_api_key_env_var(self) -> str: - return self.api_key_env_var - - def validate_environment( - self, - *, - headers: dict[str, object], - model: str, - api_key: str | None, - api_base: str | None, - litellm_params: dict[str, object], - ) -> dict[str, object]: - return {"Authorization": f"Bearer {api_key}", **headers} - - def get_complete_url( - self, - *, - api_base: str | None, - model: str, - optional_params: dict[str, object], - litellm_params: dict[str, object], - ) -> str: - return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" - - -def build_prepared_request( +def build_request( *, logging_obj: RecordingLogging | None = None, - provider_config: FakeOCRConfig | None = None, model: str = "mistral-ocr-latest", document: dict[str, object] = DOCUMENT, api_key: str | None = "sk-test", api_base: str | None = None, - custom_llm_provider: str = "mistral", + custom_llm_provider: str | None = "mistral", extra_headers: dict[str, object] | None = None, optional_params: dict[str, object] | None = None, litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = 12.5, -) -> Any: - return ocr_main._PreparedOCRRequest( +) -> rust_bridge.LiteLLMOcrRequest: + return rust_bridge.LiteLLMOcrRequest( model=model, document=document, api_key=api_key, api_base=api_base, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, - provider_config=provider_config or FakeOCRConfig(), - optional_params=optional_params or {}, - litellm_params=litellm_params or {}, - effective_timeout=timeout, - litellm_logging_obj=logging_obj or RecordingLogging(), + timeout=timeout, + kwargs={ + **(optional_params or {}), + **(litellm_params or {}), + "litellm_logging_obj": logging_obj or RecordingLogging(), + }, ) @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge.set_rust_ocr(ocr=None, aocr=None) + rust_bridge._OCR.reset() + rust_bridge._AOCR.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - rust_bridge.set_rust_ocr(ocr=None, aocr=None) + rust_bridge._OCR.reset() + rust_bridge._AOCR.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -229,7 +214,7 @@ def fake_bridge(): """Enable the Rust path with an injected recording bridge (no native wheel).""" bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) return bridge @@ -238,34 +223,14 @@ def fake_async_bridge(): """Enable the async Rust path with an injected recording bridge.""" bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._AOCR.override(bridge) return bridge -def test_rust_toggles_flag(): - assert rust_bridge.rust_ocr_enabled() is False - litellm.rust(True) - assert rust_bridge.rust_ocr_enabled() is True - litellm.rust(False) - assert rust_bridge.rust_ocr_enabled() is False - - -def test_env_var_enables_rust_ocr(monkeypatch): - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert rust_bridge.rust_ocr_enabled() is True - - -def test_explicit_false_overrides_process_enable(): - litellm.rust(True) - - assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False - - def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) assert rust_bridge.load_rust_ocr() is bridge @@ -329,7 +294,7 @@ def test_native_bridge_available_reflects_loader(monkeypatch): def test_load_rust_aocr_returns_injected_impl(): bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._AOCR.override(bridge) assert rust_bridge.load_rust_aocr() is bridge @@ -338,7 +303,8 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) + rust_bridge._OCR.override(bridge) + rust_bridge._AOCR.override(async_bridge) litellm.rust(False) assert rust_bridge.load_rust_ocr() is bridge @@ -350,16 +316,18 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): def test_explicit_ocr_none_clears_injected_impl(monkeypatch): monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + rust_bridge_bindings, "get_native_bridge", lambda: None, ) bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) + rust_bridge._OCR.override(bridge) + rust_bridge._AOCR.override(async_bridge) - rust_bridge.set_rust_ocr(ocr=None, aocr=None) + rust_bridge._OCR.override(None) + rust_bridge._AOCR.override(None) assert rust_bridge.load_rust_ocr() is None assert rust_bridge.load_rust_aocr() is None @@ -368,7 +336,7 @@ def test_load_rust_ocr_none_when_extension_absent(monkeypatch): """With no injected impl and no compiled wheel, the loader returns None so the caller degrades to the Python path instead of raising ImportError.""" monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + rust_bridge_bindings, "get_native_bridge", lambda: None, ) @@ -385,7 +353,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + rust_bridge_bindings, "get_native_bridge", lambda: fake_module, ) @@ -406,7 +374,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) response = rust_bridge.ocr( model="mistral-ocr-latest", document=DOCUMENT, @@ -431,6 +399,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): "x-trace-id": "trace-1", }, "optional_params": {"include_image_base64": True, "pages": [0]}, + "input_sources": {}, "timeout_seconds": 12.5, } @@ -441,7 +410,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._AOCR.override(bridge) response = await rust_bridge.aocr( model="mistral-ocr-maas", document=DOCUMENT, @@ -462,6 +431,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): "custom_llm_provider": "vertex_ai", "extra_headers": None, "optional_params": {"vertex_project": "project-1"}, + "input_sources": {}, "timeout_seconds": 42.0, } @@ -470,10 +440,10 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): bridge = RecordingBridge() logging_obj = RecordingLogging() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) response = ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( logging_obj=logging_obj, api_base="https://proxy.internal", extra_headers={"x-trace-id": "trace-1"}, @@ -492,21 +462,35 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): "api_base": "https://proxy.internal", "custom_llm_provider": "mistral", "extra_headers": { - "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", }, "optional_params": {"include_image_base64": True}, + "input_sources": {}, "timeout_seconds": 12.5, } +def test_rust_upstream_error_uses_ocr_provider_error_mapping(): + error = RustUpstreamError(400, '{"message":"invalid model"}') + + mapped = ocr_main._map_rust_ocr_error( + error, + build_request(), + (RuntimeError, RustUpstreamError), + ) + + assert isinstance(mapped, BaseLLMException) + assert mapped.status_code == 400 + assert mapped.message == '{"message":"invalid model"}' + + def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request(api_key=None, timeout=None), + request=build_request(api_key=None, timeout=None), resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, ) @@ -516,13 +500,13 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): def test_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) def _resolver(name: str) -> str | None: raise AssertionError(f"resolver should not be called for {name}") ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( api_key="sk-explicit", timeout=None, ), @@ -532,37 +516,36 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver(): assert bridge.calls[0]["api_key"] == "sk-explicit" -def test_run_rust_ocr_uses_provider_api_key_env_var(): +def test_run_rust_ocr_uses_mistral_secret_manager_without_provider_config(): bridge = RecordingBridge() resolver_calls = [] litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) def _resolver(name): resolver_calls.append(name) return "sk-provider-env" ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( - provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), - model="provider-ocr-model", + request=build_request( + model="mistral-ocr-latest", api_key=None, timeout=None, ), resolve_api_key=_resolver, ) - assert resolver_calls == ["PROVIDER_OCR_API_KEY"] + assert resolver_calls == ["MISTRAL_API_KEY"] assert bridge.calls[0]["api_key"] == "sk-provider-env" def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", litellm_params={ @@ -580,22 +563,24 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): "include_image_base64": True, "vertex_project": "project-1", "vertex_location": "us-central1", + "vertex_credentials": "redacted", } def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) def _resolver(name: str) -> str | None: return { "VERTEXAI_PROJECT": "project-from-secret", "VERTEXAI_LOCATION": "us-east5", + "VERTEXAI_CREDENTIALS": "credentials-from-secret", }.get(name) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", timeout=None, @@ -605,54 +590,220 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret" assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" + assert bridge.calls[0]["optional_params"]["vertex_credentials"] == "credentials-from-secret" -def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): +def test_prepare_rust_ocr_call_defers_azure_environment_resolution_to_rust(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( custom_llm_provider="azure_ai", model="pixtral-12b-2409", + api_key=None, api_base=None, timeout=None, ), - resolve_api_key=lambda name: "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None, + resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), ) - assert bridge.calls[0]["api_base"] == "https://azure.example.com" + assert bridge.calls[0]["api_base"] is None + assert bridge.calls[0]["api_key"] is None + assert bridge.calls[0]["extra_headers"] is None -def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): +def test_prepare_rust_ocr_call_defers_document_intelligence_environment_to_rust(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( custom_llm_provider="azure_ai", model="doc-intelligence/prebuilt-layout", api_base=None, timeout=None, ), - resolve_api_key=lambda name: ( - "https://document-intelligence.example.com" if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" else None - ), + resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), ) - assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com" + assert bridge.calls[0]["api_base"] is None + + +def test_prepare_rust_ocr_call_forwards_raw_azure_auth_inputs(): + bridge = RecordingBridge() + litellm.rust(True) + rust_bridge._OCR.override(bridge) + + ocr_main._run_rust_ocr( + request=build_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_key=None, + api_base="https://azure.example.com", + extra_headers={"x-trace-id": "trace-1"}, + litellm_params={ + "azure_ad_token": "entra-token", + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_scope": "scope", + "azure_authority_host": "https://login.example.com", + "azure_credential": "ClientSecretCredential", + "azure_federated_token_file": "/token", + }, + timeout=None, + ), + resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), + ) + + call = bridge.calls[0] + assert call["api_key"] is None + assert call["api_base"] == "https://azure.example.com" + assert call["extra_headers"] == {"x-trace-id": "trace-1"} + assert call["optional_params"] == { + "azure_ad_token": "entra-token", + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_scope": "scope", + "azure_authority_host": "https://login.example.com", + "azure_credential": "ClientSecretCredential", + "azure_federated_token_file": "/token", + } + assert call["input_sources"] == {} + + +def test_prepare_rust_ocr_call_preserves_proxy_input_sources(): + bridge = RecordingBridge() + litellm.rust(True) + rust_bridge._OCR.override(bridge) + request_values = { + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_authority_host": "https://login.example.com", + "api_base": "https://azure.example.com", + } + + ocr_main._run_rust_ocr( + request=build_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_key="request-key", + api_base="https://azure.example.com", + litellm_params={ + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_authority_host": "https://login.example.com", + "proxy_server_request": {"body": request_values, "credential_fields": ("api_key",)}, + }, + ), + resolve_api_key=lambda _name: None, + ) + + assert bridge.calls[0]["input_sources"] == { + **{name: "request" for name in request_values}, + "api_key": "request", + } + + marshaled = rust_bridge._marshal( + build_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_key="request-key", + api_base="https://azure.example.com", + litellm_params={ + "proxy_server_request": { + "body": {"api_base": "https://azure.example.com"}, + "credential_fields": ("api_key",), + } + }, + ), + lambda _name: None, + lambda document: document, + ) + assert marshaled.input_sources == {"api_base": "request", "api_key": "request"} + + +def test_rust_ocr_logging_redacts_azure_credentials(): + bridge = RecordingBridge() + logging_obj = RecordingLogging() + litellm.rust(True) + rust_bridge._OCR.override(bridge) + + ocr_main._run_rust_ocr( + request=build_request( + logging_obj=logging_obj, + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_key=None, + litellm_params={"azure_ad_token": "token", "client_secret": "secret"}, + ), + resolve_api_key=lambda _name: None, + ) + + assert logging_obj.update_kwargs["optional_params"] == { + "azure_ad_token": "****", + "client_secret": "****", + } + assert logging_obj.pre_call_kwargs is not None + additional_args = logging_obj.pre_call_kwargs["additional_args"] + assert isinstance(additional_args, dict) + complete_input = additional_args["complete_input_dict"] + assert isinstance(complete_input, dict) + assert complete_input["azure_ad_token"] == "****" + assert complete_input["client_secret"] == "****" + + +def test_rust_eligibility_rejects_python_only_azure_auth_modes(): + for params in ( + {"azure_ad_token_provider": lambda: "token"}, + {"azure_username": "user"}, + {"azure_password": "password"}, + ): + assert not ocr_main._rust_ocr_supported( + build_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + litellm_params=params, + ) + ) + + +def test_prepare_rust_ocr_call_forwards_global_azure_refresh(monkeypatch: pytest.MonkeyPatch): + bridge = RecordingBridge() + litellm.rust(True) + rust_bridge._OCR.override(bridge) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + + ocr_main._run_rust_ocr( + request=build_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_key=None, + api_base="https://azure.example.com", + litellm_params={"proxy_server_request": {"body": {"enable_azure_ad_token_refresh": True}}}, + timeout=None, + ), + resolve_api_key=lambda _name: None, + ) + + assert bridge.calls[0]["optional_params"] == {"enable_azure_ad_token_refresh": True} + assert bridge.calls[0]["input_sources"] == {"enable_azure_ad_token_refresh": "deployment"} def test_run_rust_ocr_runs_pre_call_logging(): logging_obj = RecordingLogging() bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( logging_obj=logging_obj, api_base="https://api.mistral.ai/v1", extra_headers={"x-trace-id": "trace-1"}, @@ -668,9 +819,8 @@ def test_run_rust_ocr_runs_pre_call_logging(): complete_input = additional_args["complete_input_dict"] assert complete_input["document"] == DOCUMENT assert complete_input["include_image_base64"] is True - assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" + assert additional_args["api_base"] == "https://api.mistral.ai/v1" assert additional_args["headers"] == { - "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", } @@ -688,12 +838,11 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): assert response.pages[0].markdown == "hello world" assert len(fake_bridge.calls) == 1 call = fake_bridge.calls[0] - assert call["model"] == "mistral-ocr-latest" + assert call["model"] == MODEL assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" - assert call["custom_llm_provider"] == "mistral" + assert call["custom_llm_provider"] is None assert call["extra_headers"] == { - "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", } assert call["optional_params"].get("include_image_base64") is True @@ -709,8 +858,29 @@ def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): assert isinstance(response, OCRResponse) assert len(fake_bridge.calls) == 1 - assert fake_bridge.calls[0]["model"] == "pixtral-12b-2409" - assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai" + assert fake_bridge.calls[0]["model"] == "azure_ai/pixtral-12b-2409" + assert fake_bridge.calls[0]["custom_llm_provider"] is None + assert fake_bridge.calls[0]["extra_headers"] is None + + +def test_ocr_routes_azure_entra_inputs_to_rust_without_python_auth(fake_bridge): + response = litellm.ocr( + model="azure_ai/pixtral-12b-2409", + document=DOCUMENT, + api_base="https://example.services.ai.azure.com", + azure_ad_token="entra-token", + tenant_id="tenant", + client_id="client", + ) + + assert isinstance(response, OCRResponse) + assert fake_bridge.calls[0]["api_key"] is None + assert fake_bridge.calls[0]["extra_headers"] is None + assert fake_bridge.calls[0]["optional_params"] == { + "azure_ad_token": "entra-token", + "tenant_id": "tenant", + "client_id": "client", + } def test_ocr_rust_path_converts_file_document_before_bridge(fake_bridge): @@ -737,7 +907,7 @@ def test_ocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=RaisingBridge()) + rust_bridge._OCR.override(RaisingBridge()) with pytest.raises(CapturedException): litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -760,12 +930,11 @@ async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge): assert response.pages[0].markdown == "hello world" assert len(fake_async_bridge.calls) == 1 call = fake_async_bridge.calls[0] - assert call["model"] == "mistral-ocr-latest" + assert call["model"] == MODEL assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" - assert call["custom_llm_provider"] == "mistral" + assert call["custom_llm_provider"] is None assert call["extra_headers"] == { - "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", } assert call["optional_params"].get("include_image_base64") is True @@ -783,7 +952,7 @@ async def test_aocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=RaisingAsyncBridge()) + rust_bridge._AOCR.override(RaisingAsyncBridge()) with pytest.raises(CapturedException): await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -812,9 +981,7 @@ def test_ocr_does_not_route_to_rust_when_disabled(): """With the flag off, the bridge must not be consulted even if an impl exists.""" bridge = RecordingBridge() litellm.rust(False) - rust_bridge.set_rust_ocr(ocr=bridge) - - assert rust_bridge.rust_ocr_enabled() is False + rust_bridge._OCR.override(bridge) # The impl stays available for injection, but the disabled flag gates usage, # so ocr() never reaches the Rust path (asserted via the enabled-path test). assert bridge.calls == [] @@ -858,3 +1025,137 @@ def test_ocr_provider_configs_expose_api_key_env_vars(): assert AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_rust_receives_unmapped_azure_options(asynchronous, fake_bridge, fake_async_bridge): + from typing import Final + + arguments: Final = { + "model": "azure_ai/doc-intelligence/prebuilt-layout", + "document": DOCUMENT, + "api_key": "test-key", + "pages": [0, 2], + "features": ["languages", "style"], + "provider_extension": {"enabled": True}, + } + if asynchronous: + await litellm.aocr(**arguments) + else: + litellm.ocr(**arguments) + call: Final = (fake_async_bridge if asynchronous else fake_bridge).calls[0] + assert call["model"] == arguments["model"] + assert call["custom_llm_provider"] is None + assert call["extra_headers"] is None + assert call["optional_params"] == { + "pages": [0, 2], + "features": ["languages", "style"], + "provider_extension": {"enabled": True}, + } + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.asyncio +async def test_python_fallback_maps_original_options_once(enabled, monkeypatch): + from io import BytesIO + from typing import Final + + class PythonHandler: + def __init__(self): + self.calls = [] + + def ocr(self, **kwargs): + self.calls.append(kwargs) + return OCRResponse(pages=[], model=kwargs["model"]) + + handler: Final = PythonHandler() + monkeypatch.setattr(ocr_main, "base_llm_http_handler", handler) + litellm.rust(enabled) + rust_bridge._OCR.override(None) + rust_bridge._AOCR.override(None) + for asynchronous in (False, True): + file: Final = BytesIO(b"test document") + arguments: Final = { + "model": "azure_ai/doc-intelligence/prebuilt-layout", + "document": {"type": "file", "file": file}, + "api_key": "test-key", + "pages": [0, 2], + } + if asynchronous: + await litellm.aocr(**arguments) + else: + litellm.ocr(**arguments) + assert handler.calls[-1]["optional_params"]["pages"] == "1,3" + assert handler.calls[-1]["document"]["document_url"].endswith("dGVzdCBkb2N1bWVudA==") + assert len(handler.calls) == 2 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/doc-intelligence/prebuilt-read"]) +@pytest.mark.asyncio +async def test_native_public_ocr_matches_python(model, asynchronous): + import json + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + from threading import Thread + from typing import Final + from urllib.parse import parse_qsl, urlsplit + + native: Final = rust_bridge_loader.get_native_bridge() + if native is None: + pytest.skip("requires the compiled Rust extension") + calls: Final = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + body: Final = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + target: Final = urlsplit(self.path) + calls.append( + ( + target.path, + parse_qsl(target.query), + self.headers.get("Authorization"), + self.headers.get("Ocp-Apim-Subscription-Key"), + body, + ) + ) + payload: Final = ( + {"status": "succeeded", "analyzeResult": {"pages": []}} + if "doc-intelligence" in model + else {"pages": [{"index": 0, "markdown": "hello"}]} + ) + encoded: Final = json.dumps(payload).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, *_args): + pass + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + responses: Final = [] + try: + for enabled in (False, True): + litellm.rust(enabled) + arguments: Final = { + "model": model, + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{server.server_port}", + "pages": [0, 2], + "timeout": 3.0, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + responses.append(response.model_dump()) + assert len(calls) == 2 + assert calls[0] == calls[1] + for key in ("model", "pages", "object"): + assert responses[0][key] == responses[1][key] + finally: + server.shutdown() + server.server_close() + thread.join(timeout=3) diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 1950c37a12e..546cff18b5d 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -873,3 +873,118 @@ def test_llm_passthrough_route_propagates_allm_passthrough_route_to_logging_obj( assert captured_litellm_params.get("allm_passthrough_route") is True assert LitellmLogging._is_sync_litellm_request(captured_litellm_params) is False + + +FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" + + +def _foundry_parse_response() -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"id":"parse-1","pages":[]}', + request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), + ) + + +def test_azure_ai_relay_reaches_the_deployment_with_its_own_credential(): + """ + Regression for LIT-7022: azure_ai had no passthrough config, so every + /azure_ai// relay raised "Provider azure_ai not found" + before a request was built. + """ + client = HTTPHandler() + + with patch.object(client.client, "send", return_value=_foundry_parse_response()) as mock_send: + response = llm_passthrough_route( + model="azure_ai/Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + method="POST", + custom_llm_provider="azure_ai", + api_base=FOUNDRY_BASE, + api_key="deployment-key", + json={"model": "Cohere-parse-v5", "document": {"type": "image_url", "image_url": "https://x/y.png"}}, + client=client, + litellm_logging_obj=MagicMock(), + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "Cohere-parse-v5" + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_router_relays_azure_ai_model_through_the_deployment_api_base(): + router = litellm.Router( + model_list=[ + { + "model_name": "foundry-parse", + "litellm_params": { + "model": "azure_ai/Cohere-parse-v5", + "api_base": FOUNDRY_BASE, + "api_key": "deployment-key", + }, + } + ] + ) + async_client = AsyncHTTPHandler() + + with patch.object(async_client.client, "send", AsyncMock(return_value=_foundry_parse_response())) as mock_send: + response = await router.allm_passthrough_route( + model="foundry-parse", + method="POST", + endpoint="foundry-parse/providers/cohere/v2/parse", + json={"model": "foundry-parse", "document": {"type": "image_url", "image_url": "https://x/y.png"}}, + client=async_client, + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "Cohere-parse-v5" + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_router_relays_an_openai_model_on_a_foundry_base_as_azure_ai(monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") + router = litellm.Router( + model_list=[ + { + "model_name": "foundry-gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": FOUNDRY_BASE, + "api_key": "deployment-key", + }, + } + ] + ) + async_client = AsyncHTTPHandler() + upstream = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=( + b'{"id":"chatcmpl-1","object":"chat.completion","model":"gpt-5.4-mini",' + b'"choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"hi"}}],' + b'"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' + ), + request=httpx.Request("POST", f"{FOUNDRY_BASE}/models/chat/completions"), + ) + + with patch.object(async_client.client, "send", AsyncMock(return_value=upstream)) as mock_send: + await router.allm_passthrough_route( + model="foundry-gpt", + method="POST", + endpoint="foundry-gpt/models/chat/completions", + request_query_params={"api-version": "2024-05-01-preview"}, + json={"model": "foundry-gpt", "messages": [{"role": "user", "content": "hi"}]}, + client=async_client, + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "gpt-5.4-mini" diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py index a88b0ef0c4b..5e13db9439b 100644 --- a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py +++ b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py @@ -34,6 +34,34 @@ class _ImmediateExecutor: fn(*args, **kwargs) +class _RecordingCollector: + def __init__(self) -> None: + self.chunks: List[bytes] = [] + + def add(self, chunk: bytes) -> None: + self.chunks.append(chunk) + + def build_logged_response(self, litellm_logging_obj: MagicMock) -> bytes: + return b"".join(self.chunks) + + +class _FailingCollector(_RecordingCollector): + def add(self, chunk: bytes) -> None: + raise ValueError("bad frame") + + +def _provider_config(collector: _RecordingCollector) -> MagicMock: + provider_config = MagicMock() + provider_config.create_stream_collector.return_value = collector + return provider_config + + +def _spend_payload(flush_mock: MagicMock) -> bytes: + flush_mock.assert_called_once() + collector = flush_mock.call_args.kwargs["collector"] + return collector.build_logged_response(litellm_logging_obj=MagicMock()) + + @pytest.mark.asyncio async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): from litellm.passthrough.main import AsyncPassthroughStreamingResponse @@ -48,13 +76,12 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): return mock_response mock_logging_obj = _make_logging_obj() - provider_config = MagicMock() received = [] received_response = AsyncPassthroughStreamingResponse( response=response_coro(), litellm_logging_obj=mock_logging_obj, - provider_config=provider_config, + provider_config=_provider_config(_RecordingCollector()), ) async for chunk in received_response: @@ -67,12 +94,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): await asyncio.sleep(0) - mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once() - call_kwargs = ( - mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs - ) - assert call_kwargs["raw_bytes"] == chunks - assert call_kwargs["provider_config"] is provider_config + assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == b"".join(chunks) @pytest.mark.asyncio @@ -93,12 +115,11 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_client_disconnect(): return mock_response mock_logging_obj = _make_logging_obj() - provider_config = MagicMock() gen = AsyncPassthroughStreamingResponse( response=response_coro(), litellm_logging_obj=mock_logging_obj, - provider_config=provider_config, + provider_config=_provider_config(_RecordingCollector()), ) received = [await gen.__anext__()] @@ -108,11 +129,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_client_disconnect(): await asyncio.sleep(0) - mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once() - call_kwargs = ( - mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs - ) - assert call_kwargs["raw_bytes"] == [chunks[0]] + assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == chunks[0] @pytest.mark.asyncio @@ -178,14 +195,13 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_w return mock_response mock_logging_obj = _make_logging_obj() - provider_config = MagicMock() received = [] async def _drain(): async for chunk in AsyncPassthroughStreamingResponse( response=response_coro(), litellm_logging_obj=mock_logging_obj, - provider_config=provider_config, + provider_config=_provider_config(_RecordingCollector()), ): received.append(chunk) @@ -196,11 +212,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_w await asyncio.sleep(0) - mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once() - call_kwargs = ( - mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs - ) - assert call_kwargs["raw_bytes"] == partial_chunks + assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == b"".join(partial_chunks) def test_passthroughstreamingresponse_flushes_on_normal_completion(): @@ -221,12 +233,11 @@ def test_passthroughstreamingresponse_flushes_on_normal_completion(): mock_logging_obj = MagicMock() mock_logging_obj.flush_passthrough_collected_chunks = MagicMock() - provider_config = MagicMock() received_responce = PassthroughStreamingResponse( response=mock_response, litellm_logging_obj=mock_logging_obj, - provider_config=provider_config, + provider_config=_provider_config(_RecordingCollector()), ) with patch("litellm.utils.executor", _ImmediateExecutor()): @@ -237,7 +248,7 @@ def test_passthroughstreamingresponse_flushes_on_normal_completion(): assert received_responce.headers["content-type"] == "application/octet-stream" assert received_responce.headers["x-request-id"] == "req-123" - mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once() + assert _spend_payload(mock_logging_obj.flush_passthrough_collected_chunks) == b"".join(chunks) def test_passthroughstreamingresponse_flushes_on_early_close(): @@ -258,19 +269,66 @@ def test_passthroughstreamingresponse_flushes_on_early_close(): mock_logging_obj = MagicMock() mock_logging_obj.flush_passthrough_collected_chunks = MagicMock() - provider_config = MagicMock() with patch("litellm.utils.executor", _ImmediateExecutor()): gen = PassthroughStreamingResponse( response=mock_response, litellm_logging_obj=mock_logging_obj, - provider_config=provider_config, + provider_config=_provider_config(_RecordingCollector()), ) first = next(gen) gen.close() assert first == chunks[0] - mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once() - call_kwargs = mock_logging_obj.flush_passthrough_collected_chunks.call_args.kwargs - assert call_kwargs["raw_bytes"] == [chunks[0]] + assert _spend_payload(mock_logging_obj.flush_passthrough_collected_chunks) == chunks[0] + + +@pytest.mark.asyncio +async def test_asyncpassthroughstreamingresponse_relays_the_stream_when_spend_parsing_fails(): + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + + chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] + mock_response = _make_streaming_response(chunks) + + async def response_coro(): + return mock_response + + mock_logging_obj = _make_logging_obj() + + received = [ + chunk + async for chunk in AsyncPassthroughStreamingResponse( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=_provider_config(_FailingCollector()), + ) + ] + await asyncio.sleep(0) + + assert received == chunks + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called() + + +def test_passthroughstreamingresponse_relays_the_stream_when_spend_parsing_fails(): + from litellm.passthrough.main import PassthroughStreamingResponse + + chunks = [b"a", b"b", b"c"] + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream"}) + mock_response.iter_bytes = lambda: iter(chunks) + + mock_logging_obj = MagicMock() + mock_logging_obj.flush_passthrough_collected_chunks = MagicMock() + + received = list( + PassthroughStreamingResponse( + response=mock_response, + litellm_logging_obj=mock_logging_obj, + provider_config=_provider_config(_FailingCollector()), + ) + ) + + assert received == chunks + mock_logging_obj.flush_passthrough_collected_chunks.assert_not_called() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index f1e299802fb..e6c8d4ee039 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -11,6 +11,7 @@ from starlette.datastructures import Headers from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + UnloadableEntitlementError, _is_mcp_admitted_user_subject, ) from litellm.proxy._types import ( @@ -4396,6 +4397,147 @@ class TestAgentMCPPermissions: ) assert sorted(result) == ["tool_a", "tool_b"] + def _agent_object_permission(self, *, toolset_ids, servers=(), tool_permissions=None): + agent_object_permission = MagicMock() + agent_object_permission.mcp_servers = list(servers) + agent_object_permission.mcp_access_groups = [] + agent_object_permission.mcp_tool_permissions = tool_permissions + agent_object_permission.mcp_toolsets = list(toolset_ids) + return agent_object_permission + + def _mock_manager_with_toolsets(self, toolset_perms): + mock_manager = MagicMock() + mock_manager.expand_permission_list = MagicMock(side_effect=lambda servers: list(servers)) + mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {}) + mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value=toolset_perms) + return mock_manager + + def _agent_toolset_patches(self, agent_object_permission, mock_manager): + return ( + patch.object( # test-quality-ok: stub the agent perm loader; the resolver reads module globals with no injection seam + MCPRequestHandler, "_get_agent_object_permission", AsyncMock(return_value=agent_object_permission) + ), + patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling toolset tests + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here + MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[]) + ), + ) + + async def test_get_allowed_mcp_servers_for_agent_includes_toolset_servers(self): + """An agent granted only mcp_toolsets reaches the toolset's servers, exactly as a + key, team, or org granted only toolsets does""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"], servers=["server-direct"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + result = await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) + + assert sorted(result) == ["server-a", "server-direct"] + mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with(toolset_ids=["toolset-1"]) + + async def test_get_allowed_mcp_servers_toolset_only_agent_caps_key_servers(self): + """Regression: an agent whose only grant is a toolset used to resolve to [] and place + no ceiling at all, so a key bound to it kept every server the key itself granted""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + stack.enter_context( + patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[]) + ) + ) + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert result == ["server-a"] + + async def test_get_allowed_mcp_servers_agent_dangling_toolset_denies(self): + """An agent toolset that resolves to nothing is a known restriction with unknown + contents: deny, never fall through to the key's own servers""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-gone"]) + mock_manager = self._mock_manager_with_toolsets({}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + with pytest.raises(UnloadableEntitlementError): + await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) + stack.enter_context( + patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[]) + ) + ) + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert result == [] + + async def test_get_agent_tool_permissions_for_server_unions_direct_and_toolset_tools(self): + """The agent's tool ceiling on a server is its direct tool grants plus the tools its + toolsets grant there, and None only when neither names the server""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission( + toolset_ids=["toolset-1"], tool_permissions={"server-a": ["tool_direct"]} + ) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_via_toolset"], "server-b": ["tool_b"]}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-a", user_api_key_auth) + server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-b", user_api_key_auth) + server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-c", user_api_key_auth) + + assert sorted(server_a_tools) == ["tool_direct", "tool_via_toolset"] + assert server_b_tools == ["tool_b"] + assert server_c_tools is None + + async def test_get_allowed_tools_for_server_toolset_only_agent_caps_key_tools(self): + """Regression: a key allowing [tool_a, tool_b] bound to an agent whose toolset grants + only tool_a on the server ends with [tool_a]; the toolset used to be ignored""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_a"]}) + key_perm = MagicMock() + key_perm.mcp_tool_permissions = {"server-a": ["tool_a", "tool_b"]} + key_perm.mcp_toolsets = [] + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + stack.enter_context( + patch.object( # test-quality-ok: stub the key perm loader; the resolver reads module globals with no injection seam + MCPRequestHandler, "_get_key_object_permission", return_value=key_perm + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: team resolution has its own tests; pin it absent here + MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None) + ) + ) + result = await MCPRequestHandler.get_allowed_tools_for_server("server-a", user_api_key_auth) + + assert result == ["tool_a"] + async def test_get_agent_object_permission_uses_shared_helper(self): """``_get_agent_object_permission`` must resolve the agent's ``object_permission_id`` and then defer to the shared @@ -7051,14 +7193,13 @@ class TestAggregateGatewayDcrChallenge: www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f"Bearer {self._EXPECTED_RESOURCE_METADATA}" - async def test_per_server_challenge_for_gateway_managed_oauth2(self): - """Anonymous request to a per-server path whose single target is a gateway-managed - oauth2 server: 401 plus the RFC 9728 challenge advertising the PER-SERVER - protected-resource metadata in the same URL spelling the request used, so a keyless - DCR client configured with either per-server spelling discovers the gateway as the - authorization server (LIT-4864). Covers interactive and M2M, which the gateway can - both serve end to end.""" - from litellm.types.mcp import MCPAuth + @pytest.mark.parametrize( + "auth_type", + (None, "none", "api_key", "bearer_token", "basic", "aws_sigv4", "authorization", "token", "oauth2"), + ) + @pytest.mark.parametrize("bearer_presented", (False, True)) + async def test_per_server_challenge_for_gateway_owned_auth(self, auth_type, bearer_presented): + """Gateway admission challenges are independent of upstream authentication.""" from litellm.types.mcp_server.mcp_server_manager import MCPServer server = MCPServer( @@ -7067,7 +7208,7 @@ class TestAggregateGatewayDcrChallenge: server_name="github", url="https://upstream.example/mcp", transport="http", - auth_type=MCPAuth.oauth2, + auth_type=auth_type, ) for path, expected_metadata_path in ( ("/mcp/github", "/.well-known/oauth-protected-resource/mcp/github"), @@ -7081,10 +7222,16 @@ class TestAggregateGatewayDcrChallenge: ): mock_mgr.get_mcp_server_by_name.return_value = server with pytest.raises(HTTPException) as exc_info: - await MCPRequestHandler.process_mcp_request(self._scope(path=path)) + await MCPRequestHandler.process_mcp_request( + self._scope( + path=path, + extra_headers=((b"authorization", b"Bearer invalid-key"),) if bearer_presented else (), + ) + ) assert exc_info.value.status_code == 401 www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] - assert www_authenticate == f'Bearer resource_metadata="http://testserver{expected_metadata_path}"' + error = 'error="invalid_token", ' if bearer_presented else "" + assert www_authenticate == f'Bearer {error}resource_metadata="http://testserver{expected_metadata_path}"' async def test_per_server_challenge_keeps_spelling_under_server_root_path(self): """On a sub-path deployment the challenge must still advertise the spelling the client @@ -7161,10 +7308,7 @@ class TestAggregateGatewayDcrChallenge: ) def test_challenge_target_excludes_every_non_gateway_managed_mode(self): - """Unit pin of the challenge-target owner: only a resolved gateway-managed oauth2 - target (interactive or M2M) yields a per-server challenge; delegate-auth oauth2 - (whose keyless flow is upstream PKCE via the relay), every client-forwarded auth - type, OBO, api_key, unknown names, and CSV paths yield None (LIT-4864).""" + """Gateway challenges exclude unresolved, delegated, and client-forwarded targets.""" from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( _gateway_dcr_challenge_target, ) @@ -7184,13 +7328,18 @@ 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), (_server(MCPAuth.true_passthrough), None), (_server(MCPAuth.oauth_delegate), None), (_server(MCPAuth.oauth_delegate, dcr_bridge=True), None), - (_server(MCPAuth.api_key), None), + (_server(MCPAuth.api_key), "srv"), + (_server(MCPAuth.none, extra_headers=["Authorization"]), None), + (_server(None, extra_headers=["X-API-Key"]), None), + (_server(MCPAuth.none, extra_headers=["Authorization"], oauth_passthrough=True), None), + (_server(MCPAuth.oauth2_id_jag), None), (None, None), ] for resolved, expected in cases: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py index 2ccba2b2055..9a66f130d24 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -2,15 +2,15 @@ import os import pytest -from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, -) - @pytest.fixture(autouse=True) def _hermetic_mcp_server_registry(): """Restore the singleton ``global_mcp_server_manager``'s registry state around every test, so entries seeded by one test never leak into another on a shared shard.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + saved_registry = dict(global_mcp_server_manager.registry) saved_config_servers = dict(global_mcp_server_manager.config_mcp_servers) saved_tool_mapping = dict(global_mcp_server_manager.tool_name_to_mcp_server_name_mapping) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_classify.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_classify.py index dc20d664a53..30107db4055 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_classify.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_classify.py @@ -1,6 +1,9 @@ """Classification matrix for upstream OAuth/DCR rejections: who is blamed depends only on the §5.2 code and whose credentials the gateway presented, never on the upstream's HTTP status.""" +from typing import Final + +import pytest import httpx from litellm.proxy._experimental.mcp_server.faults.classify import ( @@ -12,6 +15,7 @@ from litellm.proxy._experimental.mcp_server.faults.types import ( GatewayRejected, UpstreamProtocolFault, UpstreamReportedFault, + UpstreamRegistrationRefused, ) @@ -145,3 +149,28 @@ def test_dcr_server_error_code_is_not_blamed_on_caller(): log_context="srv", ) assert isinstance(fault, UpstreamReportedFault) + + +@pytest.mark.parametrize("status_code", [401, 403]) +@pytest.mark.parametrize("body", ["Forbidden", 'private upstream details', '{"error": ""}', '{"error": 12}']) +def test_dcr_access_refusal_without_oauth_error(status_code: int, body: str) -> None: + fault: Final = classify_upstream_dcr_rejection(_response(status_code, text_body=body), log_context="srv") + assert isinstance(fault, UpstreamRegistrationRefused) + assert fault.status_code == status_code + + +@pytest.mark.parametrize("status_code", [401, 403]) +def test_dcr_access_refusal_preserves_oauth_error(status_code: int) -> None: + fault: Final = classify_upstream_dcr_rejection( + _response(status_code, json_body={"error": "invalid_redirect_uri", "error_description": "not allowed"}), + log_context="srv", + ) + assert fault == CallerRejected(code="invalid_redirect_uri", description="not allowed") + + +@pytest.mark.parametrize("status_code", [401, 403]) +def test_token_access_refusal_remains_protocol_fault(status_code: int) -> None: + fault: Final = classify_upstream_token_rejection( + _response(status_code, text_body="Forbidden"), credential_source="gateway_stored", log_context="srv" + ) + assert isinstance(fault, UpstreamProtocolFault) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py index 78513e315a7..a6807ae1454 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_render_oauth.py @@ -2,6 +2,9 @@ code can never ship on a server-fault status and gateway-side faults never carry provider prose.""" import json +from typing import Final, Literal + +import pytest from litellm.proxy._experimental.mcp_server.faults.render_oauth import ( dcr_fault_detail, @@ -12,6 +15,7 @@ from litellm.proxy._experimental.mcp_server.faults.types import ( GatewayRejected, UpstreamProtocolFault, UpstreamReportedFault, + UpstreamRegistrationRefused, ) @@ -94,3 +98,17 @@ def test_dcr_upstream_reported_fault_maps_to_5xx(): status_code, detail = dcr_fault_detail(UpstreamReportedFault(code="server_error")) assert status_code == 502 assert "internal error" in detail + + +@pytest.mark.parametrize("upstream_status", [401, 403]) +def test_registration_refusal_gives_configuration_guidance(upstream_status: Literal[401, 403]) -> None: + fault: Final = UpstreamRegistrationRefused(status_code=upstream_status) + status, detail = dcr_fault_detail(fault) + assert status == 403 + assert f"HTTP {upstream_status}" in detail + assert "may require a pre-registered OAuth client" in detail + assert "client_id" in detail and "client_secret" in detail + response: Final = render_token_fault(fault) + assert response.status_code == 400 + assert json.loads(response.body) == {"error": "unauthorized_client", "error_description": detail} + assert response.headers["cache-control"] == "no-store" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index d67d0df4d0e..1b003e11993 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -476,16 +476,47 @@ def test_raise_public_plain_unauthorized_has_no_challenge(): @pytest.mark.parametrize( - "root_path, expected_prefix", + "root_path", [ - ("/", ""), # "/" means no prefix - ("", ""), # empty means no prefix - ("/api/v1", "/api/v1"), # a real root path is prepended verbatim + "/", # "/" means no prefix + "", # empty means no prefix ], ) -def test_oauth_protected_resource_path_honors_root_path(root_path, expected_prefix): +def test_oauth_protected_resource_path_no_prefix(root_path, monkeypatch): + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) path = oauth_protected_resource_path(root_path, _server(alias="my-srv")) - assert path == f"/.well-known/oauth-protected-resource{expected_prefix}/mcp/my-srv" + assert path == "/.well-known/oauth-protected-resource/mcp/my-srv" + + +def test_oauth_protected_resource_path_scalar_prefix_uses_rfc8414_insertion(monkeypatch): + # A scalar SERVER_ROOT_PATH deployment registers the well-known routes with + # the prefix inserted (via well_known_root_suffix at import time). The URL + # must match that insertion or a client fetching it 404s. + monkeypatch.setenv("SERVER_ROOT_PATH", "/api/v1") + path = oauth_protected_resource_path("/api/v1", _server(alias="my-srv")) + assert path == "/.well-known/oauth-protected-resource/api/v1/mcp/my-srv" + + +def test_oauth_protected_resource_path_per_request_prefix_goes_before_wellknown(monkeypatch): + # Per-request deployment: SERVER_ROOT_PATHS matched /tenant-a for this + # request but the scalar SERVER_ROOT_PATH is unset. Routes were registered + # without the well-known insertion, so the URL must place the prefix + # *before* .well-known — PerRequestRootPathMiddleware strips it and the + # router matches the un-inserted route. + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + path = oauth_protected_resource_path("/tenant-a", _server(alias="my-srv")) + assert path == "/tenant-a/.well-known/oauth-protected-resource/mcp/my-srv" + + +def test_oauth_protected_resource_path_dynamic_prefix_wins_over_scalar(monkeypatch): + # Both env vars configured: the middleware matched a SERVER_ROOT_PATHS + # prefix (/tenant-a) that differs from the scalar (/legacy). The URL must + # advertise /tenant-a — the prefix the client called — with no /legacy + # segment stacked onto it. Same review-fix invariant get_custom_url pins. + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + path = oauth_protected_resource_path("/tenant-a", _server(alias="my-srv")) + assert path == "/tenant-a/.well-known/oauth-protected-resource/mcp/my-srv" + assert "/legacy" not in path @pytest.mark.parametrize( @@ -510,7 +541,11 @@ def test_raise_user_oauth_challenge_points_at_per_server_prm(): ) -def test_raise_user_oauth_challenge_includes_server_root_path(): +def test_raise_user_oauth_challenge_includes_server_root_path(monkeypatch): + # The scalar deployment: routes are registered with the prefix inserted + # (via well_known_root_suffix at import time), so the challenge URL uses + # the RFC 8414 §3 insertion form. + monkeypatch.setenv("SERVER_ROOT_PATH", "/api/v1") with pytest.raises(HTTPException) as exc_info: raise_user_oauth_challenge(_server(alias="my-srv"), root_path="/api/v1") assert ( @@ -519,6 +554,20 @@ def test_raise_user_oauth_challenge_includes_server_root_path(): ) +def test_raise_user_oauth_challenge_per_request_prefix_is_routable(monkeypatch): + # Per-request deployment (SERVER_ROOT_PATHS matched /tenant-a): the + # challenge URL must place /tenant-a before .well-known so the client's + # discovery fetch routes through the same middleware strip the original + # request went through. The scalar-inserted form would 404 here. + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + with pytest.raises(HTTPException) as exc_info: + raise_user_oauth_challenge(_server(alias="my-srv"), root_path="/tenant-a") + assert ( + exc_info.value.headers["WWW-Authenticate"] + == 'Bearer resource_metadata="/tenant-a/.well-known/oauth-protected-resource/mcp/my-srv"' + ) + + def test_raise_token_exchange_challenge_is_rfc9728_invalid_token(): from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( raise_token_exchange_challenge, @@ -535,17 +584,30 @@ def test_raise_token_exchange_challenge_is_rfc9728_invalid_token(): assert "error_description=" in www -def test_raise_token_exchange_challenge_includes_server_root_path(): +def test_raise_token_exchange_challenge_includes_server_root_path(monkeypatch): from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( raise_token_exchange_challenge, ) + monkeypatch.setenv("SERVER_ROOT_PATH", "/api/v1") with pytest.raises(HTTPException) as exc_info: raise_token_exchange_challenge(_server(alias="obo-srv"), root_path="/api/v1") www = exc_info.value.headers["WWW-Authenticate"] assert 'resource_metadata="/.well-known/oauth-protected-resource/api/v1/mcp/obo-srv"' in www +def test_raise_token_exchange_challenge_per_request_prefix_is_routable(monkeypatch): + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_token_exchange_challenge, + ) + + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + with pytest.raises(HTTPException) as exc_info: + raise_token_exchange_challenge(_server(alias="obo-srv"), root_path="/tenant-a") + www = exc_info.value.headers["WWW-Authenticate"] + assert 'resource_metadata="/tenant-a/.well-known/oauth-protected-resource/mcp/obo-srv"' in www + + def test_raise_token_exchange_challenge_static_form_is_unchanged_without_step_up(): from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( raise_token_exchange_challenge, @@ -589,9 +651,7 @@ def test_id_jag_client_secret_maps_to_config(): # ID-JAG asserts the user's id_token; the access_token default maps to id_token. assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:id_token" assert isinstance(spec.config.client_auth, ClientSecretAuth) - assert spec.config.client_auth.client_secret.get_secret_value() == ( - "litellm-client-secret" - ) + assert spec.config.client_auth.client_secret.get_secret_value() == ("litellm-client-secret") def test_id_jag_private_key_maps_to_private_key_jwt_auth(): @@ -617,9 +677,7 @@ def test_id_jag_private_key_wins_over_client_secret(): def test_id_jag_honors_explicit_subject_token_type(): - spec = to_server_spec( - _id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2") - ) + spec = to_server_spec(_id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2")) assert spec is not None and isinstance(spec.config, IdJagConfig) assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:saml2" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py index bb2f2ff8b02..df068b60338 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py @@ -22,6 +22,7 @@ class _Server: server_id="srv", configured_token_url=None, ): + self.oauth_identity_binding = None self.token_url = token_url self.configured_token_url = configured_token_url self.client_id = client_id @@ -287,3 +288,40 @@ async def test_refresh_uses_admin_entered_token_url_when_issuer_yield_empties_re assert token is not None assert token.access_token == "new-at" assert posted[0][0] == "https://idp.example.com/token" + + +@pytest.mark.asyncio +async def test_identity_rejection_never_persists_or_returns_refreshed_token(): + from unittest.mock import AsyncMock + + from fastapi import HTTPException + + validator = AsyncMock(side_effect=HTTPException(status_code=403, detail="oauth_principal_mismatch")) + persist = AsyncMock() + refresher = AuthorizationCodeRefresher( + _lookup(_Server()), + _endpoint({"access_token": "foreign-token", "id_token": "foreign-identity"}), + persist, + identity_validator=validator, + ) + assert await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) is None + persist.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_verified_refresh_preserves_binding_proof_in_storage(): + from unittest.mock import AsyncMock + + validator = AsyncMock(return_value="verified-binding") + persist = AsyncMock() + refresher = AuthorizationCodeRefresher( + _lookup(_Server()), + _endpoint({"access_token": "new", "refresh_token": "rotated"}), + persist, + identity_validator=validator, + ) + token = await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) + assert token.access_token == "new" + assert token.refresh_token == "rotated" + assert token.identity_binding_proof == "verified-binding" + assert persist.await_args.kwargs["identity_binding_proof"] == "verified-binding" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py index 010e7e14d39..774cd022703 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py @@ -478,3 +478,47 @@ async def test_bearer_auth_advertises_the_header_it_will_occupy(): assert ClientCredentialsBearerAuth("t", refetch, ClientCredentialsConfig()).header_name == "Authorization" default_carrier = ClientCredentialsConfig(header_name="esb-oauth") assert ClientCredentialsBearerAuth("t", refetch, default_carrier).header_name == "esb-oauth" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["denied", "invalid", "missing", "success", "timeout", "connect", "cancel"]) +async def test_token_exchange_failure_diagnostics(mode, monkeypatch, caplog): + import asyncio + import logging + from litellm.llms.custom_httpx import http_handler + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import post_client_credentials_grant + + class Poster: + async def post(self, url, headers, data): + request = httpx.Request("POST", url, headers=headers, data=data) + if mode == "timeout": + raise httpx.ReadTimeout("private-transport-message", request=request) + if mode == "connect": + raise httpx.ConnectError("private-transport-message", request=request) + if mode == "cancel": + raise asyncio.CancelledError + response = httpx.Response(401 if mode == "denied" else 200, request=request, + content=b"not-json-private" if mode == "invalid" else None, + json=None if mode == "invalid" else {"error": "invalid_client", "client_secret":"first second", **({"access_token":"private-token"} if mode == "success" else {})}) + response.raise_for_status() + return response + + monkeypatch.setattr(http_handler, "get_async_httpx_client", lambda **kwargs: Poster()) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + if mode == "cancel": + with pytest.raises(asyncio.CancelledError): + await post_client_credentials_grant("https://idp/token", {}, {}) + assert not caplog.text + return + result = await post_client_credentials_grant("https://idp/token?key=query-secret", {"client_secret":"first second"}, {"X-Custom":"header-secret"}) + for secret in ("first", "second", "query-secret", "header-secret", "private-token", "not-json-private", "private-transport-message"): + assert secret not in caplog.text + if mode == "success": + assert isinstance(result, TokenEndpointSuccess) and result.body["access_token"] == "private-token" + assert not caplog.text + elif mode in {"timeout", "connect"}: + assert isinstance(result, TokenEndpointUnreachable) + assert "POST https://idp/ failed" in caplog.text + else: + assert "POST https://idp/ -> HTTP" in caplog.text + assert {"denied":"denied", "invalid":"invalid response", "missing":"no access token"}[mode] in caplog.text diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py index ca32cf2bb8d..e7308884060 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py @@ -246,3 +246,94 @@ async def test_lazy_store_invalidate_works_after_redis_chain_is_built() -> None: assert build_calls == 1 assert redis_store.invalidations == [("u", "s")] + + +@pytest.mark.asyncio +async def test_enforcement_invalidates_cached_legacy_credentials_before_use(): + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + server = MCPServer( + server_id="srv", + name="srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", + issuer="https://idp.example.com", + audiences=["client"], + ), + ) + cached = _RecordingStore("belongs-to-bob") + + store = LazyPerUserOAuthTokenStore( + lambda server_id: server, + store_builder=lambda lookup: (cached, False), + redis_available=lambda: False, + ) + assert await store.fetch("alice", "srv") is None + assert cached.calls == [("alice", "srv")] + assert cached.invalidations == [("alice", "srv")] + + +@pytest.mark.asyncio +async def test_enforced_cache_hit_avoids_credential_read_and_rejects_changed_policy(monkeypatch): + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof + from litellm.proxy._experimental.mcp_server.outbound_credentials import per_user_oauth_store as module + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", + name="srv", + transport="http", + auth_type="oauth2", + oauth_identity_binding={ + "mode": "enforce", + "issuer": "https://idp.example", + "audiences": ["client"], + "caller_field": "user_id", + "principal_claim": "sub", + }, + ) + proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv") + read = AsyncMock(return_value={"access_token": "alice-token", "identity_binding_proof": proof}) + monkeypatch.setattr(module, "_read_credential", read) + monkeypatch.setattr(module, "_runtime_backend_and_coordinator", lambda: (None, None, False)) + store = LazyPerUserOAuthTokenStore(lambda _: server, redis_available=lambda: False) + assert (await store.fetch("alice", "srv")).access_token == "alice-token" + assert (await store.fetch("alice", "srv")).access_token == "alice-token" + read.assert_awaited_once_with("alice", "srv") + server.oauth_identity_binding = server.oauth_identity_binding.model_copy(update={"audiences": ["changed"]}) + assert await store.fetch("alice", "srv") is None + read.assert_awaited_once() + assert await store.fetch("alice", "srv") is None + assert read.await_count == 2 + + +@pytest.mark.asyncio +async def test_expired_unverified_credential_never_reaches_refresh(monkeypatch): + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.outbound_credentials import per_user_oauth_store as module + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", + name="srv", + transport="http", + auth_type="oauth2", + token_url="https://idp.example/token", + oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"]}, + ) + read = AsyncMock( + return_value={"access_token": "bob", "refresh_token": "bob-refresh", "expires_at": "2000-01-01T00:00:00Z"} + ) + post = AsyncMock() + monkeypatch.setattr(module, "_read_credential", read) + monkeypatch.setattr(module, "_post_token_endpoint", post) + monkeypatch.setattr(module, "_runtime_backend_and_coordinator", lambda: (None, None, False)) + store = LazyPerUserOAuthTokenStore(lambda _: server, redis_available=lambda: False) + assert await store.fetch("alice", "srv") is None + post.assert_not_awaited() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 6b3098d9e60..5fab4ceec72 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -19,6 +19,7 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials import ( ApiKeyConfig, + AuthConfig, AuthorizationCodeConfig, AwsSigV4Config, Byok, @@ -1203,3 +1204,71 @@ async def test_passthrough_ignores_the_carrier_and_keeps_the_callers_slot(): assert isinstance(result, Ok) headers, _ = await _emitted_async(result.ok) assert headers["Authorization"] == "caller-token" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("config", "subject", "expected_source", "expected_header"), + [ + (NoneConfig(), _SUBJECT, "no-auth", None), + (PassthroughConfig(), _SUBJECT, "no-auth", None), + (PassthroughConfig(), _with_inbound("Bearer caller-token"), "oauth2-passthrough", "Bearer caller-token"), + (ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))), _SUBJECT, "static-token", "Bearer static-key"), + (AuthorizationCodeConfig(), Subject(tenant_id="", subject_id="alice"), "stored-user-token", "Bearer stored-alice"), + ], +) +async def test_resolved_source_matches_the_credential_sent_upstream( + config: AuthConfig, subject: Subject, expected_source: str, expected_header: str | None +) -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="stored-alice")}) + provider = UpstreamCredentialProvider(oauth_token_store=store) + result = await resolve_credentials_with_source(provider, subject, _spec(config)) + assert isinstance(result, Ok) + assert result.ok.source.value == expected_source + assert _emitted(result.ok.auth).get("Authorization") == expected_header + assert "stored-alice" not in repr(result.ok) + assert "static-key" not in repr(result.ok) + + +@pytest.mark.asyncio +async def test_resolved_source_preserves_missing_user_token_error() -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + result = await resolve_credentials_with_source(UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig())) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + + +@pytest.mark.asyncio +async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) + m2m = await resolve_credentials_with_source(UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M)) + assert isinstance(m2m, Ok) + headers, _ = await _emitted_async(m2m.ok.auth) + assert headers["Authorization"] == "Bearer m2m-at" + assert m2m.ok.source.value == "m2m-client-credentials" + assert source.gets == ["s"] + + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) + exchanged = await resolve_credentials_with_source(UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO)) + assert isinstance(exchanged, Ok) + assert _emitted(exchanged.ok.auth)["Authorization"] == "Bearer exchanged-at" + assert exchanged.ok.source.value == "token-exchange" + assert len(exchanger.calls) == 1 + + +@pytest.mark.asyncio +async def test_id_jag_source_describes_final_token_after_both_exchanges() -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + endpoint = _FakeTokenEndpoint(_two_leg_ok("resource-token")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await resolve_credentials_with_source(provider, _with_inbound("identity-token"), _spec(_id_jag_config())) + assert isinstance(result, Ok) + assert result.ok.source.value == "id-jag" + assert _emitted(result.ok.auth)["Authorization"] == "Bearer resource-token" + assert len(endpoint.calls) == 2 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py index 17a7e13af03..12a04c4b4dc 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py @@ -52,3 +52,17 @@ def test_undecryptable_blob_is_a_miss(): def test_empty_plaintext_is_a_miss(): codec = OAuthTokenCacheCodec(encrypt=lambda s: s, decrypt=lambda b: b) assert codec.decode("") is None + + +def test_bound_token_round_trip_preserves_proof_without_refresh_secret(): + codec = _wrapping_codec() + blob = codec.encode(OAuthToken("alice-token", refresh_token="private-refresh", identity_binding_proof="proof")) + assert "private-refresh" not in blob + decoded = codec.decode(blob) + assert decoded == OAuthToken("alice-token", identity_binding_proof="proof") + + +def test_malformed_bound_entries_fail_closed(): + codec = _wrapping_codec() + for payload in ("not-json", "{}", '{"access_token":"at"}', '{"access_token":1,"identity_binding_proof":"p"}'): + assert codec.decode("enc:litellm-bound-oauth-v1:" + payload) is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 4d9142ad4c5..60a5e1a22bb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -12,28 +12,39 @@ import base64 import json from datetime import datetime, timedelta, timezone from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest +from prisma.models import LiteLLM_MCPServerTable as PrismaMCPServer from litellm.proxy._experimental.mcp_server.db import ( _decode_user_credential, _prepare_mcp_server_data, + create_mcp_server, decrypt_credentials, encrypt_credentials, + get_all_mcp_servers, + get_mcp_servers, + get_mcp_submissions, get_user_credential, get_user_oauth_credential, is_oauth_credential_expired, list_user_oauth_credentials, resolve_valid_user_oauth_token, + rotate_mcp_server_credentials_master_key, rotate_mcp_user_credentials_master_key, rotate_mcp_user_env_vars_master_key, store_user_credential, store_user_oauth_credential, + update_mcp_server, ) -from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest +from litellm.proxy._types import LiteLLM_MCPServerTable, NewMCPServerRequest, UpdateMCPServerRequest from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + SecretMapDecodeError, + decode_secret_map, decrypt_value_helper, + encrypt_secret_map, encrypt_value_helper, ) from litellm.types.mcp import MCPAuth, MCPTransport @@ -44,6 +55,7 @@ SALT_KEY = "test-salt-key-for-byok-credential-tests-1234" @pytest.fixture(autouse=True) def _set_salt_key(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": "xsalsa20-poly1305"}) def _make_prisma_with_existing(row): @@ -368,6 +380,169 @@ def test_client_private_key_encrypted_at_rest(): assert decrypted["client_secret"] == "shh" +@pytest.fixture(params=["xsalsa20-poly1305", "aes-256-gcm"]) +def map_algorithm(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> str: + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": request.param}) + return request.param + + +def _prisma_map_row(data: dict[str, object], quoted: bool = False) -> PrismaMCPServer: + return PrismaMCPServer.model_validate({ + "transport": "http", "mcp_access_groups": [], "allowed_tools": [], "extra_headers": [], "args": [], + "allow_all_keys": False, "available_on_public_internet": True, "delegate_auth_to_upstream": False, + "oauth_passthrough": False, "per_server_oauth_discovery": False, "is_byok": False, "byok_description": [], + **data, + **{field: json.dumps(data[field]) for field in ("static_headers", "env") if quoted and data.get(field)}, + }) + + +class _MapTable: + def __init__(self, *rows: dict[str, object], quoted: bool = False) -> None: + self.rows = {row["server_id"]: row for row in rows} + self.quoted = quoted + + async def create(self, *, data: dict[str, object]) -> PrismaMCPServer: + self.rows = {**self.rows, data["server_id"]: dict(data)} + return _prisma_map_row(data, self.quoted) + + async def update(self, *, where: dict[str, str], data: dict[str, object]) -> PrismaMCPServer: + return await self.create(data={**self.rows[where["server_id"]], **data}) + + async def find_many(self, where: object = None) -> list[PrismaMCPServer]: + return [_prisma_map_row(row, self.quoted) for row in self.rows.values()] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["static_headers", "env"]) +@pytest.mark.parametrize("quoted", [False, True]) +async def test_secret_maps_create_update_round_trip(map_algorithm: str, field: str, quoted: bool) -> None: + table: Final = _MapTable(quoted=quoted) + prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + original: Final = {"TOKEN": " sensitive-secret\n", "PREFIX": "v2:gcm:literal", "TEMPLATE": "Bearer ${TOKEN}"} + create: Final = NewMCPServerRequest.model_validate({ + "server_id": "srv-map", "transport": "http", "url": "https://up.example.com/mcp", field: original, + }) + created: Final = await create_mcp_server(prisma, create, touched_by="test") + first: Final = table.rows["srv-map"][field] + assert isinstance(first, str) and isinstance(json.loads(first), str) + assert json.loads(first).startswith("v2:gcm:") is (map_algorithm == "aes-256-gcm") + assert "sensitive-secret" not in first and "TEMPLATE" not in first + assert getattr(created, field) == original == getattr(create, field) + assert decode_secret_map(first, key=field) == original + replacement: Final = {**original, "TOKEN": "updated-sensitive-secret"} + update: Final = UpdateMCPServerRequest.model_validate({"server_id": "srv-map", field: replacement}) + updated: Final = await update_mcp_server(prisma, update, touched_by="test") + second: Final = table.rows["srv-map"][field] + assert second != first and "updated-sensitive-secret" not in second + assert decode_secret_map(second, key=field) == replacement + assert getattr(updated, field) == replacement == getattr(update, field) + assert original["TOKEN"] == " sensitive-secret\n" + omitted: Final = await update_mcp_server(prisma, UpdateMCPServerRequest(server_id="srv-map"), touched_by="test") + assert table.rows["srv-map"][field] == second and getattr(omitted, field) == replacement + cleared: Final = await update_mcp_server( + prisma, UpdateMCPServerRequest.model_validate({"server_id": "srv-map", field: {}}), touched_by="test" + ) + assert table.rows["srv-map"][field] == "{}" and getattr(cleared, field) == {} + + +@pytest.mark.parametrize("field", ["static_headers", "env"]) +@pytest.mark.parametrize("as_json", [False, True]) +def test_secret_map_legacy_model_read_preserves_exact_values(field: str, as_json: bool) -> None: + original: Final = {"PREFIX": "v2:gcm:literal", "SPACE": " secret\n", "TEMPLATE": "${TOKEN}", "B64": "YWJjZA=="} + incoming: Final = { + "server_id": "srv-map", "transport": "http", field: json.dumps(original) if as_json else original, + } + snapshot: Final = json.dumps(incoming) + parsed: Final = LiteLLM_MCPServerTable.model_validate(incoming) + assert getattr(parsed, field) == original + assert json.dumps(incoming) == snapshot + assert LiteLLM_MCPServerTable.model_validate(parsed.model_dump()).model_dump() == parsed.model_dump() + empty: Final = LiteLLM_MCPServerTable.model_validate({"server_id": "srv-map", "transport": "http", field: None}) + assert getattr(empty, field) == ({} if field == "env" else None) + + +@pytest.mark.parametrize("field", ["static_headers", "env"]) +@pytest.mark.parametrize("failure", ["wrong-key", "corrupt", "invalid-values", "invalid-shape", "invalid-json"]) +def test_secret_map_model_read_fails_closed(map_algorithm: str, field: str, failure: str) -> None: + plaintext: Final = {"invalid-values": '{"TOKEN": ["sensitive-secret"]}', "invalid-shape": '["sensitive-secret"]', + "invalid-json": "sensitive-secret"}.get(failure, '{"TOKEN": "sensitive-secret"}') + ciphertext: Final = encrypt_value_helper( + plaintext, new_encryption_key="wrong-map-key" if failure == "wrong-key" else None + ) + stored: Final = json.dumps(ciphertext[:-8] if failure == "corrupt" else ciphertext) + with pytest.raises(SecretMapDecodeError) as exc: + LiteLLM_MCPServerTable.model_validate({"server_id": "srv-map", "transport": "http", field: stored}) + assert field in str(exc.value) and "LITELLM_SALT_KEY" in str(exc.value) + assert all(secret not in str(exc.value) for secret in (plaintext, ciphertext, "sensitive-secret")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field,other", [("static_headers", "env"), ("env", "static_headers")]) +async def test_secret_map_rotation_migrates_rekeys_and_preserves_corrupt( + map_algorithm: str, field: str, other: str, monkeypatch: pytest.MonkeyPatch +) -> None: + values: Final = {"TOKEN": "rotation-sensitive-secret", "TEMPLATE": "Bearer ${TOKEN}"} + old: Final = encrypt_secret_map(values) + corrupt: Final = json.dumps(json.loads(old)[:-8]) + table: Final = _MapTable( + {"server_id": "broken", field: corrupt, other: old}, + {"server_id": "legacy", field: json.dumps(values), other: "{}"}, + {"server_id": "encrypted", field: old, other: None}, + ) + prisma: Final = SimpleNamespace(db=SimpleNamespace( + litellm_mcpservertable=table, litellm_mcpserveroauthclient=SimpleNamespace(find_many=AsyncMock(return_value=[])) + )) + await rotate_mcp_server_credentials_master_key(prisma, touched_by="test", new_master_key="rotated-map-key") + assert table.rows["broken"][field] == corrupt + assert table.rows["legacy"][other] == "{}" and table.rows["encrypted"][other] is None + for server_id, map_field in (("broken", other), ("legacy", field), ("encrypted", field)): + stored: Final = table.rows[server_id][map_field] + assert isinstance(json.loads(stored), str) and stored != old and "rotation-sensitive-secret" not in stored + with pytest.raises(SecretMapDecodeError): + decode_secret_map(stored, key=map_field) + monkeypatch.setenv("LITELLM_SALT_KEY", "rotated-map-key") + for server_id, map_field in (("broken", other), ("legacy", field), ("encrypted", field)): + assert decode_secret_map(table.rows[server_id][map_field], key=map_field) == values + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reader", [get_all_mcp_servers, get_mcp_servers, get_mcp_submissions]) +@pytest.mark.parametrize("field", ["static_headers", "env"]) +async def test_bulk_reads_isolate_corrupt_secret_maps(reader, field, map_algorithm, caplog): + secret = {"TOKEN": "bulk-sensitive-secret"} + encrypted = encrypt_secret_map(secret) + corrupt = encrypt_secret_map(secret, new_encryption_key="wrong-bulk-key") + rows = [ + _prisma_map_row({"server_id": "broken", field: corrupt, "approval_status": "pending_review"}), + _prisma_map_row({"server_id": "healthy", field: encrypted, "approval_status": "active"}), + ] + snapshot = [row.model_dump() for row in rows] + table = SimpleNamespace(find_many=AsyncMock(return_value=rows)) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + result = await reader(prisma, ["broken", "healthy"]) if reader is get_mcp_servers else await reader(prisma) + items = result.items if reader is get_mcp_submissions else result + assert [row.server_id for row in items] == ["healthy"] + assert getattr(items[0], field) == secret + assert [row.model_dump() for row in rows] == snapshot + assert "broken" in caplog.text + assert all(value not in caplog.text for value in ("bulk-sensitive-secret", corrupt, encrypted)) + if reader is get_mcp_submissions: + assert (result.total, result.pending_review, result.active, result.rejected) == (1, 0, 1, 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reader", [get_all_mcp_servers, get_mcp_servers, get_mcp_submissions]) +async def test_bulk_reads_do_not_swallow_unrelated_validation_errors(reader): + from pydantic import ValidationError + + row = _prisma_map_row({"server_id": "invalid", "transport": "unsupported"}) + table = SimpleNamespace(find_many=AsyncMock(return_value=[row])) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + request = reader(prisma, ["invalid"]) if reader is get_mcp_servers else reader(prisma) + with pytest.raises(ValidationError, match="transport"): + await request + + # ── BYOK round-trip ─────────────────────────────────────────────────────────── @@ -453,6 +628,7 @@ async def test_oauth_round_trip_returns_payload(): access_token, refresh_token="rfr-xyz", scopes=["a", "b"], + identity_binding_proof="verified-proof", ) stored = _stored_value(prisma) @@ -467,6 +643,7 @@ async def test_oauth_round_trip_returns_payload(): assert result["access_token"] == access_token assert result["refresh_token"] == "rfr-xyz" assert result["scopes"] == ["a", "b"] + assert result["identity_binding_proof"] == "verified-proof" @pytest.mark.asyncio @@ -1009,6 +1186,7 @@ class _RefreshResponse: def _refresh_server(**overrides): base = dict( + oauth_identity_binding=None, token_url="https://idp.example.com/token", server_id="srv-1", client_id="cid", @@ -1134,7 +1312,7 @@ async def test_refresh_user_oauth_token_uses_client_secret_basic(monkeypatch): sends HTTP Basic and keeps the secret out of the body.""" import litellm.proxy._experimental.mcp_server.db as db_mod - server = MagicMock() + server = MagicMock(oauth_identity_binding=None) server.token_url = "https://idp.example.com/oauth2/token" server.server_id = "srv" server.client_id = "cid" @@ -1173,7 +1351,7 @@ async def test_refresh_user_oauth_token_defaults_to_client_secret_post(monkeypat the body (client_secret_post) and sends no Authorization header.""" import litellm.proxy._experimental.mcp_server.db as db_mod - server = MagicMock() + server = MagicMock(oauth_identity_binding=None) server.token_url = "https://idp.example.com/oauth2/token" server.server_id = "srv" server.client_id = "cid" @@ -1362,3 +1540,167 @@ 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 + + +@pytest.mark.asyncio +async def test_enforcement_rejects_preexisting_unverified_credential(): + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + server = MCPServer( + server_id="srv-1", name="srv-1", url="https://mcp.example.com", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", issuer="https://idp.example.com", audiences=["client"], + ), + ) + result = await resolve_valid_user_oauth_token( + user_id="alice", server=server, + cred={"access_token": "belongs-to-bob", "refresh_token": "bobs-refresh-token"}, + ) + assert result is None + + +@pytest.mark.asyncio +async def test_refresh_identity_rejection_returns_reauthentication_without_persisting(monkeypatch): + from fastapi import HTTPException + from litellm.proxy._experimental.mcp_server import db as module + + validator = AsyncMock(side_effect=HTTPException(status_code=403, detail="oauth_principal_mismatch")) + monkeypatch.setattr(module, "enforce_oauth_identity_binding", validator) + result, captured = await _run_refresh( + monkeypatch, _refresh_server(), {"access_token": "bob", "refresh_token": "rotated"} + ) + assert result is None + assert captured["data"]["grant_type"] == "refresh_token" + module.store_user_oauth_credential.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_verified_legacy_cache_reads_avoid_database_and_reject_policy_changes(monkeypatch): + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db as module + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", + name="srv", + transport="http", + auth_type="oauth2", + oauth_identity_binding={ + "mode": "enforce", + "issuer": "https://idp.example", + "audiences": ["client"], + "caller_field": "user_id", + "principal_claim": "sub", + }, + ) + proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv") + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + read = AsyncMock(return_value=None) + monkeypatch.setattr(module, "get_user_oauth_credential", read) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + await mcp_per_user_token_cache.set("alice", "srv", "alice-token", 60, identity_binding_proof=proof) + assert await module.resolve_user_oauth_access_token("alice", server) == "alice-token" + assert await module.resolve_user_oauth_access_token("alice", server) == "alice-token" + read.assert_not_awaited() + server.oauth_identity_binding = server.oauth_identity_binding.model_copy(update={"audiences": ["changed"]}) + assert await module.resolve_user_oauth_access_token("alice", server) is None + read.assert_awaited_once() + assert await mcp_per_user_token_cache.get_token("alice", "srv") is None + + +@pytest.mark.asyncio +async def test_unverified_legacy_cache_cannot_bypass_enforcement(monkeypatch): + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db as module + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", + name="srv", + transport="http", + auth_type="oauth2", + oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"]}, + ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(module, "get_user_oauth_credential", AsyncMock(return_value={"access_token": "bob"})) + await mcp_per_user_token_cache.set("alice", "srv", "bob", 60) + assert await module.resolve_user_oauth_access_token("alice", server) is None + assert await mcp_per_user_token_cache.get("alice", "srv") is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 775f6e5f3b8..1b7e3d6a1c3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4,6 +4,7 @@ import hashlib import json import time from base64 import urlsafe_b64encode +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -18,6 +19,57 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer +def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None): + credential = {"type": "oauth2", "access_token": access_token} + if refresh_token is not None: + credential["refresh_token"] = refresh_token + if expires_in_seconds is not None: + credential["expires_at"] = (datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)).isoformat() + if expires_at is not None: + credential["expires_at"] = expires_at + return credential + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("fields", "egress_has_token"), + [ + (None, False), + ({"access_token": "", "refresh_token": "refresh-token"}, False), + ({}, True), + ({"expires_at": "never"}, True), + ({"expires_in_seconds": 600}, True), + ({"expires_in_seconds": 30}, False), + ({"expires_in_seconds": -300}, False), + ({"expires_in_seconds": -300, "refresh_token": ""}, False), + ({"expires_in_seconds": 30, "refresh_token": "refresh-token"}, True), + ({"expires_in_seconds": -300, "refresh_token": "refresh-token"}, True), + ], +) +async def test_vendor_credential_state_agrees_with_egress_token_resolution(monkeypatch, fields, egress_has_token): + from litellm.proxy._experimental.mcp_server import db as mcp_db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + + monkeypatch.setattr(mcp_db, "MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", 60) + credential = _stored_grant(**fields) if fields is not None else None + read = AsyncMock(return_value=credential) + refresh = AsyncMock(return_value=_stored_grant(access_token="fresh-token", expires_in_seconds=3600)) + prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr(mcp_db, "get_user_oauth_credential", read) + monkeypatch.setattr(mcp_db, "refresh_user_oauth_token", refresh) + + connect = await discoverable_endpoints._vendor_credential_state("user-1", "server-1") + read.assert_awaited_once_with(prisma, "user-1", "server-1") + refresh.assert_not_awaited() + egress = await mcp_db.resolve_valid_user_oauth_token( + user_id="user-1", server=MagicMock(), cred=credential, prisma_client=prisma + ) + + assert (egress is not None) is egress_has_token + assert connect == ("present" if egress_has_token else "absent") + + # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control @pytest.fixture(autouse=True) @@ -3342,12 +3394,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 +3416,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 ) @@ -7088,7 +7150,7 @@ async def test_extract_user_id_rehydrates_cross_replica_dict_cache(proxy_globals key = "sk-alice-key" cache = UserApiKeyCache() - cache.in_memory_cache.set_cache(hash_token(key), {"token": hash_token(key), "user_id": "alice"}) + cache.set_cache(hash_token(key), {"token": hash_token(key), "user_id": "alice"}) proxy_globals.user_api_key_cache = cache proxy_globals.prisma_client = object() @@ -7589,12 +7651,6 @@ async def test_token_endpoint_client_secret_basic_without_secret_returns_400(): assert exc_info.value.status_code == 400 -# ------------------------------------------------------------------- -# Non-oauth2 (auth_type=none, access-group gated) servers must not be -# driven through the gateway OAuth authorize/token/register/discovery -# flow, and must not be advertised as OAuth-protected in discovery docs. -# ------------------------------------------------------------------- - def _access_group_none_server(server_name="access_group_server"): """A non-oauth2, access-group gated MCP server: no client_id, no OAuth.""" @@ -7732,35 +7788,38 @@ async def test_register_client_rejects_non_oauth2_server(): @pytest.mark.asyncio -async def test_oauth_protected_resource_404_for_non_oauth2_server(): - """Discovery must not advertise a none-auth server as an OAuth-protected resource.""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _build_oauth_protected_resource_response, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - except ImportError: - pytest.skip("MCP discoverable endpoints not available") +@pytest.mark.parametrize( + "auth_type", (None, "none", "api_key", "bearer_token", "basic", "aws_sigv4", "authorization", "token") +) +@pytest.mark.parametrize("use_standard_pattern", (False, True)) +async def test_oauth_protected_resource_for_gateway_owned_auth(auth_type, use_standard_pattern): + from starlette.requests import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_protected_resource_response, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = _access_group_none_server().model_copy(update={"auth_type": auth_type}) + request = Request( + {"type": "http", "scheme": "https", "path": "/", "headers": [(b"host", b"litellm.example.com")]} + ) global_mcp_server_manager.registry.clear() - server = _access_group_none_server() global_mcp_server_manager.registry[server.server_id] = server - - mock_request = MagicMock() - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - try: - with pytest.raises(HTTPException) as exc_info: - await _build_oauth_protected_resource_response( - request=mock_request, - mcp_server_name="access_group_server", - use_standard_pattern=False, - ) - assert exc_info.value.status_code == 404 - assert "not an OAuth-protected resource" in str(exc_info.value.detail) + response = await _build_oauth_protected_resource_response( + request=request, + mcp_server_name="access_group_server", + use_standard_pattern=use_standard_pattern, + ) + resource_path = "/mcp/access_group_server" if use_standard_pattern else "/access_group_server/mcp" + assert response == { + "resource": f"https://litellm.example.com{resource_path}", + "authorization_servers": ["https://litellm.example.com/mcp"], + "scopes_supported": [], + } finally: global_mcp_server_manager.registry.clear() @@ -7852,9 +7911,7 @@ async def test_oauth_protected_resource_passthrough_none_auth_not_404(): @pytest.mark.asyncio async def test_oauth_protected_resource_404_for_unknown_server_name(): - """A discovery request for an unknown server name returns the same 404 as a non-oauth2 - server (not a 200 metadata doc with broken URLs), so the well-known paths cannot be used - to enumerate non-OAuth server names.""" + """Unknown server names must not produce metadata advertising nonexistent resources.""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _build_oauth_protected_resource_response, @@ -8066,6 +8123,137 @@ async def test_token_exchange_pairs_client_secret_with_server_client_id(): assert "client_secret" not in sent +@pytest.mark.asyncio +async def test_token_exchange_refresh_passes_presented_refresh_ownership(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import RefreshTokenPresented + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + server = MCPServer( + server_id="srv-1", + name="srv-1", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + token_url="https://provider.example/token", + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", + issuer="https://provider.example", + audiences=["cid"], + ), + ) + request = MagicMock(spec=Request) + request.base_url = "https://litellm.example.com/" + request.headers = {} + response = MagicMock() + response.json.return_value = {"access_token": "at"} + response.raise_for_status = MagicMock() + client = MagicMock() + client.post = AsyncMock(return_value=response) + enforce = AsyncMock() + + with ( + patch( # test-quality-ok: no injection seam exists for the exchange's HTTP and identity collaborators + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=client, + ), + patch( # test-quality-ok: no injection seam exists for request identity extraction + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request", + new=AsyncMock(return_value="user-a"), + ), + patch( # test-quality-ok: captures the ownership value at the exchange boundary + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.enforce_oauth_identity_binding", + new=enforce, + ), + ): + await exchange_token_with_server( + request=request, + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="cid", + client_secret=None, + code_verifier=None, + refresh_token="rt-1", + ) + + ownership = enforce.await_args.kwargs["refresh_ownership"] + assert isinstance(ownership, RefreshTokenPresented) + assert ownership.refresh_token == "rt-1" + + +@pytest.mark.asyncio +async def test_token_exchange_authorization_code_passes_no_refresh_ownership(monkeypatch): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + seal_bridge_authorization_code, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + monkeypatch.setenv("LITELLM_SALT_KEY", "identity-binding-test-salt") + server = MCPServer( + server_id="srv-1", + name="srv-1", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + token_url="https://provider.example/token", + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", + issuer="https://provider.example", + audiences=["cid"], + ), + ) + request = MagicMock(spec=Request) + request.base_url = "https://litellm.example.com/" + request.headers = {} + response = MagicMock() + response.json.return_value = {"access_token": "at"} + response.raise_for_status = MagicMock() + client = MagicMock() + client.post = AsyncMock(return_value=response) + enforce = AsyncMock() + + with ( + patch( # test-quality-ok: no injection seam exists for the exchange's HTTP and identity collaborators + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=client, + ), + patch( # test-quality-ok: no injection seam exists for request identity extraction + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request", + new=AsyncMock(return_value="user-a"), + ), + patch( # test-quality-ok: captures the ownership value at the exchange boundary + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.enforce_oauth_identity_binding", + new=enforce, + ), + ): + result = await exchange_token_with_server( + request=request, + mcp_server=server, + grant_type="authorization_code", + code=seal_bridge_authorization_code("auth-code", "user-a", "srv-1", "login-nonce"), + redirect_uri="https://litellm.example.com/callback", + client_id="cid", + client_secret=None, + code_verifier="test-verifier", + ) + + assert result.status_code == 200 + assert json.loads(result.body)["access_token"] == "at" + assert enforce.await_args.kwargs["refresh_ownership"] is None + assert enforce.await_args.kwargs["expected_nonce"] == "login-nonce" + + def _upstream_token_response(status_code: int, *, json_body: object = None, text_body: str = "") -> "httpx.Response": import httpx @@ -10332,6 +10520,230 @@ async def test_upstream_resource_sent_on_dcr_bridge_relay_authorize(): assert query["client_id"] == ["caller-client"] +# --------------------------------------------------------------------------- +# Per-request root_path (SERVER_ROOT_PATHS / PerRequestRootPathMiddleware): +# one app fronting several client-visible URL path prefixes, each prefix's +# discovery documents emitting URLs under the prefix the client called +# (RFC 9728 §3 exact-match). A scalar PROXY_BASE_URL / SERVER_ROOT_PATH can +# encode at most one prefix per pod; these tests pin the N-prefix case. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _no_proxy_base_url(monkeypatch): + """Discovery must derive URLs from the request in these tests, so the + scalar env overrides are cleared.""" + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + + +@pytest.fixture +def _isolated_mcp_registry(): + """Fixture-owned registry state: snapshot the shared registry, hand the + test an empty one, restore afterwards so nothing leaks between cases.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + saved = dict(global_mcp_server_manager.registry) + global_mcp_server_manager.registry.clear() + try: + yield global_mcp_server_manager.registry + finally: + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry.update(saved) + + +def _prefixed_discovery_client(prefixes): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy.middleware.per_request_root_path_middleware import ( + PerRequestRootPathMiddleware, + ) + + app = FastAPI() + app.include_router(router) + app.add_middleware(PerRequestRootPathMiddleware, root_paths=prefixes) + return TestClient(app) + + +class TestPerRequestRootPathDiscovery: + def test_prefixed_wellknown_not_routable_without_middleware(self, _isolated_mcp_registry): + """Control: on a plain app (the only shape a scalar root_path can + express), a prefixed well-known request 404s before any discovery + builder runs — the routing gap this feature exists to close.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + + server = _create_oauth2_server(server_id="srv_a", name="server_a", server_name="server_a", alias="server_a") + _isolated_mcp_registry[server.server_id] = server + app = FastAPI() + app.include_router(router) + client = TestClient(app) + resp = client.get("/tenant-a/.well-known/oauth-protected-resource/mcp/server_a") + assert resp.status_code == 404 + + def test_two_prefixes_one_app_each_resource_matches_the_called_url( + self, _no_proxy_base_url, _isolated_mcp_registry + ): + """The multi-origin case itself: two prefixes served by the same app, + each per-server document's ``resource`` equal to the URL its client + called — including the prefix.""" + for sid, name in (("srv_a", "server_a"), ("srv_b", "server_b")): + _isolated_mcp_registry[sid] = _create_oauth2_server(server_id=sid, name=name, server_name=name, alias=name) + client = _prefixed_discovery_client(["/tenant-a", "/tenant-b"]) + + resp_a = client.get("/tenant-a/.well-known/oauth-protected-resource/mcp/server_a") + resp_b = client.get("/tenant-b/.well-known/oauth-protected-resource/mcp/server_b") + + assert resp_a.status_code == 200 + assert resp_a.json()["resource"] == "http://testserver/tenant-a/mcp/server_a" + assert resp_b.status_code == 200 + assert resp_b.json()["resource"] == "http://testserver/tenant-b/mcp/server_b" + + # Every URL the document advertises stays under the request's + # prefix, so it resolves on this same app. + for auth_server in resp_a.json()["authorization_servers"]: + assert auth_server.startswith("http://testserver/tenant-a/") + + def test_unprefixed_requests_unchanged_on_the_same_app(self, _no_proxy_base_url, _isolated_mcp_registry): + """Backward compat on the very same app: a root request emits the + document byte-identical to a deployment without the middleware.""" + server = _create_oauth2_server(server_id="srv_a", name="server_a", server_name="server_a", alias="server_a") + _isolated_mcp_registry[server.server_id] = server + client = _prefixed_discovery_client(["/tenant-a"]) + resp = client.get("/.well-known/oauth-protected-resource/mcp/server_a") + assert resp.status_code == 200 + assert resp.json()["resource"] == "http://testserver/mcp/server_a" + + def test_unlisted_prefix_404s(self, _no_proxy_base_url): + client = _prefixed_discovery_client(["/tenant-a"]) + assert client.get("/tenant-c/.well-known/oauth-protected-resource/mcp/server_a").status_code == 404 + + def test_aggregate_documents_and_as_endpoints_under_prefix(self, _no_proxy_base_url, _isolated_mcp_registry): + """Aggregate PRM/AS documents carry the prefix, and the advertised + authorize endpoint actually resolves under it — the 404 trap that + invalidated prefixing discovery URLs without per-request routing + (#35226 review round 1).""" + client = _prefixed_discovery_client(["/tenant-a"]) + + prm = client.get("/tenant-a/.well-known/oauth-protected-resource/mcp") + asm = client.get("/tenant-a/.well-known/oauth-authorization-server/mcp") + + assert prm.status_code == 200 + assert prm.json()["resource"] == "http://testserver/tenant-a/mcp" + assert prm.json()["authorization_servers"] == ["http://testserver/tenant-a/mcp"] + + assert asm.status_code == 200 + assert asm.json()["issuer"] == "http://testserver/tenant-a/mcp" + assert asm.json()["authorization_endpoint"] == "http://testserver/tenant-a/authorize" + + # The prefixed authorize URL routes to the real handler (not 404): + # under per-request root_path the whole app is reachable per-prefix, + # so discovery may advertise prefixed AS endpoints safely. + assert client.get("/tenant-a/authorize").status_code != 404 + + def test_passthrough_challenge_metadata_url_carries_prefix(self, _no_proxy_base_url): + """The WWW-Authenticate resource_metadata URL a 401 advertises must + land under the request's prefix, or the client is bounced to a + document whose ``resource`` cannot match the URL it called.""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + get_passthrough_resource_metadata_url, + ) + + def _scope(path, root_path=None): + scope = { + "type": "http", + "method": "POST", + "path": path, + "headers": [], + "scheme": "http", + "server": ("testserver", 80), + "client": ("1.2.3.4", 4444), + } + if root_path is not None: + scope["root_path"] = root_path + return scope + + prefixed = get_passthrough_resource_metadata_url( + scope=_scope("/tenant-a/mcp/github", root_path="/tenant-a"), + server_name="github", + ) + assert prefixed == "http://testserver/tenant-a/.well-known/oauth-protected-resource/mcp/github" + + # Regression guard: no root_path → today's URL, unchanged. + bare = get_passthrough_resource_metadata_url( + scope=_scope("/mcp/github"), + server_name="github", + ) + assert bare == "http://testserver/.well-known/oauth-protected-resource/mcp/github" + + def test_user_oauth_challenge_url_routes_and_resource_matches_client_url( + self, _no_proxy_base_url, _isolated_mcp_registry + ): + """The reviewer's expected end-state, pinned end-to-end: an MCP endpoint + raising ``raise_user_oauth_challenge`` under a per-request prefix must + emit a resource_metadata URL the client can actually fetch, and the + document it returns must carry the same prefix the client originally + called. If either half breaks the client's discovery is dead.""" + import re + + from fastapi import FastAPI, HTTPException, Request + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_user_oauth_challenge, + ) + from litellm.proxy.middleware.per_request_root_path_middleware import ( + PerRequestRootPathMiddleware, + get_request_root_path, + ) + + server = _create_oauth2_server( + server_id="srv_a", name="server_a", server_name="server_a", alias="server_a" + ) + _isolated_mcp_registry[server.server_id] = server + + app = FastAPI() + app.include_router(router) + + @app.post("/mcp/{name}") + def _mcp(name: str, request: Request): + try: + raise_user_oauth_challenge(server, root_path=get_request_root_path()) + except HTTPException as exc: + return {"www_authenticate": exc.headers["WWW-Authenticate"]} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a", "/tenant-b"]) + client = TestClient(app) + + for prefix, mcp_url in ( + ("/tenant-a", "http://testserver/tenant-a/mcp/server_a"), + ("/tenant-b", "http://testserver/tenant-b/mcp/server_a"), + ("", "http://testserver/mcp/server_a"), + ): + call = client.post(f"{prefix}/mcp/server_a") + assert call.status_code == 200, call.text + www = call.json()["www_authenticate"] + match = re.search(r'resource_metadata="([^"]+)"', www) + assert match, www + discovery = client.get(match.group(1)) + # The challenge URL must route (a client that can't fetch it has + # no way to reach the resource metadata). + assert discovery.status_code == 200, ( + f"challenge URL {match.group(1)} for prefix {prefix!r} 404s; " + "the client can't reach the resource metadata." + ) + # And the doc's `resource` must equal the URL the client called + # (RFC 9728 §3 exact match): a mismatch bounces a strict client. + assert discovery.json()["resource"] == mcp_url, discovery.json() + + def _s256(verifier: str) -> str: return urlsafe_b64encode(hashlib.sha256(verifier.encode("ascii")).digest()).rstrip(b"=").decode("ascii") @@ -10635,3 +11047,142 @@ def test_introspect_route_answers_for_authenticated_caller(monkeypatch): assert active.status_code == 200 assert active.json()["active"] is True assert active.json()["sub"] == "u1" + + +@pytest.mark.asyncio +async def test_identity_bound_authorization_carries_nonce_and_caller_through_callback(monkeypatch): + from http.cookies import SimpleCookie + from urllib.parse import parse_qs, urlparse + from fastapi import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _oauth_state_cookie_name, authorize_with_server, callback, open_bridge_authorization_code, + ) + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + monkeypatch.setenv("LITELLM_SALT_KEY", "identity-binding-test-salt") + server = MCPServer( + server_id="srv", name="srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + client_id="client", authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", issuer="https://idp.example.com", audiences=["client"], + ), + ) + request = Request({"type": "http", "scheme": "https", "server": ("proxy.example.com", 443), + "path": "/authorize", "query_string": b"", "headers": []}) + with ( + patch( # test-quality-ok: isolate authenticated request resolution from the real encrypted OAuth round trip + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request", + new=AsyncMock(return_value="alice")), + patch( # test-quality-ok: isolate user access lookup while testing nonce and caller preservation + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._bridge_authorize_access_denial", + new=AsyncMock(return_value=None)), + ): + authorized = await authorize_with_server( + request, server, "client", "http://127.0.0.1:6274/callback", state="client-state", + code_challenge="pkce-challenge", code_challenge_method="S256", + ) + query = parse_qs(urlparse(authorized.headers["location"]).query) + assert len(query["nonce"][0]) >= 32 + cookies = SimpleCookie() + cookies.load(authorized.headers["set-cookie"]) + name = _oauth_state_cookie_name(query["state"][0]) + callback_request = Request({**request.scope, "path": "/callback", + "headers": [(b"cookie", f"{name}={cookies[name].value}".encode())]}) + completed = await callback(callback_request, code="upstream-code", state=query["state"][0]) + returned = parse_qs(urlparse(completed.headers["location"]).query) + sealed = open_bridge_authorization_code(returned["code"][0]) + assert sealed.litellm_user_id == "alice" + assert sealed.mcp_server_id == "srv" + assert sealed.upstream_code == "upstream-code" + assert sealed.oauth_nonce == query["nonce"][0] + assert returned["state"] == ["client-state"] + + +@pytest.mark.asyncio +async def test_enforced_login_warms_verified_token_readable_without_database_lookup(monkeypatch): + from types import SimpleNamespace + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db, mcp_server_manager + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _store_per_user_token_server_side + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof + from litellm.proxy._experimental.mcp_server.outbound_credentials.dual_cache_token_backend import DualCacheTokenCacheBackend + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import CachedOAuthTokenStore + from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import LazyPerUserOAuthTokenStore + from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import V2PerUserTokenStore + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", name="srv", transport="http", auth_type="oauth2", + oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"], + "caller_field": "user_id", "principal_claim": "sub"}, + ) + proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv") + cache = DualCache() + monkeypatch.setenv("LITELLM_SALT_KEY", "test-cache-warm-encryption-salt") + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", SimpleNamespace(invalidate_user_oauth_token_cache=AsyncMock())) + monkeypatch.setattr(db, "store_user_oauth_credential", AsyncMock()) + await _store_per_user_token_server_side( + server=server, user_id="alice", token_response={"access_token": "alice-token", "refresh_token": "private", "expires_in": 3600}, + identity_binding_proof=proof, + ) + read = AsyncMock(return_value=None) + cached = CachedOAuthTokenStore( + V2PerUserTokenStore(read), default_ttl_seconds=300, backend=DualCacheTokenCacheBackend(cache, mcp_per_user_token_cache._codec()), + ) + store = LazyPerUserOAuthTokenStore(lambda _: server, store_builder=lambda _: (cached, True), redis_available=lambda: True) + token = await store.fetch("alice", "srv") + assert token is not None and token.access_token == "alice-token" + assert token.identity_binding_proof == proof + assert token.refresh_token is None + read.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("upstream_status", [401, 403]) +@pytest.mark.parametrize("auth_type", [MCPAuth.true_passthrough, MCPAuth.oauth_delegate]) +@pytest.mark.parametrize("dcr_bridge", [False, True]) +@pytest.mark.parametrize("flow", ["register", "mint"]) +async def test_dcr_refusal_is_actionable_without_upstream_body( + upstream_status: int, auth_type: MCPAuth, dcr_bridge: bool, flow: str, monkeypatch: pytest.MonkeyPatch +) -> None: + import httpx + from typing import Final + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + mint_ephemeral_dcr_client, + register_client_with_server, + ) + + server: Final = _bridge_server( + auth_type=auth_type, dcr_bridge=dcr_bridge, server_id=f"refused-{auth_type}-{dcr_bridge}-{flow}-{upstream_status}", + client_id=None, + ) + import respx + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + with respx.mock as upstream: + registration: Final = upstream.post(server.registration_url).mock( + return_value=httpx.Response(upstream_status, text="Forbidden private upstream details") + ) + operation: Final = ( + mint_ephemeral_dcr_client(_bridge_mock_request(), server) + if flow == "mint" + else register_client_with_server( + request=_bridge_mock_request(), mcp_server=server, client_name="Test client", + grant_types=None, response_types=None, token_endpoint_auth_method=None, + client_redirect_uris=["http://localhost:9999/callback"], + ) + ) + with pytest.raises(HTTPException) as exc: + await operation + assert registration.call_count == 1 + assert exc.value.status_code == 403 + assert f"HTTP {upstream_status}" in str(exc.value.detail) + assert "pre-registered OAuth client" in str(exc.value.detail) + assert "private upstream details" not in str(exc.value.detail) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 1670370f082..7c80ee77cd7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -6,6 +6,7 @@ import re from base64 import urlsafe_b64encode from datetime import datetime, timedelta, timezone from http.cookies import SimpleCookie +from typing import Final from urllib.parse import parse_qs, urlparse import pytest @@ -18,6 +19,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( GATEWAY_AUTH_CODE_PREFIX, GATEWAY_AUTH_CODE_TTL_SECONDS, MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS, + MAX_CLIENT_ID_LENGTH, ConsentTeam, MintedProxyCredential, _GatewayAuthCode, @@ -26,6 +28,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( aggregate_authorize, aggregate_token, complete_connect_flow, + describe_connect_flow, introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, @@ -52,6 +55,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i MASTER_KEY = "sk-gateway-dcr-flow-tests" REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback" +VSCODE_REDIRECT_URIS: Final = ( + "https://insiders.vscode.dev/redirect", + "https://vscode.dev/redirect", + "http://127.0.0.1/", + "http://127.0.0.1:33418/", +) +MAX_LENGTH_REDIRECT_URIS: Final = tuple(f"https://client.example/{index}/".ljust(256, "a") for index in range(4)) CODE_VERIFIER = "verifier-" + "v" * 43 CODE_CHALLENGE = urlsafe_b64encode(hashlib.sha256(CODE_VERIFIER.encode("ascii")).digest()).rstrip(b"=").decode("ascii") @@ -103,6 +113,58 @@ async def test_register_mints_stateless_public_client(): assert record.redirect_uris == (REDIRECT_URI,) +@pytest.mark.asyncio +@pytest.mark.parametrize("redirect_uris", [VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS]) +async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[str, ...]) -> None: + response: Final = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={ + "client_name": "Visual Studio Code", + "client_uri": "https://code.visualstudio.com", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "redirect_uris": list(redirect_uris), + "token_endpoint_auth_method": "none", + "application_type": "native", + }, + ) + assert response.status_code == 201 + body: Final = json.loads(response.body) + assert body["redirect_uris"] == list(redirect_uris) + assert body["token_endpoint_auth_method"] == "none" + assert "client_secret" not in body + assert len(body["client_id"]) <= MAX_CLIENT_ID_LENGTH + record: Final = open_gateway_dcr_client(body["client_id"]) + assert record is not None + assert record.redirect_uris == redirect_uris + + +@pytest.mark.asyncio +async def test_register_rejects_five_valid_callbacks() -> None: + response: Final = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": [*VSCODE_REDIRECT_URIS, "http://127.0.0.1:33419/"]}, + ) + assert response.status_code == 400 + assert json.loads(response.body) == { + "error": "invalid_redirect_uri", + "error_description": "redirect_uris must be a list of 1 to 4 URIs", + } + + +@pytest.mark.asyncio +async def test_register_four_callbacks_preserves_encoded_size_guard() -> None: + response: Final = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": [f"https://client.example/{index}/".ljust(256, "é") for index in range(4)]}, + ) + assert response.status_code == 400 + assert json.loads(response.body) == { + "error": "invalid_client_metadata", + "error_description": "registered metadata is too large", + } + + @pytest.mark.asyncio async def test_register_allows_loopback_http_for_dev_clients(): body = await _register(["http://localhost:6274/oauth/callback"]) @@ -161,7 +223,6 @@ async def test_register_rejects_userinfo_spoofed_origin(): ["https://claude.ai/cb#fragment"], ["ftp://claude.ai/cb"], ["https://a.example.com/" + "p" * 300], - ["https://a.example.com/1", "https://a.example.com/2", "https://a.example.com/3", "https://a.example.com/4"], [12345], ], ) @@ -230,7 +291,7 @@ async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_co assert location.path == "/ui/connect" params = parse_qs(location.query) handle = params["connect_flow"][0] - assert params["connect_client"] == ["https://claude.ai"] + assert set(params) == {"connect_flow"} set_cookie = response.headers["set-cookie"] assert f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" in set_cookie assert "HttpOnly" in set_cookie @@ -247,12 +308,14 @@ def _flow_cookie_from(response) -> tuple: @pytest.mark.asyncio -async def test_full_walk_register_authorize_complete_token_and_replay(): +@pytest.mark.parametrize("redirect_uris", [(REDIRECT_URI,), VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS]) +async def test_full_walk_register_authorize_complete_token_and_replay(redirect_uris: tuple[str, ...]): """The whole front door on one deterministic walk: register -> authorize -> complete -> token, then the security edges on the same artifacts (user mismatch, PKCE mismatch, single-use replay, refresh rotation, cross-client refresh).""" - client_id = (await _register([REDIRECT_URI]))["client_id"] - authorize_response = _authorize(client_id, session_user_id="u1") + redirect_uri: Final = redirect_uris[-1] + client_id = (await _register(list(redirect_uris)))["client_id"] + authorize_response = _authorize(client_id, session_user_id="u1", redirect_uri=redirect_uri) handle, cookies = _flow_cookie_from(authorize_response) denied = await complete_connect_flow( @@ -279,7 +342,7 @@ async def test_full_walk_register_authorize_complete_token_and_replay(): ) assert completed.status_code == 303 redirect = urlparse(completed.headers["location"]) - assert f"{redirect.scheme}://{redirect.netloc}{redirect.path}" == REDIRECT_URI + assert f"{redirect.scheme}://{redirect.netloc}{redirect.path}" == redirect_uri params = parse_qs(redirect.query) assert params["state"] == ["client-state-123"] code = params["code"][0] @@ -292,7 +355,7 @@ async def test_full_walk_register_authorize_complete_token_and_replay(): "request": _request("/token", method="POST"), "grant_type": "authorization_code", "code": code, - "redirect_uri": REDIRECT_URI, + "redirect_uri": redirect_uri, "client_id": client_id, "code_verifier": CODE_VERIFIER, "refresh_token": None, @@ -832,7 +895,7 @@ async def test_manual_delivery_page_renders_the_url_as_data_never_as_a_shell_com assert 'value="' in body -def _scoped_mcp_server(name="github", **kw): +def _scoped_mcp_server(name="github", auth_type="oauth2", **kw): from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -843,7 +906,7 @@ def _scoped_mcp_server(name="github", **kw): alias=name, url="https://upstream.example/mcp", transport="http", - auth_type=MCPAuth.oauth2, + auth_type=MCPAuth(auth_type) if auth_type is not None else None, **kw, ) @@ -889,14 +952,60 @@ def _opened_principal(payload): return admitted.principal -async def _finish_connect_page(response): +class _VendorCredential: + def __init__(self, state="present"): + self.calls = [] + self.state = state + + async def __call__(self, user_id, server_id): + self.calls.append((user_id, server_id)) + return self.state + + +class _ServerReachability: + def __init__(self, reachable=True): + self.calls = [] + self.reachable = reachable + + async def __call__(self, user_id, server_id): + self.calls.append((user_id, server_id)) + return self.reachable + + +async def _complete_page(response, scoped_server=None, vendor=None, reachable=None, cache=None, **overrides): + from unittest.mock import patch + handle, cookies = _flow_cookie_from(response) - completed = await complete_connect_flow( - request=_request("/authorize/complete", cookies=cookies, method="POST"), - flow_handle=handle, - session_user_id="u1", - cache=DualCache(), - ) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.return_value = scoped_server + return await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=cache or DualCache(), + lookup_vendor_credential=vendor or _VendorCredential(), + lookup_server_reachability=reachable or _ServerReachability(), + **overrides, + ) + + +async def _describe_page(response, scoped_server=None, vendor=None, reachable=None, session_user_id="u1", cookies=None): + from unittest.mock import patch + + handle, flow_cookies = _flow_cookie_from(response) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.return_value = scoped_server + return await describe_connect_flow( + request=_request("/authorize/flow", cookies=flow_cookies if cookies is None else cookies), + flow_handle=handle, + session_user_id=session_user_id, + lookup_vendor_credential=vendor or _VendorCredential(), + lookup_server_reachability=reachable or _ServerReachability(), + ) + + +async def _finish_connect_page(response, scoped_server=None): + completed = await _complete_page(response, scoped_server=scoped_server) return parse_qs(urlparse(completed.headers["location"]).query)["code"][0] @@ -910,23 +1019,43 @@ def _sealed_wire_json(sealed, prefix, debug_key): @pytest.mark.asyncio async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): - """LIT-4917: a per-server RFC 8707 resource naming a gateway-managed oauth2 server - seals that server into the flow. The connect page interlude runs exactly as before - (the scope restricts, it never skips consent), and the code minted at the finish step - and the session pair it redeems for are both scoped.""" + """LIT-4917 plus LIT-7075: a per-server RFC 8707 resource naming a gateway-managed oauth2 + server seals that server into the flow. The connect URL carries only the handle; the page + learns the scoped server and its vendor state from describe_connect_flow, and the finish + step refuses to mint a scoped code until that vendor credential exists, without burning + the flow. The code minted afterwards and the session pair it redeems for are both scoped.""" from unittest.mock import patch client_id = (await _register([REDIRECT_URI]))["client_id"] + github = _scoped_mcp_server() with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + manager.get_mcp_server_by_name.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) assert response.status_code == 303 - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) assert ( _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == "github-id" ) - code = await _finish_connect_page(response) + described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("absent")) + assert json.loads(described.body) == { + "state": "interactive", + "client_origin": "https://claude.ai", + "server_id": "github-id", + "server_name": "github", + "connected": False, + } + cache = DualCache() + premature = await _complete_page(response, scoped_server=github, vendor=_VendorCredential("absent"), cache=cache) + assert premature.status_code == 400 + assert "authorize the requested MCP server" in json.loads(premature.body)["error_description"] + present = _VendorCredential("present") + completed = await _complete_page(response, scoped_server=github, vendor=present, cache=cache) + assert completed.status_code == 303 + assert present.calls == [("u1", "github-id")] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] assert ( _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"] == "github-id" @@ -952,9 +1081,11 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): ) async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, resolves): """Every resource shape outside 'exactly one gateway-managed server' keeps today's flow: - connect page interlude, and NONE of the minted artifacts carry the scope key on the - wire, not the flow cookie, not the code, not the session JWT, so an unscoped flow - started on a new pod completes on a pod whose strict models predate the claim.""" + the generic connect grid (describe names no server, the finish step never consults the + vendor credential), and NONE of the minted + artifacts carry the scope key on the wire, not the flow cookie, not the code, not the + session JWT, so an unscoped flow started on a new pod completes on a pod whose strict + models predate the claim.""" import base64 from unittest.mock import patch @@ -963,10 +1094,18 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, manager.get_mcp_server_by_name.return_value = None if resolves is None else _scoped_mcp_server() response = _scoped_authorize(client_id, resource) assert response.status_code == 303 - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) assert "resource_server_id" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") - code = await _finish_connect_page(response) + vendor = _VendorCredential("absent") + described = await _describe_page(response, vendor=vendor) + assert json.loads(described.body)["state"] == "unscoped" + assert json.loads(described.body)["server_id"] is None + completed = await _complete_page(response, vendor=vendor) + assert vendor.calls == [] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] assert "resource_server_id" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code") token_response = await _redeem(code, client_id) payload = json.loads(token_response.body) @@ -979,19 +1118,150 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, @pytest.mark.asyncio async def test_scoped_authorize_delegate_server_stays_unscoped(): """A delegate-auth oauth2 server is outside the gateway-managed set (its keyless flow is - upstream PKCE via the relay), so a resource naming it never scopes the gateway flow.""" + upstream PKCE via the relay), so a resource naming it never scopes the gateway flow and + never narrows the connect page to it.""" from unittest.mock import patch client_id = (await _register([REDIRECT_URI]))["client_id"] with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True) response = _scoped_authorize(client_id, SCOPED_RESOURCE) - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} + assert json.loads((await _describe_page(response)).body)["server_id"] is None code = await _finish_connect_page(response) token_response = await _redeem(code, client_id) assert _opened_principal(json.loads(token_response.body)).resource_server_id is None +@pytest.mark.asyncio +async def test_m2m_scoped_flow_mints_without_a_user_credential(): + """A client-credentials server is already authorized by its gateway service credential, so + a resource-scoped flow finishes without consulting the per-user vault.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + m2m = _scoped_mcp_server(oauth2_flow="client_credentials") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = m2m + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + vendor = _VendorCredential("unavailable") + described = await _describe_page(response, scoped_server=m2m, vendor=vendor) + assert json.loads(described.body)["state"] == "m2m" + assert json.loads(described.body)["connected"] is True + assert vendor.calls == [] + completed = await _complete_page(response, scoped_server=m2m, vendor=vendor) + assert completed.status_code == 303 + assert vendor.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("oauth2_flow", ["authorization_code", "client_credentials"]) +async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow): + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server(oauth2_flow=oauth2_flow) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + reachable = _ServerReachability(False) + vendor = _VendorCredential("present") + described = await _describe_page(response, scoped_server=server, reachable=reachable, vendor=vendor) + assert json.loads(described.body) == { + "state": "stale", + "client_origin": "https://claude.ai", + "server_id": None, + "server_name": None, + "connected": None, + } + cache = DualCache() + refused = await _complete_page(response, scoped_server=server, reachable=reachable, vendor=vendor, cache=cache) + assert refused.status_code == 400 + assert vendor.calls == [] + assert reachable.calls == [("u1", "github-id"), ("u1", "github-id")] + completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert completed.status_code == 303 + + +@pytest.mark.asyncio +async def test_stale_scoped_flow_remains_distinct_from_unscoped(): + """A server removed after authorize stays a stale scoped flow, so the page cannot offer a + broader unscoped grant or report a misleading Finish action.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + described = await _describe_page(response, scoped_server=None) + assert json.loads(described.body)["state"] == "stale" + assert json.loads(described.body)["connected"] is None + stale = await _complete_page(response, scoped_server=None) + assert stale.status_code == 400 + assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available" + + +@pytest.mark.asyncio +async def test_scoped_flow_deny_and_stale_server_never_need_the_vendor_credential(): + """Cancel is the escape hatch: a scoped user who cannot finish the vendor step still ends + the flow with access_denied and no credential lookup. A scoped server that is no longer + gateway-managed refuses to mint (nothing could serve that code) but also burns nothing.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + github = _scoped_mcp_server() + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = github + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + cache = DualCache() + skipped_reachability = _ServerReachability(False) + stale = await _complete_page(response, scoped_server=None, reachable=skipped_reachability, cache=cache) + assert stale.status_code == 400 + assert skipped_reachability.calls == [] + assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available" + described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("unavailable")) + assert described.status_code == 503 + vendor = _VendorCredential("absent") + deny_reachability = _ServerReachability(False) + denied = await _complete_page( + response, + scoped_server=github, + vendor=vendor, + reachable=deny_reachability, + cache=cache, + decision="deny", + ) + assert denied.status_code == 303 + assert parse_qs(urlparse(denied.headers["location"]).query)["error"] == ["access_denied"] + assert vendor.calls == [] + assert deny_reachability.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "session_user_id, cookies, expected_status, expected_error", + [ + ("u1", {}, 400, "invalid_request"), + ("u1", {"mcp_connect_flow_wrong": "garbage"}, 400, "invalid_request"), + (None, None, 401, "login_required"), + ("u2", None, 403, "access_denied"), + ], +) +async def test_describe_connect_flow_refuses_exactly_like_the_finish_step( + session_user_id, cookies, expected_status, expected_error +): + """The page's read of the flow is gated the same way minting is: the HttpOnly cookie for + that handle must open and the signed-in user must be the sealed one. A lure link with a + made-up handle therefore learns nothing and starts nothing.""" + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _authorize(client_id, session_user_id="u1") + described = await _describe_page(response, session_user_id=session_user_id, cookies=cookies) + assert described.status_code == expected_status + assert json.loads(described.body)["error"] == expected_error + + @pytest.mark.asyncio async def test_token_rejects_resource_conflicting_with_sealed_scope(): """RFC 8707 section 2.2: redeeming a scoped code (or rotating a scoped refresh token) @@ -1005,7 +1275,7 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope(): with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) - code = await _finish_connect_page(response) + code = await _finish_connect_page(response, scoped_server=github) with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = linear @@ -1836,3 +2106,45 @@ async def test_introspect_fails_closed_on_dead_user_and_503s_on_outage(): status, body = await _introspect(minted.token.get_secret_value(), master_key=None) assert (status, body["error"]) == (500, "server_error") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_type", [None, "none", "api_key", "bearer_token", "basic", "authorization", "token", "aws_sigv4"] +) +@pytest.mark.parametrize("resource", ["https://llm.example.com/mcp/github", "https://llm.example.com/github/mcp"]) +async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(auth_type, resource): + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server(auth_type=auth_type) + vendor = _VendorCredential("absent") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize(client_id, resource) + described = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(described.body) == { + "state": "m2m", + "client_origin": "https://claude.ai", + "server_id": "github-id", + "server_name": "github", + "connected": True, + } + unreachable = await _complete_page(response, scoped_server=server, reachable=_ServerReachability(False)) + assert unreachable.status_code == 400 + cache = DualCache() + completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert completed.status_code == 303 + assert vendor.calls == [] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + redeemed = await _redeem(code, client_id, cache=cache, resource=resource) + assert redeemed.status_code == 200 + payload = json.loads(redeemed.body) + assert _opened_principal(payload).resource_server_id == "github-id" + renewed = await _redeem( + None, client_id, cache=cache, grant_type="refresh_token", refresh_token=payload["refresh_token"] + ) + assert renewed.status_code == 200 + assert _opened_principal(json.loads(renewed.body)).resource_server_id == "github-id" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index d299239f68e..b6535e6326a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -3,11 +3,21 @@ Tests for MCPDebug — MCP OAuth2 debug response headers. """ import asyncio -from unittest.mock import MagicMock +from typing import Final + +import pytest +from starlette.types import Message + +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution + +import httpx from litellm.proxy._experimental.mcp_server.mcp_debug import ( MCP_DEBUG_REQUEST_HEADER, MCPDebug, + describe_upstream_http_failure, + + MCPAuthDiagnostics, ) @@ -166,77 +176,6 @@ class TestBuildDebugHeaders: assert set(headers.keys()) == expected_keys -class TestResolveAuthResolution: - def _make_server(self, **kwargs): - server = MagicMock() - server.alias = kwargs.get("alias", "test") - server.server_name = kwargs.get("server_name", "test") - server.has_client_credentials = kwargs.get("has_client_credentials", False) - server.authentication_token = kwargs.get("authentication_token", None) - server.auth_type = kwargs.get("auth_type", None) - return server - - def test_per_request_header(self): - server = self._make_server() - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header="Bearer xxx", - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "per-request-header" - - def test_server_specific_header(self): - server = self._make_server(alias="atlas") - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers={"atlas": {"Authorization": "Bearer xxx"}}, - oauth2_headers=None, - ) - assert result == "per-request-header" - - def test_m2m(self): - server = self._make_server(has_client_credentials=True) - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "m2m-client-credentials" - - def test_static_token(self): - server = self._make_server(authentication_token="static-tok") - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "static-token" - - def test_oauth2_passthrough(self): - server = self._make_server(auth_type="oauth2") - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers={"Authorization": "Bearer eyJ..."}, - ) - assert result == "oauth2-passthrough" - - def test_no_auth(self): - server = self._make_server() - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "no-auth" - - class TestWrapSendWithDebugHeaders: def test_injects_headers(self): captured = [] @@ -269,3 +208,386 @@ class TestWrapSendWithDebugHeaders: asyncio.run(wrapped(body_msg)) assert captured[0] == body_msg + + +class TestDescribeUpstreamHttpFailure: + @staticmethod + def _status_error(*, body: bytes, response_body: bytes | None = None) -> httpx.HTTPStatusError: + request = httpx.Request( + "POST", + "https://upstream.example/apis/mcp", + headers={"Authorization": "Bearer secret-token-abcdef0123456789", "Content-Type": "application/json" if body.startswith(b"{") else "application/x-www-form-urlencoded"}, + content=body, + ) + response = ( + httpx.Response(500, request=request, content=response_body) + if response_body is not None + else httpx.Response(500, request=request, stream=httpx.ByteStream(b'{"error":"boom"}')) + ) + return httpx.HTTPStatusError("500", request=request, response=response) + + def test_includes_method_url_status_and_request_body(self): + exc = self._status_error( + body=b'{"method":"initialize","jsonrpc":"2.0","id":0}', + response_body=b'{"error":"boom"}', + ) + described = describe_upstream_http_failure(exc) + assert described is not None + assert "POST https://upstream.example/ -> HTTP 500" in described + assert '{"method":"initialize"' in described + assert 'response body: {"error":"boom"}' in described + + def test_masks_authorization_header_and_secret_body_fields(self): + exc = self._status_error( + body=b"grant_type=client_credentials&client_id=abc&client_secret=super-secret-value-1234", + response_body=b"{}", + ) + described = describe_upstream_http_failure(exc) + assert described is not None + assert "secret-token-abcdef0123456789" not in described + assert "super-secret-value-1234" not in described + assert "client_id=abc" in described + assert "client_secret=" in described + + def test_reports_unread_streamed_response_body(self): + described = describe_upstream_http_failure(self._status_error(body=b"{}")) + assert described is not None + assert "response body: (not read)" in described + + def test_finds_response_behind_cause_chain(self): + wrapper = RuntimeError("token minting failed") + wrapper.__cause__ = self._status_error(body=b"{}", response_body=b'{"error":"invalid_client"}') + described = describe_upstream_http_failure(wrapper) + assert described is not None + assert "invalid_client" in described + + def test_returns_none_without_http_response(self): + assert describe_upstream_http_failure(ConnectionError("refused")) is None + + +@pytest.mark.parametrize("body", [ + b'{"password":"first second","token":"demo-secret"}', + b'{"nested":[{"access_token":"first,second"}]}', + b'client%5Fsecret=first+second&token=demo-secret', +]) +def test_failure_log_fully_redacts_structured_secrets(body): + request = httpx.Request("POST", "https://upstream/mcp?credential=query-secret", + headers={"X-Custom-Credential": "custom-secret"}, content=body) + response = httpx.Response(500, request=request, content=body) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response)) + assert detail is not None + for secret in ("first", "second", "demo-secret", "custom-secret", "query-secret"): + assert secret not in detail + + +def test_failure_log_omits_unstructured_body(): + request = httpx.Request("POST", "https://upstream/mcp", content=b"arbitrary-secret") + response = httpx.Response(500, request=request, content=b"arbitrary-secret") + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response)) + assert detail is not None + assert "arbitrary-secret" not in detail + assert "omitted" in detail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["error", "empty", "large", "timeout", "read_failure", "closed", "success", "cancel"]) +async def test_error_capture_is_bounded_and_preserves_success_and_cancellation(mode): + from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response + + class Stream(httpx.AsyncByteStream): + def __init__(self): + self.reads = 0 + + async def __aiter__(self): + self.reads += 1 + if mode == "timeout": + await asyncio.sleep(10) + if mode == "closed": + raise httpx.StreamClosed() + if mode == "read_failure": + raise httpx.ReadError("private-read-error") + if mode == "cancel": + raise asyncio.CancelledError + yield b"" if mode == "empty" else b'{"error":"missing_scope","password":"first second"}' if mode != "large" else b"x" * 20000 + + stream = Stream() + request = httpx.Request("POST", "https://upstream/mcp") + response = httpx.Response(200 if mode == "success" else 500, request=request, stream=stream) + if mode == "cancel": + with pytest.raises(asyncio.CancelledError): + await capture_upstream_error_response(response) + return + await capture_upstream_error_response(response) + if mode == "success": + assert stream.reads == 0 + assert await response.aread() == b'{"error":"missing_scope","password":"first second"}' + return + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response)) + assert detail is not None + assert "first" not in detail and "second" not in detail and "private-read-error" not in detail + expected = {"empty": "(empty)", "error": "missing_scope", "large": "capture limit", "timeout": "read failed", "read_failure": "read failed", "closed":"read failed"} + assert expected[mode] in detail + if mode == "error": + assert await response.aread() == b'{"error":"missing_scope","password":"first second"}' + + +@pytest.mark.parametrize("body", [b"", b'"scalar"', b'{"hint":"line1\\nline2"}', b'{"hint":"' + b'x' * 600 + b'"}']) +def test_failure_preview_handles_empty_scalar_control_and_long_bodies(body): + request = httpx.Request("POST", "https://user:secret@upstream/mcp?key=private#private", content=body) + response = httpx.Response(500, request=request, content=body) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response)) + assert detail is not None + assert "private" not in detail and "user:secret" not in detail and "\n" not in detail + if not body: + assert "(empty)" in detail + elif body.startswith(b'"'): + assert "omitted" in detail + elif len(body) > 512: + assert "truncated" in detail and len(detail) < 1300 + else: + assert "line1\\nline2" in detail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("slow_error", [False, True]) +async def test_error_capture_preserves_httpx_auth_retry(slow_error): + from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response + + class RetryAuth(httpx.Auth): + def auth_flow(self, request): + response = yield request + if response.status_code == 401: + request.headers["Authorization"] = "Bearer refreshed" + yield request + + class SlowStream(httpx.AsyncByteStream): + async def __aiter__(self): + await asyncio.sleep(10) + yield b'{"error":"expired_token"}' + + def upstream(request): + if request.headers.get("Authorization"): + return httpx.Response(200, json={"ok": True}) + return httpx.Response(401, stream=SlowStream()) if slow_error else httpx.Response(401, json={"error":"expired_token"}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(upstream), auth=RetryAuth(), + event_hooks={"response":[capture_upstream_error_response]}) as client: + response = await client.get("https://upstream/mcp") + assert response.status_code == 200 and response.json() == {"ok":True} + if slow_error: + assert response.history[0].content == b"" + else: + assert response.history[0].json() == {"error":"expired_token"} + + +def test_failure_diagnostics_without_request_and_with_streamed_request(): + response = httpx.Response(503) + exc = httpx.HTTPStatusError("failed", request=httpx.Request("GET", "https://upstream"), response=response) + assert describe_upstream_http_failure(exc) == "HTTP 503 | request unavailable" + request = httpx.Request("POST", "https://upstream", content=iter((b"private-body",))) + response = httpx.Response(503, request=request) + described = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response)) + assert described is not None and "streamed, not captured" in described and "private-body" not in described + + + +def test_deep_error_body_is_bounded_without_exposing_nested_values(): + body = b'{"nested":' * 18 + b'{"password":"hidden-value"}' + b'}' * 18 + request = httpx.Request("POST", "https://upstream/mcp", content=body) + response = httpx.Response(500, request=request, content=body) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response)) + assert detail is not None and "hidden-value" not in detail + assert "nested" in detail and "REDACTED" in detail + + +@pytest.mark.parametrize("body", [b'client%5Fsecret=first+second&client_id=visible', b'client_secret=first%26second&client_id=visible']) +def test_encoded_form_credentials_are_decoded_before_redaction(body): + request = httpx.Request("POST", "https://upstream/token", content=body, + headers={"Content-Type":"application/x-www-form-urlencoded"}) + response = httpx.Response(400, request=request, content=body, + headers={"Content-Type":"application/x-www-form-urlencoded"}) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failure", request=request, response=response)) + assert detail is not None and "client_id=visible" in detail + assert "first" not in detail and "second" not in detail + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", tuple(AuthResolution)) +@pytest.mark.parametrize("method", ("GET", "DELETE", "POST")) +async def test_debug_defers_resolution_until_first_frame_only_for_post(source: AuthResolution, method: str) -> None: + captured: Final[list[Message]] = [] + diagnostics: Final = MCPAuthDiagnostics() + + async def send(message: Message) -> None: + captured.append(message) + + wrapped: Final = MCPDebug.wrap_send_with_debug_headers( + send, diagnostics.headers(), diagnostics.headers, request_method=method + ) + await wrapped({"type": "http.response.start", "status": 200, "headers": []}) + assert len(captured) == (0 if method == "POST" else 1) + diagnostics.record("s1", source) + body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True} + await wrapped(body) + assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == ( + source.value.encode() if method == "POST" else b"unresolved" + ) + assert captured[1] == body + + +@pytest.mark.asyncio +async def test_early_stream_frame_reports_unresolved_without_waiting() -> None: + captured: Final[list[Message]] = [] + diagnostics: Final = MCPAuthDiagnostics() + + async def send(message: Message) -> None: + captured.append(message) + + wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers, request_method="POST") + await wrapped({"type": "http.response.start", "status": 200, "headers": []}) + await wrapped({"type": "http.response.body", "body": b": ping\n\n", "more_body": True}) + diagnostics.record("s1", AuthResolution.stored_user_token) + await wrapped({"type": "http.response.body", "body": b"data: pong\n\n", "more_body": False}) + assert len(captured) == 3 + assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == b"unresolved" + + +def test_diagnostics_keep_requests_separate_and_do_not_collapse_multiple_servers() -> None: + alice: Final = MCPAuthDiagnostics() + bob: Final = MCPAuthDiagnostics() + alice.record("s1", AuthResolution.stored_user_token) + assert bob.resolution() == "unresolved" + alice.record("s1", AuthResolution.token_exchange) + assert alice.resolution() == "token-exchange" + alice.record("s2", AuthResolution.static_token) + assert alice.resolution() == "multiple" + assert alice.headers()["x-mcp-debug-auth-resolutions"] == '{"s1":"token-exchange","s2":"static-token"}' + + +@pytest.mark.asyncio +async def test_concurrent_mcp_messages_record_on_their_own_http_scope() -> None: + from unittest.mock import MagicMock + + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_debug import ( + MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, + record_auth_resolution, + ) + + session: Final = MagicMock() + first: Final = MCPAuthDiagnostics() + second: Final = MCPAuthDiagnostics() + + async def record(diagnostics: MCPAuthDiagnostics, source: AuthResolution) -> None: + context: Final = RequestContext( + request_id=1, meta=None, session=session, lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + ) + token: Final = request_ctx.set(context) + try: + await asyncio.sleep(0) + record_auth_resolution("same-server", source) + finally: + request_ctx.reset(token) + + await asyncio.gather(record(first, AuthResolution.stored_user_token), record(second, AuthResolution.per_request_header)) + assert first.resolution() == "stored-user-token" + assert second.resolution() == "per-request-header" + + +@pytest.mark.parametrize("source", ["header", "bearer", "basic", "cookie", "query", "form", "json"]) +def test_reflected_credentials_are_removed_from_normal_response_fields(source): + import base64 + + secret = "generic-credential-123" + headers = {"X-Custom":secret} if source == "header" else {"Authorization":"Bearer " + secret} if source == "bearer" else {"Authorization":"Basic " + base64.b64encode(("client:" + secret).encode()).decode()} if source == "basic" else {"Cookie":"session=" + secret} if source == "cookie" else {} + request = httpx.Request("POST", "https://upstream/token" + ("?credential=" + secret if source == "query" else ""), + headers=headers, data={"client_secret":secret} if source == "form" else None, + json={"nested":{"client_secret":secret}} if source == "json" else None) + response = httpx.Response(401, request=request, json={"error":"invalid_client", "error_description":"Rejected " + secret}) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response)) + assert detail is not None and "invalid_client" in detail + assert secret not in detail and "REDACTED" in detail + + + +@pytest.mark.parametrize("secret", ['value"with\ncharacters€', "R"]) +def test_reflected_values_are_redacted_before_truncation_without_expanding_replacements(secret): + request = httpx.Request("POST", "https://upstream/token", json={"client_secret":secret}) + response = httpx.Response(401, request=request, json={"error":"invalid_client", "detail":"x" * 460 + secret}) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response)) + assert detail is not None and "invalid_client" in detail + assert "value" not in detail and "characters" not in detail and len(detail) < 1400 + + +@pytest.mark.parametrize("headers", [{"Authorization":"Basic !!!"}, {"Cookie":"bad@key=opaque"}]) +def test_malformed_auth_headers_do_not_break_failure_diagnostics(headers): + request = httpx.Request("POST", "https://upstream/token", headers=headers) + response = httpx.Response(401, request=request, json={"error":"invalid_client"}) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response)) + assert detail is not None and "invalid_client" in detail + assert "!!!" not in detail and "opaque" not in detail + + +def test_oversized_request_omits_potentially_reflected_response_credentials(): + request = httpx.Request("POST", "https://upstream/token", content=b"x" * 17000) + response = httpx.Response(401, request=request, json={"error_description":"unknown-secret"}) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response)) + assert detail is not None and "capture limit" in detail and "credentials unavailable" in detail + assert "unknown-secret" not in detail + + + +@pytest.mark.asyncio +async def test_streamed_error_redacts_reflected_credentials_before_capture(): + import json + from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response + + secret = "generic-credential-123" + request = httpx.Request("POST", "https://upstream/token", data={"client_secret":secret}) + raw = json.dumps({"error":"invalid_client", "error_description":"Rejected " + secret}).encode() + response = httpx.Response(401, request=request, stream=httpx.ByteStream(raw)) + await capture_upstream_error_response(response) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response)) + assert detail is not None and "invalid_client" in detail and "Rejected" in detail + assert secret not in detail and "REDACTED" in detail + assert await response.aread() == raw + + +@pytest.mark.parametrize("path", ["/credential-path-value/mcp", "/oauth/credential-path-value/token"]) +def test_failure_diagnostics_omit_credential_bearing_url_paths(path): + request = httpx.Request("POST", "https://upstream.example" + path) + response = httpx.Response(401, request=request, json={"error": "access_denied"}) + error = httpx.HTTPStatusError("denied", request=request, response=response) + diagnostic = describe_upstream_http_failure(error) + assert diagnostic is not None + assert "credential-path-value" not in diagnostic + assert "POST https://upstream.example/ -> HTTP 401" in diagnostic + assert "access_denied" in diagnostic + + +def test_deep_request_omits_response_when_credentials_cannot_be_inspected(): + from litellm.proxy._experimental.mcp_server.utils import MAX_STRUCTURED_CONTENT_SCAN_DEPTH + + raw = "[" * (MAX_STRUCTURED_CONTENT_SCAN_DEPTH + 1) + '{"client_secret":"nested-credential"}' + "]" * (MAX_STRUCTURED_CONTENT_SCAN_DEPTH + 1) + request = httpx.Request("POST", "https://upstream/token", content=raw, headers={"Content-Type": "application/json"}) + response = httpx.Response(401, request=request, json={"error_description": "Rejected nested-credential"}) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response)) + assert detail is not None and "HTTP 401" in detail + assert "response body: (omitted: request credentials unavailable)" in detail + assert "nested-credential" not in detail + + +@pytest.mark.parametrize("field", ["accessToken", "refreshToken", "clientSecret", "apikey", "CLIENTASSERTION", "cost_token"]) +@pytest.mark.parametrize("encoding", ["json", "form"]) +def test_compact_credential_fields_and_reflected_values_are_redacted(field, encoding): + secret = "generic-private-value" + fields = {field: secret} + request = httpx.Request("POST", "https://upstream/token", json=fields if encoding == "json" else None, + data=fields if encoding == "form" else None) + response = httpx.Response(401, request=request, json={field: secret, "error": "invalid_client", "detail": "Rejected " + secret}) + detail = describe_upstream_http_failure(httpx.HTTPStatusError("failed", request=request, response=response)) + assert detail is not None and "invalid_client" in detail + assert "REDACTED" in detail and secret not in detail diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py index 76cc235f7eb..36b545ad031 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py @@ -861,6 +861,7 @@ _SALT_KEY = "test-salt-key-for-env-vars-tests-1234" @pytest.fixture def env_vars_salt_key(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", _SALT_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": "xsalsa20-poly1305"}) def _mock_env_vars_prisma(row=None): @@ -1518,9 +1519,16 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri assert "s3cr3t-p@ss" not in encrypted_env_vars_str def _prisma_row_with_json_string_env_vars(): - row = MagicMock() - row.env_vars = encrypted_env_vars_str - return row + import json + + from prisma.models import LiteLLM_MCPServerTable + + return LiteLLM_MCPServerTable.model_validate({ + "server_id": "srv-returned", "transport": "http", "mcp_access_groups": [], "allowed_tools": [], + "extra_headers": [], "args": [], "allow_all_keys": False, "available_on_public_internet": True, + "delegate_auth_to_upstream": False, "oauth_passthrough": False, "per_server_oauth_discovery": False, + "is_byok": False, "byok_description": [], "env_vars": json.dumps(encrypted_env_vars_str), + }) mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.create = AsyncMock( @@ -1537,7 +1545,9 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri touched_by="test-user", ) assert isinstance(created.env_vars, list) - assert created.env_vars[0]["value"] == "s3cr3t-p@ss" + assert created.env_vars[0].value == "s3cr3t-p@ss" + assert created.env_vars[0].name == "DB_PASSWORD" + assert created.env == {} mock_prisma_upd = MagicMock() mock_prisma_upd.db.litellm_mcpservertable.update = AsyncMock( @@ -1549,7 +1559,9 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri touched_by="test-user", ) assert isinstance(updated.env_vars, list) - assert updated.env_vars[0]["value"] == "s3cr3t-p@ss" + assert updated.env_vars[0].value == "s3cr3t-p@ss" + assert updated.env_vars[0].name == "DB_PASSWORD" + assert updated.env == {} def test_reencrypt_global_env_var_values_handles_json_string(env_vars_salt_key): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index 9667224de98..3f5d4ad83ea 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -1,5 +1,6 @@ """Unit tests for MCP OAuth passthrough tool-fetch behavior.""" +import logging import sys from unittest.mock import AsyncMock, MagicMock @@ -11,7 +12,7 @@ if sys.version_info < (3, 11): from exceptiongroup import ExceptionGroup -from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.exceptions import MCPServerListError, MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _extract_upstream_auth_failure, @@ -434,3 +435,43 @@ async def test_aggregate_with_single_accessible_server_still_absorbs(): assert listing.tools == [] assert listing.outcomes["delegate_docs"].tag == "auth_required" + + +@pytest.mark.asyncio +async def test_fetch_tools_logs_upstream_request_details_on_500(caplog): + manager = MCPServerManager() + request = httpx.Request( + "POST", + "https://upstream/apis/mcp", + headers={"Authorization": "Bearer upstream-token-0123456789"}, + content=b'{"method":"initialize","jsonrpc":"2.0","id":0}', + ) + response = httpx.Response(500, request=request) + mock_client = MagicMock() + mock_client.list_tools = AsyncMock( + side_effect=httpx.HTTPStatusError("500", request=request, response=response) + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + with pytest.raises(MCPServerListError): + await manager._fetch_tools_with_timeout(mock_client, "sample_docs") + + assert "POST https://upstream/ -> HTTP 500" in caplog.text + assert '"method":"initialize"' in caplog.text + assert "upstream-token-0123456789" not in caplog.text + + + +@pytest.mark.asyncio +async def test_client_creation_failure_logs_sanitized_exchange(monkeypatch, caplog): + manager = MCPServerManager() + server = MCPServer(server_id="sample", name="sample", url="https://upstream/mcp", transport=MCPTransport.http, auth_type=MCPAuth.none) + request = httpx.Request("POST", "https://upstream/mcp?credential=query-secret") + response = httpx.Response(500, request=request, json={"error":"missing_scope"}) + error = httpx.HTTPStatusError("query-secret", request=request, response=response) + monkeypatch.setattr(manager, "_create_mcp_client", AsyncMock(side_effect=error)) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + with pytest.raises(MCPServerListError): + await manager._get_tools_from_server(server) + assert "POST https://upstream/ -> HTTP 500" in caplog.text + assert "missing_scope" in caplog.text and "query-secret" not in caplog.text diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index f6bd79c5d2d..669e094fee4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -1,17 +1,18 @@ """ -Tests for partial-update semantics of PUT /v1/mcp/server. +Tests for partial-update semantics of PUT /v1/mcp/server and PUT /v1/mcp/toolset. A partial update must only write the fields the caller explicitly provided. Omitting a field must NOT reset it to its Pydantic schema default (e.g. ``transport=sse``, ``mcp_access_groups=[]``, ``allow_all_keys=False``), which -would silently overwrite the existing DB row. +would silently overwrite the existing DB row, and a field the caller sent as null +must be cleared rather than left at its stored value. """ import json from unittest.mock import AsyncMock, MagicMock import pytest -from prisma import Json +from prisma import Json, models from litellm.proxy._experimental.mcp_server.db import ( create_mcp_server, @@ -28,8 +29,11 @@ def _credentials_cleared(value) -> bool: def _mock_prisma(): mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable = AsyncMock() - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) - mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=MagicMock()) + row = models.LiteLLM_MCPServerTable.model_construct( + server_id="test-server", transport="http", env={}, env_vars=[] + ) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=row) + mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=row) return mock_prisma @@ -847,3 +851,69 @@ async def test_cf_pair_switch_does_not_clear_dcr_bridge(): data = UpdateMCPServerRequest(server_id="s", auth_type="oauth_delegate") data_dict = await _run_update_with_existing(data, existing_auth_type="true_passthrough") assert "dcr_bridge" not in data_dict + + +def _mock_toolset_prisma(): + """A prisma double whose update answers with a row the reader can expand, so the + call under test returns instead of failing inside the row mapper.""" + updated_row = MagicMock() + updated_row.model_dump.return_value = { + "toolset_id": "ts-1", + "toolset_name": "ops", + "description": None, + "tools": "[]", + } + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcptoolsettable = AsyncMock() + mock_prisma.db.litellm_mcptoolsettable.update = AsyncMock(return_value=updated_row) + return mock_prisma + + +async def _run_toolset_update(payload: dict) -> dict: + """The columns PUT /v1/mcp/toolset writes for this payload, minus the audit stamp + every write carries. The prisma double is injected, so nothing is patched.""" + from litellm.proxy._experimental.mcp_server.toolset_db import update_mcp_toolset + from litellm.types.mcp_server.mcp_toolset import UpdateMCPToolsetRequest + + mock_prisma = _mock_toolset_prisma() + await update_mcp_toolset(mock_prisma, UpdateMCPToolsetRequest.model_validate(payload), "test-user") + written = dict(mock_prisma.db.litellm_mcptoolsettable.update.call_args[1]["data"]) + assert written["updated_by"] == "test-user" + return {name: value for name, value in written.items() if name != "updated_by"} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_clears_description_on_explicit_null(): + """The dump used to drop None, so a null description could never clear the stored + one: the toolset kept a description its owner had deleted.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "description": None}) == {"description": None} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_omits_the_fields_the_caller_left_out(): + tools = [{"server_id": "s1", "tool_name": "alpha"}] + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": tools}) == {"tools": json.dumps(tools)} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_ignores_null_tools_rather_than_revoking_them(): + """A client that sends tools=null means "leave the selection alone", so the grants + survive. Clearing them is an explicit [], which cannot be confused with an omitted + field; treating null as a clear would silently revoke every tool the toolset grants.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": None, "description": "kept"}) == { + "description": "kept" + } + + +@pytest.mark.asyncio +async def test_toolset_partial_update_empties_the_selection_on_an_explicit_empty_list(): + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": []}) == {"tools": "[]"} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_ignores_a_null_name(): + """A toolset always has a name, so a null toolset_name is a no-op, not a clear + that would write a NOT NULL column to null.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "toolset_name": None, "description": "kept"}) == { + "description": "kept" + } diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py new file mode 100644 index 00000000000..67b7c5a3414 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py @@ -0,0 +1,118 @@ +import json +from datetime import datetime + +import pytest +from fastapi import HTTPException +from mcp.shared.exceptions import McpError +from pydantic import AnyUrl + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._experimental.mcp_server import server +from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth + +AUTH = UserAPIKeyAuth(api_key="key") + + +@pytest.fixture +def proxy_mode(): + token = _mcp_proxy_mode.set(True) + try: + yield + finally: + _mcp_proxy_mode.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("proxy_mode") +async def test_proxy_call_rejects_non_proxy_tool_names() -> None: + result = await server._dispatch_virtual_mcp_tool( + name="math_stdio-add", arguments={"a": 1, "b": 2}, user_api_key_auth=AUTH, client_ip=None + ) + + assert result is not None + assert result.isError is True + assert "unavailable on /mcp/proxy" in result.content[0].text + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("proxy_mode") +async def test_proxy_rejects_non_tool_protocol_operations() -> None: + options = server.server.create_initialization_options() + assert options.capabilities.prompts is None + assert options.capabilities.resources is None + assert options.capabilities.tools is not None + + with pytest.raises(McpError): + await server.list_prompts() + with pytest.raises(McpError): + await server.get_prompt("prompt", {}) + with pytest.raises(McpError): + await server.list_resources() + with pytest.raises(McpError): + await server.list_resource_templates() + with pytest.raises(McpError): + await server.read_resource(AnyUrl("https://example.com/resource")) + + +class FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: list[tuple[str, str]] = [] + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.events.append(("failure", json.dumps(kwargs.get("standard_logging_object"), default=str))) + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.events.append(("success", json.dumps(kwargs.get("standard_logging_object"), default=str))) + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + self.events.append(("post_failure", json.dumps(request_data, default=str))) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("proxy_mode") +async def test_proxy_scope_exception_emits_failure_log(monkeypatch: pytest.MonkeyPatch) -> None: + recorder = FailureRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + auth = UserAPIKeyAuth( + api_key="scope-denial-key-hash", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="denied", mcp_servers=["no-mcp-servers"]), + ) + arguments = {"tool_id": "denied-scope", "arguments": {}} + + with pytest.raises(HTTPException) as denied: + await server._dispatch_virtual_mcp_tool( + name="call_tool", + arguments=arguments, + user_api_key_auth=auth, + client_ip=None, + mcp_servers=["ungranted"], + raw_headers={"authorization": "Bearer raw-scope-secret", "x-litellm-call-id": "scope-denial"}, + ) + + assert denied.value.status_code == 403 + assert denied.value.detail == {"error": "The key is not allowed to access the requested MCP servers: ungranted"} + assert [kind for kind, _ in recorder.events] == ["failure", "post_failure"] + payload = json.loads(recorder.events[0][1]) + assert payload["id"] == "scope-denial" + assert payload["call_type"] == "call_mcp_tool" + assert payload["status"] == "failure" + assert payload["response_cost"] == 0 + assert "ungranted" in payload["error_str"] + hook_payload = json.loads(recorder.events[1][1]) + assert hook_payload["standard_logging_object"] == payload + assert hook_payload["arguments"] == arguments + assert "raw_headers" not in hook_payload + assert "raw-scope-secret" not in recorder.events[1][1] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 086ab854e36..7eb4396e60d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2,6 +2,8 @@ import asyncio import contextvars import os from datetime import datetime, timedelta +from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1368,6 +1370,295 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): mock_logger.info.assert_any_call("Successfully fetched %s tools total from all MCP servers", 0) +def _denied_scope_manager(known_server_names_to_ids: dict[str, str]) -> MagicMock: + servers = {name: MagicMock(server_id=server_id) for name, server_id in known_server_names_to_ids.items()} + manager = MagicMock() + manager.get_mcp_server_by_name = lambda name, client_ip=None: servers.get(name) + return manager + + +def _scope_resolver(resolved_without_agent: dict[str, str], access_groups: tuple[str, ...] = ()) -> AsyncMock: + async def resolve(user_api_key_auth, mcp_servers, client_ip=None): + if user_api_key_auth is not None and user_api_key_auth.agent_id: + return [] + return [ + SimpleNamespace( + server_id=server_id, + server_name=server_name, + alias=None, + short_prefix=None, + access_groups=list(access_groups), + ) + for server_name, server_id in resolved_without_agent.items() + ] + + return AsyncMock(side_effect=resolve) + + +async def _denied_scoped_list( + user_api_key_auth: UserAPIKeyAuth, + mcp_servers: list[str], + mock_manager: MagicMock, + resolver: AsyncMock, +) -> HTTPException: + from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers + + with ( + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + resolver, + ), + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + mock_manager, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=None, + mcp_servers=mcp_servers, + ) + return exc_info.value + + +@pytest.mark.asyncio +async def test_scoped_list_denied_by_agent_binding_raises_403_naming_agent(): + """The agent-binding veto must raise a 403 naming the agent, never a silent 200 with no tools.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + resolver = _scope_resolver(resolved_without_agent={"github": "srv-github"}) + + denial = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "MCP server 'github'" in message + assert "agent 'agent-123'" in message + assert "mcp_servers" in message + rerun_auth = resolver.await_args_list[1].kwargs["user_api_key_auth"] + assert rerun_auth.agent_id is None + assert rerun_auth.user_id == "test_user" + assert resolver.await_args_list[1].kwargs["mcp_servers"] == ["github"] + + +@pytest.mark.asyncio +async def test_empty_scope_lists_nothing_instead_of_raising_a_nameless_denial(): + """An empty ``x-mcp-servers`` header scopes to no servers; that is an empty listing, not a 403.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + resolver = AsyncMock(return_value=[]) + + with ( + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + resolver, + ), + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + _denied_scope_manager({"github": "srv-github"}), + ), + ): + listing = await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=[] + ) + + assert listing.tools == [] + resolver.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scoped_list_denied_for_non_agent_key_raises_generic_403(): + """A denial for a key with no agent binding stays generic and skips the agent-stripped rerun.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + resolver = AsyncMock(return_value=[]) + + denial = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "github" in message + assert "agent" not in message + resolver.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scoped_list_unknown_name_raises_same_generic_403_as_unauthorized(): + """Unknown and registered-but-unauthorized names raise byte-identical generic 403s, so a + caller cannot probe which server names exist.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + unknown = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({}), _scope_resolver(resolved_without_agent={}) + ) + unauthorized = await _denied_scoped_list( + user_api_key_auth, + ["github"], + _denied_scope_manager({"github": "srv-github"}), + _scope_resolver(resolved_without_agent={}), + ) + + assert unknown.status_code == unauthorized.status_code == 403 + assert unknown.detail["error"] == unauthorized.detail["error"] + assert "github" in unknown.detail["error"] + assert "agent" not in unknown.detail["error"] + + +@pytest.mark.asyncio +async def test_scoped_list_access_group_vetoed_by_agent_names_agent_and_group(): + """An access-group scope vetoed by the agent binding raises the 403 naming the agent and the + group instead of the silent empty list.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + denial = await _denied_scoped_list( + user_api_key_auth, + ["prod-group"], + _denied_scope_manager({}), + _scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)), + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "access group 'prod-group'" in message + assert "agent 'agent-123'" in message + assert "mcp_access_groups" in message + + +@pytest.mark.asyncio +async def test_scoped_list_mixed_unknown_and_vetoed_group_names_the_group_that_resolved(): + """With an unknown name ahead of the agent-vetoed group in the scope, the 403 must name the group + whose servers the key can reach, never the unknown name, or the admin is told to grant a group + that does not exist.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + denial = await _denied_scoped_list( + user_api_key_auth, + ["no-such-group", "prod-group"], + _denied_scope_manager({}), + _scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)), + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "access group 'prod-group'" in message + assert "no-such-group" not in message + assert "agent 'agent-123'" in message + + +@pytest.mark.asyncio +async def test_scoped_list_agent_key_denied_by_key_grants_raises_generic_403(): + """When the agent-stripped rerun still resolves nothing, the 403 stays generic instead of + blaming the agent binding.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + resolver = _scope_resolver(resolved_without_agent={}) + + denial = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "github" in message + assert "agent" not in message + assert resolver.await_count == 2 + + +@pytest.mark.asyncio +async def test_scoped_list_agent_veto_attributed_for_differently_cased_server_name(): + """The scope filter matches `/mcp/GitHub` to a server named `github` case-insensitively, so the + agent-attributed 403 must match the same way instead of falling back to the generic denial.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + denial = await _denied_scoped_list( + user_api_key_auth, + ["GitHub"], + _denied_scope_manager({"github": "srv-github"}), + _scope_resolver(resolved_without_agent={"github": "srv-github"}), + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "MCP server 'GitHub'" in message + assert "agent 'agent-123'" in message + + +@pytest.mark.asyncio +async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(): + """The MCP protocol handler surfaces a permission HTTPException as a clean JSON-RPC error + (McpError, INVALID_REQUEST) carrying the denial message, instead of a raw 500.""" + try: + from litellm.proxy._experimental.mcp_server.server import handle_list_tools + except ImportError: + pytest.skip("MCP server not available") + + from mcp.shared.exceptions import McpError + from mcp.types import INVALID_REQUEST + + denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'" + denial = HTTPException(status_code=403, detail={"error": denial_message}) + + with ( + patch( # test-quality-ok: the protocol handler reads auth from module context; no injection seam + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", + new=AsyncMock(return_value=(None, None, None, None, None, None, None)), + ), + patch( # test-quality-ok: the listing helper is the handler's only collaborator; the suite's seam + "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + new=AsyncMock(side_effect=denial), + ), + ): + with pytest.raises(McpError) as exc_info: + await handle_list_tools() + + assert exc_info.value.error.code == INVALID_REQUEST + assert exc_info.value.error.message == denial_message + + +@pytest.mark.asyncio +async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict(): + try: + from litellm.proxy._experimental.mcp_server.server import mcp_server_tool_call + except ImportError: + pytest.skip("MCP server not available") + + denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'" + denial = HTTPException(status_code=403, detail={"error": denial_message}) + + with ( + patch( # test-quality-ok: the protocol handler reads auth from module context; no injection seam + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", + new=AsyncMock(return_value=(None, None, None, None, None, None, None)), + ), + patch( # test-quality-ok: the tool-call helper is the handler's only collaborator; the suite's seam + "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + new=AsyncMock(side_effect=denial), + ), + ): + result = await mcp_server_tool_call("github-search_issues", {}) + + assert result.isError is True + assert result.content[0].text == f"Error: {denial_message}" + + @pytest.mark.asyncio async def test_mcp_server_tool_call_body_with_none_arguments(): """Test that proxy_server_request body handles None arguments correctly""" @@ -1576,96 +1867,85 @@ async def test_streamable_http_session_manager_is_stateless(): @pytest.mark.asyncio -async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(): - """ - Test that routing correctly sends: - - initialize (no mcp-session-id) → stateful manager (so client gets mcp-session-id) - - tools/list (no mcp-session-id) → stateless manager (curl, Inspector) - """ - try: - from litellm.proxy._experimental.mcp_server.server import ( - handle_streamable_http_mcp, - session_manager_stateful, - session_manager_stateless, +@pytest.mark.parametrize("debug", (False, True)) +@pytest.mark.parametrize( + ("method", "request_body", "stateful"), + ( + ("POST", b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', True), + ("POST", b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}', False), + ("GET", b"", False), + ("DELETE", b"", False), + ), +) +async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless( + debug: bool, method: str, request_body: bytes, stateful: bool +) -> None: + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + from starlette.types import Message, Receive, Scope, Send + + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.mcp_debug import record_auth_resolution + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution + + scope: Final[Scope] = {"type": "http", "method": method, "path": "/mcp", "headers": []} + receive: Final = AsyncMock(return_value={"type": "http.request", "body": request_body, "more_body": False}) + send: Final = AsyncMock() + observe_start: Final = AsyncMock() + body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True} + + async def handle_request(request_scope: Scope, receive: Receive, outgoing: Send) -> None: + await outgoing({"type": "http.response.start", "status": 200, "headers": []}) + await observe_start(send.await_count) + context: Final = RequestContext( + request_id=1, meta=None, session=MagicMock(), lifespan_context=None, request=Request(request_scope) ) - except ImportError: - pytest.skip("MCP server not available") + token: Final = request_ctx.set(context) + try: + record_auth_resolution("s1", AuthResolution.stored_user_token) + finally: + request_ctx.reset(token) + await outgoing(body) - async def make_request(method_body: bytes, path: str = "/mcp/progress_test"): - scope = { - "type": "http", - "method": "POST", - "path": path, - "headers": [ - (b"content-type", b"application/json"), - (b"authorization", b"Bearer test-key"), - ], - } - receive = AsyncMock( - return_value={ - "type": "http.request", - "body": method_body, - "more_body": False, - } - ) - send = AsyncMock() - - stateless_called = [] - stateful_called = [] - - async def stateless_handle(s, r, se): - stateless_called.append(1) - - async def stateful_handle(s, r, se): - stateful_called.append(1) - - with ( - patch( - "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", - new_callable=AsyncMock, - return_value=(MagicMock(), None, ["progress_test"], None, None, None), + stateless_handle: Final = AsyncMock(side_effect=handle_request) + stateful_handle: Final = AsyncMock(side_effect=handle_request) + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=( + UserAPIKeyAuth(user_id="debug-user"), + None, + None, + None, + None, + {"x-litellm-mcp-debug": "true"} if debug else {}, ), - patch( - "litellm.proxy._experimental.mcp_server.server.set_auth_context", - ), - patch( - "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", - True, - ), - patch.object( - session_manager_stateless, - "handle_request", - side_effect=stateless_handle, - ), - patch.object( - session_manager_stateful, - "handle_request", - side_effect=stateful_handle, - ), - patch.object( - session_manager_stateless, - "_server_instances", - {}, - ), - patch.object( - session_manager_stateful, - "_server_instances", - {}, - ), - ): - await handle_streamable_http_mcp(scope, receive, send) + ), + patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) - return bool(stateless_called), bool(stateful_called) - - # initialize → stateful - init_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}' - stateless_called, stateful_called = await make_request(init_body) - assert stateful_called and not stateless_called, "initialize (no session) should route to stateful, not stateless" - - # tools/list → stateless - tools_body = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}' - stateless_called, stateful_called = await make_request(tools_body) - assert stateless_called and not stateful_called, "tools/list (no session) should route to stateless, not stateful" + assert stateful_handle.await_count == (1 if stateful else 0) + assert stateless_handle.await_count == (0 if stateful else 1) + observe_start.assert_awaited_once_with(0 if debug and method == "POST" else 1) + assert send.await_count == 2 + assert send.call_args_list[0].args[0]["status"] == 200 + assert send.call_args_list[1].args[0] == body + headers: Final = dict(send.call_args_list[0].args[0]["headers"]) + if debug: + assert headers[b"x-mcp-debug-auth-resolution"] == (b"stored-user-token" if method == "POST" else b"unresolved") + else: + assert not any(name.startswith(b"x-mcp-debug") for name in headers) @pytest.mark.asyncio @@ -1723,6 +2003,11 @@ async def test_mcp_routing_chunked_initialize_to_stateful(): patch( "litellm.proxy._experimental.mcp_server.server.set_auth_context", ), + patch( # test-quality-ok: registry is empty in unit tests; key owns one server + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[MagicMock()], + ), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, @@ -2208,6 +2493,11 @@ async def test_initialize_request_tracks_active_session_after_response_header(): new_callable=AsyncMock, return_value=(owner_auth, None, None, None, None, None), ), + patch( # test-quality-ok: registry is empty in unit tests; key owns one server + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[MagicMock()], + ), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, @@ -2320,6 +2610,11 @@ async def test_initialize_request_with_existing_session_tracks_new_session(): {"x-new-header": "new"}, ), ), + patch( # test-quality-ok: registry is empty in unit tests; key owns one server + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[MagicMock()], + ), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, @@ -3518,6 +3813,35 @@ async def test_call_mcp_tool_user_unauthorized_access(): assert "User not allowed to call this tool" in exc_info.value.detail +@pytest.mark.asyncio +async def test_call_mcp_tool_scoped_denial_names_the_binding_agent(): + from litellm.proxy._experimental.mcp_server.server import call_mcp_tool + + agent_bound_key = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-123") + + with ( + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_allowed_mcp_servers", + AsyncMock(return_value=[]), + ), + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + _scope_resolver({"github": "srv-github"}), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await call_mcp_tool( + name="github-search_issues", + arguments={}, + user_api_key_auth=agent_bound_key, + mcp_servers=["github"], + ) + + assert exc_info.value.status_code == 403 + assert "MCP server 'github'" in exc_info.value.detail["error"] + assert "agent 'agent-123'" in exc_info.value.detail["error"] + + @pytest.mark.asyncio async def test_call_mcp_tool_unauthorized_403_does_not_leak_server_credentials(): """Regression for LIT-4703 / GH #29936. @@ -5305,6 +5629,78 @@ class TestGatewayCreateInitializationOptions: assert server.create_initialization_options().server_name == "litellm-mcp-server" + @pytest.mark.asyncio + async def test_initialize_with_no_granted_servers_returns_403(self): + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + ) + from litellm.proxy._types import UserAPIKeyAuth + + with patch( # test-quality-ok: grant resolution is the input under test + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[], + ): + with pytest.raises(HTTPException) as exc_info: + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=UserAPIKeyAuth(api_key="sk-no-mcp"), + mcp_servers=None, + client_ip=None, + is_initialize=True, + ): + pytest.fail("initialize must not proceed when the key grants no MCP servers") + + assert exc_info.value.status_code == 403 + assert "no MCP servers granted" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_initialize_with_no_granted_scoped_servers_returns_scoped_denial(self): + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + ) + from litellm.proxy._types import UserAPIKeyAuth + + with patch( # test-quality-ok: grant resolution is the input under test + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[], + ): + with pytest.raises(HTTPException) as exc_info: + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=UserAPIKeyAuth(api_key="sk-no-mcp"), + mcp_servers=["grafana"], + client_ip=None, + is_initialize=True, + ): + pytest.fail("scoped initialize must not proceed when nothing resolves") + + assert exc_info.value.status_code == 403 + assert "grafana" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_non_initialize_request_with_no_granted_servers_is_not_rejected_here(self): + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + _mcp_gateway_initialize_instructions, + ) + from litellm.proxy._types import UserAPIKeyAuth + + with patch( # test-quality-ok: grant resolution is the input under test + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[], + ): + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=UserAPIKeyAuth(api_key="sk-no-mcp"), + mcp_servers=None, + client_ip=None, + ): + assert _mcp_gateway_initialize_instructions.get() is None + @pytest.mark.asyncio async def test_sse_handler_scopes_server_name_from_single_server_path(self): try: @@ -7319,6 +7715,75 @@ class TestMCPMetaTraceCarrier: assert _mcp_meta_trace_carrier(SimpleNamespace(meta=only_progress)) is None +@pytest.mark.asyncio +async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations() -> None: + from types import SimpleNamespace + + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + + from litellm.integrations.otel.model.destination import OtelDestination + from litellm.integrations.otel.plumbing.context import ( + request_destinations, + reset_request_destinations, + set_request_destinations, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.server import ( + _MCP_DESTINATIONS_SCOPE_KEY, + mcp_server_tool_call, + set_auth_context, + ) + from litellm.proxy._experimental.mcp_server.tool_registry import ( + global_mcp_tool_registry, + ) + + initialized_destination = OtelDestination(endpoint="https://initialize.example", callback_name="langfuse_otel") + current_destination = OtelDestination(endpoint="https://current.example", callback_name="arize") + server = MCPServer( + server_id="otel-context-test", + name="otelcontext", + transport=MCPTransport.http, + allow_all_keys=True, + ) + + async def observe_destinations() -> str: + assert request_destinations() == (current_destination,) + return "ok" + + global_mcp_server_manager.registry[server.server_id] = server + global_mcp_server_manager.tool_name_to_mcp_server_name_mapping["otelcontext-observe"] = server.name + global_mcp_tool_registry.register_tool( + name="otelcontext-observe", + description="Observe request destinations", + input_schema={"type": "object"}, + handler=observe_destinations, + ) + set_auth_context(None, raw_headers={}) + destinations_token = set_request_destinations((initialized_destination,)) + scope = {_MCP_DESTINATIONS_SCOPE_KEY: (current_destination,)} + current_request_context = RequestContext( + request_id=1, + meta=None, + session=SimpleNamespace(), + lifespan_context=None, + request=SimpleNamespace(scope=scope), + ) + request_token = request_ctx.set(current_request_context) + try: + result = await mcp_server_tool_call("otelcontext-observe", {}) + assert result.isError is False + assert request_destinations() == (initialized_destination,) + finally: + request_ctx.reset(request_token) + reset_request_destinations(destinations_token) + global_mcp_tool_registry.tools.pop("otelcontext-observe", None) + global_mcp_server_manager.registry.pop(server.server_id, None) + global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.pop("otelcontext-observe", None) + + @pytest.mark.asyncio async def test_get_allowed_mcp_servers_includes_active_servers_submitted_by_user(): """BYOM submitters can see approved servers they submitted without allow_all_keys.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e3ed48713aa..b637586d7ff 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5,7 +5,7 @@ import logging import os import sys from datetime import datetime -from typing import Any, Dict, Final, Optional +from typing import Any, Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -58,6 +58,11 @@ from litellm.proxy._types import ( from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer +from litellm.caching.caching import DualCache +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks def _reload_mcp_manager_module(): @@ -900,6 +905,68 @@ class TestMCPServerManager: assert retry_slot is not None assert retry_slot.generation > old_generation + @pytest.mark.asyncio + @pytest.mark.parametrize("corrupt_column", ("static_headers", "env")) + async def test_database_reload_drops_cached_server_whose_secret_map_stops_decoding( + self, monkeypatch, caplog, corrupt_column + ): + from types import SimpleNamespace + + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_secret_map + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-reload-secret-map-salt") + monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": "aes-256-gcm"}) + headers = {"Authorization": "Bearer dummy-header-secret-4f1c"} + env = {"UPSTREAM_TOKEN": "dummy-env-secret-9a2b"} + stamp = datetime.now() + cached = MCPServer( + server_id="cached-server", + name="cached_server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + static_headers=dict(headers), + env=dict(env), + updated_at=stamp, + ) + manager = MCPServerManager() + manager.registry[cached.server_id] = cached + stored = {"static_headers": encrypt_secret_map(headers), "env": encrypt_secret_map(env)} + corrupted = {**stored, corrupt_column: stored[corrupt_column][:-6] + 'AAAAA"'} + + def _row(server_id, maps): + row = MagicMock() + row.server_id = server_id + row.alias = server_id + row.model_dump.return_value = { + "server_id": server_id, + "alias": server_id, + "server_name": server_id, + "url": "https://up.example.com/mcp", + "transport": MCPTransport.http, + "updated_at": stamp, + **maps, + } + return row + + table = SimpleNamespace( + find_many=AsyncMock(return_value=[_row(cached.server_id, corrupted), _row("healthy-sibling", stored)]) + ) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + await manager.reload_servers_from_database() + + assert set(manager.registry) == {"healthy-sibling"} + sibling = manager.registry["healthy-sibling"] + assert dict(sibling.static_headers) == headers + assert dict(sibling.env) == env + logged = "\n".join(caplog.messages) + assert cached.server_id in logged + for secret in (*headers.values(), *env.values(), *stored.values(), corrupted[corrupt_column]): + assert secret not in logged + @pytest.mark.asyncio async def test_lazy_oauth_discovery_preserves_manual_authorization_url_gate(self): with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "false"}): @@ -1229,6 +1296,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() @@ -4362,6 +4473,116 @@ class TestMCPServerManager: assert len(result) == 1 assert result[0].name == "github_tool_1" + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.api_key, MCPAuth.oauth2]) + @pytest.mark.parametrize("is_byok", [False, True]) + @pytest.mark.parametrize("scheme", ["http", "https"]) + async def test_openapi_health_loads_spec_without_mcp_handshake(self, respx_mock, monkeypatch, auth_type, is_byok, scheme): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="openapi-health", + name="openapi-health", + transport=MCPTransport.http, + url="https://rest.example.com", + spec_path=f"{scheme}://93.184.216.34/openapi.json", + auth_type=auth_type, + is_byok=is_byok, + authentication_token=None if is_byok else "shared-secret", + static_headers={"Authorization": "Bearer static-secret"}, + ) + manager.registry = {server.server_id: server} + route = respx_mock.get(server.spec_path).respond(200, json={"openapi": "3.0.0", "paths": {}}) + result = await manager.health_check_server(server.server_id, mcp_auth_header="caller-secret") + assert result.status == "healthy" + assert result.health_check_error is None + assert result.last_health_check is not None + assert result.spec_path == server.spec_path + assert route.call_count == 1 + assert "authorization" not in route.calls[0].request.headers + assert "x-api-key" not in route.calls[0].request.headers + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token]) + @pytest.mark.parametrize("spec_path", ["/config/openapi.json", "relative/openapi.json"]) + async def test_openapi_local_spec_health_is_unknown(self, respx_mock, auth_type, spec_path): + manager = MCPServerManager() + server = MCPServer( + server_id="local-openapi-health", + name="local-openapi-health", + transport=MCPTransport.http, + url="https://rest.example.com", + spec_path=spec_path, + auth_type=auth_type, + is_byok=True, + ) + manager.registry = {server.server_id: server} + result = await manager.health_check_server(server.server_id) + assert result.status == "unknown" + assert result.health_check_error == "OpenAPI servers have no protocol-level health probe" + assert result.last_health_check is not None + assert not respx_mock.calls + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("failure", "expected_status", "expected_error"), + [ + (httpx.Response(401, text="secret response content"), "unhealthy", "OpenAPI specification request failed (HTTP 401)"), + (httpx.Response(404), "unhealthy", "OpenAPI specification request failed (HTTP 404)"), + (httpx.Response(500), "unhealthy", "OpenAPI specification request failed (HTTP 500)"), + (httpx.ConnectError("secret network details"), "unhealthy", "OpenAPI specification could not be loaded (ConnectError)"), + (httpx.Response(200, text="secret invalid JSON body"), "unhealthy", "OpenAPI specification could not be loaded (JSONDecodeError)"), + ], + ) + async def test_openapi_health_reports_safe_failures(self, respx_mock, monkeypatch, failure, expected_status, expected_error): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="failed-openapi-health", + name="failed-openapi-health", + transport=MCPTransport.http, + url="https://rest.example.com", + spec_path="https://93.184.216.34/key-secret?token=query-secret", + auth_type=MCPAuth.bearer_token, + is_byok=True, + ) + manager.registry = {server.server_id: server} + route = respx_mock.get(server.spec_path).mock(side_effect=[failure]) + result = await manager.health_check_server(server.server_id) + assert result.status == expected_status + assert result.health_check_error == expected_error + assert result.last_health_check is not None + assert route.call_count == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("cancel", [False, True]) + async def test_openapi_health_timeout_and_cancellation_cleanup(self, respx_mock, monkeypatch, cancel): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_spec_health + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + started = asyncio.Event() + cancelled = asyncio.Event() + + async def slow_load(request): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + respx_mock.get("https://93.184.216.34/slow.json").mock(side_effect=slow_load) + task = asyncio.create_task(_openapi_spec_health("https://93.184.216.34/slow.json", timeout=0.1)) + await asyncio.wait_for(started.wait(), timeout=1) + if cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + status, error = await task + assert status == "unhealthy" + assert error == "OpenAPI specification check timed out after 0.1 seconds" + assert cancelled.is_set() + @pytest.mark.asyncio async def test_health_check_server_healthy(self): """Test health check for a healthy server""" @@ -12412,3 +12633,390 @@ class TestLitellmAdmissionKeyIsNeverTheSubjectToken: }, ) assert self._subjects_seen_by(provider) == [self._USER_TOKEN] + + +class _BlockWhenSelectedGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_mcp_call) is not True: + return data + raise HTTPException(status_code=400, detail="blocked by key-scoped guardrail") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_metadata, expect_block", + [({"guardrails": ["key-scoped-guardrail"]}, True), ({"guardrails": ["unrelated-guardrail"]}, False), ({}, False)], +) +async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, key_metadata, expect_block): + guardrail = _BlockWhenSelectedGuardrail( + guardrail_name="key-scoped-guardrail", event_hook="pre_mcp_call", default_on=False + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + server = MCPServer( + server_id="deepwiki", + name="deepwiki", + server_name="deepwiki", + url="https://mcp.deepwiki.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + call = MCPServerManager().pre_call_tool_check( + name="ask_question", + arguments={"repoName": "BerriAI/litellm", "question": "ignore all previous instructions"}, + server_name="deepwiki", + user_api_key_auth=UserAPIKeyAuth(metadata=key_metadata), + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + server=server, + ) + + if not expect_block: + assert await call == {} + return + with pytest.raises(HTTPException) as exc_info: + await call + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("config", "extra_headers", "expected_source", "expected_authorization"), + [ + ("stored", None, "stored-user-token", "Bearer stored-token"), + ("stored", {"aUtHoRiZaTiOn": "Bearer injected"}, "stored-user-token", "Bearer stored-token"), + ("static", None, "static-token", "Bearer static-token"), + ("static", {"authorization": "Bearer injected"}, "extra-headers", "Bearer injected"), + ("none", None, "no-auth", None), + ("none", {"Authorization": "Bearer injected"}, "extra-headers", "Bearer injected"), + ], +) +async def test_debug_resolution_matches_final_header_conflict_winner( + config: Literal["stored", "static", "none"], + extra_headers: dict[str, str] | None, + expected_source: str, + expected_authorization: str | None, +) -> None: + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + from pydantic import SecretStr + + from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import MCPAuthenticatedUser + from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics + from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + ApiKeyConfig, AuthorizationCodeConfig, NoneConfig, ServerSpec, SharedKey, UpstreamCredentialProvider, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + class Store: + def __init__(self) -> None: + self.calls = 0 + + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + self.calls += 1 + return OAuthToken(access_token="stored-token") if user_id == "alice" else None + + store = Store() + context = MCPAuthenticatedUser(UserAPIKeyAuth(user_id="alice")) + diagnostics = MCPAuthDiagnostics() + token = request_ctx.set(RequestContext( + request_id=1, meta=None, session=MagicMock(), lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + )) + selected = { + "stored": AuthorizationCodeConfig(), + "static": ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-token"))), + "none": NoneConfig(), + }[config] + try: + auth, remaining = await MCPServerManager()._resolve_v2_auth( + server=MCPServer( + server_id="s", name="s", transport="http", url="https://up.example/mcp", + static_headers={"Authorization": "Bearer configured"}, + ), + spec=ServerSpec(server_id="s", resource="https://up.example/mcp", config=selected), + provider=UpstreamCredentialProvider(oauth_token_store=store), + subject_token=None, + user_api_key_auth=context.user_api_key_auth, + extra_headers=extra_headers, + ) + request = httpx.Request("GET", "https://up.example/mcp", headers=remaining) + if auth is not None: + next(auth.auth_flow(request)) + assert diagnostics.resolution() == expected_source + assert request.headers.get("Authorization") == expected_authorization + assert store.calls == (1 if config == "stored" else 0) + finally: + request_ctx.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "stdio"]) +async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Literal["http", "stdio"]) -> None: + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + diagnostics = MCPAuthDiagnostics() + token = request_ctx.set(RequestContext( + request_id=1, meta=None, session=MagicMock(), lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + )) + try: + server = MCPServer( + server_id="signed", name="signed", transport=transport, + url="https://up.example/mcp", auth_type="aws_sigv4", + aws_access_key_id="AKIDEXAMPLE", aws_secret_access_key="test-signing-secret", + aws_region_name="us-east-1", aws_service_name="execute-api", + command="python", args=["-c", "pass"], + ) + client = await MCPServerManager()._create_mcp_client(server) + if transport == "stdio": + assert diagnostics.resolution() == "not-applicable" + else: + assert diagnostics.resolution() == "aws-sigv4" + request = httpx.Request("POST", "https://up.example/mcp", content=b"{}") + next(client._aws_auth.auth_flow(request)) + assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256 ") + assert "Credential=AKIDEXAMPLE/" in request.headers["Authorization"] + finally: + request_ctx.reset(token) + + +@pytest.mark.asyncio +async def test_temporary_server_discovery_reuses_resolved_metadata_without_publishing() -> None: + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="temporary-oauth-discovery", name="temporary", url="https://idp.example.com/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough, + ) + manager._set_oauth_discovery_deferred(server.server_id, True) + metadata: Final = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + ) + with patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery: + resolved: Final = await manager.ensure_oauth_metadata_discovered(server) + repeated: Final = await manager.ensure_oauth_metadata_discovered(server) + assert resolved.authorization_url == metadata.authorization_url + assert resolved.token_url == metadata.token_url + assert resolved.registration_url == metadata.registration_url + assert repeated is resolved + assert server.server_id not in manager.registry + assert server.server_id not in manager.config_mcp_servers + discovery.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type", [MCPAuth.oauth2, MCPAuth.true_passthrough]) +async def test_repeated_stale_oauth_discovery_is_bounded(auth_type: MCPAuth) -> None: + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="repeated-stale", name="stale", url="https://idp.example.com/mcp", + transport=MCPTransport.http, auth_type=auth_type, oauth2_flow="authorization_code", + ) + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + metadata: Final = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + ) + with ( + patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery, + patch.object(manager, "_publish_resolved_oauth_server", return_value=None), + ): + if auth_type == MCPAuth.true_passthrough: + assert await manager.ensure_oauth_metadata_discovered(server) is server + else: + with pytest.raises(HTTPException) as exc: + await manager.ensure_oauth_metadata_discovered(server) + assert exc.value.status_code == 503 + assert "changed repeatedly" in str(exc.value.detail) + assert discovery.await_count == 2 + + +@pytest.mark.asyncio +async def test_stale_discovery_falls_back_to_resolved_registered_server() -> None: + manager: Final = MCPServerManager() + original: Final = MCPServer( + server_id="resolved-replacement", name="replacement", url="https://old.example.com/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", + ) + replacement: Final = original.model_copy(update={ + "url": "https://new.example.com/mcp", "authorization_url": "https://new.example.com/authorize", + "token_url": "https://new.example.com/token", + }) + manager.registry[original.server_id] = replacement + assert await manager._rejoin_oauth_metadata_discovery(original, retry_stale=False) is replacement + + +def test_stale_discovery_cannot_overwrite_new_registered_server() -> None: + manager: Final = MCPServerManager() + original: Final = MCPServer( + server_id="stale-publication", name="publication", url="https://old.example.com/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + ) + manager._set_oauth_discovery_deferred(original.server_id, True) + original_slot: Final = manager._oauth_discovery_slot(original.server_id) + assert original_slot is not None + replacement: Final = original.model_copy(update={"url": "https://new.example.com/mcp"}) + manager.registry[original.server_id] = replacement + manager._set_oauth_discovery_deferred(original.server_id, True) + assert manager._publish_resolved_oauth_server(original, original_slot.generation) is None + assert manager.registry[original.server_id] is replacement + + +@pytest.mark.asyncio +async def test_temporary_oauth_discovery_expires_without_more_requests() -> None: + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="expiring-session", name="temporary", url="https://idp.example.com/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough, + authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + ) + manager._set_oauth_discovery_deferred(server.server_id, True) + resolved: Final = await manager.ensure_oauth_metadata_discovered(server) + assert manager._oauth_discovery_slot(server.server_id) is not None + loop: Final = asyncio.get_running_loop() + expired: Final = loop.create_future() + with patch.object(loop, "time", return_value=loop.time() + 301): + loop.call_later(0, expired.set_result, None) + await expired + assert resolved.authorization_url == server.authorization_url + assert manager._oauth_discovery_slot(server.server_id) is None + + +def test_old_temporary_discovery_expiry_preserves_replacement() -> None: + manager: Final = MCPServerManager() + manager._set_oauth_discovery_deferred("reused-session", True) + old_slot: Final = manager._oauth_discovery_slot("reused-session") + assert old_slot is not None + manager._set_oauth_discovery_deferred("reused-session", True) + replacement: Final = manager._oauth_discovery_slot("reused-session") + manager._expire_temporary_oauth_discovery("reused-session", old_slot.generation) + assert manager._oauth_discovery_slot("reused-session") is replacement + assert replacement is not None + manager._expire_temporary_oauth_discovery("reused-session", replacement.generation) + assert manager._oauth_discovery_slot("reused-session") is None + manager._expire_temporary_oauth_discovery("reused-session", replacement.generation) + assert manager._oauth_discovery_slot("reused-session") is None + + +@pytest.mark.asyncio +async def test_openapi_health_coalesces_concurrent_checks_and_reuses_results(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="coalesced", + name="coalesced", + transport=MCPTransport.http, + spec_path="https://93.184.216.34/coalesced.json", + auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} + started = asyncio.Event() + release = asyncio.Event() + + async def serve(request): + started.set() + await release.wait() + return httpx.Response(200, json={"paths": {}}) + + route = respx_mock.get(server.spec_path).mock(side_effect=serve) + tasks = [asyncio.create_task(manager.health_check_server(server.server_id)) for _ in range(4)] + await asyncio.wait_for(started.wait(), timeout=1) + release.set() + results = await asyncio.gather(*tasks) + cached = await manager.health_check_server(server.server_id) + assert [result.status for result in results] == ["healthy"] * 4 + assert cached.status == "healthy" + assert {result.last_health_check for result in [*results, cached]} == {results[0].last_health_check} + assert route.call_count == 1 + + +@pytest.mark.asyncio +async def test_openapi_health_cache_expires_at_thirty_seconds(respx_mock, monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _OpenAPIHealthProbe + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + clock = iter([0.0, 29.0, 30.0, 30.0]) + probe = _OpenAPIHealthProbe("https://93.184.216.34/expiry.json", clock=clock.__next__) + route = respx_mock.get(probe.spec_path).mock( + side_effect=[ + httpx.Response(200, json={"paths": {}}), + httpx.Response(503), + ] + ) + first = await probe.check() + assert first[0] == "healthy" + assert await probe.check() == first + refreshed = await probe.check() + assert refreshed[0] == "unhealthy" + assert refreshed[1] == "OpenAPI specification request failed (HTTP 503)" + assert refreshed[2] >= first[2] + assert route.call_count == 2 + + +@pytest.mark.asyncio +async def test_openapi_health_reports_size_limit_as_unknown_and_caches_failure(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="oversized", + name="oversized", + transport=MCPTransport.http, + spec_path="https://93.184.216.34/large.json", + auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} + route = respx_mock.get(server.spec_path).respond(200, headers={"content-length": str(12 * 1024 * 1024)}) + result = await manager.health_check_server(server.server_id) + cached = await manager.health_check_server(server.server_id) + assert result.status == "unknown" + assert result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit" + assert cached.health_check_error == result.health_check_error + assert cached.last_health_check == result.last_health_check + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("already_waiting", [False, True]) +async def test_openapi_health_cancellation_does_not_poison_cache(respx_mock, monkeypatch, already_waiting): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="cancelled-cache", name="cancelled-cache", transport=MCPTransport.http, + spec_path="https://93.184.216.34/cancelled-cache.json", auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} + started = asyncio.Event() + attempts = [] + + async def serve(request): + attempts.append(request.url) + if not started.is_set(): + started.set() + await asyncio.Event().wait() + return httpx.Response(200, json={"paths": {}}) + + route = respx_mock.get(server.spec_path).mock(side_effect=serve) + leader = asyncio.create_task(manager.health_check_server(server.server_id)) + await asyncio.wait_for(started.wait(), timeout=1) + follower = asyncio.create_task(manager.health_check_server(server.server_id)) if already_waiting else None + await asyncio.sleep(0) + leader.cancel() + cancelled = await leader + assert cancelled.status == "unknown" + assert cancelled.health_check_error == "OpenAPI specification check was cancelled" + recovered = await follower if follower is not None else await manager.health_check_server(server.server_id) + assert recovered.status == "healthy" + assert recovered.health_check_error is None + cached = await manager.health_check_server(server.server_id) + assert cached.last_health_check == recovered.last_health_check + assert cached.status == "healthy" + assert len(attempts) == 2 + assert route.call_count == 1 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index f6b61c1d9f7..e814425c9a2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -15,6 +15,11 @@ import httpx from litellm.experimental_mcp_client.client import MCPSigV4Auth, MCPClient from litellm.types.mcp import MCPAuth, MCPTransport +from prisma import models + + +def _updated_row() -> models.LiteLLM_MCPServerTable: + return models.LiteLLM_MCPServerTable.model_construct(server_id="test-server", transport="http", env={}, env_vars=[]) class TestMCPSigV4Auth: @@ -600,7 +605,7 @@ class TestCredentialMergeOnUpdate: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -639,7 +644,7 @@ class TestCredentialMergeOnUpdate: from litellm.proxy._types import UpdateMCPServerRequest mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -667,7 +672,7 @@ class TestCredentialMergeOnUpdate: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -709,7 +714,7 @@ class TestCredentialMergeOnUpdate: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -752,7 +757,7 @@ class TestCredentialMergeOnUpdate: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -1083,7 +1088,7 @@ class TestAuthTypeSwitchClearsCredentials: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index feab179570b..9420eecd222 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -750,8 +750,7 @@ async def test_admitted_subject_missing_stored_token_challenged_with_resource_me challenge = exc_info.value.headers["www-authenticate"] assert "authorization_uri=" not in challenge assert challenge == ( - 'Bearer resource_metadata="http://localhost:8000' - '/.well-known/oauth-protected-resource/mcp/repro_oauth_server"' + 'Bearer resource_metadata="http://localhost:8000/.well-known/oauth-protected-resource/mcp/repro_oauth_server"' ) @@ -938,6 +937,11 @@ async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_cha "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", return_value=delegated_server, ), + patch( # test-quality-ok: registry is empty in unit tests; key owns the delegated server + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[delegated_server], + ), patch.object( session_manager_stateful, "handle_request", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 239f89ebd90..b8935d07774 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -11,6 +11,7 @@ Covers: import json from collections.abc import Sequence +from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -24,6 +25,7 @@ from litellm.proxy._experimental.mcp_server.tool_search import ( AGENT_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, SemanticToolRanker, ToolSearchResult, coerce_top_k, @@ -272,8 +274,8 @@ class TestSearchTools: class TestGetVirtualToolDefinitions: - def test_returns_three_tools(self) -> None: - assert len(get_virtual_tool_definitions()) == 3 + def test_returns_four_tools(self) -> None: + assert len(get_virtual_tool_definitions()) == 4 def test_agent_search_schema_requires_query(self) -> None: tools = get_virtual_tool_definitions() @@ -330,6 +332,7 @@ class TestGetVirtualToolDefinitions: MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, } @@ -364,7 +367,12 @@ class TestListToolRestApiWithToolSearch: assert result["error"] is None tool_names = [t["name"] for t in result["tools"]] - assert set(tool_names) == {MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME} + assert set(tool_names) == { + MCP_TOOL_SEARCH_TOOL_NAME, + MCP_TOOL_CALL_TOOL_NAME, + AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, + } @pytest.mark.asyncio async def test_returns_full_catalog_when_flag_disabled(self) -> None: @@ -737,6 +745,31 @@ class TestCallToolRestApiVirtualTools: assert mock_search.await_args.kwargs["top_k"] == 1 assert mock_search.await_args.kwargs["agents"] == (translator,) + @pytest.mark.asyncio + async def test_skill_search_call_tolerates_malformed_top_k(self) -> None: + """Regression: a caller-supplied non-numeric top_k must be coerced to the default, + the same as agent_search, instead of raising a pydantic ValidationError that the + endpoint's catch-all turns into an HTTP 500.""" + from mcp.types import CallToolResult, TextContent + + from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K + + user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) + request = self._make_request( + {"name": SKILL_SEARCH_TOOL_NAME, "arguments": {"query": "translate a document", "top_k": "not-a-number"}} + ) + fake_result = CallToolResult(content=[TextContent(type="text", text="[]")], isError=False) + with patch( # test-quality-ok: the embedding router only resolves via proxy_server globals, no injection seam + "litellm.proxy._experimental.mcp_server.tool_search.handle_skill_search", + new_callable=AsyncMock, + return_value=fake_result, + ) as mock_search: + result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict) + + assert result.isError is False + assert mock_search.await_args.kwargs["top_k"] == DEFAULT_SKILL_SEARCH_TOP_K + assert mock_search.await_args.kwargs["query"] == "translate a document" + @pytest.mark.asyncio async def test_agent_search_call_reports_missing_embedding_model_as_tool_error(self) -> None: from litellm.proxy.agent_endpoints.agent_search import AgentSearchNotConfigured @@ -782,10 +815,15 @@ class TestCallToolRestApiVirtualTools: router = MagicMock() router.aembedding = AsyncMock(side_effect=fake_aembedding) + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) with ( patch( # test-quality-ok: the proxy's router is a module global; the handler reaches it the way production does "litellm.proxy.proxy_server.llm_router", router ), + patch( # test-quality-ok: the proxy's key-limit hooks are a module global; the embedding call runs them like /embeddings does + "litellm.proxy.proxy_server.proxy_logging_obj", key_limits + ), patch( # test-quality-ok: the authorized catalog is the seam every virtual tool shares; the ranking under test stays real "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", new_callable=AsyncMock, @@ -795,6 +833,8 @@ class TestCallToolRestApiVirtualTools: result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) assert mock_list.await_args.kwargs["user_api_key_auth"] is user_api_key_dict + assert key_limits.pre_call_hook.await_args.kwargs["call_type"] == "aembedding" + assert key_limits.pre_call_hook.await_args.kwargs["data"]["model"] == "emb" assert result.isError is False assert [t["name"] for t in json.loads(result.content[0].text)] == [FX_TOOL.name] @@ -810,7 +850,9 @@ class TestCallToolRestApiVirtualTools: assert "mcp_tool_search.embedding_model" in result.content[0].text @pytest.mark.asyncio - async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(self, monkeypatch: pytest.MonkeyPatch) -> None: + async def test_mcp_tool_search_reports_invalid_settings_as_tool_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setattr(litellm, "mcp_tool_search", {"top_k": 0}) user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) @@ -1196,6 +1238,7 @@ class TestHandleListToolsVirtual: MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, } @@ -1229,3 +1272,37 @@ class TestMcpServerToolCallErrorHandling: assert result.isError is True assert "User not allowed to call this tool" in result.content[0].text + + +@pytest.mark.asyncio +async def test_handle_mcp_tool_call_scoped_denial_names_the_binding_agent() -> None: + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.tool_search import handle_mcp_tool_call + + agent_bound_key = UserAPIKeyAuth(api_key="test_key", agent_id="agent-123") + + async def resolve(user_api_key_auth, mcp_servers, client_ip=None): + if user_api_key_auth.agent_id: + return [] + return [ + SimpleNamespace( + server_id="srv-github", server_name="github", alias=None, short_prefix=None, access_groups=[] + ) + ] + + with patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(side_effect=resolve), + ): + with pytest.raises(HTTPException) as exc_info: + await handle_mcp_tool_call( + tool_name="github-create_issue", + arguments={}, + user_api_key_dict=agent_bound_key, + mcp_servers=["github"], + ) + + assert exc_info.value.status_code == 403 + assert "MCP server 'github'" in exc_info.value.detail["error"] + assert "agent 'agent-123'" in exc_info.value.detail["error"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py new file mode 100644 index 00000000000..0036035f448 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py @@ -0,0 +1,667 @@ +import time +from collections.abc import Mapping +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch + +import jwt +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import HTTPException +from pydantic import ValidationError + +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + RefreshOwnershipProven, + RefreshTokenPresented, + VerifiedRefreshToken, + _discover_jwks_url, + _fetch_issuer_jwks, + _load_caller_principal, + _load_stored_refresh_token, + _select_signing_key, + current_binding_proof, + enforce_oauth_identity_binding, +) +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + +ISSUER: Final = "https://idp.example.com" +AUDIENCE: Final = "litellm-client" +KID: Final = "test-key" + +_PRIVATE_KEY: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) +_PRIVATE_PEM: Final = _PRIVATE_KEY.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), +) +_PUBLIC_JWK: Final = { + **jwt.algorithms.RSAAlgorithm.to_jwk(_PRIVATE_KEY.public_key(), as_dict=True), + "kid": KID, + "alg": "RS256", + "use": "sig", +} + + +def _sign_id_token(claims: Mapping[str, object]) -> str: + payload: Final = { + "iss": ISSUER, + "aud": AUDIENCE, + "exp": int(time.time()) + 300, + "iat": int(time.time()), + "nonce": "test-nonce", + "sub": "upstream-user", + **claims, + } + return jwt.encode(payload, _PRIVATE_PEM, algorithm="RS256", headers={"kid": KID}) + + +async def _jwks_fetcher(_binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]: + return [_PUBLIC_JWK] + + +def _caller_loader(email: str | None): + async def load(_user_id: str, _binding: MCPOAuthIdentityBinding) -> str | None: + return email + + return load + + +def _stored_refresh_token_loader(refresh_token: str | None): + async def load(_user_id: str, _server_id: str, _binding: MCPOAuthIdentityBinding) -> VerifiedRefreshToken | None: + return VerifiedRefreshToken(refresh_token, "verified-binding") if refresh_token else None + + return load + + +def _server(mode: str = "enforce", **binding_overrides: object) -> MCPServer: + return MCPServer( + server_id="srv-1", + name="srv-1", + url="https://mcp.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth_identity_binding=MCPOAuthIdentityBinding( + mode=mode, + issuer=ISSUER, + audiences=[AUDIENCE], + **binding_overrides, + ), + ) + + +@pytest.mark.asyncio +async def test_fetch_issuer_jwks_fetches_and_caches_keys(): + binding: Final = _server(jwks_url="https://idp.example.com/jwks-cache").oauth_identity_binding + assert binding is not None + response: Final = MagicMock() + response.json.return_value = {"keys": [_PUBLIC_JWK]} + client: Final = MagicMock() + client.get = AsyncMock(return_value=response) + + with patch( # test-quality-ok: no HTTP dependency injection seam exists for JWKS fetching + "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", + return_value=client, + ): + first: Final = await _fetch_issuer_jwks(binding) + second: Final = await _fetch_issuer_jwks(binding) + + assert first == second == [_PUBLIC_JWK] + client.get.assert_awaited_once_with("https://idp.example.com/jwks-cache") + + +@pytest.mark.asyncio +async def test_fetch_issuer_jwks_rejects_malformed_document(): + binding: Final = _server(jwks_url="https://idp.example.com/jwks-invalid").oauth_identity_binding + assert binding is not None + response: Final = MagicMock() + response.json.return_value = {} + client: Final = MagicMock() + client.get = AsyncMock(return_value=response) + + with patch( # test-quality-ok: no HTTP dependency injection seam exists for JWKS fetching + "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", + return_value=client, + ): + with pytest.raises(TypeError, match="has no 'keys' array"): + await _fetch_issuer_jwks(binding) + + +@pytest.mark.asyncio +async def test_discover_jwks_url_returns_provider_uri(): + response: Final = MagicMock() + response.json.return_value = {"jwks_uri": "https://idp.example.com/jwks"} + client: Final = MagicMock() + client.get = AsyncMock(return_value=response) + + with patch( # test-quality-ok: no HTTP dependency injection seam exists for OIDC discovery + "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", + return_value=client, + ): + result: Final = await _discover_jwks_url("https://idp.example.com/") + + assert result == "https://idp.example.com/jwks" + client.get.assert_awaited_once_with("https://idp.example.com/.well-known/openid-configuration") + + +@pytest.mark.asyncio +async def test_discover_jwks_url_rejects_missing_provider_uri(): + response: Final = MagicMock() + response.json.return_value = {} + client: Final = MagicMock() + client.get = AsyncMock(return_value=response) + + with patch( # test-quality-ok: no HTTP dependency injection seam exists for OIDC discovery + "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", + return_value=client, + ): + with pytest.raises(ValueError, match="returned no jwks_uri"): + await _discover_jwks_url("https://idp.example.com") + + +def test_select_signing_key_returns_matching_key_or_rejection(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + + selected: Final = _select_signing_key(token, [_PUBLIC_JWK]) + rejected: Final = _select_signing_key(token, []) + + assert selected.__class__.__name__ == "PyJWK" + assert rejected.code == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_load_caller_principal_supports_user_id_and_database_email(): + user_id_binding: Final = _server(caller_field="user_id").oauth_identity_binding + assert user_id_binding is not None + assert await _load_caller_principal("user-a", user_id_binding) == "user-a" + + with patch( # test-quality-ok: caller loading is a lazy database boundary without injection + "litellm.proxy._experimental.mcp_server.bridge_token_flow.load_active_user_by_id", + new=AsyncMock(side_effect=["no_active_key", SimpleNamespace(user_email="alice@example.com")]), + ): + assert await _load_caller_principal("user-a", _server().oauth_identity_binding) is None + assert await _load_caller_principal("user-a", _server().oauth_identity_binding) == "alice@example.com" + + +@pytest.mark.asyncio +async def test_load_stored_refresh_token_returns_credential_and_fails_closed(): + binding: Final = _server(caller_field="user_id", principal_claim="sub").oauth_identity_binding + proof: Final = await current_binding_proof(binding, "user-a", "srv-1") + get_credential: Final = AsyncMock(return_value={"refresh_token": "rt-1", "identity_binding_proof": proof}) + with ( + patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection + "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", + new=get_credential, + ), + patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value="prisma", + ), + ): + assert await _load_stored_refresh_token("user-a", "srv-1", binding) == VerifiedRefreshToken("rt-1", proof) + get_credential.return_value = {"refresh_token": "rt-1"} + assert await _load_stored_refresh_token("user-a", "srv-1", binding) is None + + with patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection + "litellm.proxy.utils.get_prisma_client_or_throw", + side_effect=RuntimeError("database unavailable"), + ): + assert await _load_stored_refresh_token("user-a", "srv-1", binding) is None + + +@pytest.mark.asyncio +async def test_matching_principal_passes(): + token: Final = _sign_id_token({"email": "Alice@Example.com", "email_verified": True}) + result: Final = await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_mismatched_principal_rejected(): + token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_principal_mismatch" + assert exc_info.value.detail["credential_stored"] is False + + +@pytest.mark.asyncio +async def test_missing_upstream_principal_rejected(): + token: Final = _sign_id_token({"email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "no usable" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_jwks_fetch_failure_is_rejected(): + async def fail(_binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]: + raise RuntimeError("jwks unavailable") + + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=fail, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "jwks unavailable" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_missing_signing_key_is_rejected(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + + async def no_keys(_binding: MCPOAuthIdentityBinding) -> tuple[Mapping[str, object], ...]: + return () + + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=no_keys, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "signing key" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_missing_caller_principal_is_rejected(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader(None), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "has no" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_requires_litellm_identity_for_presented_token(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id=None, + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-1"), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "no resolvable LiteLLM user identity" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_user_id_principal_matching_uses_exact_comparison(): + token: Final = _sign_id_token({"sub": "user-a"}) + result: Final = await enforce_oauth_identity_binding( + server=_server(principal_claim="sub", caller_field="user_id"), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("user-a"), + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_missing_id_token_rejected_on_authorization_code(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_allowed_when_presented_token_matches_stored_credential(): + result: Final = await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-1"), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_rejects_different_presented_token(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-stolen"), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_rejects_missing_stored_token(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-1"), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader(None), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_identity_envelope_does_not_prove_upstream_binding(): + with pytest.raises(HTTPException) as error: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshOwnershipProven(), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert error.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_rejects_without_ownership_proof(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_refresh_with_mismatched_id_token_rejected(): + token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True}) + with pytest.raises(HTTPException): + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + + +@pytest.mark.asyncio +async def test_audit_mode_logs_but_does_not_reject(caplog): + token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True}) + result: Final = await enforce_oauth_identity_binding( + server=_server(mode="audit"), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert result is None + assert "oauth_principal_mismatch" in caplog.text + assert "nonce" not in caplog.text + + +@pytest.mark.asyncio +async def test_unverified_email_rejected(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": False}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_wrong_issuer_rejected(): + payload: Final = { + "iss": "https://evil.example.com", + "aud": AUDIENCE, + "exp": int(time.time()) + 300, + "email": "alice@example.com", + "email_verified": True, + } + token: Final = jwt.encode(payload, _PRIVATE_PEM, algorithm="RS256", headers={"kid": KID}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_no_litellm_identity_rejected(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id=None, + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_disabled_binding_is_noop(): + result: Final = await enforce_oauth_identity_binding( + server=_server(mode="disabled"), + token_response={"access_token": "at"}, + litellm_user_id=None, + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader(None), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_no_binding_is_noop(): + server: Final = MCPServer( + server_id="srv-2", + name="srv-2", + url="https://mcp.example.com", + transport=MCPTransport.http, + ) + result: Final = await enforce_oauth_identity_binding( + server=server, + token_response={"access_token": "at"}, + litellm_user_id=None, + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader(None), + ) + assert result is None + + +def test_identity_binding_requires_non_empty_audiences(): + with pytest.raises(ValidationError): + MCPOAuthIdentityBinding(mode="enforce", issuer=ISSUER, audiences=[]) + with pytest.raises(ValidationError): + MCPOAuthIdentityBinding(mode="enforce", issuer=ISSUER) + + +@pytest.mark.asyncio +async def test_wrong_audience_rejected(): + token: Final = _sign_id_token({"aud": "other-client", "email": "alice@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + expected_nonce="test-nonce", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nonce", [None, "another-login"]) +async def test_authorization_code_rejects_missing_or_foreign_nonce(nonce): + token = _sign_id_token({"email": "alice@example.com", "email_verified": True, "nonce": nonce}) + with pytest.raises(HTTPException) as error: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + expected_nonce="this-login", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert error.value.status_code == 403 + assert error.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_binding_proof_rejects_changed_user_or_policy(): + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import credential_binding_matches + + binding = _server(caller_field="user_id", principal_claim="sub").oauth_identity_binding + proof = await current_binding_proof(binding, "alice", "srv-1") + credential = {"identity_binding_proof": proof} + assert await credential_binding_matches(binding, "alice", "srv-1", credential) + assert not await credential_binding_matches(binding, "bob", "srv-1", credential) + assert not await credential_binding_matches(binding, "alice", "other-server", credential) + changed = binding.model_copy(update={"audiences": ["different-client"]}) + assert not await credential_binding_matches(changed, "alice", "srv-1", credential) + + +@pytest.mark.parametrize("auth_type", [MCPAuth.oauth_delegate, MCPAuth.true_passthrough]) +def test_identity_binding_rejects_modes_without_gateway_credential_custody(auth_type): + with pytest.raises(ValidationError, match="gateway-managed per-user"): + MCPServer( + server_id="srv", + name="srv", + transport=MCPTransport.http, + auth_type=auth_type, + oauth_identity_binding=_server().oauth_identity_binding, + ) + + +@pytest.mark.asyncio +async def test_audit_matching_login_without_nonce_does_not_report_failure(caplog): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + result: Final = await enforce_oauth_identity_binding( + server=_server(mode="audit"), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert result is None + assert "oauth_identity_binding audit" not in caplog.text diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index e59616e53c1..5fa202224e3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -1378,3 +1378,83 @@ class TestUpstreamStatusIsClassified: assert exc.value.status_code == status_code assert secret_body not in str(exc.value) assert str(exc.value) == f"upstream returned HTTP {status_code}" + +class TestBoundedOpenAPISpecLoading: + @pytest.mark.asyncio + @pytest.mark.parametrize("max_bytes", [12, 13]) + async def test_exact_size_and_smaller_specs_load(self, respx_mock, monkeypatch, max_bytes): + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + route = respx_mock.get("https://93.184.216.34/spec.json").respond(200, content=b'{"paths":{}}') + assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=max_bytes) == {"paths": {}} + assert route.calls[0].request.headers["accept-encoding"] == "identity" + + @pytest.mark.asyncio + @pytest.mark.parametrize("headers", [{"content-length": "1000000"}, {"content-encoding": "gzip"}]) + async def test_unsafe_response_headers_reject_before_reading(self, respx_mock, monkeypatch, headers): + import httpx + from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + load_openapi_spec_async, + ) + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + closed = [] + + class UnreadableStream(httpx.AsyncByteStream): + async def __aiter__(self): + pytest.fail("Oversized or compressed response must not be consumed") + yield b"" + + async def aclose(self): + closed.append(True) + + respx_mock.get("https://93.184.216.34/spec.json").respond(200, headers=headers, stream=UnreadableStream()) + with pytest.raises(HTTPResponseLimitError): + await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=12) + assert closed == [True] + + @pytest.mark.asyncio + async def test_chunked_response_is_bounded_and_closed(self, respx_mock, monkeypatch): + import httpx + from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + load_openapi_spec_async, + ) + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + consumed = [] + closed = [] + + class ChunkedStream(httpx.AsyncByteStream): + async def __aiter__(self): + for index in range(10): + consumed.append(index) + yield b"x" * 65536 + + async def aclose(self): + closed.append(True) + + respx_mock.get("https://93.184.216.34/spec.json").respond(200, stream=ChunkedStream()) + with pytest.raises(HTTPResponseLimitError, match="size limit"): + await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=65536) + assert consumed == [0, 1] + assert closed == [True] + + @pytest.mark.asyncio + @pytest.mark.parametrize("target", ["https://93.184.216.35/final.json", "http://127.0.0.1/private.json"]) + async def test_bounded_spec_redirects_preserve_ssrf_protection(self, respx_mock, monkeypatch, target): + from litellm.litellm_core_utils.url_utils import SSRFError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + respx_mock.get("https://93.184.216.34/spec.json").respond(302, headers={"location": target}) + destination = respx_mock.get(target).respond(200, json={"paths": {}}) + if "127.0.0.1" in target: + with pytest.raises(SSRFError): + await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) + assert not destination.called + else: + assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}} + assert destination.call_count == 1 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index c007d22117f..31ccd5c9817 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -3,7 +3,7 @@ import inspect import json import sys from datetime import datetime -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 @@ -87,6 +87,125 @@ def _route_has_dependency(route, dependency) -> bool: class TestExecuteWithMcpClient: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("auth_type", "auth_value", "expected_auth"), + ( + (MCPAuth.none, None, {}), + (MCPAuth.basic, "preview:correct", {"Authorization": "Basic cHJldmlldzpjb3JyZWN0"}), + (MCPAuth.basic, None, {"Authorization": "Basic cHJldmlldzpzdG9yZWQ="}), + (MCPAuth.bearer_token, "edited", {"Authorization": "Bearer edited"}), + (MCPAuth.api_key, "edited", {"X-API-Key": "edited"}), + (MCPAuth.token, "edited", {"Authorization": "token edited"}), + (MCPAuth.authorization, "Custom edited", {"Authorization": "Custom edited"}), + ), + ) + async def test_static_preview_uses_edited_connection_instead_of_registered_server( + self, + monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuth, + auth_value: str | None, + expected_auth: dict[str, str], + ) -> None: + from starlette.datastructures import Headers + + from litellm.experimental_mcp_client.client import MCPClient + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy.management_endpoints import mcp_management_endpoints + + saved: Final = MCPServer( + server_id="saved-preview-server", + name="saved", + url="https://stored.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + authentication_token="preview:stored", + ) + manager: Final = MCPServerManager() + manager.registry = {saved.server_id: saved} + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) + payload: Final = NewMCPServerRequest( + server_id=saved.server_id, + server_name="edited", + url="https://stored.example/corrected-mcp", + transport=MCPTransport.sse, + auth_type=auth_type, + credentials={"auth_value": auth_value} if auth_value is not None else None, + static_headers={"X-Preview": "edited"}, + ) + staged: Final = rest_endpoints._stage_server_test(payload, Headers()) + + async def inspect_connection(client: MCPClient) -> dict[str, object]: + return {"url": client.server_url, "transport": client.transport_type, "headers": client._get_auth_headers()} + + result: Final = await rest_endpoints._execute_with_mcp_client( + staged.request, + inspect_connection, + mcp_auth_header=staged.mcp_auth_header, + oauth2_headers=staged.oauth2_headers, + ) + assert result == { + "url": "https://stored.example/corrected-mcp", + "transport": MCPTransport.sse, + "headers": {"X-Preview": "edited", **expected_auth}, + } + assert manager.get_mcp_server_by_id(saved.server_id) is saved + assert saved.url == "https://stored.example/mcp" + + @pytest.mark.parametrize( + ("saved_url", "url", "same_origin"), + ( + ("https://stored.example/mcp", "https://other.example/mcp", False), + ("https://stored.example/mcp", "http://stored.example/mcp", False), + ("https://stored.example/mcp", "https://stored.example:8443/mcp", False), + ("https://stored.example/mcp", "https://stored.example:443/mcp", True), + ("http://stored.example/mcp", "http://stored.example:80/edited", True), + ("https://stored.example/mcp", "HTTPS://STORED.EXAMPLE/edited", True), + ("https://[::1]/mcp", "https://[::1]/edited", True), + ("https://[::1]/mcp", "https://[::1]:443/edited", True), + ("https://[::1]/mcp", "https://[::2]/edited", False), + ("https://stored.example/mcp", "https://stored.example:invalid/mcp", False), + ), + ) + @pytest.mark.parametrize("explicit_credential", (None, "preview:explicit")) + def test_static_preview_respects_origin_when_inheriting_credentials( + self, + monkeypatch: pytest.MonkeyPatch, + saved_url: str, + url: str, + same_origin: bool, + explicit_credential: str | None, + ) -> None: + from starlette.datastructures import Headers + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy.management_endpoints import mcp_management_endpoints + + saved: Final = MCPServer( + server_id="saved-preview-server", + name="saved", + url=saved_url, + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + authentication_token="preview:stored", + ) + manager: Final = MCPServerManager() + manager.registry = {saved.server_id: saved} + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) + payload: Final = NewMCPServerRequest( + server_id=saved.server_id, + url=url, + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + credentials={"auth_value": explicit_credential} if explicit_credential else None, + ) + staged: Final = rest_endpoints._stage_server_test(payload, Headers()) + expected: Final = explicit_credential or ("preview:stored" if same_origin else None) + assert staged.mcp_auth_header == expected + assert staged.request.credentials == ({"auth_value": expected} if expected else None) + @pytest.mark.asyncio async def test_redacts_stack_trace(self, monkeypatch): async def fake_create_client(*args, **kwargs): @@ -113,7 +232,7 @@ class TestExecuteWithMcpClient: assert "stack_trace" not in result @pytest.mark.asyncio - async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch): + async def test_timeout_caps_hanging_operation_and_names_origin(self, monkeypatch): async def fake_create_client(*args, **kwargs): return object() @@ -138,7 +257,7 @@ class TestExecuteWithMcpClient: ) assert result["error"] is True - assert "https://mcp.example.com/mcp/" in result["message"] + assert "https://mcp.example.com" in result["message"] @pytest.mark.asyncio async def test_timeout_covers_client_creation(self, monkeypatch): @@ -166,15 +285,15 @@ class TestExecuteWithMcpClient: ) assert result["error"] is True - assert "https://mcp.example.com/mcp/" in result["message"] + assert "https://mcp.example.com" in result["message"] def test_timeout_defaults_to_tool_listing_timeout(self): default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default assert default == MCP_TOOL_LISTING_TIMEOUT - def test_connection_error_message_timeout_names_url_and_budget(self): + def test_connection_error_message_timeout_names_origin_and_budget(self): message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0) - assert "https://api.example.com/mcp/" in message + assert "https://api.example.com" in message assert "30s" in message def test_connection_error_message_hides_arbitrary_http_exception_detail(self): @@ -592,7 +711,7 @@ class TestExecuteWithMcpClient: assert result["status"] == "error" assert result["error"] is True - assert "Failed to connect to MCP server" in result["message"] + assert "reference" in result["message"] # Error message must not leak raw exception details assert "cancel scope" not in result["message"] @@ -602,6 +721,140 @@ class TestTestConnection: route = _get_route("/mcp-rest/test/connection", "POST") assert _route_has_dependency(route, user_api_key_auth) + @staticmethod + def _capture_execute(monkeypatch) -> dict: + captured: dict = {} + + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): + captured["request"] = request + captured["mcp_auth_header"] = mcp_auth_header + captured["oauth2_headers"] = oauth2_headers + return {"status": "ok"} + + monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False) + return captured + + @staticmethod + def _oauth2_authorization_code_payload(**overrides) -> NewMCPServerRequest: + return NewMCPServerRequest( + server_name="github_mcp", + url="https://api.githubcopilot.com/mcp/", + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + authorization_url="https://github.com/login/oauth/authorize", + token_url="https://github.com/login/oauth/access_token", + **overrides, + ) + + @pytest.mark.asyncio + async def test_forwards_staged_oauth2_bearer(self, monkeypatch): + """The just-authorized upstream token rides the request's Authorization header, exactly + as /test/tools/list receives it; dropping it makes every authorization_code server fail + the connection test that its tools preview passes.""" + from litellm.proxy._types import LitellmUserRoles + + captured = self._capture_execute(monkeypatch) + request = _build_request( + {"x-litellm-api-key": "sk-admin-session", "authorization": "Bearer upstream-oauth-token"}, + path="/mcp-rest/test/connection", + ) + + result = await rest_endpoints.test_connection( + request, + self._oauth2_authorization_code_payload(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result == {"status": "ok"} + assert captured["oauth2_headers"] == {"Authorization": "Bearer upstream-oauth-token"} + assert captured["mcp_auth_header"] is None + + @pytest.mark.asyncio + async def test_forwards_staged_auth_value(self, monkeypatch): + from litellm.proxy._types import LitellmUserRoles + + captured = self._capture_execute(monkeypatch) + request = _build_request({"x-litellm-api-key": "sk-admin-session"}, path="/mcp-rest/test/connection") + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com/mcp", + auth_type=MCPAuth.bearer_token, + credentials={"auth_value": "upstream-static-token"}, + ) + + result = await rest_endpoints.test_connection( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result == {"status": "ok"} + assert captured["mcp_auth_header"] == "upstream-static-token" + assert captured["oauth2_headers"] is None + + @pytest.mark.asyncio + async def test_inherits_stored_credentials_of_saved_server(self, monkeypatch): + """The edit form resends a saved server without its masked credential; the stored one + must be used, as /test/tools/list already does.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + captured = self._capture_execute(monkeypatch) + saved = MCPServer( + server_id="saved-server-id", + name="example", + url="https://example.com/mcp", + transport="http", + auth_type=MCPAuth.bearer_token, + authentication_token="stored-upstream-token", + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: saved if server_id == "saved-server-id" else None, + ) + request = _build_request({"x-litellm-api-key": "sk-admin-session"}, path="/mcp-rest/test/connection") + payload = NewMCPServerRequest( + server_id="saved-server-id", + server_name="example", + url="https://example.com/mcp", + auth_type=MCPAuth.bearer_token, + ) + + result = await rest_endpoints.test_connection( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result == {"status": "ok"} + assert captured["mcp_auth_header"] == "stored-upstream-token" + assert captured["request"].credentials == {"auth_value": "stored-upstream-token"} + + @pytest.mark.asyncio + async def test_does_not_forward_authorization_that_satisfied_admission(self, monkeypatch): + """With no x-litellm-api-key, the Authorization value is the caller's LiteLLM key and + must never reach the upstream.""" + from litellm.proxy._types import LitellmUserRoles + + captured = self._capture_execute(monkeypatch) + request = _build_request({"authorization": "Bearer sk-litellm-admission-key"}, path="/mcp-rest/test/connection") + + result = await rest_endpoints.test_connection( + request, + self._oauth2_authorization_code_payload(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result == {"status": "ok"} + assert captured["oauth2_headers"] is None + class TestTestToolsList: pytestmark = pytest.mark.asyncio @@ -2342,6 +2595,79 @@ class TestCallToolRestAPI: assert result == masked_result + async def test_success_logging_start_time_excludes_pre_call_processing(self, monkeypatch): + """Pre-call hook latency (guardrails, header resolution) must not inflate the tool call's + logged duration on success.""" + from litellm.proxy import proxy_server + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + server_id = "server-1" + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = None + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + auth_type = None + + stub_server = StubServer() + + async def fake_add_litellm_data_to_request(**kwargs): + return kwargs.get("data", {}) + + pre_call_finished_at = {} + + async def slow_pre_call_hook(user_api_key_dict, data, call_type): + await asyncio.sleep(0.05) + pre_call_finished_at["value"] = datetime.now() + return data + + captured = {} + + async def fake_execute_mcp_tool(**kwargs): + captured.update(kwargs) + return {"result": "ok"} + + fire_logging = AsyncMock(return_value={"result": "ok"}) + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + proxy_server, "add_litellm_data_to_request", fake_add_litellm_data_to_request, raising=False + ) + monkeypatch.setattr(proxy_server, "proxy_config", {}, raising=False) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "pre_call_hook", slow_pre_call_hook) + monkeypatch.setattr(rest_endpoints, "execute_mcp_tool", fake_execute_mcp_tool, raising=False) + monkeypatch.setattr(rest_endpoints, "_fire_mcp_tool_call_logging", fire_logging, raising=False) + + request = _build_request( + path="/mcp-rest/tools/call", + method="POST", + json_body={"server_id": "server-1", "name": "demo-tool", "arguments": {}}, + ) + + await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=UserAPIKeyAuth()) + + logged_start_time = fire_logging.await_args.args[2] + assert captured["start_time"] >= pre_call_finished_at["value"] + assert logged_start_time == captured["start_time"] + async def test_success_logging_guardrail_rejection_propagates(self, monkeypatch): """A guardrail rejecting the tool result must not be swallowed as a logging failure, otherwise the unguarded result would still be returned to the caller.""" @@ -2512,6 +2838,139 @@ class TestCallToolRestAPI: info_messages = [_rendered_log_message(c) for c in mock_logger.info.call_args_list if c.args] assert not any("relaying upstream" in m for m in info_messages) + @pytest.mark.parametrize("raise_site", ["pre_call_hook", "execute_mcp_tool"]) + async def test_guardrail_block_runs_failure_logging_before_http_translation(self, monkeypatch, raise_site): + """A pre_mcp_call guardrail block, whether raised by the pre-call hook or from inside + execute_mcp_tool, must reach proxy_logging_obj.post_call_failure_hook (the only path that + writes the failure spend-log row) with the logging object's failure payload already built, + and the REST caller must still get the same 400 it got before.""" + from litellm.proxy import proxy_server + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + server_id = "server-1" + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = None + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + auth_type = None + + async def fake_add_litellm_data_to_request(**kwargs): + return kwargs.get("data", {}) + + guardrail_error = HTTPException( + status_code=400, + detail={"error": "Content blocked: keyword 'confidential' detected", "keyword": "confidential"}, + ) + + async def passthrough_pre_call_hook(user_api_key_dict, data, call_type): + return data + + async def blocking_pre_call_hook(user_api_key_dict, data, call_type): + raise guardrail_error + + async def fake_execute_mcp_tool(**kwargs): + raise guardrail_error + + async def passthrough_execute_mcp_tool(**kwargs): + return [] + + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: StubServer() if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + proxy_server, "add_litellm_data_to_request", fake_add_litellm_data_to_request, raising=False + ) + monkeypatch.setattr(proxy_server, "proxy_config", {}, raising=False) + monkeypatch.setattr( + proxy_server.proxy_logging_obj, + "pre_call_hook", + blocking_pre_call_hook if raise_site == "pre_call_hook" else passthrough_pre_call_hook, + ) + monkeypatch.setattr( + rest_endpoints, + "execute_mcp_tool", + fake_execute_mcp_tool if raise_site == "execute_mcp_tool" else passthrough_execute_mcp_tool, + raising=False, + ) + post_call_failure_hook = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "post_call_failure_hook", post_call_failure_hook) + + user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key", request_route="/mcp-rest/tools/call") + request = _build_request( + headers={"x-mcp-deepwiki-authorization": "Bearer upstream-secret"}, + path="/mcp-rest/tools/call", + method="POST", + json_body={"server_id": "server-1", "name": "demo-tool", "arguments": {"q": "confidential"}}, + ) + + with pytest.raises(HTTPException) as exc_info: + await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=user_api_key_dict) + + assert exc_info.value is guardrail_error + + post_call_failure_hook.assert_awaited_once() + hook_kwargs = post_call_failure_hook.await_args.kwargs + assert hook_kwargs["original_exception"] is guardrail_error + assert hook_kwargs["user_api_key_dict"] is user_api_key_dict + assert hook_kwargs["route"] == "/mcp/call_tool" + request_data = hook_kwargs["request_data"] + assert "raw_headers" not in request_data + assert "mcp_server_auth_headers" not in request_data + standard_logging_object = request_data["litellm_logging_obj"].model_call_details["standard_logging_object"] + assert standard_logging_object["status"] == "failure" + assert standard_logging_object["error_str"] == str(guardrail_error) + + async def test_failure_logging_error_does_not_replace_guardrail_error(self, monkeypatch): + from litellm.proxy import proxy_server + + guardrail_error = HTTPException(status_code=400, detail={"error": "Content blocked"}) + + async def fake_add_litellm_data_to_request(**kwargs): + return kwargs.get("data", {}) + + async def blocking_pre_call_hook(user_api_key_dict, data, call_type): + raise guardrail_error + + failure_logging = AsyncMock(side_effect=RuntimeError("spend log db down")) + monkeypatch.setattr( + proxy_server, "add_litellm_data_to_request", fake_add_litellm_data_to_request, raising=False + ) + monkeypatch.setattr(proxy_server, "proxy_config", {}, raising=False) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "pre_call_hook", blocking_pre_call_hook) + monkeypatch.setattr(rest_endpoints, "fire_mcp_tool_call_failure_logging", failure_logging, raising=False) + + request = _build_request( + path="/mcp-rest/tools/call", + method="POST", + json_body={"server_id": "server-1", "name": "demo-tool", "arguments": {"q": "confidential"}}, + ) + + with pytest.raises(HTTPException) as exc_info: + await rest_endpoints.call_tool_rest_api( + request, user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", request_route="/mcp-rest/tools/call") + ) + + assert exc_info.value is guardrail_error + failure_logging.assert_awaited_once() + async def test_success_logging_cancellation_propagates(self, monkeypatch): fire_logging = AsyncMock(side_effect=asyncio.CancelledError()) monkeypatch.setattr( @@ -2548,7 +3007,7 @@ class TestCallToolRestAPI: class _FakePreCall: def __init__(self, data): - pass + self.data = data async def common_processing_pre_call_logic(self, **kwargs): return None, MagicMock() @@ -2582,6 +3041,58 @@ class TestCallToolRestAPI: assert exc_info.value.headers is not None assert exc_info.value.headers.get("www-authenticate") == challenge + async def test_virtual_mcp_tool_call_guardrail_block_runs_failure_logging(self, monkeypatch): + """A pre_mcp_call guardrail block on the virtual mcp_tool_call branch must write a failure + spend log, same as the direct tool call branch, and still raise the original error.""" + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + guardrail_error = HTTPException(status_code=400, detail={"error": "Content blocked"}) + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_add_litellm_data_to_request(**kwargs): + return kwargs.get("data", {}) + + async def blocking_pre_call_hook(user_api_key_dict, data, call_type): + raise guardrail_error + + failure_logging = AsyncMock() + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + proxy_server, "add_litellm_data_to_request", fake_add_litellm_data_to_request, raising=False + ) + monkeypatch.setattr(proxy_server, "proxy_config", {}, raising=False) + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "pre_call_hook", blocking_pre_call_hook) + monkeypatch.setattr(rest_endpoints, "fire_mcp_tool_call_failure_logging", failure_logging, raising=False) + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + request_route="/mcp-rest/tools/call", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="search-scope", + mcp_tool_search_enabled=True, + ), + ) + request = _build_request( + path="/mcp-rest/tools/call", + method="POST", + json_body={"name": "mcp_tool_call", "arguments": {"tool_name": "x", "arguments": {"q": "confidential"}}}, + ) + + with pytest.raises(HTTPException) as exc_info: + await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=user_api_key_dict) + + assert exc_info.value is guardrail_error + failure_logging.assert_awaited_once() + logging_obj, exception, _start_time, user_api_key_auth, request_data = failure_logging.await_args.args + assert exception is guardrail_error + assert user_api_key_auth is user_api_key_dict + assert logging_obj is request_data.get("litellm_logging_obj") + assert logging_obj is not None + class TestGetToolsForSingleServer: """Test _get_tools_for_single_server with object_permission filtering""" @@ -3293,10 +3804,191 @@ class TestConnectionErrorMessage: message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "503" in message + @pytest.mark.parametrize( + "error_type", [httpx.ReadError, httpx.WriteError, httpx.RemoteProtocolError, ConnectionResetError] + ) + def test_interrupted_connection_message_is_safe(self, error_type: type[Exception]) -> None: + message: Final = rest_endpoints._connection_error_message( + error_type("secret-transport-detail"), "https://example.com/?token=secret-query", 30 + ) + assert "connection was interrupted" in message + assert "secret" not in message + + def test_closed_connection_explains_incomplete_request(self) -> None: + from mcp import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=-32000, message="Connection closed", data="secret-data")), None, 30 + ) + assert "connection was closed before the request completed" in message + assert "secret" not in message + + def test_timeout_does_not_claim_the_server_sent_nothing(self) -> None: + message: Final = rest_endpoints._connection_error_message(TimeoutError(), None, 30) + assert "no valid MCP response received" in message + + @pytest.mark.asyncio + @pytest.mark.parametrize("sdk_timeout", [True, False]) + @pytest.mark.parametrize("read_timeout", [0, 1]) + async def test_timeout_message_uses_the_deadline_that_expired(self, sdk_timeout: bool, read_timeout: int) -> None: + from mcp import McpError + from mcp.types import ErrorData + + async def operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + try: + raise TimeoutError("secret-timeout") + except TimeoutError as elapsed: + if not sdk_timeout: + raise + try: + raise McpError(ErrorData(code=408, message="secret-sdk-timeout")) from elapsed + except McpError as sdk_error: + raise TimeoutError() from sdk_error + + payload: Final = NewMCPServerRequest( + server_name="timeout", url="https://example.com", auth_type=MCPAuth.none, timeout=read_timeout + ) + result: Final = await rest_endpoints._execute_with_mcp_client(payload, operation, timeout_seconds=30) + assert (f"within {read_timeout}s" if sdk_timeout else "within 30s") in result["message"] + assert "secret" not in result["message"] + def test_unknown_error_falls_back_to_generic(self): message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message - assert "proxy logs" in message.lower() + assert "reference" in message.lower() + + def test_sdk_session_terminated_explains_endpoint_and_retry(self) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=32600, message="Session terminated")), "https://example.com/mcp", 30.0 + ) + + assert "session was terminated" in message + assert "MCP endpoint" in message + assert "transport" in message + assert "retry" in message + assert "404" not in message + + @pytest.mark.parametrize("code", [-32700, -32601, -32602, -32603, -32000, 32600, 408]) + def test_rpc_errors_include_code_without_echoing_upstream_data(self, code: int) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=code, message="secret-message", data={"token": "secret-data"})), + "https://example.com/secret-path?token=secret-query", + 30.0, + ) + + assert f"JSON-RPC code {code}" in message + assert "secret" not in message + assert "timed out" not in message + assert "session was terminated" not in message + + @pytest.mark.parametrize("status_code", [401, 403, 404, 405, 429, 503]) + def test_wrapped_http_failures_preserve_status(self, status_code: int) -> None: + response: Final = httpx.Response(status_code, text="secret-body") + upstream: Final = httpx.HTTPStatusError( + "secret-exception", + request=httpx.Request("POST", "https://example.com/?token=secret-query"), + response=response, + ) + wrapped: Final = BaseExceptionGroup( + "secret-group", [asyncio.CancelledError(), BaseExceptionGroup("nested", [upstream])] + ) + + message: Final = rest_endpoints._connection_error_message(wrapped, "https://example.com", 30.0) + + assert f"HTTP {status_code}" in message + assert "secret" not in message + + def test_explicit_cause_is_classified_before_incidental_context(self) -> None: + wrapped: Final = RuntimeError("secret-wrapper") + wrapped.__cause__ = httpx.ConnectError("secret-cause") + wrapped.__context__ = TimeoutError("secret-context") + + message: Final = rest_endpoints._connection_error_message(wrapped, "https://example.com", 30.0) + + assert "unreachable" in message + assert "secret" not in message + + def test_timeout_url_redacts_credentials_path_query_and_fragment(self) -> None: + message: Final = rest_endpoints._connection_error_message( + TimeoutError("secret-error"), + "https://secret-user:secret-pass@example.com:8443/secret-path?token=secret-query#secret-fragment", + 30.0, + ) + + assert "https://example.com:8443" in message + assert "30s" in message + assert "secret" not in message + + def test_unknown_failure_reference_matches_safe_diagnostics(self, caplog: pytest.LogCaptureFixture) -> None: + import re + + try: + raise RuntimeError("secret-exception-body") + except RuntimeError as exc: + message: Final = rest_endpoints._connection_error_message( + exc, "https://secret-user:secret-password@example.com/secret-path?token=secret-query", 30.0 + ) + + reference: Final = re.search(r"reference ([a-f0-9]{32})", message) + assert reference is not None + diagnostics: Final = tuple( + record for record in caplog.records if "MCP connection test failed" in record.message + ) + assert len(diagnostics) == 1 + assert reference.group(1) in diagnostics[0].message + assert "RuntimeError" in diagnostics[0].message + assert "test_unknown_failure_reference_matches_safe_diagnostics" in diagnostics[0].message + assert diagnostics[0].exc_info is None + assert "secret" not in message + diagnostics[0].message + + @pytest.mark.parametrize("exc", [ValueError("secret-config"), HTTPException(500, "secret-detail")]) + def test_unrelated_errors_are_not_misreported_as_invalid_mcp(self, exc: Exception) -> None: + message: Final = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) + + assert "reference" in message + assert "invalid MCP response" not in message + assert "secret" not in message + + def test_configuration_validation_error_uses_unknown_fallback(self) -> None: + from pydantic import ValidationError + + with pytest.raises(ValidationError) as caught: + NewMCPServerRequest.model_validate({"server_name": "example", "transport": "secret-invalid-transport"}) + + message: Final = rest_endpoints._connection_error_message(caught.value, "https://example.com", 30.0) + assert "reference" in message + assert "invalid MCP response" not in message + assert "secret" not in message + + @pytest.mark.asyncio + async def test_connection_test_preserves_cancellation(self) -> None: + async def cancelled_operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + raise asyncio.CancelledError + + payload: Final = NewMCPServerRequest(server_name="cancelled", url="https://example.com", auth_type=MCPAuth.none) + with pytest.raises(asyncio.CancelledError): + await rest_endpoints._execute_with_mcp_client(payload, cancelled_operation) + + @pytest.mark.asyncio + async def test_unknown_failure_preserves_response_contract(self) -> None: + async def failing_operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + raise RuntimeError("secret-operation") + + payload: Final = NewMCPServerRequest(server_name="unknown", url="https://example.com", auth_type=MCPAuth.none) + result: Final = await rest_endpoints._execute_with_mcp_client(payload, failing_operation) + + assert result["error"] is True + assert result["status"] == "error" + assert "reference" in result["message"] + assert "secret" not in result["message"] + assert "stack_trace" not in result class TestGetServerAuthHeaderGroupDefault: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index bf0df17fafb..f0b4e94f72f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -2216,3 +2216,26 @@ async def test_top_k_above_router_default_is_respected(): assert len(filtered) == 6 print("✅ Configured top_k above the semantic-router default of 5 is honored") + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves(): + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + gateway_reference = {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"} + external_tool = { + "type": "mcp", + "server_label": "zapier", + "server_url": "https://mcp.zapier.com/api/mcp/mcp", + "allowed_tools": ["zapier_send_email"], + } + + async def served_names(names): + assert names == {"mcp"} + return frozenset() + + narrowed = await SemanticToolFilterHook._narrow_mcp_references( + [gateway_reference, external_tool], ["srv-tool_1"], served_names=served_names + ) + + assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool] diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 43034f889f6..e0476361074 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -7,11 +7,24 @@ Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request import json import socket import sys -from contextlib import ExitStack +from collections.abc import Awaitable, Callable, Mapping +from contextlib import AbstractContextManager, ExitStack +from dataclasses import dataclass +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.proxy._types import UserAPIKeyAuth + +AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]] + + +@dataclass(frozen=True, slots=True) +class CapturedAgentCall: + request_id: object + agent_extra_headers: dict[str, str] | None + @pytest.mark.asyncio async def test_invoke_agent_a2a_adds_litellm_data(): @@ -364,7 +377,7 @@ def _make_agent_mock(url: str = "http://backend-agent:10001") -> MagicMock: def _make_request_mock( - method: str, params: dict, request_id: object = "req-1" + method: str, params: Mapping[str, object], request_id: object = "req-1" ) -> MagicMock: req = MagicMock() req.headers = {} @@ -379,7 +392,9 @@ def _make_request_mock( return req -def _base_patches(agent: MagicMock): +def _base_patches( + agent: MagicMock, add_litellm_data: AddLiteLLMData | None = None +) -> list[AbstractContextManager[object]]: return [ patch( "litellm.proxy.agent_endpoints.a2a_endpoints._get_agent", @@ -391,7 +406,7 @@ def _base_patches(agent: MagicMock): ), patch( "litellm.proxy.common_request_processing.add_litellm_data_to_request", - new=AsyncMock(side_effect=_add_proxy_data), + new=AsyncMock(side_effect=add_litellm_data or _add_proxy_data), ), patch("litellm.proxy.proxy_server.general_settings", {}), patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), @@ -399,84 +414,67 @@ def _base_patches(agent: MagicMock): ] -async def _add_proxy_data(data, **kwargs): - data["proxy_server_request"] = { - "url": "http://localhost:4000", - "method": "POST", - "headers": {}, - "body": {}, +async def _add_proxy_data(data: dict[str, object], **kwargs: object) -> dict[str, object]: + return { + **data, + "proxy_server_request": {"url": "http://localhost:4000", "method": "POST", "headers": {}, "body": {}}, + "metadata": data.get("metadata", {}), } - data.setdefault("metadata", {}) - return data -@pytest.mark.asyncio -@pytest.mark.parametrize("method", ["message/send", "message/stream"]) -async def test_message_methods_preserve_numeric_zero_request_id(method: str): +_HELLO_MESSAGE_PARAMS = { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-123", + } +} + + +async def _invoke_message_method( + method: str, + mock_request: MagicMock, + user_api_key_dict: UserAPIKeyAuth, + add_litellm_data: AddLiteLLMData | None = None, +) -> CapturedAgentCall: from fastapi.responses import JSONResponse - from litellm.proxy._types import UserAPIKeyAuth class MessageSendParams: - def __init__(self, **kwargs): + def __init__(self, **kwargs: object) -> None: self.__dict__.update(kwargs) class SendMessageRequest: - def __init__(self, **kwargs): + def __init__(self, **kwargs: object) -> None: self.__dict__.update(kwargs) - agent = _make_agent_mock() - params = { - "message": { - "role": "user", - "parts": [{"kind": "text", "text": "Hello"}], - "messageId": "msg-123", - } - } - mock_request = _make_request_mock(method, params, request_id=0) - user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") - captured = {} - - async def capture_asend_message(request, **kwargs): - captured["request_id"] = request.id - response = MagicMock() + async def fake_asend_message(request: SendMessageRequest, **kwargs: object) -> MagicMock: + response: Final = MagicMock() response.model_dump.return_value = { "jsonrpc": "2.0", - "id": request.id, + "id": request.__dict__["id"], "result": {"status": "success"}, } return response - async def capture_stream_message(**kwargs): - captured["request_id"] = kwargs["request_id"] - return JSONResponse({"jsonrpc": "2.0", "id": kwargs["request_id"]}) + async def fake_stream_message(request_id: object, **kwargs: object) -> JSONResponse: + return JSONResponse({"jsonrpc": "2.0", "id": request_id}) - mock_a2a_types = MagicMock() + mock_a2a_types: Final = MagicMock() mock_a2a_types.MessageSendParams = MessageSendParams mock_a2a_types.SendMessageRequest = SendMessageRequest + is_send: Final = method == "message/send" + downstream: Final = AsyncMock(side_effect=fake_asend_message if is_send else fake_stream_message) with ExitStack() as stack: - for p in _base_patches(agent): + for p in _base_patches(_make_agent_mock(), add_litellm_data): stack.enter_context(p) stack.enter_context(patch("litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", True)) - if method == "message/send": - stack.enter_context( - patch.dict( - sys.modules, - {"a2a": MagicMock(), "a2a.types": mock_a2a_types}, - ) - ) - stack.enter_context( - patch( - "litellm.a2a_protocol.asend_message", - new=AsyncMock(side_effect=capture_asend_message), - ) - ) + if is_send: + stack.enter_context(patch.dict(sys.modules, {"a2a": MagicMock(), "a2a.types": mock_a2a_types})) + stack.enter_context(patch("litellm.a2a_protocol.asend_message", new=downstream)) else: stack.enter_context( - patch( - "litellm.proxy.agent_endpoints.a2a_endpoints._handle_stream_message", - new=AsyncMock(side_effect=capture_stream_message), - ) + patch("litellm.proxy.agent_endpoints.a2a_endpoints._handle_stream_message", new=downstream) ) from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a @@ -488,7 +486,82 @@ async def test_message_methods_preserve_numeric_zero_request_id(method: str): user_api_key_dict=user_api_key_dict, ) - assert captured["request_id"] == 0 + kwargs: Final = downstream.call_args.kwargs + request_id: Final = kwargs["request"].__dict__["id"] if is_send else kwargs["request_id"] + return CapturedAgentCall(request_id=request_id, agent_extra_headers=kwargs.get("agent_extra_headers")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["message/send", "message/stream"]) +async def test_message_methods_preserve_numeric_zero_request_id(method: str): + mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS, request_id=0) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") + + captured = await _invoke_message_method(method, mock_request, user_api_key_dict) + + assert captured.request_id == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["message/send", "message/stream"]) +async def test_message_methods_forward_caller_identity_headers(method: str): + mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="user-abc", team_id="team-xyz") + + captured = await _invoke_message_method(method, mock_request, user_api_key_dict) + + forwarded_headers = captured.agent_extra_headers or {} + assert forwarded_headers.get("X-LiteLLM-User-Id") == "user-abc" + assert forwarded_headers.get("X-LiteLLM-Team-Id") == "team-xyz" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["message/send", "message/stream"]) +async def test_message_methods_caller_identity_headers_cannot_be_spoofed(method: str): + mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS) + mock_request.headers = { + "x-a2a-test-agent-x-litellm-user-id": "attacker-user", + "x-a2a-test-agent-x-litellm-team-id": "attacker-team", + } + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="real-user", team_id="real-team") + + captured = await _invoke_message_method(method, mock_request, user_api_key_dict) + + forwarded_headers = captured.agent_extra_headers or {} + assert ( + forwarded_headers.get("X-LiteLLM-User-Id") == "real-user" + ), "authenticated user id must not be overridden by forwarded client headers" + assert ( + forwarded_headers.get("X-LiteLLM-Team-Id") == "real-team" + ), "authenticated team id must not be overridden by forwarded client headers" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["message/send", "message/stream"]) +async def test_message_methods_forward_key_bound_identity_not_pre_call_rewrite(method: str): + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS) + mock_request.headers = {"X-OpenWebUI-User-Id": "header-mapped-user"} + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="key-user", team_id="key-team") + general_settings: Final = { + "user_header_mappings": [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}] + } + + async def apply_user_header_mapping(data: dict[str, object], **kwargs: object) -> dict[str, object]: + LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + general_settings, user_api_key_dict, dict(mock_request.headers) + ) + return await _add_proxy_data(data, **kwargs) + + captured = await _invoke_message_method( + method, mock_request, user_api_key_dict, add_litellm_data=apply_user_header_mapping + ) + + assert user_api_key_dict.user_id == "header-mapped-user", "precondition: pre-call rewrite ran" + forwarded_headers = captured.agent_extra_headers or {} + assert forwarded_headers.get("X-LiteLLM-User-Id") == "key-user" + assert forwarded_headers.get("X-LiteLLM-Team-Id") == "key-team" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py index 15864417489..e894f4ad69a 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py @@ -223,7 +223,7 @@ async def test_static_overrides_dynamic(): @pytest.mark.asyncio async def test_no_headers(): - """When no headers are configured, agent_extra_headers is None and behaviour is unchanged.""" + """When no headers are configured, only the caller identity is forwarded.""" mock_agent = _make_mock_agent() # no static_headers or extra_headers mock_request = _make_mock_request() @@ -231,7 +231,7 @@ async def test_no_headers(): call_kwargs = mock_asend.call_args.kwargs headers = call_kwargs.get("agent_extra_headers") - assert headers is None + assert headers == {"X-LiteLLM-User-Id": "u1"} # --------------------------------------------------------------------------- @@ -303,7 +303,7 @@ async def test_convention_unrelated_prefix_not_forwarded(): mock_asend = await _invoke(mock_agent, mock_request, None) headers = mock_asend.call_args.kwargs.get("agent_extra_headers") - assert headers is None + assert headers == {"X-LiteLLM-User-Id": "u1"} # --------------------------------------------------------------------------- @@ -393,7 +393,7 @@ async def test_non_databricks_agent_skips_oauth_resolution(): mock_resolve.assert_not_called() headers = mock_asend.call_args.kwargs.get("agent_extra_headers") - assert headers == {"x-custom": "v"} + assert headers == {"x-custom": "v", "X-LiteLLM-User-Id": "u1"} assert "Authorization" not in headers @@ -477,7 +477,7 @@ async def test_convention_header_blocked_by_case_variant_static(): headers = mock_asend.call_args.kwargs.get("agent_extra_headers") assert headers is not None - assert headers == {"Authorization": "Bearer admin-token"} + assert headers == {"Authorization": "Bearer admin-token", "X-LiteLLM-User-Id": "u1"} assert "authorization" not in headers diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py index daca244c0a1..c02ed1f37e5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py @@ -211,11 +211,24 @@ class TestAgentSearchIndex: assert isinstance(outcome, AgentSearchEmbeddingFailed) +def _pass_through_key_limits() -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + return limits + + class TestSearchAgents: @pytest.mark.asyncio async def test_no_embedding_model_is_not_configured(self) -> None: outcome = await search_agents( - "q", AGENTS, 5, router=MagicMock(), embedding_model=None, index=AgentSearchIndex(), user_api_key_dict=CALLER + "q", + AGENTS, + 5, + router=MagicMock(), + embedding_model=None, + index=AgentSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchNotConfigured) assert "agent_search_embedding_model" in outcome.reason @@ -223,7 +236,14 @@ class TestSearchAgents: @pytest.mark.asyncio async def test_no_router_is_not_configured(self) -> None: outcome = await search_agents( - "q", AGENTS, 5, router=None, embedding_model="m", index=AgentSearchIndex(), user_api_key_dict=CALLER + "q", + AGENTS, + 5, + router=None, + embedding_model="m", + index=AgentSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchNotConfigured) @@ -244,6 +264,7 @@ class TestSearchAgents: embedding_model="text-embedding-3-small", index=AgentSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchHits) assert [hit.agent.agent_id for hit in outcome.hits] == ["translator"] @@ -266,6 +287,7 @@ class TestSearchAgents: embedding_model="text-embedding-3-small", index=AgentSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) metadata = router.aembedding.await_args.kwargs["metadata"] assert metadata["user_api_key"] == "hashed-caller-key" @@ -302,6 +324,7 @@ def embedding_router(monkeypatch: pytest.MonkeyPatch) -> MagicMock: ) ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", _pass_through_key_limits()) monkeypatch.setattr(litellm, "agent_search_embedding_model", "text-embedding-3-small") return router diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py index 99b09f48fe5..a585666743f 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -1,7 +1,7 @@ """ Unit tests for claude_code_marketplace.py source validation. -Covers the git-subdir source type added alongside the existing github and url types. +Covers the git-subdir and archive source types added alongside the existing github and url types. """ import json @@ -11,17 +11,20 @@ from fastapi import HTTPException from unittest.mock import AsyncMock, MagicMock import litellm -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, ProxyException, UserAPIKeyAuth from litellm.proxy.proxy_server import LitellmUserRoles from litellm.types.proxy.claude_code_endpoints import ( RegisterPluginRequest, UpdatePluginRequest, ) +from litellm.proxy.anthropic_endpoints.claude_code_endpoints import claude_code_marketplace from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( delete_plugin, disable_plugin, enable_plugin, get_marketplace, + get_plugin, + list_plugins, register_plugin, update_plugin, ) @@ -38,11 +41,15 @@ def _make_mock_prisma(): async def _find_unique(where): return store.get(where.get("name")) + def _matches(record, where) -> bool: + if "OR" in where: + return any(_matches(record, clause) for clause in where["OR"]) + if "enabled" in where and record.enabled != where["enabled"]: + return False + return "name" not in where or record.name in where["name"]["in"] + async def _find_many(where=None): - records = list(store.values()) - if where and "enabled" in where: - return [r for r in records if r.enabled == where["enabled"]] - return records + return [r for r in store.values() if _matches(r, where or {})] async def _create(data): record = MagicMock() @@ -52,6 +59,8 @@ def _make_mock_prisma(): record.description = data.get("description") record.manifest_json = data.get("manifest_json", "{}") record.enabled = data.get("enabled", True) + record.created_at = data.get("created_at") + record.updated_at = data.get("updated_at") store[data["name"]] = record return record @@ -87,6 +96,12 @@ _GIT_SUBDIR_SOURCE = { "path": "plugins/my-plugin", } +_ARCHIVE_SOURCE = { + "source": "archive", + "url": "https://skills-bucket.s3.us-east-1.amazonaws.com/plugins/s3-skill-1.0.0.zip", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", +} + @pytest.fixture(autouse=True) def _patch_proxy_globals(monkeypatch): @@ -265,13 +280,89 @@ async def test_get_marketplace_skips_plugin_with_null_manifest(): table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable await table.create(data={"name": "null-manifest-plugin", "manifest_json": None, "enabled": True}) - response = await get_marketplace() + response = await get_marketplace(request=MagicMock()) assert response.status_code == 200 body = json.loads(response.body) assert [plugin["name"] for plugin in body["plugins"]] == ["good-plugin"] +def _granted_user(skills: list[str]) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-granted", + user_id="granted-user", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="perm-1", skills=skills), + ) + + +async def _register_public_and_private_plugins() -> None: + for name, enabled in (("public-skill", True), ("private-skill", False)): + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), + user_api_key_dict=_USER, + ) + if not enabled: + await disable_plugin(plugin_name=name, user_api_key_dict=_USER) + + +async def _listed_names(user: UserAPIKeyAuth) -> set[str]: + response = await list_plugins(user_api_key_dict=user) + return {plugin.name for plugin in response.plugins} + + +@pytest.mark.asyncio +async def test_list_plugins_shows_disabled_plugin_only_to_granted_key_or_admin(): + await _register_public_and_private_plugins() + + assert await _listed_names(_NON_ADMIN_USER) == {"public-skill"} + assert await _listed_names(_granted_user(["other-skill"])) == {"public-skill"} + assert await _listed_names(_granted_user(["private-skill"])) == {"public-skill", "private-skill"} + assert await _listed_names(_USER) == {"public-skill", "private-skill"} + + +@pytest.mark.asyncio +async def test_get_plugin_returns_403_for_disabled_plugin_the_key_is_not_granted(): + await _register_public_and_private_plugins() + + with pytest.raises(HTTPException) as exc_info: + await get_plugin(plugin_name="private-skill", user_api_key_dict=_NON_ADMIN_USER) + assert exc_info.value.status_code == 403 + + assert (await get_plugin(plugin_name="public-skill", user_api_key_dict=_NON_ADMIN_USER))["name"] == "public-skill" + granted = await get_plugin(plugin_name="private-skill", user_api_key_dict=_granted_user(["private-skill"])) + assert granted["name"] == "private-skill" + assert granted["enabled"] is False + + +async def _marketplace_names(key: str | None) -> list[str]: + response = await get_marketplace(request=MagicMock(), key=key) + assert response.status_code == 200 + return sorted(plugin["name"] for plugin in json.loads(response.body)["plugins"]) + + +@pytest.mark.asyncio +async def test_get_marketplace_key_query_param_adds_granted_disabled_plugins(monkeypatch): + await _register_public_and_private_plugins() + keys = {"sk-granted": _granted_user(["private-skill"]), "sk-plain": _NON_ADMIN_USER} + + async def _fake_auth(request, api_key: str) -> UserAPIKeyAuth: + token = api_key.removeprefix("Bearer ") + if token not in keys: + raise ProxyException(message="invalid key", type="auth_error", param="key", code=401) + return keys[token] + + monkeypatch.setattr(claude_code_marketplace, "user_api_key_auth", _fake_auth) + + assert await _marketplace_names(None) == ["public-skill"] + assert await _marketplace_names("sk-plain") == ["public-skill"] + assert await _marketplace_names("sk-granted") == ["private-skill", "public-skill"] + + with pytest.raises(ProxyException) as exc_info: + await get_marketplace(request=MagicMock(), key="sk-bogus") + assert exc_info.value.code == "401" + + @pytest.mark.asyncio async def test_register_plugin_git_subdir_missing_url(): """git-subdir without url field raises HTTP 400.""" @@ -377,6 +468,75 @@ async def test_register_plugin_unknown_source_type(): assert exc_info.value.status_code == 400 assert "git-subdir" in exc_info.value.detail["error"] + assert "archive" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_archive_source_registers_and_is_served_verbatim_in_marketplace(): + response = await register_plugin( + request=RegisterPluginRequest(name="s3-skill", source=_ARCHIVE_SOURCE), + user_api_key_dict=_USER, + ) + + assert response.action == "created" + assert response.plugin.source == _ARCHIVE_SOURCE + + marketplace = json.loads((await get_marketplace(request=MagicMock())).body) + assert marketplace["plugins"] == [{"name": "s3-skill", "source": _ARCHIVE_SOURCE, "version": "1.0.0"}] + + +@pytest.mark.asyncio +async def test_archive_source_without_sha256_is_accepted(): + source = {"source": "archive", "url": "https://artifacts.example.com/plugin.zip"} + + response = await register_plugin( + request=RegisterPluginRequest(name="unpinned-skill", source=source), + user_api_key_dict=_USER, + ) + + assert response.plugin.source == source + + +@pytest.mark.asyncio +async def test_update_plugin_to_archive_source(): + name = "my-monorepo-plugin" + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE), + user_api_key_dict=_USER, + ) + + response = await update_plugin( + plugin_name=name, + request=UpdatePluginRequest(source=_ARCHIVE_SOURCE), + user_api_key_dict=_USER, + ) + + assert response.action == "updated" + assert (await _read_stored_manifest(name))["source"] == _ARCHIVE_SOURCE + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "source, expected_fragment", + [ + ({"source": "archive"}, "url"), + ({"source": "archive", "url": ""}, "url"), + ({"source": "archive", "url": "http://artifacts.example.com/plugin.zip"}, "https"), + ({"source": "archive", "url": "s3://skills-bucket/plugin.zip"}, "https"), + ({"source": "archive", "url": "https://"}, "https"), + ({"source": "archive", "url": "https:///plugin.zip"}, "https"), + ({"source": "archive", "url": "https://[::1/plugin.zip"}, "https"), + ({"source": "archive", "url": "https://artifacts.example.com/plugin.zip", "sha256": "a" * 63}, "sha256"), + ({"source": "archive", "url": "https://artifacts.example.com/plugin.zip", "sha256": "a" * 65}, "sha256"), + ({"source": "archive", "url": "https://artifacts.example.com/plugin.zip", "sha256": "g" * 64}, "sha256"), + ], +) +async def test_register_plugin_archive_rejects_malformed_source(source, expected_fragment): + with pytest.raises(HTTPException) as exc_info: + await register_plugin(request=RegisterPluginRequest(name="bad-plugin", source=source), user_api_key_dict=_USER) + + assert exc_info.value.status_code == 400 + assert expected_fragment in exc_info.value.detail["error"] @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_skill_access.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_skill_access.py new file mode 100644 index 00000000000..dc4016d6db0 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_skill_access.py @@ -0,0 +1,79 @@ +"""Unit tests for claude_code_skill_access.py: key/team grant resolution.""" + +from unittest.mock import MagicMock + +import pytest + +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( + granted_skills, + skill_visibility, +) + + +def _perm(skills: list[str] | None) -> LiteLLM_ObjectPermissionTable: + return LiteLLM_ObjectPermissionTable(object_permission_id="perm", skills=skills) + + +def _key(key_skills: list[str] | None, team_skills: list[str] | None) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-user", + user_role=LitellmUserRoles.INTERNAL_USER, + object_permission=_perm(key_skills) if key_skills is not None else None, + team_object_permission=_perm(team_skills) if team_skills is not None else None, + ) + + +def _plugin(name: str, enabled: bool) -> MagicMock: + record = MagicMock() + record.name = name + record.enabled = enabled + return record + + +@pytest.mark.parametrize( + ("key_skills", "team_skills", "expected"), + [ + (None, None, frozenset()), + ([], [], frozenset()), + (["a", "b"], None, frozenset({"a", "b"})), + (None, ["a", "b"], frozenset({"a", "b"})), + (["a", "b"], ["b", "c"], frozenset({"b"})), + (["a"], ["c"], frozenset()), + (["a", "b"], [], frozenset({"a", "b"})), + ([], ["a", "b"], frozenset({"a", "b"})), + ], +) +def test_granted_skills_intersects_key_with_team(key_skills, team_skills, expected): + assert granted_skills(_key(key_skills, team_skills)) == expected + + +def test_visibility_enabled_plugin_is_public_for_everyone(): + public = _plugin("public-skill", enabled=True) + + assert skill_visibility(None).allows(public) + assert skill_visibility(_key(None, None)).allows(public) + assert skill_visibility(_key([], ["other"])).allows(public) + + +def test_visibility_disabled_plugin_needs_grant_or_admin(): + private = _plugin("private-skill", enabled=False) + admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + + assert not skill_visibility(None).allows(private) + assert not skill_visibility(_key(None, None)).allows(private) + assert not skill_visibility(_key(["other-skill"], None)).allows(private) + assert not skill_visibility(_key(["private-skill"], ["other-skill"])).allows(private) + assert skill_visibility(_key(["private-skill"], None)).allows(private) + assert skill_visibility(_key(None, ["private-skill"])).allows(private) + assert skill_visibility(admin).allows(private) + + +def test_where_clause_bounds_the_plugin_query_to_what_the_caller_may_see(): + admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + + assert skill_visibility(None).where() == {"enabled": True} + assert skill_visibility(_key(None, None)).where() == {"enabled": True} + assert skill_visibility(_key(["b", "a"], None)).where() == {"OR": [{"enabled": True}, {"name": {"in": ["a", "b"]}}]} + assert skill_visibility(_key(["a", "b"], ["b"])).where() == {"OR": [{"enabled": True}, {"name": {"in": ["b"]}}]} + assert skill_visibility(admin).where() == {} diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py new file mode 100644 index 00000000000..ae6b1471c93 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py @@ -0,0 +1,175 @@ +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from openai import APIConnectionError + +import litellm +from litellm.llms.litellm_proxy.skills.skill_search import skill_search_text +from litellm.proxy._types import LiteLLM_SkillsTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.skills_endpoints import router +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + +TRANSLATOR: Final = LiteLLM_SkillsTable( + skill_id="translate-file", + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", +) +SQL_ANALYST: Final = LiteLLM_SkillsTable( + skill_id="warehouse-sql-analyst", + display_title="Warehouse SQL Analyst", + description="Runs SQL against the inventory database", +) +TRIP_PLANNER: Final = LiteLLM_SkillsTable( + skill_id="trip-planner", + display_title="Trip Planner", + description="Books flights and hotels", +) +SKILLS: Final = (TRANSLATOR, SQL_ANALYST, TRIP_PLANNER) + +VECTORS: Final = MappingProxyType( + { + "language translation": (1.0, 0.0, 0.0), + skill_search_text(TRANSLATOR): (0.9, 0.1, 0.0), + skill_search_text(SQL_ANALYST): (0.0, 1.0, 0.0), + skill_search_text(TRIP_PLANNER): (0.3, 0.0, 1.0), + } +) + + +def _client(role: LitellmUserRoles) -> TestClient: + app = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="u", user_role=role) + return TestClient(app) + + +@pytest.fixture +def accessible_skills(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + list_for_search = AsyncMock(return_value=list(SKILLS)) + monkeypatch.setattr( + "litellm.llms.litellm_proxy.skills.handler.LiteLLMSkillsHandler.list_skills_for_search", list_for_search + ) + return list_for_search + + +@pytest.fixture +def key_limits(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", limits) + return limits + + +@pytest.fixture +def embedding_router(monkeypatch: pytest.MonkeyPatch, key_limits: MagicMock) -> MagicMock: + embedding_router = MagicMock() + embedding_router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], + ) + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", embedding_router) + monkeypatch.setattr(litellm, "skill_search_embedding_model", "text-embedding-3-small") + return embedding_router + + +class TestGetSkillsQuery: + def test_query_ranks_and_scores_and_truncates( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation", "top_k": 2}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 200 + body = response.json()["data"] + assert [skill["id"] for skill in body] == ["translate-file", "trip-planner"] + assert body[0]["search_score"] > body[1]["search_score"] + assert embedding_router.aembedding.await_args.kwargs["metadata"]["user_api_key_user_id"] == "u" + + def test_restricted_key_only_ranks_the_skills_it_can_access( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + accessible_skills.return_value = [SQL_ANALYST] + response = _client(LitellmUserRoles.INTERNAL_USER).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 200 + assert [skill["id"] for skill in response.json()["data"]] == ["warehouse-sql-analyst"] + + def test_no_accessible_skills_is_a_no_match_empty_result( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + accessible_skills.return_value = [] + response = _client(LitellmUserRoles.INTERNAL_USER).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 200 + assert response.json()["data"] == [] + embedding_router.aembedding.assert_not_awaited() + + def test_query_is_unsupported_for_the_anthropic_passthrough_provider( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", params={"query": "anything"}, headers={"Authorization": "Bearer k"} + ) + assert response.status_code == 400 + assert response.json()["detail"]["error"] == "skill_search_unsupported_provider" + accessible_skills.assert_not_awaited() + + def test_missing_embedding_model_is_a_400( + self, accessible_skills: AsyncMock, embedding_router: MagicMock, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "skill_search_embedding_model", None) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 400 + assert response.json()["detail"]["error"] == "skill_search_not_configured" + + def test_embedding_provider_failure_is_a_503( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + embedding_router.aembedding = AsyncMock(side_effect=APIConnectionError(request=MagicMock())) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 503 + assert response.json()["detail"]["error"] == "skill_search_unavailable" + + def test_a_key_over_its_rate_limit_gets_a_429_without_embedding( + self, accessible_skills: AsyncMock, embedding_router: MagicMock, key_limits: MagicMock + ) -> None: + key_limits.pre_call_hook = AsyncMock(side_effect=ProxyRateLimitError(detail="rpm exceeded")) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 429 + embedding_router.aembedding.assert_not_awaited() + + def test_top_k_is_validated(self, accessible_skills: AsyncMock, embedding_router: MagicMock) -> None: + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything", "top_k": 0}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 422 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 2284a05b2e9..8777e24e209 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,7 +1,7 @@ import asyncio import json from types import SimpleNamespace -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch if TYPE_CHECKING: @@ -37,6 +37,7 @@ from litellm.proxy.auth.auth_checks import ( _can_object_call_vector_stores, _check_end_user_budget, _check_team_member_budget, + _fetch_key_object_from_db_with_reconnect, _get_fuzzy_user_object, _get_team_db_check, _log_budget_lookup_failure, @@ -55,6 +56,7 @@ from litellm.caching.redis_cache import RedisCache from litellm.constants import ( DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + PROXY_DB_LOOKUP_MAX_CONCURRENCY, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, ) @@ -564,6 +566,43 @@ async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_e assert mock_prisma_client.get_data.await_count == 1 +class _InFlightCountingPrisma: + def __init__(self) -> None: + self.in_flight = 0 + self.max_in_flight = 0 + + async def get_data( + self, token: str, table_name: str, parent_otel_span: None, proxy_logging_obj: None + ) -> UserAPIKeyAuth: + self.in_flight += 1 + self.max_in_flight = max(self.max_in_flight, self.in_flight) + await asyncio.sleep(0.001) + self.in_flight -= 1 + return UserAPIKeyAuth(token=token) + + +@pytest.mark.asyncio +async def test_fetch_key_object_from_db_bounds_in_flight_prisma_requests(): + prisma: Final = _InFlightCountingPrisma() + burst: Final = PROXY_DB_LOOKUP_MAX_CONCURRENCY * 5 + + results: Final = await asyncio.gather( + *( + _fetch_key_object_from_db_with_reconnect( + hashed_token=f"hashed-token-{i}", + prisma_client=prisma, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + parent_otel_span=None, + proxy_logging_obj=None, + ) + for i in range(burst) + ) + ) + + assert len(results) == burst + assert {r.token for r in results if r is not None} == {f"hashed-token-{i}" for i in range(burst)} + assert prisma.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY + + def _fake_redis_cache(): fake_redis = MagicMock() fake_redis.async_get_cache = AsyncMock(return_value=None) @@ -2374,6 +2413,44 @@ def _mock_prisma_for_team_lookup(find_unique): return mock_prisma_client +_TEAM_ALIAS_TABLE_ROW = {"id": 1, "model_aliases": '{"fast": "gpt-4o"}', "created_by": "admin", "updated_by": "admin"} + + +def _prisma_team_row(include): + """Mimics Prisma: the `litellm_model_table` relation rides on the row only when the query `include`s it.""" + columns = {"team_id": "team-aliases", "team_alias": "aliases", "models": ["gpt-4o"]} + row = ( + {**columns, "litellm_model_table": _TEAM_ALIAS_TABLE_ROW} + if (include or {}).get("litellm_model_table") + else columns + ) + return SimpleNamespace(dict=lambda: row, model_dump=lambda: row) + + +@pytest.mark.asyncio +async def test_get_team_object_loads_model_aliases_relation(): + """LIT-5858: the auth path read teams without `include`ing `litellm_model_table`, so every JWT + team came back with `model_aliases=None` and alias requests 403'd.""" + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.auth.team_grants import team_model_aliases + + async def find_unique(where, include=None): + return _prisma_team_row(include) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + team = await get_team_object( + team_id="team-aliases", + prisma_client=_mock_prisma_for_team_lookup(AsyncMock(side_effect=find_unique)), + user_api_key_cache=mock_cache, + check_db_only=True, + ) + + assert team_model_aliases(team) == {"fast": "gpt-4o"} + + @pytest.mark.asyncio async def test_get_team_object_distinguishes_absent_team_from_unreadable_row(): """A deleted team and a database that would not answer both surface as a 404, @@ -6195,6 +6272,32 @@ async def test_get_team_object_by_alias_db_fetch_returns_cached_obj(): assert result.models == ["gpt-4"] +@pytest.mark.asyncio +async def test_get_team_object_by_alias_loads_model_aliases_relation(): + """LIT-5858: same regression as `test_get_team_object_loads_model_aliases_relation`, for the + `team_alias_jwt_field` lookup.""" + from litellm.proxy.auth.auth_checks import get_team_object_by_alias + from litellm.proxy.auth.team_grants import team_model_aliases + + async def find_many(where, include=None): + return [_prisma_team_row(include)] + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + team = await get_team_object_by_alias( + team_alias="aliases", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + assert team_model_aliases(team) == {"fast": "gpt-4o"} + + @pytest.mark.asyncio async def test_get_org_object_by_alias_db_fetch_returns_validated_org(): from litellm.proxy._types import LiteLLM_OrganizationTable diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index f513f397b64..cdf1f897707 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -3,6 +3,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext """ import base64 +import logging from typing import Optional from unittest.mock import MagicMock, patch @@ -15,6 +16,7 @@ from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, check_complete_credentials, custom_auth_common_checks_warning, + log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, get_end_user_id_from_request_body, get_key_mcp_rpm_limit, @@ -101,6 +103,41 @@ class TestWarnOnceIfCustomAuthSkipsCommonChecks: assert logger.warning.call_count == 0 +class TestLogOnceIfBudgetReservationDisabled: + @pytest.fixture(autouse=True) + def _reset_sentinel(self, monkeypatch): + monkeypatch.setattr( + "litellm.constants.budget_reservation_disabled_info_emitted", + False, + ) + + def test_logs_info_only_once_when_enabled(self, caplog): + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + log_once_if_budget_reservation_disabled(disabled=False) + assert not any( + "disable_budget_reservation is enabled" in record.message + for record in caplog.records + ) + for _ in range(3): + log_once_if_budget_reservation_disabled(disabled=True) + + records = [ + record + for record in caplog.records + if "disable_budget_reservation is enabled" in record.message + ] + assert len(records) == 1 + assert records[0].levelno == logging.INFO + + def test_logs_to_injected_logger_only_once(self): + logger = MagicMock() + log_once_if_budget_reservation_disabled(disabled=False, logger=logger) + for _ in range(3): + log_once_if_budget_reservation_disabled(disabled=True, logger=logger) + assert logger.info.call_count == 1 + assert "disable_budget_reservation is enabled" in logger.info.call_args[0][0] + + class TestGetKeyModelRpmLimit: """Tests for get_key_model_rpm_limit function.""" @@ -428,6 +465,138 @@ def test_get_model_from_request_openai_deployment_route_still_works(): ) +def test_get_model_from_request_bedrock_converse_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/converse", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_invoke_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_v2_converse_stream_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/v2/model/us.anthropic.claude-sonnet-4-6/converse-stream", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_model_id_with_slashes(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/aws/anthropic/model-name/invoke", + ) + == "aws/anthropic/model-name" + ) + + +def test_get_model_from_request_bedrock_unparseable_endpoint_returns_none(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/agents/some-agent-route", + ) + is None + ) + + +def test_get_model_from_request_bedrock_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/converse", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_invoke_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/invoke", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_count_tokens_uses_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/v1/messages/count_tokens", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_uppercase_count_tokens_segment_is_not_count_tokens(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-haiku-4-5-20251001-v1:0"}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke/COUNT_TOKENS", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/agents/some-agent-route", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def _azure_relay_router(): + from litellm.router import Router + + return Router( + model_list=[ + { + "model_name": "gpt", + "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://a.services.ai.azure.com", "api_key": "k"}, + }, + { + "model_name": "other-group", + "litellm_params": {"model": "azure/gpt-5.4", "api_base": "https://b.openai.azure.com", "api_key": "k"}, + }, + ] + ) + + +@pytest.mark.parametrize( + "route, request_data, expected", + [ + ("/azure_ai/other-group/openai/deployments/other-group/chat/completions", {"model": "gpt"}, "other-group"), + ("/azure_ai/other-group/models/chat/completions", {}, "other-group"), + ("/azure/openai/deployments/gpt/chat/completions", {"model": "other-group"}, "gpt"), + ("/azure/openai/deployments/gpt/chat/completions", {}, "gpt"), + ("/azure/openai/deployments/my-azure-deployment/chat/completions", {"model": "gpt"}, "gpt"), + ("/azure_ai/gpt", {"model": "other-group"}, "other-group"), + ], +) +def test_get_model_from_request_azure_relay_routes_use_the_model_group_in_the_path(route, request_data, expected): + assert get_model_from_request(request_data=request_data, route=route, llm_router=_azure_relay_router()) == expected + + def test_get_model_from_request_includes_file_endpoint_header_model(): assert ( get_model_from_request( @@ -3360,7 +3529,7 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors: @pytest.mark.parametrize( "selector", - ["aws_profile_name", "aws_session_name", "aws_external_id"], + ["aws_profile_name", "aws_session_name", "aws_external_id", "aws_session_tags"], ) def test_aws_identity_selector_in_batch_body_is_rejected(self, selector): with pytest.raises(ValueError, match=selector): @@ -3379,7 +3548,7 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors: @pytest.mark.parametrize( "selector", - ["aws_profile_name", "aws_session_name", "aws_external_id"], + ["aws_profile_name", "aws_session_name", "aws_external_id", "aws_session_tags"], ) def test_aws_identity_selector_under_extra_body_is_rejected(self, selector): with pytest.raises(ValueError, match=selector): diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 99a0a4c0a8b..94226b5404d 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -13,6 +13,7 @@ from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, JWTLiteLLMRoleMap, LiteLLM_JWTAuth, + LiteLLM_ModelTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, @@ -1255,6 +1256,57 @@ async def test_find_team_with_model_access_model_group(monkeypatch): assert team_obj.team_id == "team-1" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_aliases", + ['{"fast": "gpt-4o"}', {"fast": "gpt-4o"}], + ids=["json-string", "dict"], +) +async def test_find_team_with_model_access_resolves_team_model_alias(monkeypatch, model_aliases): + """LIT-5858: a JWT team that grants `gpt-4o` under the alias `fast` must resolve a request + for `fast`. The JWT path used to pass `team_model_aliases=None`, so every alias request 403'd.""" + import sys + import types + + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + router = Router(model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}]) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + team = LiteLLM_TeamTable( + team_id="team-aliases", + models=["gpt-4o"], + litellm_model_table=LiteLLM_ModelTable(model_aliases=model_aliases, created_by="admin", updated_by="admin"), + ) + + async def mock_get_team_object(*args, **kwargs): + return team + + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + user_api_key_cache = DualCache() + + team_id, team_obj = await JWTAuthManager.find_team_with_model_access( + team_ids={"team-aliases"}, + requested_model="fast", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), + ) + + assert team_id == "team-aliases" + assert team_obj is team + + @pytest.mark.asyncio async def test_find_team_with_model_access_v1_messages_default_routes(monkeypatch): """Regression for #31189: a single-team JWT that grants the requested model diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 8d93d801bfd..e209a491b0a 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -23,6 +23,7 @@ from litellm.proxy.auth.login_utils import ( LoginResult, authenticate_user, get_ui_credentials, + is_env_credential_login_enabled, ) @@ -185,6 +186,7 @@ async def test_authenticate_user_invalid_credentials(): assert exc_info.value.type == ProxyErrorTypes.auth_error assert exc_info.value.code == "401" assert "Invalid credentials" in exc_info.value.message + assert "UI_USERNAME" in exc_info.value.message @pytest.mark.asyncio @@ -799,3 +801,158 @@ class TestDisablePasswordLoginWhenSSOEnabled: assert isinstance(result, LoginResult) assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + +class TestDisableEnvCredentialLogin: + """`disable_env_credential_login` must reject a login with the env + credentials (UI_USERNAME/UI_PASSWORD, or the master-key fallback when + UI_PASSWORD is unset) while leaving database-user password logins + untouched, so admins with real accounts keep a way in.""" + + @pytest.mark.asyncio + async def test_rejects_correct_env_credentials_when_disabled(self): + master_key = "sk-1234" + ui_username = "admin" + ui_password = "env-only-password" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password}): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username=ui_username, + password=ui_password, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "401" + assert "UI_USERNAME" not in exc_info.value.message + assert "UI_PASSWORD" not in exc_info.value.message + + @pytest.mark.asyncio + async def test_rejects_master_key_fallback_when_disabled(self): + """With UI_PASSWORD unset, the master key IS the env password, so the + setting must reject it too or it protects nothing by default.""" + master_key = "sk-1234" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": "admin"}, clear=True): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username="admin", + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert exc_info.value.code == "401" + + @pytest.mark.asyncio + async def test_db_user_login_still_works_when_disabled(self): + master_key = "sk-1234" + user_email = "admin@example.com" + password = "Str0ng!Passw0rd" + + mock_user = LiteLLM_UserTable( + user_id="db-admin-1", + user_email=user_email, + password=hash_token(token=password), + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=mock_user) + + with patch.dict( + os.environ, + { + "UI_USERNAME": "admin", + "UI_PASSWORD": "env-password", + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + stack.enter_context( + patch( # test-quality-ok: internal orchestration, no HTTP boundary; matches pre-existing tests + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "db-user-token"}, + ) + ) + result = await authenticate_user( + username=user_email, + password=password, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == "db-admin-1" + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + @pytest.mark.asyncio + async def test_env_login_still_works_when_setting_absent(self): + """Env-credential login is the bootstrap path on a fresh install and + must stay on by default.""" + master_key = "sk-1234" + ui_username = "admin" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": master_key, + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + _patch_successful_admin_login_deps(stack) + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + +class TestIsEnvCredentialLoginEnabled: + """Drives the Admin UI warning banner: it must be True exactly when a + login with the env credentials could actually succeed.""" + + def test_enabled_by_default(self): + assert is_env_credential_login_enabled({}) is True + + def test_disabled_by_dedicated_setting(self): + assert is_env_credential_login_enabled({"disable_env_credential_login": True}) is False + + def test_explicit_false_keeps_it_enabled(self): + assert is_env_credential_login_enabled({"disable_env_credential_login": False}) is True + + def test_disabled_when_sso_gate_blocks_all_password_logins(self): + """`disable_password_login_when_sso_enabled` with SSO configured + rejects every username/password login before the env comparison runs, + so the banner must not nag about an already-unreachable path.""" + with ExitStack() as stack: + _patch_sso_configured(stack, configured=True) + assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is False + + def test_enabled_when_sso_gate_is_set_but_sso_not_configured(self): + with ExitStack() as stack: + _patch_sso_configured(stack, configured=False) + assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is True diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index d58683fd1e5..36bfc4c5dd3 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 48926cb7bc2..c8b3d789665 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2033,6 +2033,82 @@ def test_proxy_admin_viewer_can_access_logs_page_endpoints(route): ) +@pytest.mark.parametrize( + "user_role", + [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY], +) +def test_internal_user_can_access_logs_drawer_detail_route(user_role): + """ + The Logs drawer detail fetch (GET /spend/logs/ui/{request_id}) must pass + route_checks for plain internal users, not just admins — the handler + itself already self-authorizes row ownership via + _assert_user_can_view_request_id. + """ + route = "/spend/logs/ui/abc-request-id" + user_obj = LiteLLM_UserTable( + user_id="internal_user", + user_email="user@example.com", + user_role=user_role.value, + ) + valid_token = UserAPIKeyAuth( + user_id="internal_user", + user_role=user_role.value, + ) + request = MagicMock(spec=Request) + request.query_params = {} + + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=user_role.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + except Exception as e: + pytest.fail(f"{user_role.value} should be able to access {route}. Got error: {str(e)}") + + +@pytest.mark.parametrize( + "route_group_name", + [ + "spend_tracking_routes", + "internal_user_routes", + "internal_user_view_only_routes", + "admin_viewer_routes", + "org_admin_allowed_routes", + ], +) +def test_logs_drawer_detail_route_in_every_route_group(route_group_name): + """ + /spend/logs/ui/{request_id} must be reachable through + RouteChecks.check_route_access under each role's own route group, so a + partial revert (removing the route from `spend_tracking_routes` while + leaving `non_proxy_admin_allowed_routes_check` alone) is also caught. + """ + from litellm.proxy._types import LiteLLMRoutes + + allowed_routes = getattr(LiteLLMRoutes, route_group_name).value + assert RouteChecks.check_route_access( + route="/spend/logs/ui/req-34099", allowed_routes=allowed_routes + ) + + +def test_logs_drawer_detail_route_allowed_for_scoped_virtual_key(): + """ + A virtual key scoped to `allowed_routes=["spend_tracking_routes"]` must be + able to reach the Logs drawer detail route. + """ + valid_token = UserAPIKeyAuth( + user_id="scoped_key_user", + allowed_routes=["spend_tracking_routes"], + ) + assert RouteChecks.is_virtual_key_allowed_to_call_route( + route="/spend/logs/ui/req-34099", valid_token=valid_token + ) + + @pytest.mark.parametrize("route", ADMIN_VIEWER_LOGS_PAGE_ROUTES) def test_internal_user_blocked_from_admin_viewer_logs_routes(route): """ @@ -3217,6 +3293,65 @@ def test_internal_user_blocked_from_search_tool_writes(route): assert "Your role=internal_user" in str(exc_info.value) +@pytest.mark.parametrize( + "user_role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ], +) +def test_non_admin_can_open_vector_store_details(user_role): + """Regression for LIT-7132: the dashboard lists a vector store via /vector_store/list + (an LLM API route) but opened it via /vector_store/info, which no non-admin allowlist + granted, so the route gate 401'd before the handler's per-store access check ran.""" + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="user@example.com", + user_role=user_role, + ) + valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role) + request = MagicMock(spec=Request) + request.query_params = {} + + granted = RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=user_role, + route="/vector_store/info", + request=request, + valid_token=valid_token, + request_data={}, + ) + assert granted is None + + +@pytest.mark.parametrize( + "route", + ["/vector_store/new", "/vector_store/update", "/vector_store/delete"], +) +def test_internal_user_blocked_from_vector_store_writes(route): + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="user@example.com", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + valid_token = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + request = MagicMock(spec=Request) + request.query_params = {} + + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update"): + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + + def test_proxy_admin_viewer_can_read_another_users_info(): """Admin Viewer has read parity with Proxy Admin, so the /user/info key-ownership gate must not apply to it — the Users page reads every row.""" @@ -3579,3 +3714,230 @@ def test_agent_registry_route_gate_open_to_non_admin_roles(user_role, method, ro valid_token=valid_token, request_data={}, ) +TEAM_CALLBACK_ROUTES = ( + "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback", + "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback/langfuse", + # the routes register team_id with the :path converter, so a team id may + # contain a slash + "/team/tenant/06bda574/callback", + "/team/tenant/06bda574/callback/langfuse", + # team_id is a free-form string, so it may also contain a colon + "/team/tenant:06bda574/callback", + "/team/tenant:06bda574/callback/langfuse", + # or both, which is the shape neither a "[^:]+" nor a "[^/]+" expansion + # of the placeholder reaches on its own + "/team/tenant:acme/prod/callback", + "/team/tenant:acme/prod/callback/langfuse", +) + + +def _gate(route, role) -> str: + """Drive the real route gate for a non-proxy-admin caller. + + Reports "allowed" when the gate lets the request through to its handler, and + the denial message otherwise, so a caller asserts the verdict as a value + instead of on whether an exception escaped. + """ + user_obj = LiteLLM_UserTable( + user_id="team_admin_user", + user_email="team-admin@example.com", + user_role=role, + ) + request = MagicMock(spec=Request) + request.query_params = {} + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=role, + route=route, + request=request, + valid_token=UserAPIKeyAuth(user_id="team_admin_user", user_role=role), + request_data={}, + ) + except Exception as exc: + return f"denied: {exc}" + return "allowed" + + +def test_team_callback_routes_are_self_managed(): + """The grant has to come from self_managed_routes specifically. + + That list is the one whose entries carry no role predicate, so the handler + decides. Granting the same paths through internal_user_routes instead would + look identical for an internal_user while silently denying the org admins and + view-only roles that list does not cover. + """ + for template in ( + "/team/{team_id:path}/callback", + "/team/{team_id:path}/callback/{callback_name}", + ): + assert template in LiteLLMRoutes.self_managed_routes.value + + +@pytest.mark.parametrize("route", TEAM_CALLBACK_ROUTES) +@pytest.mark.parametrize( + "role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + LitellmUserRoles.ORG_ADMIN.value, + ], +) +def test_team_callback_routes_reach_their_handler_for_non_admins(route, role): + """A team admin manages their own team's logging callbacks, so the route gate + must let a non-proxy-admin through to the handler. + + The handler is what authorizes: every team callback endpoint calls + _verify_team_access, which admits only a proxy admin, an org admin for the + team, or an admin of that team, and 403s everyone else. Before this, the gate + rejected the team admin with a 401 naming proxy admin, so the handler's own + check was unreachable for them. + """ + assert _gate(route, role) == "allowed" + + +@pytest.mark.parametrize( + "pattern, route, matches", + [ + # a :path placeholder takes what the router's path converter takes + ("/team/{team_id:path}/callback", "/team/plain/callback", True), + ("/team/{team_id:path}/callback", "/team/tenant/acme/callback", True), + ("/team/{team_id:path}/callback", "/team/tenant:acme/callback", True), + ("/team/{team_id:path}/callback", "/team/tenant:acme/prod/callback", True), + # and still has to reach the template's own suffix + ("/team/{team_id:path}/callback", "/team/tenant:acme/disable_logging", False), + # a template with a ":" literal after the placeholder keeps the suffix + ( + "/v1beta/models/{model_name:path}:generateContent", + "/v1beta/models/gemini-2.5-flash:generateContent", + True, + ), + ( + "/v1beta/models/{model_name:path}:generateContent", + "/v1beta/models/publishers/google/gemini-2.5-flash:generateContent", + True, + ), + # the value must not swallow that suffix and match a different verb + ( + "/v1beta/models/{model_name:path}:generateContent", + "/v1beta/models/gemini-2.5-flash:countTokens", + False, + ), + # a %0A in the value reaches the handler through the path converter, so + # the gate has to see it too or DISABLE_ADMIN_ENDPOINTS is bypassable + ("/v1/mcp/server/{path:path}", "/v1/mcp/server/abc\ndef", True), + ("/team/{team_id:path}/callback", "/team/ten\nant/callback", True), + ("/v1beta/models/{model_name:path}:generateContent", "/v1beta/models/gem\nini:generateContent", True), + # an ordinary placeholder stays one segment + ("/team/{team_id}/members/me", "/team/abc/members/me", True), + ("/team/{team_id}/members/me", "/team/tenant/abc/members/me", False), + ("/team/{team_id}/members/me", "/team/ab\nc/members/me", True), + ], +) +def test_path_placeholder_matches_what_the_router_accepts(pattern, route, matches): + """The gate's placeholder expansion has to agree with the router's. + + A team id may carry a slash, a colon, or both, and the router mounted these + paths with the same :path converter, so an id the router routes must not be + an id the gate fails to recognize. The one narrowing that stays is a template + whose own suffix begins with a colon: there the value stops before it, or + ":generateContent" would also match a ":countTokens" request. + """ + assert RouteChecks._route_matches_pattern(route=route, pattern=pattern) is matches + + +# Every other route the proxy mounts under /team/{team_id}, spelled the way it +# is registered. None of them takes a path converter, so none can be reached by +# a URL that ends in the callback suffix. +PROTECTED_TEAM_ROUTES = ( + "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112", + "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/disable_logging", + "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/members/me", + "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/member/u-1/reset_spend", + # the same routes with the callback suffix spliced in, which is the shape a + # caller would craft to make a protected route look self-managed + "/team/06bda574/callback/disable_logging/x", + "/team/06bda574/callback/member/u-1/reset_spend", + "/team/06bda574/callback/members/me", +) + + +@pytest.mark.parametrize("route", PROTECTED_TEAM_ROUTES) +def test_the_callback_grant_does_not_reach_another_team_route(route): + """Widening the callback templates must not hand out any neighbouring route. + + The grant is two templates ending in the callback suffix. Every other team + route registers an ordinary single-segment placeholder, so no URL the router + sends to one of them can end in "/callback" or "/callback/" -- and the + gate must agree, or a crafted team id would carry a caller into a handler + the grant never covered. + """ + for template in ( + "/team/{team_id:path}/callback", + "/team/{team_id:path}/callback/{callback_name}", + ): + assert RouteChecks._route_matches_pattern(route=route, pattern=template) is False + + +def test_team_disable_logging_stays_proxy_admin_only(): + """disable_logging was left out of the grant, so it must still be rejected at + the gate. It is the one team callback route a team admin cannot reach.""" + verdict = _gate( + "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/disable_logging", + LitellmUserRoles.INTERNAL_USER.value, + ) + + assert "Only proxy admin" in verdict + assert "disable_logging" in verdict + + +@pytest.mark.parametrize( + "route", + [ + "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112", + "/team/update", + "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/model/add", + ], +) +def test_neighbouring_team_routes_stay_closed(route): + """The grant is the callback paths and nothing else on the team namespace.""" + assert "Only proxy admin" in _gate(route, LitellmUserRoles.INTERNAL_USER.value) + + +@pytest.mark.parametrize( + "route", + [ + "/claude-code/marketplace.json", + "/claude-code/plugins", + "/claude-code/plugins/my-skill", + ], +) +def test_claude_code_marketplace_routes_open_to_internal_users(route): + """Per-skill visibility is enforced inside the handler, so the route gate must let non-admins through.""" + assert RouteChecks.is_llm_api_route(route) is True + assert _gate(route, LitellmUserRoles.INTERNAL_USER.value) == "allowed" + + +@pytest.mark.parametrize("user_role", [None, LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value]) +def test_auto_router_session_is_reachable_by_any_key_but_benchmarks_stays_admin_only(user_role): + valid_token = UserAPIKeyAuth(api_key="hash-of-caller", user_role=user_role) + request = MagicMock(spec=Request) + request.query_params = {"session_id": "sess-1"} + + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=user_role, + route="/auto_router/session", + request=request, + valid_token=valid_token, + request_data={}, + ) + with pytest.raises(Exception, match="Only proxy admin"): + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=user_role, + route="/auto_router/benchmarks", + request=request, + valid_token=valid_token, + request_data={}, + ) diff --git a/tests/test_litellm/proxy/auth/test_team_grants.py b/tests/test_litellm/proxy/auth/test_team_grants.py new file mode 100644 index 00000000000..447fc1c93a1 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_team_grants.py @@ -0,0 +1,129 @@ +import pytest + +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_VerificationTokenView, + Member, + UserAPIKeyAuth, +) +from litellm.models.team import LiteLLM_ModelTable +from litellm.proxy.auth.team_grants import team_grants, team_model_aliases + +TEAM_ID = "team-grants" +USER_ID = "user-in-team" +ALIASES = {"fast": "gpt-4o-mini", "smart": "gpt-4o"} + + +def _alias_table(model_aliases) -> LiteLLM_ModelTable: + return LiteLLM_ModelTable(model_aliases=model_aliases, created_by="admin", updated_by="admin") + + +def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=TEAM_ID, + team_alias="grants-team", + tpm_limit=1000, + rpm_limit=10, + max_budget=50.0, + soft_budget=25.0, + spend=12.5, + models=["gpt-4o", "gpt-4o-mini"], + blocked=True, + metadata={"tier": "gold"}, + litellm_model_table=_alias_table(model_aliases), + object_permission_id="op-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", mcp_servers=["mcp-a"]), + members_with_roles=[ + Member(user_id="someone-else", role="user"), + Member(user_id=USER_ID, role="admin"), + ], + ) + + +def _membership() -> LiteLLM_TeamMembership: + return LiteLLM_TeamMembership( + user_id=USER_ID, + team_id=TEAM_ID, + spend=3.25, + litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=500, rpm_limit=5), + ) + + +def test_team_grants_cover_every_team_field_the_key_path_gets(): + """Class guard for LIT-5858 and its siblings: every ``team_*`` column the combined-view SQL hands the + virtual-key path must come out of the projection too, with the team's actual value, so adding a column + to ``LiteLLM_VerificationTokenView`` without teaching ``team_grants`` fails here instead of in prod.""" + team = _full_team() + grants = team_grants(team_object=team, team_membership=_membership(), user_id=USER_ID) + token = UserAPIKeyAuth(team_id=TEAM_ID, **grants) + + view_team_fields = {name for name in LiteLLM_VerificationTokenView.model_fields if name.startswith("team_")} + assert view_team_fields - {"team_id"} <= set(grants) + assert all(grants[name] is not None for name in view_team_fields - {"team_id"}) + + assert token.team_alias == "grants-team" + assert token.team_tpm_limit == 1000 + assert token.team_rpm_limit == 10 + assert token.team_max_budget == 50.0 + assert token.team_soft_budget == 25.0 + assert token.team_spend == 12.5 + assert token.team_models == ["gpt-4o", "gpt-4o-mini"] + assert token.team_blocked is True + assert token.team_metadata == {"tier": "gold"} + assert token.team_model_aliases == ALIASES + assert token.team_object_permission_id == "op-1" + assert token.team_object_permission is not None + assert token.team_object_permission.mcp_servers == ["mcp-a"] + assert token.team_member == Member(user_id=USER_ID, role="admin") + assert token.team_member_spend == 3.25 + assert token.team_member_tpm_limit == 500 + assert token.team_member_rpm_limit == 5 + + +def test_team_grants_without_team_leave_token_defaults(): + token = UserAPIKeyAuth(**team_grants(team_object=None, team_membership=None, user_id=USER_ID)) + assert token == UserAPIKeyAuth() + + +@pytest.mark.parametrize( + "stored_aliases", + [ALIASES, '{"fast": "gpt-4o-mini", "smart": "gpt-4o"}'], + ids=["json-object", "json-string-as-written-by-team-new"], +) +def test_team_model_aliases_decode_both_storage_shapes(stored_aliases): + team = _full_team(model_aliases=stored_aliases) + assert team_model_aliases(team) == ALIASES + assert team_grants(team_object=team, team_membership=None, user_id=None)["team_model_aliases"] == ALIASES + + +@pytest.mark.parametrize("stored_aliases", [None, "not json", '["a", "b"]', {"fast": 3}], ids=str) +def test_team_model_aliases_treat_unusable_column_as_no_aliases(stored_aliases): + team = _full_team(model_aliases=stored_aliases) + assert team_model_aliases(team) is None + assert team_grants(team_object=team, team_membership=None, user_id=None)["team_model_aliases"] is None + + +def test_team_model_aliases_none_without_relation_loaded(): + team = _full_team() + team.litellm_model_table = None + assert team_model_aliases(team) is None + assert team_model_aliases(None) is None + + +def test_team_member_is_the_callers_row_only(): + team = _full_team() + assert team_grants(team_object=team, team_membership=None, user_id="someone-else")["team_member"] == Member( + user_id="someone-else", role="user" + ) + assert team_grants(team_object=team, team_membership=None, user_id="stranger")["team_member"] is None + assert team_grants(team_object=team, team_membership=None, user_id=None)["team_member"] is None + + +def test_membership_limits_absent_without_membership_row(): + grants = team_grants(team_object=_full_team(), team_membership=None, user_id=USER_ID) + assert grants["team_member_spend"] is None + assert grants["team_member_tpm_limit"] is None + assert grants["team_member_rpm_limit"] is None diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index d44f96d95bf..6cce6d0316b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1,7 +1,13 @@ import asyncio import json +import logging +import os +import subprocess +import sys from contextlib import contextmanager from datetime import datetime, timedelta +from pathlib import Path +from textwrap import dedent from types import SimpleNamespace from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ -146,6 +152,35 @@ async def test_disable_budget_reservation_skips_reservation(): assert user_api_key_auth_obj.budget_reservation is None +@pytest.mark.asyncio +async def test_disable_budget_reservation_does_not_log_per_request(caplog): + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + for _ in range(3): + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={"disable_budget_reservation": True}, + ) + + records = [ + record + for record in caplog.records + if "disable_budget_reservation is enabled" in record.message + ] + assert records == [] + assert user_api_key_auth_obj.budget_reservation is None + + @pytest.mark.asyncio async def test_budget_reservation_runs_when_not_disabled(): """Control for #27639: with the flag absent, the reservation still runs and is stored.""" @@ -6737,6 +6772,109 @@ async def test_temp_budget_increase_applied_for_cached_key(): assert cached_after.max_budget == 2.0 +@pytest.mark.asyncio +@pytest.mark.parametrize( + "team_member_spend, expect_blocked", + [ + (2.4, True), + (2.4000000000000004, True), + (2.39, False), + ], +) +async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spend, expect_blocked): + """A team member counter sitting exactly at the cap (where a resized reservation + lands it) must be rejected by the cached-key auth path like every other budget check.""" + from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj + from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key + from litellm.proxy.utils import hash_token + + api_key = "sk-team-member-exact-cap" + hashed_token = hash_token(api_key) + team_id = "team-exact-cap" + user_id = "user-exact-cap" + max_budget = 2.4 + + user_api_key_cache = DualCache() + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=UserAPIKeyAuth( + token=hashed_token, + team_id=team_id, + user_id=user_id, + team_member_spend=team_member_spend, + ), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=None, + ) + await user_api_key_cache.async_set_cache( + key=f"team_id:{team_id}", + value=LiteLLM_TeamTableCachedObj(team_id=team_id), + ) + await user_api_key_cache.async_set_cache( + key=user_id, + value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER), + ) + await user_api_key_cache.async_set_cache( + key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=team_member_spend, + budget_id="budget-exact-cap", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=max_budget), + ), + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/messages" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {api_key}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + async def _auth(): + return await _user_api_key_auth_builder( + request=mock_request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]}, + ) + + with ( + patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam + "litellm.proxy.proxy_server.general_settings", {"disable_budget_reservation": True} + ), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state + patch( # test-quality-ok: seed the cached key, team and membership without a DB + "litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache + ), + patch( # test-quality-ok: module-global proxy state + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ), + patch( # test-quality-ok: the live counter needs Redis or a DB; pin the spend the check compares + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=team_member_spend), + ), + ): + if not expect_blocked: + result = await _auth() + assert result.team_member_spend == team_member_spend + return + with pytest.raises(ProxyException) as exc_info: + await _auth() + + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message + + async def _proxy_exception_for_key( api_key: str, general_settings: dict[str, bool], @@ -6872,3 +7010,285 @@ class TestLitellmReceivedAtStamping: assert result == earlier assert request.state.litellm_received_at == earlier + + +_RECORDING_DDTRACE = dedent( + ''' + import functools + import inspect + + + class _Span: + def __enter__(self): + return self + + def __exit__(self, *exc): + return None + + + class _Tracer: + def __init__(self): + self.spans = [] + + def wrap(self, name=None, **kwargs): + def decorator(f): + span_name = name or f"{f.__module__}.{f.__name__}" + if inspect.iscoroutinefunction(f): + + @functools.wraps(f) + async def async_wrapped(*args, **kw): + self.spans.append(span_name) + return await f(*args, **kw) + + return async_wrapped + + @functools.wraps(f) + def wrapped(*args, **kw): + self.spans.append(span_name) + return f(*args, **kw) + + return wrapped + + return decorator + + def trace(self, name, **kwargs): + return _Span() + + def current_span(self): + return None + + def current_root_span(self): + return None + + + tracer = _Tracer() + ''' +) + +_DDTRACE_AUTH_PROBE = dedent( + ''' + import asyncio + import json + + import ddtrace + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + proxy_server.master_key = "sk-probe" + + + async def auth(api_key): + request = Request(scope={"type": "http", "headers": [], "method": "POST", "path": "/chat/completions"}) + request._url = URL(url="/chat/completions") + try: + await user_api_key_auth( + request=request, + api_key=api_key, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + custom_litellm_key_header=None, + ) + return "accepted" + except ProxyException: + return "rejected" + + + async def main(): + outcomes = [await auth("Bearer sk-probe"), await auth("Bearer sk-wrong")] + print(json.dumps({"outcomes": outcomes, "spans": ddtrace.tracer.spans})) + + + asyncio.run(main()) + ''' +) + + +def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(tmp_path: Path): + stub_root = tmp_path / "site" + (stub_root / "ddtrace").mkdir(parents=True) + (stub_root / "ddtrace" / "__init__.py").write_text(_RECORDING_DDTRACE) + probe = tmp_path / "probe.py" + probe.write_text(_DDTRACE_AUTH_PROBE) + repo_root = Path(litellm.__file__).resolve().parent.parent + env = { + **os.environ, + "USE_DDTRACE": "true", + "PYTHONPATH": os.pathsep.join( + [str(stub_root), str(repo_root)] + [p for p in (os.environ.get("PYTHONPATH"),) if p] + ), + } + + result = subprocess.run( + [sys.executable, str(probe)], env=env, cwd=repo_root, capture_output=True, text=True, check=False + ) + + assert result.returncode == 0, result.stderr[-4000:] + report = json.loads(result.stdout.strip().splitlines()[-1]) + assert report["outcomes"] == ["accepted", "rejected"] + auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" + assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard-return", "proxy-admin-return"]) +async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_admin): + """LIT-5858: the team-based JWT path hand-built ``UserAPIKeyAuth`` from a short list of team fields, so the + team's model aliases (and on the admin return, its object permission) never reached the token and alias + requests 403'd. Both returns now go through ``team_grants``; pin the fields that used to be dropped.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.models.team import LiteLLM_ModelTable + from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + Member, + ) + + class _AcceptEveryJwt(JWTHandler): + def is_jwt(self, token: str) -> bool: + return True + + jwt_handler = _AcceptEveryJwt() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + team = LiteLLM_TeamTable( + team_id="team-jwt-aliases", + team_alias="jwt-aliases", + models=["gpt-4o"], + max_budget=40.0, + spend=4.0, + blocked=False, + metadata={"tier": "gold"}, + litellm_model_table=LiteLLM_ModelTable( + model_aliases='{"fast": "gpt-4o"}', created_by="admin", updated_by="admin" + ), + object_permission_id="op-jwt", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-jwt", mcp_servers=["mcp-a"]), + members_with_roles=[Member(user_id="jwt-user", role="admin")], + ) + membership = LiteLLM_TeamMembership(user_id="jwt-user", team_id="team-jwt-aliases", spend=1.5) + builder_result = { + "is_proxy_admin": is_proxy_admin, + "team_object": team, + "user_object": None, + "end_user_object": None, + "org_object": None, + "token": "jwt", + "team_id": "team-jwt-aliases", + "user_id": "jwt-user", + "user_email": "jwt-user@example.com", + "end_user_id": None, + "org_id": None, + "team_membership": membership, + "jwt_claims": {"sub": "jwt-user"}, + } + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + attrs = { + "prisma_client": MagicMock(), + "user_api_key_cache": DualCache(), + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {"enable_jwt_auth": True}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": jwt_handler, + "premium_user": True, + "litellm_proxy_admin_name": "admin", + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http", "headers": [], "method": "POST"}) + request._url = URL(url="/chat/completions") + with patch( # test-quality-ok: auth_builder is the claim-resolution seam; the regression is how its result is projected onto the token + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + return_value=builder_result, + ): + token = await _user_api_key_auth_builder( + request=request, + api_key="Bearer header.payload.signature", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert token.team_id == "team-jwt-aliases" + assert token.user_role == (LitellmUserRoles.PROXY_ADMIN if is_proxy_admin else LitellmUserRoles.INTERNAL_USER) + assert token.team_model_aliases == {"fast": "gpt-4o"} + assert token.team_object_permission is not None + assert token.team_object_permission.mcp_servers == ["mcp-a"] + assert token.team_object_permission_id == "op-jwt" + assert token.team_alias == "jwt-aliases" + assert token.team_models == ["gpt-4o"] + assert token.team_max_budget == 40.0 + assert token.team_spend == 4.0 + assert token.team_metadata == {"tier": "gold"} + assert token.team_member == Member(user_id="jwt-user", role="admin") + assert token.team_member_spend == 1.5 + assert token.jwt_claims == {"sub": "jwt-user"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/v1/messages", "/messages", "/v1/chat/completions", "/chat/completions", "/v1/responses", "/responses"]) +async def test_claude_view_normalizes_before_model_access(monkeypatch, route): + from starlette.requests import Request + from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access + + source = "foo[1m]" + encoded = "claude-router-" + source.encode().hex() + "[1m]" + router = litellm.Router(model_list=[{"model_name": source, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}]) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + data = {"model": encoded, "messages": [{"role": "user", "content": "hi"}]} + request = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) + token = UserAPIKeyAuth(models=[source]) + await _enforce_key_and_fallback_model_access(valid_token=token, request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + assert data["model"] == source + assert (await request.json())["model"] == source + assert json.loads(await request.body())["model"] == source + assert request.scope["parsed_body"][1]["model"] == source + with pytest.raises(ProxyException): + await _enforce_key_and_fallback_model_access(valid_token=UserAPIKeyAuth(models=["other"]), request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "hierarchical", "unclaimed"]) +async def test_claude_view_never_reinterprets_explicit_names(monkeypatch, layer): + from starlette.requests import Request + from litellm.proxy.auth.user_api_key_auth import _normalize_claude_model + + encoded = "claude-router-666f6f" + names = ("foo", "other", encoded) if layer == "literal" else ("foo", "other") + alias = {encoded: "other"} + router = litellm.Router(model_list=[{"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names], model_group_alias=alias if layer == "router" else None) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) + token = UserAPIKeyAuth(aliases=alias if layer == "key" else {}, router_settings={"model_group_alias": alias} if layer == "hierarchical" else None) + data = {"model": encoded} + request = Request({"type": "http", "method": "POST", "path": "/v1/messages", "headers": [], "query_string": b""}) + await _normalize_claude_model(data, token, request, "/v1/messages") + assert data["model"] == ("foo" if layer == "unclaimed" else encoded) + await _normalize_claude_model(data, token, request, "/v1/messages") + assert data["model"] == ("foo" if layer == "unclaimed" else encoded) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py index 4a3b3ef22c4..028ab58843f 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -6,6 +6,7 @@ from typing import Optional import yaml from click.testing import CliRunner +from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError from litellm.proxy.client.cli.commands.autoroute import commands as commands_module from litellm.proxy.client.cli.commands.autoroute import process as process_module from litellm.proxy.client.cli.commands.autoroute.commands import down, up @@ -159,6 +160,11 @@ class TestUpCommand: assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "fixed-master-key" assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true" assert "apiKeyHelper" not in captured["settings"] + # The ephemeral proxy serves only the autorouter, so a starting model left by + # `lite configure claude --model` or a user pin would 400 on the first message. + assert captured["settings"]["model"] == "autorouter" + assert captured["settings"]["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter" + assert captured["settings"]["statusLine"]["command"].endswith("statusline.py") assert captured["settings_mode"] == 0o600 assert terminate_calls == [99999] @@ -249,6 +255,33 @@ class TestUpCommand: assert not pid_record_path.exists() assert not backup_path.exists() + def test_a_status_line_install_failure_leaves_no_backup_behind(self, monkeypatch, tmp_path): + # The install runs before the backup is written, so a failure cannot strand a backup that + # would make every later `lite configure` / `lite autoroute up` think a session still owns settings.json + config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + claude_settings_path.write_text(json.dumps({"theme": "dark"})) + + def boom(): + raise ClaudeSettingsError("disk full") + + fake_process = FakeProcess(pid=778) + terminate_calls = [] + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "is_port_available", lambda port: True) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + monkeypatch.setattr(commands_module, "install_statusline_script", boom) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + result = self.runner.invoke(up) + + assert result.exit_code != 0 and "disk full" in result.output + assert terminate_calls == [778] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == {"theme": "dark"} + def test_up_uses_the_same_port_and_master_key_across_runs(self, monkeypatch, tmp_path): """The LIT-4607/LIT-4608 regression: a client configured against one session must keep working in the next, so consecutive runs must patch settings with an identical base URL diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py index f399b6f957a..47b5459d489 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py @@ -161,7 +161,13 @@ class TestBuildGeneratedModelList: config = _base_config(classifier=HeuristicClassifier(), semantic_matching=NoSemanticMatching()) autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") router_config = autorouter["litellm_params"]["complexity_router_config"] - assert set(router_config.keys()) == {"tiers", "default_model"} + assert set(router_config.keys()) == {"tiers", "default_model", "return_raw_model_name"} + + def test_the_generated_router_reports_the_tier_model_it_routed_to(self): + # The status line reads the routed model from the response body, which the proxy restamps to the + # requested alias unless the deployment opts out; "autorouter" on every line would tell nothing. + autorouter = next(m for m in build_generated_model_list(_base_config()) if m["model_name"] == "autorouter") + assert autorouter["litellm_params"]["complexity_router_config"]["return_raw_model_name"] is True class TestBuildGeneratedProxyConfig: diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py deleted file mode 100644 index 87a33c79a79..00000000000 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py +++ /dev/null @@ -1,63 +0,0 @@ -from litellm.proxy.client.cli.commands.autoroute.settings import ( - ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, - merge_claude_settings_static_token, -) - - -def test_preserves_unrelated_top_level_keys(): - merged = merge_claude_settings_static_token({"theme": "dark"}, "http://127.0.0.1:4000", "token-abc") - assert merged["theme"] == "dark" - - -def test_preserves_unrelated_env_keys(): - settings = {"env": {"SOME_OTHER_VAR": "value"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["SOME_OTHER_VAR"] == "value" - - -def test_sets_base_url_and_auth_token(): - merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000/", "token-abc") - assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" - assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" - - -def test_preserves_existing_tool_search(): - settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" - - -def test_drops_stray_api_key(): - settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert "ANTHROPIC_API_KEY" not in merged["env"] - - -def test_removes_existing_api_key_helper(): - settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token"} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert "apiKeyHelper" not in merged - - -def test_does_not_mutate_input(): - settings = {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} - merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert settings == {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} - - -def test_forces_all_claude_code_default_model_tiers_to_the_autorouter(): - # A bare "*" model_name deployment looks like the obvious way to catch every request - # regardless of which model Claude Code thinks it's using, but Router's auto-router - # registry is keyed by the literal requested model string with no wildcard resolution - # (litellm/router.py:10711-10717) -- so the only reliable way to make every one of Claude - # Code's own tiers hit the auto-router is to override the env vars it reads per tier. - merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000", "token-abc") - for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: - assert merged["env"][key] == "autorouter" - - -def test_overrides_a_preexisting_default_model_env_var(): - settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter" diff --git a/tests/test_litellm/proxy/client/cli/conftest.py b/tests/test_litellm/proxy/client/cli/conftest.py new file mode 100644 index 00000000000..c77f516a768 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/conftest.py @@ -0,0 +1,39 @@ +import os +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.client.cli.commands import claude_settings + +REAL_CLAUDE_SETTINGS: Final = Path(os.path.expanduser("~")) / ".claude" / "settings.json" + + +def _current_bytes() -> bytes | None: + return REAL_CLAUDE_SETTINGS.read_bytes() if REAL_CLAUDE_SETTINGS.exists() else None + + +@pytest.fixture(autouse=True) +def _statusline_script_under_tmp(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(claude_settings, "STATUSLINE_SCRIPT_PATH", tmp_path / "litellm-home" / "statusline.py") + + +@pytest.fixture(autouse=True) +def isolated_claude_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[Path]: + before: Final = _current_bytes() + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(tmp_path / ".claude")) + yield tmp_path + after: Final = _current_bytes() + if after == before: + return + if before is None: + REAL_CLAUDE_SETTINGS.unlink() + else: + REAL_CLAUDE_SETTINGS.write_bytes(before) + pytest.fail( + f"this test wrote the developer's real {REAL_CLAUDE_SETTINGS}; the original bytes were restored. " + "Resolve the Claude settings path at call time (never Path.home() at import) and point the test at tmp_path" + ) diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index a8a6659fe9a..7804435a60d 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -28,6 +28,7 @@ from litellm.proxy.client.cli.commands.agents import ( ) AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" +CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" def _agent_command(name): @@ -1202,3 +1203,53 @@ class TestAgentCommands: ) assert result.exit_code == 0, result.output assert captured["reattach_terminal"] is None + + +class TestPrepareCodex: + def test_registers_the_installed_script_as_a_session_scoped_stop_hook(self): + from litellm.proxy.client.cli.commands.agents import prepare_codex + + args = prepare_codex("http://localhost:4000", "sk-key", {}, install=lambda: "/py /home/me/.litellm/statusline.py") + assert args == ( + "-c", + 'hooks.Stop=[{hooks=[{type="command",command="/py /home/me/.litellm/statusline.py"}]}]', + ) + + def test_a_failed_install_is_an_agent_error_not_a_crash(self): + from litellm.proxy.client.cli.commands.agents import AgentRunError, prepare_codex + from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError + + def boom(): + raise ClaudeSettingsError("disk full") + + with pytest.raises(AgentRunError, match="disk full"): + prepare_codex("http://localhost:4000", "sk-key", {}, install=boom) + + def test_a_config_that_already_declares_hooks_keeps_them_and_skips_ours(self, tmp_path): + from litellm.proxy.client.cli.commands.agents import prepare_codex + + warnings = [] + env = {"CODEX_HOME": str(tmp_path)} + for body in ('[[hooks.Stop]]\nhooks = [{ type = "command", command = "mine" }]\n', 'hooks.Stop = []\n', "[hooks]\n"): + (tmp_path / "config.toml").write_text(body) + assert prepare_codex("http://localhost:4000", "sk", env, install=lambda: "/py /s.py", warn=warnings.append) == () + (tmp_path / "config.toml").write_text('model = "gpt-5.6-sol"\n[projects."/x"]\ntrust_level = "trusted"\n') + assert prepare_codex("http://localhost:4000", "sk", env, install=lambda: "/py /s.py", warn=warnings.append) != () + assert len(warnings) == 3 and "already declares hooks" in warnings[0] + + def test_a_config_that_cannot_be_read_or_decoded_still_lets_codex_launch(self, tmp_path): + # A UTF-16 config.toml (a Windows Notepad save) is Codex's problem to report at launch, not a reason + # for the hook pre-check to abort `lite codex` with a traceback before Codex ever starts. + from litellm.proxy.client.cli.commands.agents import codex_declares_stop_hooks, prepare_codex + + config = tmp_path / "config.toml" + config.write_bytes('[[hooks.Stop]]\nhooks = [{ type = "command", command = "mine" }]\n'.encode("utf-16")) + assert codex_declares_stop_hooks(config) is False + assert codex_declares_stop_hooks(tmp_path / "absent.toml") is False + args = prepare_codex("http://localhost:4000", "sk", {"CODEX_HOME": str(tmp_path)}, install=lambda: "/py /s.py") + assert args[0] == "-c" and "hooks.Stop=" in args[1] + + def test_codex_is_wired_through_the_preparer_registry(self): + from litellm.proxy.client.cli.commands.agents import _PREPARERS, prepare_codex + + assert _PREPARERS["codex"] is prepare_codex diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 821323e722c..3a7792db1fe 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -6,7 +6,6 @@ from pathlib import Path from unittest.mock import Mock, patch - import pytest from click.testing import CliRunner @@ -18,8 +17,15 @@ from litellm.litellm_core_utils.cli_keyring import ( SecretErased, SecretStored, ) -from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord, save_cli_token +from litellm.litellm_core_utils.cli_token_utils import ( + CliTokenRecord, + CredentialNotRecorded, + CredentialNotSaved, + save_cli_token, +) from litellm.proxy.client.cli import cli +from litellm.proxy.client.cli.commands import claude_settings as claude_settings_module +from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner from litellm.proxy.client.cli.commands.auth import ( get_stored_api_key, login, @@ -27,7 +33,6 @@ from litellm.proxy.client.cli.commands.auth import ( print_token, whoami, ) -from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner @pytest.fixture @@ -84,7 +89,7 @@ class TestPollingErrorSurfacing: } with patch("requests.get", return_value=mock_response) as mock_get, patch("time.sleep"): - with pytest.raises(ValueError, match='Your litellm CLI is out of date and uses a login flow') as exc_info: + with pytest.raises(ValueError, match="Your litellm CLI is out of date and uses a login flow") as exc_info: _poll_for_ready_data("http://test/sso/cli/poll/sk-legacy") assert mock_get.call_count == 1 @@ -151,7 +156,7 @@ class TestStartCliSsoFlowErrors: mock_response.status_code = 404 with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError, match='Either --base-url is wrong, or the proxy is older than') as exc_info: + with pytest.raises(ValueError, match="Either --base-url is wrong, or the proxy is older than") as exc_info: _start_cli_sso_flow("https://old-proxy.example.com") message = str(exc_info.value) @@ -167,7 +172,7 @@ class TestStartCliSsoFlowErrors: mock_response.json.return_value = {"detail": "Too many CLI login attempts. Try again later."} with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError, match='Too many CLI login attempts\\. Try again later\\.') as exc_info: + with pytest.raises(ValueError, match="Too many CLI login attempts\\. Try again later\\.") as exc_info: _start_cli_sso_flow("https://test.example.com") assert "HTTP 429" in str(exc_info.value) @@ -183,7 +188,7 @@ class TestStartCliSsoFlowErrors: mock_response.text = "Sign in to corporate VPN" with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError, match='A proxy, load balancer, or auth gateway in front of') as exc_info: + with pytest.raises(ValueError, match="A proxy, load balancer, or auth gateway in front of") as exc_info: _start_cli_sso_flow("https://test.example.com") message = str(exc_info.value) @@ -197,7 +202,7 @@ class TestStartCliSsoFlowErrors: from litellm.proxy.client.cli.commands.auth import _start_cli_sso_flow with patch("requests.post", side_effect=requests.ConnectionError("Connection refused")): - with pytest.raises(ValueError, match='Connection refused\\. Check that the proxy is running') as exc_info: + with pytest.raises(ValueError, match="Connection refused\\. Check that the proxy is running") as exc_info: _start_cli_sso_flow("https://unreachable.example.com") message = str(exc_info.value) @@ -584,13 +589,9 @@ class TestLogoutCommand: assert "could not be checked" in result.output assert DISABLE_KEYRING_ENV_VAR in result.output - def test_logout_warns_when_the_keychain_refuses_to_release_the_entry( - self, isolated_home, secret_vault_factory - ): + def test_logout_warns_when_the_keychain_refuses_to_release_the_entry(self, isolated_home, secret_vault_factory): """A locked keychain leaves a live credential behind that the user believes is gone.""" - vault = secret_vault_factory( - blob=_secret_blob("https://test.example.com", "sk-stored"), erasable=False - ) + vault = secret_vault_factory(blob=_secret_blob("https://test.example.com", "sk-stored"), erasable=False) _write_token_file(isolated_home, key=None) result = self.runner.invoke(logout, obj={"secret_vault": vault}) @@ -1210,9 +1211,7 @@ class TestKeychainBackedCommands: assert str(token_file) in result.output assert json.loads(token_file.read_text())["key"] == "sk-minted" - def test_login_points_a_user_missing_the_keyring_package_at_the_install( - self, isolated_home, secret_vault_factory - ): + def test_login_points_a_user_missing_the_keyring_package_at_the_install(self, isolated_home, secret_vault_factory): """`lite` ships with every install, the keyring package only with the cli extra. Telling that user their machine has no keychain sends them looking for a problem they do not have.""" result = self._login(secret_vault_factory(available=False, failure=KeyringNotInstalled())) @@ -1223,9 +1222,7 @@ class TestKeychainBackedCommands: assert "No OS keychain available" not in result.output assert json.loads(token_file.read_text())["key"] == "sk-minted" - def test_login_keeps_the_credential_when_the_backend_keeps_nothing( - self, isolated_home, secret_vault_factory - ): + def test_login_keeps_the_credential_when_the_backend_keeps_nothing(self, isolated_home, secret_vault_factory): """A backend that accepts writes and stores nothing must not be reported as keychain storage, because the file is then told to drop the only remaining copy.""" result = self._login(secret_vault_factory(discards=True)) @@ -1236,9 +1233,7 @@ class TestKeychainBackedCommands: assert "keyring --enable" in result.output assert json.loads(token_file.read_text())["key"] == "sk-minted" - def test_login_names_the_kill_switch_instead_of_blaming_the_machine( - self, isolated_home, secret_vault_factory - ): + def test_login_names_the_kill_switch_instead_of_blaming_the_machine(self, isolated_home, secret_vault_factory): result = self._login(secret_vault_factory(available=False, failure=KeyringDisabled())) assert result.exit_code == 0 @@ -1297,9 +1292,7 @@ class TestKeychainBackedCommands: assert "could not be read" in result.output assert "lite login" in result.output - def test_whoami_does_not_call_a_credential_it_cannot_read_authenticated( - self, isolated_home, secret_vault_factory - ): + def test_whoami_does_not_call_a_credential_it_cannot_read_authenticated(self, isolated_home, secret_vault_factory): """A login whose secret is stuck in an unreachable keychain authenticates nothing. Leading with "Authenticated" and a token age reads as a working session, and sends the user looking for the problem somewhere other than the keychain the notice underneath names.""" @@ -1316,9 +1309,7 @@ class TestKeychainBackedCommands: assert "the credential cannot be read" in result.output assert "could not be read" in result.output - def test_whoami_names_the_kill_switch_rather_than_a_missing_package( - self, isolated_home, secret_vault_factory - ): + def test_whoami_names_the_kill_switch_rather_than_a_missing_package(self, isolated_home, secret_vault_factory): """Every unreachable keychain used to be described as a locked one needing the keyring package installed. Someone who set the kill switch has the package and an unlocked keychain, so that advice sends them to fix two things that were never wrong.""" @@ -1330,9 +1321,7 @@ class TestKeychainBackedCommands: assert DISABLE_KEYRING_ENV_VAR in result.output assert "pip install" not in result.output - def test_print_token_points_an_install_without_keyring_at_the_package( - self, isolated_home, secret_vault_factory - ): + def test_print_token_points_an_install_without_keyring_at_the_package(self, isolated_home, secret_vault_factory): _write_token_file(isolated_home, key=None) vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) obj = {"base_url": "https://test.example.com", "secret_vault": vault} @@ -1398,9 +1387,24 @@ class TestLoginConfigClaude: def setup_method(self): self.runner = CliRunner() - def _run_login(self, tmp_path, args, base_url="https://test.example.com"): - settings_path = tmp_path / "claude" / "settings.json" + def _isolate_default_settings(self, tmp_path, monkeypatch): + """The default file, its `lite up` backup and its configure receipt all live under tmp_path.""" backup_path = tmp_path / "claude_settings_backup.json" + monkeypatch.setattr( + claude_settings_module, "SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup_path, "lite up", "lite down"),) + ) + monkeypatch.setattr( + claude_settings_module, "CLAUDE_SETTINGS_PATH", tmp_path / "default-home" / ".claude" / "settings.json" + ) + monkeypatch.setattr(claude_settings_module, "CONFIGURE_STATE_PATH", tmp_path / "claude_configure_state.json") + return backup_path + + def _run_login( + self, tmp_path, monkeypatch, args, base_url="https://test.example.com", *, config_dir_env=None, stored=None + ): + settings_path = tmp_path / "claude" / "settings.json" + backup_path = self._isolate_default_settings(tmp_path, monkeypatch) + env = {"CLAUDE_CONFIG_DIR": str(settings_path.parent)} if config_dir_env is None else config_dir_env poll_response = Mock() poll_response.status_code = 200 poll_response.json.return_value = { @@ -1414,57 +1418,139 @@ class TestLoginConfigClaude: patch("webbrowser.open"), patch("requests.post", return_value=_mock_cli_sso_start_response()), patch("requests.get", return_value=poll_response), - patch("litellm.proxy.client.cli.commands.auth.save_cli_token"), + patch("litellm.proxy.client.cli.commands.auth.save_cli_token", return_value=stored or SecretStored()), patch("litellm.proxy.client.cli.interface.show_commands"), - patch("litellm.proxy.client.cli.commands.auth.CLAUDE_SETTINGS_PATH", settings_path), - patch( - "litellm.proxy.client.cli.commands.auth.SETTINGS_FILE_OWNERS", - (SettingsFileOwner(backup_path, "lite up", "lite down"),), - ), - patch( - "litellm.proxy.client.cli.commands.claude_settings.shutil.which", - return_value="/usr/local/bin/lite", - ), ): - result = self.runner.invoke(login, args, obj={"base_url": base_url}) + result = self.runner.invoke(login, args, obj={"base_url": base_url}, env=env) return result, settings_path, backup_path - def test_default_login_does_not_touch_claude_settings(self, tmp_path): - result, settings_path, _backup_path = self._run_login(tmp_path, []) + def test_default_login_does_not_touch_claude_settings(self, tmp_path, monkeypatch): + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, []) assert result.exit_code == 0 assert "Login successful!" in result.output assert not settings_path.exists() assert "Configured Claude Code" not in result.output - def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path): - result, settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path, monkeypatch): + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) assert result.exit_code == 0 written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://test.example.com" assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" - assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" - assert "Configured Claude Code" in result.output + # The minted key goes in as a static token: an apiKeyHelper would make Claude Code spawn `lite` (and + # its keychain probe) on every credential refresh, which is what this flag used to write. + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" + assert "apiKeyHelper" not in written + assert f"Configured Claude Code: {settings_path} now routes through https://test.example.com." in result.output + assert "run `lite login --config-claude` again after it expires" in result.output + assert "the model Claude Code starts and resumes on" in result.output - def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path): + def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path, monkeypatch): settings_path = tmp_path / "claude" / "settings.json" settings_path.parent.mkdir(parents=True) settings_path.write_text(json.dumps({"theme": "dark", "env": {"KEEP": "me"}})) - result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + result, _settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) assert result.exit_code == 0 written = json.loads(settings_path.read_text()) assert written["theme"] == "dark" assert written["env"]["KEEP"] == "me" - def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path): + def _run_login_refused_before_the_sso_flow(self, tmp_path, monkeypatch, config_dir): + self._isolate_default_settings(tmp_path, monkeypatch).write_text("{}") + with patch("requests.post") as post, patch("webbrowser.open") as browser: + result = self.runner.invoke( + login, + ["--config-claude"], + obj={"base_url": "https://test.example.com"}, + env={"CLAUDE_CONFIG_DIR": config_dir}, + ) + assert result.exit_code != 0 + assert "not logging in" in result.output and "lite down" in result.output + assert "`lite up` is currently managing" in result.output + assert "Login successful!" not in result.output + post.assert_not_called() + browser.assert_not_called() + assert not (tmp_path / "default-home" / ".claude" / "settings.json").exists() + + def test_refuses_before_logging_in_while_lite_up_holds_the_default_settings_file(self, tmp_path, monkeypatch): + self._run_login_refused_before_the_sso_flow(tmp_path, monkeypatch, config_dir="") + + def test_refuses_before_logging_in_while_lite_up_holds_the_default_file_reached_through_a_symlink( + self, tmp_path, monkeypatch + ): + default_config_dir = tmp_path / "default-home" / ".claude" + default_config_dir.mkdir(parents=True) + alias = tmp_path / "claude-alias" + alias.symlink_to(default_config_dir, target_is_directory=True) + + self._run_login_refused_before_the_sso_flow(tmp_path, monkeypatch, config_dir=str(alias)) + + def test_flag_writes_an_alternate_config_dir_even_while_lite_up_holds_the_default_file(self, tmp_path, monkeypatch): + (tmp_path / "claude_settings_backup.json").write_text("{}") + + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) + + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" + assert f"Configured Claude Code: {settings_path} now routes through https://test.example.com." in result.output + + def test_flag_keeps_a_config_dir_receipt_apart_from_the_default_file_receipt(self, tmp_path, monkeypatch): + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) + + assert result.exit_code == 0, result.output + default_receipt = tmp_path / "claude_configure_state.json" + assert not default_receipt.exists() + receipts = list((tmp_path / "claude_configure_state").glob("*.json")) + assert len(receipts) == 1 + assert json.loads(receipts[0].read_text())["file_existed"] is False + + def test_a_second_login_replaces_the_key_and_unconfigure_still_restores_the_original(self, tmp_path, monkeypatch): + # The stored key expires daily, so the flag is re-run per login; the receipt must keep owning the + # slot across re-logins and hand back what was there before the first one. + from litellm.proxy.client.cli.commands.configure import unconfigure_claude + + settings_path = tmp_path / "claude" / "settings.json" + settings_path.parent.mkdir(parents=True) + settings_path.write_text(json.dumps({"theme": "dark", "env": {"ANTHROPIC_AUTH_TOKEN": "sk-theirs"}})) + self._run_login(tmp_path, monkeypatch, ["--config-claude"]) + first = json.loads(settings_path.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] + self._run_login(tmp_path, monkeypatch, ["--config-claude"]) + assert json.loads(settings_path.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == first != "sk-theirs" + + result = self.runner.invoke(unconfigure_claude, [], env={"CLAUDE_CONFIG_DIR": str(settings_path.parent)}) + assert result.exit_code == 0, result.output + assert json.loads(settings_path.read_text()) == {"theme": "dark", "env": {"ANTHROPIC_AUTH_TOKEN": "sk-theirs"}} + + @pytest.mark.parametrize( + "stored", + [CredentialNotSaved("read-only ~/.litellm"), CredentialNotRecorded()], + ids=["nothing-kept-it", "keychain-took-it-file-refused"], + ) + def test_claude_code_is_configured_even_when_the_cli_could_not_keep_the_credential( + self, tmp_path, monkeypatch, stored + ): + # The key is in hand either way, and --config-claude asked for exactly that key to be written into + # settings.json; whether the CLI's own token file or keychain kept a copy is a separate outcome. + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"], stored=stored) + + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" + assert f"Configured Claude Code: {settings_path}" in result.output + assert "even though the CLI itself could not keep it" in result.output + assert "You can now use the CLI without specifying --api-key" not in result.output + + def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path, monkeypatch): settings_path = tmp_path / "claude" / "settings.json" settings_path.parent.mkdir(parents=True) settings_path.write_text("not json at all {{{") - result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + result, _settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) assert result.exit_code != 0 assert "Login successful!" in result.output @@ -1853,7 +1939,10 @@ class TestPkcePrintToken: assert result.stdout == "" assert sum(len(session.posts) for session in _FakeSession.instances) == 1 assert result.output.count("Could not renew the key") == 1 - assert "Could not renew the key: token request failed with 400: the refresh token was already used" in result.output + assert ( + "Could not renew the key: token request failed with 400: the refresh token was already used" + in result.output + ) assert "Key expired. Run 'lite login --pkce' again." in result.output save.assert_not_called() diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index e5f2a9d95bd..cf52d41e963 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -1,22 +1,40 @@ import json +import os +import pathlib import shlex import stat +import sys import time +from pathlib import Path from unittest.mock import patch import pytest from click.testing import CliRunner -from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord -from litellm.proxy.client.cli import cli +from litellm.litellm_core_utils.private_json import commit_staged_json from litellm.proxy.client.cli.commands.claude_settings import ( + ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, AUTOROUTE_BACKUP_PATH, BACKUP_PATH, + CLAUDE_SETTINGS_PATH, + CONFIGURE_STATE_PATH, + OWNED_ENV_KEYS, + OWNED_TOP_LEVEL_KEYS, SETTINGS_FILE_OWNERS, ClaudeSettingsError, + KeepModel, SettingsFileOwner, - resolve_api_key_helper, - write_claude_settings, + StartOn, + StaticToken, + UnpinModel, + claude_settings_path, + configure_claude_settings, + install_statusline_script, + configure_state_path, + merge_claude_settings, + statusline_command, + unconfigure_claude_settings, + with_status_line, ) @@ -24,92 +42,36 @@ def _owners(*backup_paths): """Stand-in owners for the real `lite up` / `lite autoroute up` registry.""" return tuple(SettingsFileOwner(path, "lite up", "lite down") for path in backup_paths) + CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" -WINDOWS_LITE_EXE = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" - -CMD_METACHARACTERS = frozenset("&|<>^()") -CMD_PERCENT_GUARD = "%%cd:~,%" - - -def _through_cmd_exe(command): - """The line cmd.exe hands to CreateProcess after reading the apiKeyHelper. - - A `"` toggles cmd's quote state and the metacharacters only act outside it. cmd expands - `%VAR%` even inside quotes, so every `%` has to arrive as the `%%cd:~,%` guard: the first - `%` has no variable name and stays literal, and `%cd:~,%` is a zero length substring of `cd`. - """ - assert not any(CMD_METACHARACTERS & set(run) for run in command.split('"')[::2]), command - assert command.count("%") == 3 * command.count(CMD_PERCENT_GUARD), command - return command.replace(CMD_PERCENT_GUARD, "%") - - -def _through_c_runtime(command_line): - """argv as the Microsoft C runtime builds it for the `lite` executable. - - Outside quotes whitespace ends an argument. A `"` toggles quoting, and inside quotes `""` - is a literal quote. Backslashes are literal unless they run up to a `"`, where each pair - is one backslash and an odd one left over makes the quote literal. - """ - argv = [] - current = None - quoted = False - i = 0 - while i < len(command_line): - ch = command_line[i] - if ch in " \t" and not quoted: - if current is not None: - argv.append(current) - current = None - i += 1 - continue - if current is None: - current = "" - if ch == "\\": - run = len(command_line[i:]) - len(command_line[i:].lstrip("\\")) - before_quote = command_line[i + run : i + run + 1] == '"' - current += "\\" * (run // 2 if before_quote else run) - if before_quote and run % 2: - current += '"' - i += 1 - i += run - elif ch == '"': - if quoted and command_line[i + 1 : i + 2] == '"': - current += '"' - i += 1 - else: - quoted = not quoted - i += 1 - else: - current += ch - i += 1 - return argv if current is None else [*argv, current] - - @pytest.fixture def paths(tmp_path): return tmp_path / "claude" / "settings.json", tmp_path / "backup.json" -@pytest.fixture -def lite_on_path(): - with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"): - yield +def _static_configure(base_url, settings_path, owners, state_path=None): + """`lite configure claude --api-key`'s shape: a virtual key as a static token, no pinned model.""" + state = state_path if state_path is not None else settings_path.parent.parent / "state.json" + configure_claude_settings(base_url.rstrip("/"), StaticToken("sk-virtual-key"), KeepModel(), settings_path, state, owners) -class TestWriteClaudeSettings: - def test_creates_the_file_and_its_parent_when_missing(self, paths, lite_on_path): +class TestConfigureClaudeSettings: + def test_creates_the_file_and_its_parent_when_missing(self, paths): settings_path, backup_path = paths assert not settings_path.parent.exists() - write_claude_settings("https://proxy.example.com/", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com/", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" - assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" + assert written["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert "apiKeyHelper" not in written + assert "model" not in written and "ANTHROPIC_MODEL" not in written["env"] - def test_updates_an_existing_file_preserving_unrelated_settings(self, paths, lite_on_path): + def test_updates_an_existing_file_preserving_unrelated_settings(self, paths): settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_text( @@ -118,174 +80,85 @@ class TestWriteClaudeSettings: "theme": "dark", "permissions": {"allow": ["Bash"]}, "env": {"SOME_OTHER_VAR": "keep-me", "ANTHROPIC_BASE_URL": "https://old.example.com"}, - "apiKeyHelper": "old-helper", } ) ) - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["theme"] == "dark" assert written["permissions"] == {"allow": ["Bash"]} assert written["env"]["SOME_OTHER_VAR"] == "keep-me" assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" - assert written["apiKeyHelper"] != "old-helper" - def test_rerunning_against_a_new_proxy_refreshes_both_base_url_and_helper(self, paths, lite_on_path): - settings_path, backup_path = paths - - write_claude_settings("https://first.example.com", settings_path, _owners(backup_path)) - write_claude_settings("https://second.example.com", settings_path, _owners(backup_path)) - - written = json.loads(settings_path.read_text()) - assert written["env"]["ANTHROPIC_BASE_URL"] == "https://second.example.com" - assert "second.example.com" in written["apiKeyHelper"] - assert "first.example.com" not in written["apiKeyHelper"] - - def test_drops_a_stray_static_api_key_so_the_helper_token_wins(self, paths, lite_on_path): + def test_a_helper_left_by_an_older_lite_is_stripped_so_only_the_static_token_is_sent(self, paths): + # Older `lite` versions wrote `apiKeyHelper: lite auth print-token`; Claude Code would keep spawning + # `lite` (and its keychain probe) on every credential refresh, so configure takes the slot over. settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) - settings_path.write_text(json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked"}})) + settings_path.write_text( + json.dumps({"apiKeyHelper": "/usr/local/bin/lite auth print-token", "env": {"ANTHROPIC_API_KEY": "sk-leaked"}}) + ) - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) - assert "ANTHROPIC_API_KEY" not in json.loads(settings_path.read_text())["env"] + written = json.loads(settings_path.read_text()) + assert "apiKeyHelper" not in written + assert "ANTHROPIC_API_KEY" not in written["env"] + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" - def test_written_file_is_owner_only(self, paths, lite_on_path): + def test_written_file_is_owner_only(self, paths): settings_path, backup_path = paths - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert stat.S_IMODE(settings_path.stat().st_mode) == 0o600 - def test_refuses_while_lite_up_holds_a_backup(self, paths, lite_on_path): + def test_refuses_while_lite_up_holds_a_backup(self, paths): settings_path, backup_path = paths backup_path.write_text("{}") with pytest.raises(ClaudeSettingsError, match="lite down"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert not settings_path.exists() - def test_refuses_on_corrupt_existing_settings_without_touching_the_file(self, paths, lite_on_path): + def test_refuses_on_corrupt_existing_settings_without_touching_the_file(self, paths): settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_text("not json at all {{{") with pytest.raises(ClaudeSettingsError, match="invalid JSON"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert settings_path.read_text() == "not json at all {{{" - def test_reports_an_actionable_error_when_lite_is_not_on_path(self, paths): - settings_path, backup_path = paths - with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None): - with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) - - assert not settings_path.exists() - - def test_reports_an_actionable_error_on_a_non_utf8_file(self, paths, lite_on_path): - """Bytes that are not valid UTF-8 must not escape as UnicodeDecodeError. - - UnicodeDecodeError is a ValueError, not an OSError, so a decode-side catch - is easy to miss; login's broad `except Exception` would then relabel it as - an authentication failure and exit 0. - """ + def test_reports_an_actionable_error_on_a_non_utf8_file(self, paths): + # UnicodeDecodeError is a ValueError, not an OSError, so a decode-side catch is easy to miss. settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_bytes(b'{"theme": "\xff\xfe"}') with pytest.raises(ClaudeSettingsError, match="invalid JSON"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) - def test_reports_an_actionable_error_when_the_file_cannot_be_read(self, paths, lite_on_path): - """An unreadable settings file must not surface as "Authentication failed". - - login wraps the whole flow in a broad `except Exception`, so any OSError - escaping this function gets relabelled as an auth failure and sends the - user looking at their SSO config instead of at file permissions. - """ + def test_reports_an_actionable_error_when_the_file_cannot_be_read(self, paths): settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.mkdir() with pytest.raises(ClaudeSettingsError, match="Could not read"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) - def test_reports_an_actionable_error_when_the_file_cannot_be_written(self, paths, lite_on_path): + def test_reports_an_actionable_error_when_the_file_cannot_be_written(self, paths): settings_path, backup_path = paths - with patch( - f"{CLAUDE_SETTINGS_MODULE}.write_private_json", - side_effect=OSError("Read-only file system"), - ): - with pytest.raises(ClaudeSettingsError, match="Read-only file system"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) - - -class TestApiKeyHelperIsActuallyInvocable: - """The helper string is executed verbatim by Claude Code, so it has to parse. - - Asserting only on its text is what let a malformed command (`--base-url`, a - top-level group option, placed after the `print-token` subcommand) ship: click - rejects it with "No such option" and every Claude Code request loses its token. - """ - - def _helper_args(self, base_url): - with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"): - return shlex.split(resolve_api_key_helper(base_url))[1:] - - def test_the_generated_command_parses(self): - result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) - - assert "No such option" not in result.output - assert result.exit_code != 2 - - def test_the_generated_command_reaches_print_token(self): - with patch(f"{AUTH_MODULE}.load_cli_token", return_value=None): - result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) - - assert "Not authenticated" in result.output - - def test_the_generated_command_carries_the_base_url_through(self): - stale = CliTokenRecord( - base_url="http://other-proxy.example.com", - key="sk-stale", - timestamp=time.time(), - ) - with patch(f"{AUTH_MODULE}.load_cli_token", return_value=stale): - result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) - - assert "Not authenticated for this server" in result.output - - def _windows_argv(self, lite_exe, base_url): - with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=lite_exe): - helper = resolve_api_key_helper(base_url, platform="win32") - return _through_c_runtime(_through_cmd_exe(helper)) - - @pytest.mark.parametrize( - ("lite_exe", "base_url"), - [ - (WINDOWS_LITE_EXE, "http://localhost:4000"), - ("C:\\Program Files\\LiteLLM\\lite.EXE", "https://gateway.example.com/?a=1&b=2"), - ("C:\\Users\\u\\Scripts\\lite.EXE", "https://gateway.example.com/team%20a/%7Eproxy"), - ('C:\\odd "dir"\\lite.EXE', "http://localhost:4000/x\\"), - ], - ) - def test_the_windows_command_survives_cmd_exe_and_the_c_runtime(self, lite_exe, base_url): - assert self._windows_argv(lite_exe, base_url) == [lite_exe, "--base-url", base_url, "auth", "print-token"] - - def test_the_windows_command_carries_the_base_url_through_cmd_quoting(self): - stale = CliTokenRecord( - base_url="http://other-proxy.example.com", - key="sk-stale", - timestamp=time.time(), - ) - argv = self._windows_argv(WINDOWS_LITE_EXE, "http://localhost:4000") - with patch(f"{AUTH_MODULE}.load_cli_token", return_value=stale): - result = CliRunner().invoke(cli, argv[1:]) - - assert argv[0] == WINDOWS_LITE_EXE - assert "Not authenticated for this server" in result.output + settings_path.parent.mkdir(parents=True) + settings_path.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not write"): + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + finally: + settings_path.parent.chmod(0o700) + assert not settings_path.exists() class TestConflictingOwnersOfTheSettingsFile: @@ -295,7 +168,7 @@ class TestConflictingOwnersOfTheSettingsFile: write, which is the exact hazard the guard exists to prevent. """ - def test_any_owner_holding_a_backup_blocks_the_write(self, tmp_path, lite_on_path): + def test_any_owner_holding_a_backup_blocks_the_write(self, tmp_path): settings_path = tmp_path / "claude" / "settings.json" for index, owner in enumerate(SETTINGS_FILE_OWNERS): @@ -303,20 +176,20 @@ class TestConflictingOwnersOfTheSettingsFile: backup.write_text("{}") stand_in = SettingsFileOwner(backup, owner.start_command, owner.stop_command) with pytest.raises(ClaudeSettingsError, match="currently managing"): - write_claude_settings("https://proxy.example.com", settings_path, (stand_in,)) + _static_configure("https://proxy.example.com", settings_path, (stand_in,)) backup.unlink() assert not settings_path.exists() - def test_the_error_names_the_owner_that_actually_holds_the_file(self, tmp_path, lite_on_path): + def test_the_error_names_the_owner_that_actually_holds_the_file(self, tmp_path): settings_path = tmp_path / "claude" / "settings.json" backup = tmp_path / "auto.json" backup.write_text("{}") autoroute = SettingsFileOwner(backup, "lite autoroute up", "lite autoroute down") with pytest.raises(ClaudeSettingsError, match="`lite autoroute up` is currently managing"): - write_claude_settings("https://proxy.example.com", settings_path, (autoroute,)) + _static_configure("https://proxy.example.com", settings_path, (autoroute,)) with pytest.raises(ClaudeSettingsError, match="Run `lite autoroute down` first"): - write_claude_settings("https://proxy.example.com", settings_path, (autoroute,)) + _static_configure("https://proxy.example.com", settings_path, (autoroute,)) def test_the_registry_matches_the_paths_the_commands_actually_use(self): """A second definition of the autoroute dir must not drift from this one.""" @@ -328,7 +201,7 @@ class TestConflictingOwnersOfTheSettingsFile: class TestDoesNotDestroyUserOwnedStructure: - def test_writes_through_a_symlinked_settings_file(self, tmp_path, lite_on_path): + def test_writes_through_a_symlinked_settings_file(self, tmp_path): """os.replace() swaps the symlink for a regular file, detaching a dotfiles repo. There is no backup here to undo that, so the link must survive and its @@ -341,19 +214,627 @@ class TestDoesNotDestroyUserOwnedStructure: link.parent.mkdir() link.symlink_to(real) - write_claude_settings("https://proxy.example.com", link, ()) + _static_configure("https://proxy.example.com", link, ()) assert link.is_symlink() assert json.loads(real.read_text())["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" assert json.loads(real.read_text())["theme"] == "dark" - def test_refuses_rather_than_discarding_a_non_object_env(self, paths, lite_on_path): + def test_refuses_rather_than_discarding_a_non_object_env(self, paths): """merge coerces a non-dict env to {}; that is silent data loss on a persistent write.""" settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_text(json.dumps({"theme": "dark", "env": "not-an-object"})) with pytest.raises(ClaudeSettingsError, match="non-object"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert json.loads(settings_path.read_text())["env"] == "not-an-object" + + +class TestClaudeSettingsPath: + def test_defaults_to_the_home_settings_file(self): + assert claude_settings_path({}) == CLAUDE_SETTINGS_PATH + assert claude_settings_path({"CLAUDE_CONFIG_DIR": ""}) == CLAUDE_SETTINGS_PATH + + def test_follows_claude_config_dir_like_claude_code_does(self, tmp_path): + assert claude_settings_path({"CLAUDE_CONFIG_DIR": str(tmp_path)}) == tmp_path / "settings.json" + + def test_expands_a_tilde_in_claude_config_dir(self): + assert claude_settings_path({"CLAUDE_CONFIG_DIR": "~/.claude-work"}) == ( + Path.home() / ".claude-work" / "settings.json" + ) + + +class TestConfigureStatePath: + """Each settings file gets its own undo receipt: the default file keeps the long-standing path, and + a CLAUDE_CONFIG_DIR file gets one keyed by its resolved location, so `lite unconfigure claude` + under one config dir never restores the other file's history.""" + + @pytest.fixture + def default_paths(self, tmp_path): + default_settings = tmp_path / "home" / ".claude" / "settings.json" + default_state = tmp_path / "home" / ".litellm" / "claude_configure_state.json" + with ( + patch(f"{CLAUDE_SETTINGS_MODULE}.CLAUDE_SETTINGS_PATH", default_settings), + patch(f"{CLAUDE_SETTINGS_MODULE}.CONFIGURE_STATE_PATH", default_state), + ): + yield default_settings, default_state + + def test_the_default_file_keeps_the_default_receipt(self, default_paths): + default_settings, default_state = default_paths + assert configure_state_path(default_settings) == default_state + + def test_a_symlink_alias_of_the_default_file_shares_its_receipt(self, default_paths): + default_settings, default_state = default_paths + default_settings.parent.mkdir(parents=True) + alias = default_settings.parent.parent / "claude-alias" + alias.symlink_to(default_settings.parent, target_is_directory=True) + assert configure_state_path(alias / "settings.json") == default_state + + def test_another_settings_file_gets_a_receipt_of_its_own_beside_the_default_one(self, default_paths, tmp_path): + _default_settings, default_state = default_paths + work_state = configure_state_path(tmp_path / "work" / "settings.json") + play_state = configure_state_path(tmp_path / "play" / "settings.json") + assert work_state != default_state and play_state != default_state + assert work_state != play_state + assert work_state.parent == play_state.parent == default_state.parent / "claude_configure_state" + assert work_state == configure_state_path(tmp_path / "work" / "settings.json") + + def test_configure_and_unconfigure_under_a_config_dir_leave_the_default_receipt_alone( + self, default_paths, tmp_path + ): + _default_settings, default_state = default_paths + work_settings = tmp_path / "work" / "settings.json" + work_state = configure_state_path(work_settings) + configure_claude_settings( + "https://proxy.example.com", StaticToken("sk-virtual-key"), KeepModel(), work_settings, work_state, () + ) + assert work_state.exists() and not default_state.exists() + outcome = unconfigure_claude_settings(work_settings, work_state, ()) + assert outcome.file_removed and not work_settings.exists() + assert not work_state.exists() + + +class TestMergeClaudeSettings: + """One merge for every way Claude Code gets wired: `lite up`, `lite configure claude` and `lite autoroute up`.""" + + def test_a_static_token_lands_in_env_and_the_helper_slot_is_cleared(self): + settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token", "env": {"ANTHROPIC_API_KEY": "leaked"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000/", StaticToken("token-abc")) + assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" + assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert "ANTHROPIC_API_KEY" not in merged["env"] + assert "apiKeyHelper" not in merged + assert "model" not in merged + assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS) + + def test_keeps_existing_switch_values_and_unrelated_keys_without_mutating_the_input(self): + settings = {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", StaticToken("token-abc")) + assert merged["theme"] == "dark" + assert merged["env"]["SOME_OTHER_VAR"] == "value" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" + assert settings == {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}} + + def test_a_default_model_pins_the_starting_row_and_the_model_a_resumed_session_keeps(self): + # `model` is only the row Claude Code starts on: a resumed session re-sends the model its transcript + # recorded, which behind a raw-model auto-router is the tier model (403 for a key scoped to the + # router). ANTHROPIC_MODEL outranks the transcript on resume, so the pin has to land there too. + merged = merge_claude_settings( + {}, "http://127.0.0.1:4000", StaticToken("token-abc"), default_model="claude-auto" + ) + assert merged["model"] == "claude-auto" + assert merged["env"]["ANTHROPIC_MODEL"] == "claude-auto" + assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS) + + def test_without_a_default_model_neither_pin_is_written_and_a_users_own_stays(self): + settings = {"model": "mine", "env": {"ANTHROPIC_MODEL": "mine-too"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", StaticToken("token-abc")) + assert merged["model"] == "mine" and merged["env"]["ANTHROPIC_MODEL"] == "mine-too" + + def test_a_tier_model_forces_every_claude_code_tier_as_autoroute_needs(self): + # Router's auto-router registry is keyed by the literal requested model string with no + # wildcard resolution, so `lite autoroute up` overrides the env var each tier reads. + settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}} + merged = merge_claude_settings( + settings, "http://127.0.0.1:4000", StaticToken("token-abc"), tier_model="autorouter" + ) + assert {merged["env"][key] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS} == {"autorouter"} + assert "model" not in merged + + def test_touches_exactly_the_declared_owned_keys(self): + # The receipt and unconfigure restore exactly OWNED_*_KEYS, so a key the merge writes outside + # that table would be written by configure and never undone. + settings = { + "theme": "dark", + "permissions": {"allow": ["Bash"]}, + "env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "old", "ENABLE_TOOL_SEARCH": "false"}, + "apiKeyHelper": "old-helper", + "model": "old-model", + } + for credential in (StaticToken("token-abc"), StaticToken("token-rotated")): + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", credential, default_model="claude-auto") + changed_top_level = {key for key in set(settings) | set(merged) if settings.get(key) != merged.get(key)} + assert changed_top_level - {"env"} <= set(OWNED_TOP_LEVEL_KEYS) + changed_env = { + key + for key in set(settings["env"]) | set(merged["env"]) + if settings["env"].get(key) != merged["env"].get(key) + } + assert changed_env <= set(OWNED_ENV_KEYS) + assert merged["permissions"] == {"allow": ["Bash"]} + assert merged["env"]["KEEP_ME"] == "1" + + +PROXY = "http://127.0.0.1:4000" +ANTHROPIC = "https://api.anthropic.com" +RELOGIN = StaticToken("sk-fresh-login") +ORIGINAL = { + "theme": "dark", + "permissions": {"allow": ["Bash"]}, + "env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "sk-ant-mine", "ANTHROPIC_BASE_URL": ANTHROPIC}, + "apiKeyHelper": "/usr/local/bin/lite auth print-token", + "model": "claude-opus-5", +} + + +def _set(path, value): + """A user edit: set (or with `_ABSENT`, remove) the key at a dotted path in the settings file.""" + + def edit(settings): + section, _, key = path.rpartition(".") + container = settings.setdefault(section, {}) if section else settings + if value is _ABSENT: + container.pop(key, None) + else: + container[key] = value + return settings + + return edit + + +_ABSENT = object() + + +class _Rig: + """One settings file plus receipt under tmp_path, driven through the public functions only.""" + + def __init__(self, tmp_path, initial): + self.settings = tmp_path / "claude" / "settings.json" + self.state = tmp_path / "state" / "claude_configure_state.json" + if initial is not None: + self.settings.parent.mkdir(parents=True) + self.settings.write_text(json.dumps(initial)) + + def read(self): + return json.loads(self.settings.read_text()) if self.settings.exists() else None + + def configure(self, credential=StaticToken("sk-virtual-key"), model=StartOn("claude-auto"), **kwargs): + configure_claude_settings(PROXY, credential, model, self.settings, self.state, (), **kwargs) + + def edit(self, *edits): + settings = self.read() + for apply in edits: + settings = apply(settings) + self.settings.write_text(json.dumps(settings)) + + def unconfigure(self): + return unconfigure_claude_settings(self.settings, self.state, ()) + + +# Each row: initial file, steps (configure kwargs dicts or edit callables) between the first configure +# and unconfigure, the expected file afterwards, and the expected outcome fields. Sequences that used +# to be one test each; the receipt's rules are what make them all come out right. +UNDO_SCENARIOS = { + "plain round trip": (ORIGINAL, [], ORIGINAL, {"kept": ()}), + "no file before": (None, [], None, {"file_removed": True}), + "no env before": ({"theme": "dark"}, [], {"theme": "dark"}, {}), + "null env before": ({"theme": "dark", "env": None}, [], {"theme": "dark", "env": None}, {}), + "empty env before": ({"theme": "dark", "env": {}}, [], {"theme": "dark", "env": {}}, {}), + "user edits stay and are named": ( + ORIGINAL, + [_set("env.ENABLE_TOOL_SEARCH", "false"), _set("model", "claude-sonnet-4-6")], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "claude-sonnet-4-6"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}, "withheld": ()}, + ), + "user filled an env configure created": (None, [_set("env.MY_VAR", "mine")], {"env": {"MY_VAR": "mine"}}, {}), + "user deleted the file": (None, [lambda s: None], None, {"file_removed": True, "restored": (), "kept": ()}), + "user removed our key: neither restored nor kept": ( + ORIGINAL, + [_set("env.ANTHROPIC_AUTH_TOKEN", _ABSENT)], + ORIGINAL, + {"not_restored": {"env.ANTHROPIC_AUTH_TOKEN"}, "kept": ()}, + ), + "restored names only what changed": ( + {"model": "claude-opus-5"}, + [], + {"model": "claude-opus-5"}, + { + "restored": { + "env.ANTHROPIC_BASE_URL", + "env.ANTHROPIC_AUTH_TOKEN", + "env.ENABLE_TOOL_SEARCH", + "env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", + "statusLine", + }, + "kept": (), + }, + {"credential": RELOGIN, "model": KeepModel()}, + ), + "repeat across credential kinds keeps the first snapshot": ( + ORIGINAL, + [ + {"credential": RELOGIN, "model": UnpinModel()}, + {"credential": StaticToken("sk-rotated"), "model": StartOn("claude-sonnet-4-6")}, + ], + ORIGINAL, + {}, + ), + "repeat without a model lets go of our pin, user had none": ({}, [{"model": UnpinModel()}], {}, {}), + "repeat without a model lets go of our pin, user had one": ( + {"model": "claude-opus-5"}, + [{"model": UnpinModel()}], + {"model": "claude-opus-5"}, + {}, + ), + "re-login keeps our pin": (None, [{"credential": RELOGIN, "model": KeepModel()}], None, {"file_removed": True}), + "edit between configures survives an unpin repeat": ( + ORIGINAL, + [ + _set("model", "my-favourite"), + _set("env.ENABLE_TOOL_SEARCH", "false"), + {"credential": RELOGIN, "model": UnpinModel()}, + ], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, + ), + "edit between configures survives a re-login": ( + ORIGINAL, + [ + _set("model", "my-favourite"), + _set("env.ENABLE_TOOL_SEARCH", "false"), + {"credential": RELOGIN, "model": KeepModel()}, + ], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, + ), + "edit between configures: a same-model repeat displaces it, so it is what comes back": ( + ORIGINAL, + [_set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), {"credential": RELOGIN}], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH"}, "restored_includes": {"model"}}, + ), + "base URL changed since: credentials withheld, receipt kept": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")], + {**ORIGINAL, "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, "apiKeyHelper": _ABSENT}, + { + "withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC), ("apiKeyHelper", ANTHROPIC)}, + "kept": {"env.ANTHROPIC_BASE_URL"}, + "receipt_kept": True, + }, + ), + "base URL changed and back: judged against the URL the restored file holds": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", ANTHROPIC)], + ORIGINAL, + {"withheld": ()}, + ), + "credential captured beside no URL goes back only beside no URL": ( + {"env": {"ANTHROPIC_API_KEY": "sk-default-endpoint"}}, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")], + {"env": {"ANTHROPIC_BASE_URL": "http://other-proxy:4000"}}, + { + "withheld": {("env.ANTHROPIC_API_KEY", "no ANTHROPIC_BASE_URL (Anthropic's default endpoint)")}, + "receipt_kept": True, + }, + ), + "restored document empty while a credential is withheld: file goes, receipt stays": ( + None, + [ + _set("env.ANTHROPIC_API_KEY", "sk-user"), + {"credential": RELOGIN, "model": KeepModel()}, + _set("env.ANTHROPIC_BASE_URL", _ABSENT), + ], + None, + {"withheld": {("env.ANTHROPIC_API_KEY", PROXY)}, "file_removed": True, "receipt_kept": True}, + {"credential": RELOGIN, "model": KeepModel()}, + ), + "a credential the user changed is kept, never also withheld": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000"), _set("apiKeyHelper", "/opt/mine/helper")], + { + **ORIGINAL, + "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, + "apiKeyHelper": "/opt/mine/helper", + }, + { + "withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC)}, + "kept": {"env.ANTHROPIC_BASE_URL", "apiKeyHelper"}, + "receipt_kept": True, + }, + ), +} + + +def _expected_file(expected): + if expected is None: + return None + return {k: v for k, v in expected.items() if v is not _ABSENT} + + +class TestConfigureAndUnconfigure: + """`configure_claude_settings` records how to undo itself; `unconfigure_claude_settings` undoes only that.""" + + @pytest.mark.parametrize("scenario", UNDO_SCENARIOS.values(), ids=UNDO_SCENARIOS.keys()) + def test_undo_matrix(self, tmp_path, scenario): + initial, steps, expected, outcome_expectations, *first = scenario + rig = _Rig(tmp_path, initial) + rig.configure(**(first[0] if first else {})) + for step in steps: + if isinstance(step, dict): + rig.configure(**step) + elif rig.settings.exists() and step(json.loads(rig.settings.read_text())) is None: + rig.settings.unlink() + else: + rig.edit(step) + + outcome = rig.unconfigure() + + assert rig.read() == _expected_file(expected) + assert rig.state.exists() == outcome_expectations.get("receipt_kept", False) + for field, want in outcome_expectations.items(): + if field == "withheld": + assert {(item.key, item.endpoint) for item in outcome.withheld} == set(want) + elif field == "not_restored": + assert not set(want) & set(outcome.restored) and not set(want) & set(outcome.kept) + elif field == "restored_includes": + assert set(want) <= set(outcome.restored) + elif field in ("restored", "kept"): + assert set(getattr(outcome, field)) == set(want) + elif field != "receipt_kept": + assert getattr(outcome, field) == want + assert not {item.key for item in outcome.withheld} & set(outcome.kept) + + def test_configure_writes_owner_only_and_the_receipt_never_holds_the_key(self, tmp_path): + rig = _Rig(tmp_path, ORIGINAL) + rig.configure(credential=StaticToken("sk-virtual-key-never-on-disk-twice")) + configured = rig.read() + assert configured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key-never-on-disk-twice" + assert configured["env"]["ANTHROPIC_BASE_URL"] == PROXY and configured["model"] == "claude-auto" + assert "ANTHROPIC_API_KEY" not in configured["env"] and "apiKeyHelper" not in configured + assert stat.S_IMODE(rig.settings.stat().st_mode) == 0o600 == stat.S_IMODE(rig.state.stat().st_mode) + assert "sk-virtual-key-never-on-disk-twice" not in rig.state.read_text() + + def test_withheld_credentials_come_back_once_the_url_points_at_their_server_again(self, tmp_path): + # The kept receipt owns only the withheld slots: the second unconfigure restores exactly those. + rig = _Rig(tmp_path, ORIGINAL) + rig.configure() + rig.edit(_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")) + rig.unconfigure() + rig.edit(_set("env.ANTHROPIC_BASE_URL", ANTHROPIC), _set("theme", "light")) + outcome = rig.unconfigure() + assert rig.read() == {**ORIGINAL, "theme": "light"} + assert set(outcome.restored) == {"env.ANTHROPIC_API_KEY", "apiKeyHelper"} + assert outcome.kept == () and outcome.withheld == () and not rig.state.exists() + + @pytest.mark.parametrize( + ("path", "value", "repeat_credential"), + [ + ("env.ANTHROPIC_API_KEY", "sk-user-added-later", RELOGIN), + ("apiKeyHelper", "/opt/mine/helper", StaticToken("sk-rotated")), + ], + ids=["user-adds-api-key", "user-sets-own-helper"], + ) + def test_a_credential_the_user_set_between_two_configures_is_what_comes_back( + self, tmp_path, path, value, repeat_credential + ): + # The repeat's merge clears the slot, so the displaced value is snapshotted and is what returns; + # it was set while the file pointed at the proxy, so it returns once the file points there again. + rig = _Rig(tmp_path, {"theme": "dark"}) + rig.configure(credential=RELOGIN, model=KeepModel()) + rig.edit(_set(path, value)) + rig.configure(credential=repeat_credential, model=KeepModel()) + assert not _lookup(rig.read(), path) + + outcome = rig.unconfigure() + assert [(item.key, item.endpoint) for item in outcome.withheld] == [(path, PROXY)] + assert rig.read() == {"theme": "dark"} and rig.state.exists() + rig.settings.write_text(json.dumps({"theme": "dark", "env": {"ANTHROPIC_BASE_URL": PROXY}})) + outcome = rig.unconfigure() + assert _lookup(rig.read(), path) == value + assert outcome.restored == (path,) and outcome.withheld == () and not rig.state.exists() + + def test_a_receipt_commit_that_fails_leaves_no_staged_token_behind(self, tmp_path): + rig = _Rig(tmp_path, {}) + + def commit_receipt_fails(staged, path): + if path == str(rig.state): + os.unlink(staged) + raise OSError("receipt rename failed") + commit_staged_json(staged, path) + + with pytest.raises(ClaudeSettingsError, match=r"Could not write .*receipt rename failed"): + rig.configure(credential=StaticToken("sk-never-left-in-a-temp-file"), commit=commit_receipt_fails) + assert not list(rig.settings.parent.glob(".tmp-*")) and not list(rig.state.parent.glob(".tmp-*")) + assert rig.read() == {} and not rig.state.exists() + + @pytest.mark.parametrize("configured_before", [False, True], ids=["first-configure", "repeat-configure"]) + def test_a_settings_commit_that_fails_after_the_receipt_landed_puts_the_receipt_back( + self, tmp_path, configured_before + ): + # The two renames are not atomic: a settings rename that fails after the receipt landed must + # not leave a receipt describing settings that were never written. + rig = _Rig(tmp_path, ORIGINAL) + if configured_before: + rig.configure() + receipt_before = rig.state.read_text() if configured_before else None + settings_before = rig.settings.read_text() + + def commit_settings_fails(staged, path): + if path == str(rig.settings): + os.unlink(staged) + raise OSError("rename failed") + commit_staged_json(staged, path) + + with pytest.raises(ClaudeSettingsError, match="rename failed"): + rig.configure(credential=StaticToken("sk-rotated"), commit=commit_settings_fails) + assert rig.settings.read_text() == settings_before + assert (rig.state.read_text() if rig.state.exists() else None) == receipt_before + if configured_before: + rig.unconfigure() + assert rig.read() == ORIGINAL + + def test_a_failed_repeat_configure_leaves_the_earlier_undo_intact(self, tmp_path): + rig = _Rig(tmp_path, ORIGINAL) + rig.configure() + receipt_before = rig.state.read_text() + rig.settings.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not write"): + rig.configure(credential=StaticToken("sk-rotated")) + finally: + rig.settings.parent.chmod(0o700) + assert rig.state.read_text() == receipt_before and not list(rig.state.parent.glob(".tmp-*")) + assert rig.read()["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" + rig.unconfigure() + assert rig.read() == ORIGINAL + + def test_unconfigure_reports_a_receipt_it_cannot_remove_as_a_settings_error(self, tmp_path): + rig = _Rig(tmp_path, ORIGINAL) + rig.configure() + rig.state.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not remove"): + rig.unconfigure() + finally: + rig.state.parent.chmod(0o700) + + def test_configure_writes_through_a_symlinked_settings_file(self, tmp_path): + target = tmp_path / "dotfiles" / "settings.json" + target.parent.mkdir() + target.write_text(json.dumps({"theme": "dark"})) + link = tmp_path / "settings.json" + link.symlink_to(target) + configure_claude_settings(PROXY, StaticToken("sk-virtual-key"), UnpinModel(), link, tmp_path / "state.json", ()) + assert link.is_symlink() + assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" + + @pytest.mark.parametrize("operation", ["configure", "unconfigure"]) + def test_refuses_while_a_temporary_owner_holds_a_backup(self, paths, tmp_path, operation): + settings_path, backup_path = paths + backup_path.write_text("{}") + owners = _owners(backup_path) + state = tmp_path / "state.json" + attempt = ( + (lambda: configure_claude_settings(PROXY, StaticToken("k"), UnpinModel(), settings_path, state, owners)) + if operation == "configure" + else (lambda: unconfigure_claude_settings(settings_path, state, owners)) + ) + with pytest.raises(ClaudeSettingsError, match="lite down"): + attempt() + assert not settings_path.exists() + + def test_unconfigure_without_a_receipt_is_an_error_not_a_silent_no_op(self, tmp_path): + with pytest.raises(ClaudeSettingsError, match="nothing to undo"): + _Rig(tmp_path, None).unconfigure() + + +def _lookup(settings, path): + section, _, key = path.rpartition(".") + return (settings.get(section) or {}).get(key) if section else settings.get(key) + + +class TestStatusLine: + """Every configure registers the status line, and only while the slot is empty or already ours.""" + + COMMAND = "/opt/lite/bin/python /Users/me/.litellm/statusline.py" + + def test_an_empty_slot_gets_our_status_line(self): + assert with_status_line({}, self.COMMAND)["statusLine"] == {"type": "command", "command": self.COMMAND} + + def test_a_users_own_status_line_is_never_replaced(self): + theirs = {"type": "command", "command": "~/.claude/my-statusline.sh"} + assert with_status_line({"statusLine": theirs}, self.COMMAND)["statusLine"] == theirs + + def test_ours_under_an_older_interpreter_is_refreshed(self): + stale = {"type": "command", "command": "/old/python /Users/me/.litellm/statusline.py"} + assert with_status_line({"statusLine": stale}, self.COMMAND)["statusLine"]["command"] == self.COMMAND + + def test_the_merge_carries_it(self): + merged = merge_claude_settings({}, PROXY, StaticToken("tok"), status_line=self.COMMAND) + assert merged["statusLine"] == {"type": "command", "command": self.COMMAND} + + def test_the_installed_script_is_the_bundled_one_and_the_command_runs_this_interpreter(self, tmp_path): + from litellm.proxy.client.cli.commands import statusline_script + + script = tmp_path / "lite" / "statusline.py" + command = install_statusline_script(script) + assert script.read_bytes() == pathlib.Path(statusline_script.__file__).read_bytes() + assert shlex.split(command) == [sys.executable, str(script)] + assert command == statusline_command(script) + assert stat.S_IMODE(script.stat().st_mode) == 0o600 + assert stat.S_IMODE(script.parent.stat().st_mode) == 0o700 + assert install_statusline_script(script) == command + + def test_a_reinstall_replaces_the_script_in_one_step_and_a_refused_one_leaves_the_old_script_whole(self, tmp_path): + # Claude Code may be running the script at the moment `lite` reinstalls it; the file it has open + # must stay complete, and a reinstall that cannot land must not leave a truncated script behind. + from litellm.proxy.client.cli.commands import statusline_script + + script = tmp_path / "lite" / "statusline.py" + install_statusline_script(script) + bundled = pathlib.Path(statusline_script.__file__).read_bytes() + with script.open("rb") as running: + install_statusline_script(script) + assert running.read() == bundled + assert [child.name for child in script.parent.iterdir()] == ["statusline.py"] + + if os.geteuid() != 0: + script.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not install the status line script"): + install_statusline_script(script) + finally: + script.parent.chmod(0o700) + assert script.read_bytes() == bundled + + def test_configure_installs_it_and_unconfigure_removes_only_ours(self, tmp_path): + rig = _Rig(tmp_path, {"theme": "dark"}) + script = tmp_path / "statusline.py" + rig.configure(script_path=script) + assert rig.read()["statusLine"]["command"] == statusline_command(script) + assert script.exists() + + outcome = rig.unconfigure() + assert rig.read() == {"theme": "dark"} + assert "statusLine" in outcome.restored + + def test_a_status_line_the_user_replaced_after_configure_survives_unconfigure(self, tmp_path): + rig = _Rig(tmp_path, None) + rig.configure(model=KeepModel(), script_path=tmp_path / "statusline.py") + theirs = {"type": "command", "command": "~/.claude/my-statusline.sh"} + rig.edit(lambda settings: {**settings, "statusLine": theirs}) + + outcome = rig.unconfigure() + assert rig.read()["statusLine"] == theirs + assert "statusLine" in outcome.kept + + def test_a_receipt_from_before_the_status_line_existed_still_unconfigures(self, tmp_path): + # Older receipts never claimed statusLine; a key no configure wrote is never ours, so it stays. + rig = _Rig(tmp_path, None) + script = tmp_path / "statusline.py" + rig.configure(model=KeepModel(), script_path=script) + receipt = json.loads(rig.state.read_text()) + receipt["written"].pop("statusLine") + receipt["previous"].pop("statusLine") + rig.state.write_text(json.dumps(receipt)) + + outcome = rig.unconfigure() + restored = rig.read() + assert "ANTHROPIC_AUTH_TOKEN" not in restored.get("env", {}) + assert restored["statusLine"]["command"] == statusline_command(script) + assert "statusLine" not in outcome.restored diff --git a/tests/test_litellm/proxy/client/cli/test_configure_commands.py b/tests/test_litellm/proxy/client/cli/test_configure_commands.py new file mode 100644 index 00000000000..8ed188af737 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_configure_commands.py @@ -0,0 +1,421 @@ +import json +import os +import stat + +import click +import pytest +import requests +import responses +from click.testing import CliRunner + +from litellm.proxy.client.cli import cli +from litellm.proxy.client.cli.commands import claude_settings as claude_settings_module +from litellm.proxy.client.cli.commands import configure as configure_module +from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner +from litellm.proxy.client.cli.commands.configure import configure_claude, configure_group, interactive_configure + +PROXY = "http://proxy.test:4000" +VALID_KEY = "sk-virtual-key" +LISTED_MODELS = ("claude-auto", "gpt-5.6-luna") + + +def _mock_models(): + responses.get( + f"{PROXY}/v1/models", + json={"data": [{"id": model, "object": "model"} for model in LISTED_MODELS]}, + match=[responses.matchers.header_matcher({"Authorization": f"Bearer {VALID_KEY}"})], + ) + responses.get(f"{PROXY}/v1/models", status=401) + + +@pytest.fixture +def paths(monkeypatch, tmp_path): + """The default settings file, reached the way Claude Code reaches it: CLAUDE_CONFIG_DIR names its directory.""" + settings_path = tmp_path / "claude" / "settings.json" + state_path = tmp_path / "litellm" / "claude_configure_state.json" + monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(settings_path.parent)) + monkeypatch.setattr(claude_settings_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude_settings_module, "CONFIGURE_STATE_PATH", state_path) + return settings_path, state_path + + +@pytest.fixture +def lite_on_path(monkeypatch, tmp_path): + """A real `lite` executable on PATH, so the apiKeyHelper command resolves without patching.""" + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + lite = bin_dir / "lite" + lite.write_text("#!/bin/sh\nexit 0\n") + lite.chmod(lite.stat().st_mode | stat.S_IXUSR) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ.get('PATH', '')}") + return str(lite) + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def lite_up_backup(monkeypatch, tmp_path): + """A `lite up` session holding its backup, the local precondition every settings write refuses on.""" + backup = tmp_path / "claude_settings_backup.json" + backup.write_text("{}") + monkeypatch.setattr( + claude_settings_module, "SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup, "lite up", "lite down"),) + ) + return backup + + +def _configure(runner, *args): + return runner.invoke(cli, ["--base-url", PROXY, "configure", "claude", *args]) + + +class TestConfigureClaudeWithAVirtualKey: + @responses.activate + def test_writes_settings_and_reports_without_echoing_the_key(self, runner, paths): + _mock_models() + settings_path, state_path = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto") + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_BASE_URL"] == PROXY + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + assert written["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert written["model"] == "claude-auto" + assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in written["env"] + assert state_path.exists() + assert VALID_KEY not in result.output + assert written["env"]["ANTHROPIC_MODEL"] == "claude-auto" + assert "Starting model: claude-auto" in result.output + assert "1 of the proxy's 2 models" in result.output + assert "lite unconfigure claude" in result.output + assert [call.request.headers.get("x-gateway-client") for call in responses.calls] == ["claude-code"] + + @responses.activate + def test_takes_the_key_from_the_global_option_and_keeps_claude_codes_default(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = runner.invoke(cli, ["--base-url", PROXY, "--api-key", VALID_KEY, "configure", "claude"]) + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + assert "model" not in written and "ANTHROPIC_MODEL" not in written["env"] + assert "Starting model: not pinned" in result.output + + @responses.activate + def test_refuses_a_model_the_proxy_does_not_list(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-nope") + assert result.exit_code != 0 + assert "'claude-nope' is not served" in result.output + assert "claude-auto, gpt-5.6-luna" in result.output + assert not settings_path.exists() + + @responses.activate + def test_refuses_a_key_the_proxy_rejects(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = _configure(runner, "--api-key", "sk-wrong") + assert result.exit_code != 0 + assert "rejected your key (HTTP 401)" in result.output + assert not settings_path.exists() + + @responses.activate + @pytest.mark.parametrize( + ("mock", "expected", "unexpected"), + [ + ( + lambda: responses.get(f"{PROXY}/v1/models", body=requests.ConnectionError("refused")), + "Is the proxy at", + "answered", + ), + ( + lambda: responses.get(f"{PROXY}/v1/models", status=500), + "The proxy at http://proxy.test:4000 answered", + "Is the proxy at", + ), + ( + lambda: responses.get(f"{PROXY}/v1/models", body="not json"), + "answered, so check that it is a LiteLLM proxy", + "Is the proxy at", + ), + ( + lambda: responses.get(f"{PROXY}/v1/models", json={"data": []}), + "Claude Code would have nothing to run", + "Is the proxy at", + ), + ], + ids=["unreachable", "http-500", "non-json-body", "empty-list"], + ) + def test_the_listing_hint_matches_how_the_listing_failed(self, runner, paths, mock, expected, unexpected): + # Only a proxy that never answered gets the "is it running" question; a 500, a non-JSON body or an + # empty list prove it is up, and the hint says so instead. + mock() + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY) + assert result.exit_code != 0 + assert expected in result.output and unexpected not in result.output + assert not settings_path.exists() + + @responses.activate + @pytest.mark.parametrize("entry", ["virtual-key", "no-key", "interactive"]) + def test_refuses_while_lite_up_holds_a_backup_before_any_request(self, runner, paths, lite_up_backup, entry): + _mock_models() + if entry == "interactive": + ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY}) + with pytest.raises(click.ClickException, match="lite down"): + interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=lambda listed: None) + else: + args = ["--api-key", VALID_KEY] if entry == "virtual-key" else [] + result = runner.invoke(configure_claude, args, obj={"base_url": PROXY, "api_key": None}) + assert result.exit_code != 0 and "lite down" in result.output + assert len(responses.calls) == 0 + assert not paths[0].exists() + + @responses.activate + def test_says_so_when_the_key_is_written_through_a_symlink(self, runner, paths, tmp_path): + _mock_models() + settings_path, _ = paths + target = tmp_path / "dotfiles" / "settings.json" + target.parent.mkdir() + target.write_text("{}") + settings_path.parent.mkdir(parents=True) + settings_path.symlink_to(target) + result = _configure(runner, "--api-key", VALID_KEY) + assert result.exit_code == 0, result.output + assert "keep it out of version control" in result.output + assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + + +class TestConfigureClaudeWithoutAKey: + @responses.activate + def test_refuses_and_names_the_ways_to_pass_a_key_without_writing_or_logging_in(self, runner, paths): + # A `lite login` credential expires within a day; the old fallback wrote an apiKeyHelper that made + # Claude Code spawn `lite` (and its keychain probe) on every credential refresh. + _mock_models() + settings_path, state_path = paths + result = runner.invoke( + configure_claude, + ["--model", "claude-auto"], + obj={"base_url": PROXY, "api_key": "sk-login-jwt", "api_key_from_token_file": True}, + ) + assert result.exit_code != 0 + assert "--api-key" in result.output and "LITELLM_PROXY_API_KEY" in result.output + assert "apiKeyHelper" not in result.output + assert not settings_path.exists() and not state_path.exists() + assert len(responses.calls) == 0 + + @responses.activate + def test_an_explicit_key_still_wins_over_a_stored_login(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = runner.invoke( + configure_claude, + ["--api-key", VALID_KEY], + obj={"base_url": PROXY, "api_key": "sk-login-jwt", "api_key_from_token_file": True}, + ) + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY and "apiKeyHelper" not in written + + +class TestInteractiveConfigure: + @responses.activate + def test_asks_for_targets_and_a_starting_model_then_configures(self, paths): + _mock_models() + settings_path, _ = paths + asked = {} + + def pick_model(listed): + asked["listed"] = tuple(listed) + return "claude-auto" + + ctx = click.Context( + configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} + ) + interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=pick_model) + assert asked["listed"] == LISTED_MODELS + assert json.loads(settings_path.read_text())["model"] == "claude-auto" + + def test_does_nothing_when_claude_code_is_not_picked(self, paths): + settings_path, _ = paths + ctx = click.Context( + configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} + ) + interactive_configure(ctx, pick_targets=lambda: (), pick_model=lambda listed: None) + assert not settings_path.exists() + + def test_bare_configure_without_a_terminal_names_the_non_interactive_command(self, runner, paths): + result = runner.invoke(cli, ["--base-url", PROXY, "configure"]) + assert result.exit_code != 0 + assert "lite configure claude --api-key" in result.output + + +class TestUnconfigureClaude: + @responses.activate + def test_restores_the_original_file_and_removes_the_receipt(self, runner, paths): + _mock_models() + settings_path, state_path = paths + settings_path.parent.mkdir(parents=True) + original = {"theme": "dark", "model": "claude-opus-5"} + settings_path.write_text(json.dumps(original)) + assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0 + + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert json.loads(settings_path.read_text()) == original + assert not state_path.exists() + assert "Restored in" in result.output and "model" in result.output + assert "ANTHROPIC_API_KEY" not in result.output, "a key that never existed was not restored" + + @responses.activate + def test_a_file_only_configure_created_is_reported_removed_not_restored(self, runner, paths): + _mock_models() + settings_path, _ = paths + assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0 + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert not settings_path.exists() + assert "No settings file remains" in result.output and "Restored" not in result.output + + @responses.activate + def test_says_when_nothing_was_still_ours_and_names_what_it_kept(self, runner, paths): + _mock_models() + settings_path, _ = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text(json.dumps({"theme": "dark"})) + assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0 + edited = json.loads(settings_path.read_text()) + edited["env"] = {key: f"{value}-edited" for key, value in edited["env"].items()} + edited["model"] = "mine" + edited["statusLine"] = {"type": "command", "command": "~/.claude/my-statusline.sh"} + settings_path.write_text(json.dumps(edited)) + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert "Nothing in" in result.output and "was still ours to restore" in result.output + assert "Left as you changed them since:" in result.output and "model" in result.output + assert "statusLine" in result.output + + @responses.activate + def test_names_the_server_a_withheld_credential_was_captured_with_and_keeps_the_receipt(self, runner, paths): + _mock_models() + settings_path, state_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text( + json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://api.anthropic.com", "ANTHROPIC_API_KEY": "sk-ant"}}) + ) + assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0 + edited = json.loads(settings_path.read_text()) + edited["env"]["ANTHROPIC_BASE_URL"] = "http://other-proxy:4000" + settings_path.write_text(json.dumps(edited)) + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert "env.ANTHROPIC_API_KEY (captured with https://api.anthropic.com)" in result.output + assert str(state_path) in result.output and state_path.exists() + assert "sk-ant" not in result.output + + def test_refuses_while_lite_up_holds_a_backup(self, runner, paths, lite_up_backup): + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code != 0 and "lite down" in result.output + + @responses.activate + def test_a_config_dir_is_configured_and_undone_apart_from_the_default_file( + self, runner, paths, monkeypatch, tmp_path, lite_up_backup + ): + _mock_models() + default_settings, default_state = paths + work_dir = tmp_path / "claude-work" + work_dir.mkdir() + original = {"theme": "dark"} + (work_dir / "settings.json").write_text(json.dumps(original)) + monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(work_dir)) + + configured = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto") + assert configured.exit_code == 0, configured.output + assert f"Configured Claude Code: {work_dir / 'settings.json'}" in configured.output + assert json.loads((work_dir / "settings.json").read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + assert not default_settings.exists() and not default_state.exists() + + undone = runner.invoke(cli, ["unconfigure", "claude"]) + assert undone.exit_code == 0, undone.output + assert json.loads((work_dir / "settings.json").read_text()) == original + assert not default_settings.exists() and not default_state.exists() + assert runner.invoke(cli, ["unconfigure", "claude"]).exit_code != 0, "the receipt is gone with the undo" + + def test_without_a_receipt_it_fails_loudly(self, runner, paths): + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code != 0 + assert "nothing to undo" in result.output + + +class TestClaudeCodeView: + VIEW = {"anthropic-version": "2023-06-01", "x-gateway-client": "claude-code"} + + def _mock(self, rows): + responses.get( + f"{PROXY}/v1/models", + json={"data": rows}, + match=[responses.matchers.header_matcher({"Authorization": f"Bearer {VALID_KEY}", **self.VIEW})], + ) + + @responses.activate + @pytest.mark.parametrize( + "model, pinned", + [ + ("literal-claude-router-source", "emitted-literal"), + ("marked-sibling", "emitted-marked[1m]"), + ("emitted-collision", "emitted-source-priority"), + ("emitted-only", "emitted-only"), + ], + ) + def test_pins_source_identity_before_emitted_id(self, runner, paths, model, pinned): + self._mock( + [ + {"id": "emitted-collision", "source_model": "other-source"}, + {"id": "emitted-source-priority", "source_model": "emitted-collision"}, + {"id": "emitted-marked[1m]", "source_model": "marked-sibling"}, + {"id": "emitted-literal", "source_model": "literal-claude-router-source"}, + {"id": "emitted-only"}, + ] + ) + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", model) + assert result.exit_code == 0, result.output + assert json.loads(settings_path.read_text())["model"] == pinned + assert f"Starting model: {pinned}" in result.output + assert len(responses.calls) == 1 + + @responses.activate + def test_refuses_unknown_short_suffix(self, runner, paths): + self._mock([{"id": "emitted-router-source", "source_model": "literal-router-source"}]) + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", "source") + assert result.exit_code != 0 + assert "'source' is not served" in result.output + assert not settings_path.exists() + + @responses.activate + def test_interactive_picker_uses_source_names(self, paths): + self._mock([{"id": "emitted", "source_model": "source"}]) + settings_path, _ = paths + asked = {} + ctx = click.Context( + configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} + ) + + def pick_model(listed): + asked["listed"] = tuple(listed) + return "source" + + interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=pick_model) + assert asked["listed"] == ("source",) + assert json.loads(settings_path.read_text())["model"] == "emitted" + + @responses.activate + def test_counts_what_an_older_proxy_lets_the_picker_show(self, runner, paths): + _mock_models() + result = _configure(runner, "--api-key", VALID_KEY) + assert result.exit_code == 0, result.output + assert "/model will list 1 of the proxy's 2 models: Claude Code shows only ids containing" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 0dd388919a5..b73d1acc6e3 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -7,8 +7,6 @@ from unittest.mock import Mock, patch import pytest from click.testing import CliRunner - - import litellm.proxy.client.cli from litellm._version import version as litellm_version from litellm.proxy.client.cli import cli diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index 68c0ac70064..03e5d7dd197 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -4,13 +4,16 @@ import stat from concurrent.futures import ThreadPoolExecutor from pathlib import Path +import pytest import requests from litellm.proxy.client.cli.commands.pi import ( + ListingFailure, ModelLimits, PiSyncError, fetch_model_ids, fetch_model_limits, + fetch_model_listing, models_json_path, provider_block, sync_models_json, @@ -28,6 +31,10 @@ class _FakeResponse: return self._payload +def _refused(*args, **kwargs): + raise requests.ConnectionError("refused") + + class TestFetchModelIds: def test_returns_ids_in_proxy_order_deduped(self): captured = {} @@ -44,6 +51,43 @@ class TestFetchModelIds: assert captured["url"] == "http://localhost:4000/v1/models" assert captured["headers"] == {"Authorization": "Bearer sk-key"} + def test_returns_rows_with_optional_source_model_and_dedups_identical_rows(self): + result = fetch_model_listing( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse( + 200, + {"data": [{"id": "emitted", "source_model": "source"}, {"id": "emitted", "source_model": "source"}]}, + ), + ) + assert not isinstance(result, PiSyncError) + assert tuple((model.id, model.source_model) for model in result) == (("emitted", "source"),) + + @pytest.mark.parametrize( + "entry", + [ + {"id": ""}, + {"id": "emitted", "source_model": ""}, + {"id": "emitted", "source_model": 1}, + ], + ) + def test_rejects_invalid_model_identity(self, entry): + result = fetch_model_listing( + "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(200, {"data": [entry]}) + ) + assert isinstance(result, PiSyncError) and result.kind is ListingFailure.BAD_BODY + + def test_rejects_conflicting_emitted_id_mappings(self): + result = fetch_model_listing( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse( + 200, + {"data": [{"id": "emitted", "source_model": "one"}, {"id": "emitted", "source_model": "two"}]}, + ), + ) + assert isinstance(result, PiSyncError) and result.kind is ListingFailure.BAD_BODY + def test_network_error_is_a_value(self): def boom(*a, **k): raise requests.ConnectionError("refused") @@ -53,9 +97,7 @@ class TestFetchModelIds: assert "Could not list models" in result.message def test_non_200_is_a_value(self): - result = fetch_model_ids( - "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500) - ) + result = fetch_model_ids("http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500)) assert isinstance(result, PiSyncError) assert "HTTP 500" in result.message @@ -75,6 +117,22 @@ class TestFetchModelIds: ) assert isinstance(result, PiSyncError) assert "no models" in result.message + assert result.kind is ListingFailure.EMPTY + + @pytest.mark.parametrize( + ("get", "kind"), + [ + (_refused, ListingFailure.UNREACHABLE), + (lambda *a, **k: _FakeResponse(401), ListingFailure.REJECTED), + (lambda *a, **k: _FakeResponse(403), ListingFailure.REJECTED), + (lambda *a, **k: _FakeResponse(500), ListingFailure.OTHER), + (lambda *a, **k: _FakeResponse(200), ListingFailure.BAD_BODY), + ], + ids=["unreachable", "401", "403", "500", "bad-body"], + ) + def test_the_failure_kind_is_decided_where_the_response_is_classified(self, get, kind): + result = fetch_model_ids("http://localhost:4000", "sk-key", get=get) + assert isinstance(result, PiSyncError) and result.kind is kind class TestFetchModelLimits: diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py new file mode 100644 index 00000000000..691764fbef4 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -0,0 +1,404 @@ +"""The status line script is copied verbatim to the user's machine, so these drive it the way Claude Code +and Codex do: the documented stdin payload, a transcript on disk, and the proxy behind an injected fetch.""" + +import io +import json +import os +import re +import subprocess +import sys +from pathlib import Path + +import pytest + +from litellm.proxy.client.cli.commands import statusline_script +from litellm.proxy.client.cli.commands.statusline_script import ( + CACHE_TTL_SECONDS, + Credentials, + Fetched, + Session, + cache_dir_name, + cache_path, + claude_credentials, + codex_credentials, + latest_transcript_model, + load_session, + render, + run, +) + +SESSION_ID = "cf712ab8-4c7c-4d48-ba91-eed54bc2956b" +ANSI = re.compile(r"\x1b\[[0-9;]*m") +RECORDED = Session( + router_name="claude-auto", + last_model="anthropic/claude-sonnet-5", + spend=0.14, + baseline_spend=0.38, + baseline_model="anthropic/claude-opus-5", +) + + +def _assistant_line(model: str, **extra: object) -> str: + return json.dumps({"type": "assistant", "message": {"model": model, "role": "assistant"}, **extra}) + + +@pytest.fixture +def transcript(tmp_path: Path) -> Path: + path = tmp_path / "session.jsonl" + path.write_text( + "\n".join( + ( + json.dumps({"type": "user", "message": {"role": "user", "content": "hi"}}), + _assistant_line("claude-haiku-4-5"), + json.dumps({"type": "user", "message": {"role": "user", "content": "harder"}}), + _assistant_line("claude-sonnet-5"), + _assistant_line("claude-haiku-4-5", isSidechain=True), + _assistant_line("claude-haiku-4-5", agentId="agent-1"), + json.dumps({"type": "progress", "data": {}}), + ) + ) + + "\n" + ) + return path + + +@pytest.fixture +def config_dir(tmp_path: Path) -> Path: + directory = tmp_path / "claude" + (directory / "cache").mkdir(parents=True) + (directory / "cache" / "gateway-models.json").write_text( + json.dumps({"models": [{"id": "claude-opus-5", "display_name": "Claude Opus 5"}]}) + ) + return directory + + +def _payload(transcript: Path, session_id: str = SESSION_ID) -> dict: + return { + "session_id": session_id, + "transcript_path": str(transcript), + "model": {"id": "claude-auto", "display_name": "claude-auto"}, + } + + +def _env(tmp_path: Path, config_dir: Path, **extra: str) -> dict[str, str]: + return { + "TMPDIR": str(tmp_path / "tmp"), + "CLAUDE_CONFIG_DIR": str(config_dir), + "TERM": "dumb", + "ANTHROPIC_BASE_URL": "http://127.0.0.1:4000", + "ANTHROPIC_AUTH_TOKEN": "sk-virtual", + **extra, + } + + +def _run(payload: object, env: dict[str, str], fetch) -> str: + out = io.StringIO() + run(io.StringIO(json.dumps(payload)), out, env, fetch) + return out.getvalue() + + +class TestTranscript: + def test_the_latest_foreground_assistant_line_wins_over_later_sidechain_and_agent_lines(self, transcript): + assert latest_transcript_model(str(transcript)) == "claude-sonnet-5" + + def test_a_synthetic_line_is_not_a_served_model(self, tmp_path): + # Claude Code writes `` for messages it produced locally (an API error on resume, for one); + # showing "Routed to: " would name a model no proxy served. + path = tmp_path / "t.jsonl" + path.write_text(_assistant_line("claude-haiku-4-5") + "\n" + _assistant_line("") + "\n") + assert latest_transcript_model(str(path)) == "claude-haiku-4-5" + + def test_a_missing_or_empty_transcript_yields_nothing(self, tmp_path): + empty = tmp_path / "empty.jsonl" + empty.write_text("") + assert latest_transcript_model(str(tmp_path / "missing.jsonl")) == "" + assert latest_transcript_model(str(empty)) == "" + assert latest_transcript_model("") == "" + + +class TestCredentials: + MIXED = { + "ANTHROPIC_BASE_URL": "http://anthropic-side:4000", + "ANTHROPIC_AUTH_TOKEN": "sk-ant", + "OPENAI_BASE_URL": "http://openai-side:4000/v1/", + "OPENAI_API_KEY": "sk-openai", + } + + def test_each_agent_reads_the_pair_it_dials_itself(self): + # A shell that exports both families must not send Codex's hook to the Anthropic proxy. + assert claude_credentials(self.MIXED) == Credentials("http://anthropic-side:4000", "sk-ant") + assert codex_credentials(self.MIXED) == Credentials("http://openai-side:4000", "sk-openai") + + def test_lites_own_shell_variables_are_not_a_credential_either_agent_sends(self): + # A `lite login` shell exports LITELLM_PROXY_*; Claude Code and Codex never read them, so the + # status line must not query the proxy as that principal while the agent used another. + env = {"LITELLM_PROXY_URL": "http://lite:4000/", "LITELLM_PROXY_API_KEY": "sk-lite", **self.MIXED} + assert claude_credentials(env) == Credentials("http://anthropic-side:4000", "sk-ant") + assert codex_credentials(env) == Credentials("http://openai-side:4000", "sk-openai") + assert not claude_credentials({"LITELLM_PROXY_API_KEY": "sk-lite", "ANTHROPIC_BASE_URL": "http://p"}).usable + assert codex_credentials({}) == Credentials("", "") + + def test_claude_code_prefers_the_auth_token_over_a_stray_api_key(self): + env = {"ANTHROPIC_BASE_URL": "http://p", "ANTHROPIC_API_KEY": "sk-stray", "ANTHROPIC_AUTH_TOKEN": "sk-ours"} + assert claude_credentials(env).api_key == "sk-ours" + + def test_an_api_key_helper_in_settings_is_never_run(self, tmp_path, transcript, config_dir): + # `lite` once wrote `apiKeyHelper: lite auth print-token`; running it from a status line that + # refreshes every 300ms spawned `lite` (and a keychain prompt) on every refresh. Without a key + # in the env the proxy is simply not asked. + (config_dir / "settings.json").write_text(json.dumps({"apiKeyHelper": "printf sk-from-helper"})) + asked = [] + env = {k: v for k, v in _env(tmp_path, config_dir).items() if k != "ANTHROPIC_AUTH_TOKEN"} + text = _run(_payload(transcript), env, lambda c, s: asked.append(c) or Fetched(RECORDED, True)) + assert text == "Routed to: claude-sonnet-5" and asked == [] + + +class TestSessionCache: + def test_a_definite_answer_is_served_from_the_cache_within_the_ttl(self, tmp_path): + calls = [] + + def fetch(credentials, session_id): + calls.append(session_id) + return Fetched(RECORDED, definitive=True) + + clock = [100.0] + credentials = Credentials("http://p", "sk") + first = load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: clock[0]) + clock[0] = 100.0 + CACHE_TTL_SECONDS - 1 + second = load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: clock[0]) + clock[0] = 100.0 + CACHE_TTL_SECONDS + 1 + third = load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: clock[0]) + assert first == second == third == RECORDED + assert calls == [SESSION_ID, SESSION_ID] + + def test_a_404_is_cached_as_absence_but_a_transport_failure_is_retried(self, tmp_path): + outcomes = iter((Fetched(None, definitive=False), Fetched(None, definitive=True), Fetched(RECORDED, True))) + calls = [] + + def fetch(credentials, session_id): + calls.append(session_id) + return next(outcomes) + + credentials = Credentials("http://p", "sk") + assert load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: 1.0) is None + assert load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: 1.0) is None + assert load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: 1.0) is None + assert len(calls) == 2 + + def test_a_cache_directory_that_is_not_private_is_never_used(self, tmp_path): + # A shared temp root lets another user pre-create the directory; refuse it rather than write into it. + calls = [] + + def fetch(credentials, session_id): + calls.append(session_id) + return Fetched(RECORDED, definitive=True) + + shared = tmp_path / "litellm-statusline" + shared.mkdir(mode=0o755) + credentials = Credentials("http://p", "sk") + for _ in range(2): + assert load_session(credentials, SESSION_ID, shared, fetch, now=lambda: 1.0) == RECORDED + assert calls == [SESSION_ID, SESSION_ID] + assert list(shared.iterdir()) == [] + + def test_any_client_error_is_a_definite_answer_and_a_server_error_is_not(self, monkeypatch): + import urllib.error + + from litellm.proxy.client.cli.commands.statusline_script import fetch_session + + def fail_with(code): + def opener(request, timeout): + raise urllib.error.HTTPError(request.full_url, code, "x", {}, None) + + return opener + + for code, definitive in ((403, True), (401, True), (404, True), (502, False)): + monkeypatch.setattr("urllib.request.urlopen", fail_with(code)) + assert fetch_session(Credentials("http://127.0.0.1:1", "sk"), SESSION_ID) == Fetched(None, definitive) + + def test_the_cache_file_holds_the_proxy_answer_and_never_the_key(self, tmp_path): + credentials = Credentials("http://p", "sk-secret") + load_session(credentials, SESSION_ID, tmp_path, lambda c, s: Fetched(RECORDED, True)) + path = cache_path(tmp_path, credentials, SESSION_ID) + assert "sk-secret" not in written and SESSION_ID not in written if (written := path.read_text()) else False + assert "sk-secret" not in path.name + assert json.loads(written)["session"]["baseline_model"] == "anthropic/claude-opus-5" + assert (path.stat().st_mode & 0o777) == 0o600 + assert (path.parent.stat().st_mode & 0o777) == 0o700 + + def test_a_refresh_replaces_the_entry_in_one_step_so_a_concurrent_refresh_never_reads_a_torn_one(self, tmp_path): + credentials = Credentials("http://p", "sk") + load_session(credentials, SESSION_ID, tmp_path, lambda c, s: Fetched(RECORDED, True), now=lambda: 1.0) + path = cache_path(tmp_path, credentials, SESSION_ID) + first = path.read_text() + + with path.open() as concurrent_reader: + newer = RECORDED._replace(spend=0.5) + load_session(credentials, SESSION_ID, tmp_path, lambda c, s: Fetched(newer, True), now=lambda: 100.0) + assert concurrent_reader.read() == first + assert json.loads(path.read_text())["session"]["spend"] == 0.5 + assert (path.stat().st_mode & 0o777) == 0o600 + assert [child.name for child in tmp_path.iterdir()] == [path.name] + + def test_the_same_session_id_against_another_proxy_or_key_is_not_served_from_the_cache(self, tmp_path): + answers = iter((Fetched(RECORDED, True), Fetched(RECORDED._replace(spend=9.0), True))) + first = load_session(Credentials("http://p", "sk-a"), SESSION_ID, tmp_path, lambda c, s: next(answers)) + second = load_session(Credentials("http://p", "sk-b"), SESSION_ID, tmp_path, lambda c, s: next(answers)) + assert first == RECORDED and second is not None and second.spend == 9.0 + + +class TestRender: + def test_savings_header_and_bars_against_the_routers_baseline(self, config_dir): + text = render("claude-sonnet-5", RECORDED, config_dir, use_color=False, bar_width=10) + assert text.splitlines() == [ + "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5", + "LiteLLM ████░░░░░░ $0.14", + "Claude Opus 5 ██████████ $0.38", + ] + + def test_control_characters_in_any_externally_sourced_label_never_reach_the_terminal(self, tmp_path, config_dir): + # The transcript, the proxy payload and Claude Code's model cache all feed labels straight into a + # terminal, and none is under this script's control. Only the control bytes are dropped (ESC, BEL, + # C1), which is what disarms an OSC-52 clipboard write or a screen clear; the printable remainder of + # such a sequence is inert text and is kept as is. + from litellm.proxy.client.cli.commands.statusline_script import _session_from_payload, model_label + + hostile = "claude-\x1b\x07\x9bsonnet" + path = tmp_path / "t.jsonl" + path.write_text(_assistant_line(hostile) + "\n") + assert latest_transcript_model(str(path)) == "claude-sonnet" + (config_dir / "cache" / "gateway-models.json").write_text( + json.dumps({"models": [{"id": "claude-sonnet", "display_name": "Son\x1b\x07net"}]}) + ) + assert model_label("claude-sonnet", config_dir) == "Sonnet" + session = _session_from_payload( + {"router_name": "auto\x07", "last_model": hostile, "spend": 0.1, "baseline_spend": 0.2, "baseline_model": "op\x1bus"} + ) + assert session == Session("auto", "claude-sonnet", 0.1, 0.2, "opus") + assert latest_transcript_model(str(path)) == "claude-sonnet" + assert "\x1b]52;c;ZXZpbA==" not in render( + latest_transcript_model(str(path)), + _session_from_payload({"router_name": "a", "last_model": "m", "spend": 0.1, "baseline_spend": 0.2, "baseline_model": "\x1b]52;c;ZXZpbA==\x07"}), + config_dir, + use_color=False, + ) + + def test_a_session_that_cost_more_than_its_baseline_reads_as_a_plus(self, config_dir): + dearer = RECORDED._replace(spend=0.50, baseline_spend=0.40) + assert "+25% vs Claude Opus 5" in render("m", dearer, config_dir, use_color=False) + + def test_without_a_baseline_only_the_routed_line_shows(self, config_dir): + assert render("m", RECORDED._replace(baseline_model=None), config_dir, False) == "claude-auto · Routed to: m" + assert render("m", None, config_dir, False) == "Routed to: m" + + def test_color_wraps_the_same_text(self, config_dir): + colored = render("claude-sonnet-5", RECORDED, config_dir, use_color=True, bar_width=10) + assert ANSI.sub("", colored) == render("claude-sonnet-5", RECORDED, config_dir, use_color=False, bar_width=10) + + +class TestClaudeCodeMode: + def test_the_transcript_names_the_routed_model_and_the_proxy_adds_the_savings(self, tmp_path, transcript, config_dir): + seen = [] + + def fetch(credentials, session_id): + seen.append((credentials, session_id)) + return Fetched(RECORDED, definitive=True) + + text = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) + assert text.startswith("claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5\n") + assert seen == [(Credentials("http://127.0.0.1:4000", "sk-virtual"), SESSION_ID)] + + def test_an_unrecorded_session_degrades_to_the_routed_line(self, tmp_path, transcript, config_dir): + assert _run(_payload(transcript), _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) == ( + "Routed to: claude-sonnet-5" + ) + + def test_without_credentials_the_proxy_is_never_asked(self, tmp_path, transcript, config_dir): + def fetch(credentials, session_id): + raise AssertionError("must not fetch") + + env = {k: v for k, v in _env(tmp_path, config_dir).items() if k != "ANTHROPIC_AUTH_TOKEN"} + assert _run(_payload(transcript), env, fetch) == "Routed to: claude-sonnet-5" + + def test_before_the_first_response_the_payloads_display_name_shows(self, tmp_path, config_dir): + payload = _payload(tmp_path / "missing.jsonl") + assert _run(payload, _env(tmp_path, config_dir), lambda c, s: Fetched(RECORDED, True)) == "claude-auto" + + def test_a_discovered_display_name_labels_the_routed_model(self, tmp_path, config_dir): + path = tmp_path / "t.jsonl" + path.write_text(_assistant_line("anthropic/claude-opus-5") + "\n") + assert _run(_payload(path), _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) == ( + "Routed to: Claude Opus 5" + ) + + def test_the_cache_lands_under_the_platforms_temp_dir(self, tmp_path, transcript, config_dir): + env = {k: v for k, v in _env(tmp_path, config_dir).items() if k != "TMPDIR"} + env["TEMP"] = str(tmp_path / "wintemp") + _run(_payload(transcript), env, lambda c, s: Fetched(RECORDED, True)) + assert (tmp_path / "wintemp" / cache_dir_name()).is_dir() + assert cache_dir_name().endswith(str(os.getuid())) + + def test_a_crash_falls_back_to_the_model_label_claude_code_already_knows(self, tmp_path, transcript, config_dir): + def fetch(credentials, session_id): + raise RuntimeError("boom") + + assert _run(_payload(transcript), _env(tmp_path, config_dir), fetch) == "claude-auto" + + def test_garbage_on_stdin_still_prints_something(self, tmp_path, config_dir): + out = io.StringIO() + run(io.StringIO("not json"), out, _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) + assert out.getvalue() == "claude" + + +class TestCodexMode: + def test_the_stop_hook_prints_a_system_message_from_the_proxys_record(self, tmp_path, config_dir): + env = _env(tmp_path, config_dir, OPENAI_BASE_URL="http://127.0.0.1:4000/v1", OPENAI_API_KEY="sk-codex") + env = {k: v for k, v in env.items() if not k.startswith("ANTHROPIC_")} + seen = [] + + def fetch(credentials, session_id): + seen.append(credentials) + return Fetched(RECORDED, definitive=True) + + out = _run({"hook_event_name": "Stop", "session_id": SESSION_ID, "transcript_path": "/nope"}, env, fetch) + message = json.loads(out)["systemMessage"] + assert message.splitlines()[1] == "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5" + assert message.startswith("\n") + assert seen == [Credentials("http://127.0.0.1:4000", "sk-codex")] + + def test_an_unrecorded_session_prints_nothing_so_codex_shows_no_message(self, tmp_path, config_dir): + payload = {"hook_event_name": "Stop", "session_id": SESSION_ID} + assert _run(payload, _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) == "" + + def test_a_crash_prints_nothing_rather_than_text_codex_would_reject(self, tmp_path, config_dir): + def fetch(credentials, session_id): + raise RuntimeError("boom") + + env = _env(tmp_path, config_dir, OPENAI_BASE_URL="http://127.0.0.1:4000/v1", OPENAI_API_KEY="sk-codex") + assert _run({"hook_event_name": "Stop", "session_id": SESSION_ID}, env, fetch) == "" + + def test_a_turn_right_after_an_unrecorded_one_still_asks_the_proxy(self, tmp_path, config_dir): + # One hook run per turn: a miss on turn one must not be cached across turn two's fetch. + answers = iter((Fetched(None, definitive=True), Fetched(RECORDED, definitive=True))) + payload = {"hook_event_name": "Stop", "session_id": SESSION_ID} + env = _env(tmp_path, config_dir, OPENAI_BASE_URL="http://127.0.0.1:4000/v1", OPENAI_API_KEY="sk-codex") + assert _run(payload, env, lambda c, s: next(answers)) == "" + assert "Routed to: claude-sonnet-5" in json.loads(_run(payload, env, lambda c, s: next(answers)))["systemMessage"] + + +class TestStandalone: + def test_the_file_runs_under_a_bare_interpreter_with_no_litellm_on_the_path(self, tmp_path, transcript, config_dir): + # It is copied verbatim to ~/.litellm/statusline.py, so it must be self-contained. + script = tmp_path / "statusline.py" + script.write_bytes(Path(statusline_script.__file__).read_bytes()) + env = {k: v for k, v in _env(tmp_path, config_dir).items() if k != "ANTHROPIC_AUTH_TOKEN"} + completed = subprocess.run( + [sys.executable, "-I", str(script)], + input=json.dumps(_payload(transcript)), + capture_output=True, + text=True, + env=env, + check=True, + timeout=30, + ) + assert completed.stdout == "Routed to: claude-sonnet-5" diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index aead1764b0e..ddd54dd1374 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -11,16 +11,15 @@ from click.testing import CliRunner from litellm.proxy.client.cli.commands import up as up_module from litellm.proxy.client.cli.commands.agents import AgentRunError -from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError +from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError, StaticToken from litellm.proxy.client.cli.commands.up import ( BackupRecord, UpError, - _ensure_fresh_login, + ensure_fresh_login, down, load_json_or_empty, merge_claude_settings, read_backup, - resolve_api_key_helper, restore_claude_settings, up, write_backup, @@ -40,52 +39,54 @@ def _patch_paths(monkeypatch, tmp_path): class TestMergeClaudeSettings: def test_preserves_unrelated_top_level_keys(self): - merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", "helper") + merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["theme"] == "dark" def test_preserves_unrelated_env_keys(self): settings = {"env": {"SOME_OTHER_VAR": "value"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"]["SOME_OTHER_VAR"] == "value" - def test_overrides_base_url_and_helper(self): + def test_overrides_base_url_and_strips_an_old_helper(self): settings = { "env": {"ANTHROPIC_BASE_URL": "https://old.example.com"}, "apiKeyHelper": "old-helper", } - merged = merge_claude_settings(settings, "http://localhost:4000/", "new-helper") + merged = merge_claude_settings(settings, "http://localhost:4000/", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-fresh" assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" - assert merged["apiKeyHelper"] == "new-helper" + assert "apiKeyHelper" not in merged def test_preserves_existing_gateway_model_discovery(self): settings = {"env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "0"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "0" def test_preserves_existing_tool_search(self): settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" def test_drops_stray_api_key(self): settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert "ANTHROPIC_API_KEY" not in merged["env"] def test_works_from_empty_settings(self): - merged = merge_claude_settings({}, "http://localhost:4000", "helper") + merged = merge_claude_settings({}, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"] == { "ANTHROPIC_BASE_URL": "http://localhost:4000", + "ANTHROPIC_AUTH_TOKEN": "sk-fresh", "ENABLE_TOOL_SEARCH": "true", "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "1", } - assert merged["apiKeyHelper"] == "helper" + assert "apiKeyHelper" not in merged def test_does_not_mutate_input(self): settings = {"env": {"FOO": "bar"}} - merge_claude_settings(settings, "http://localhost:4000", "helper") + merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert settings == {"env": {"FOO": "bar"}} @@ -154,6 +155,25 @@ class TestBackupRoundTrip: assert restore_claude_settings() is None assert not settings_path.exists() + def test_restore_writes_owner_only_and_through_a_symlink(self, monkeypatch, tmp_path): + # The backup can hold a token the user had in the file before `up`; a plain open() would put it + # back under the umask, and would replace a dotfiles symlink with a regular file. + target = tmp_path / "dotfiles" / "settings.json" + target.parent.mkdir() + target.write_text("{}") + target.chmod(0o644) + settings_path = tmp_path / "settings.json" + settings_path.symlink_to(target) + monkeypatch.setattr(up_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(up_module, "BACKUP_PATH", tmp_path / "backup.json") + write_backup(BackupRecord(existed=True, content={"env": {"ANTHROPIC_AUTH_TOKEN": "sk-theirs"}})) + + restore_claude_settings() + + assert settings_path.is_symlink() + assert json.loads(target.read_text()) == {"env": {"ANTHROPIC_AUTH_TOKEN": "sk-theirs"}} + assert stat.S_IMODE(target.stat().st_mode) == 0o600 + def test_recreates_claude_dir_if_it_was_deleted_while_up_was_running(self, monkeypatch, tmp_path): """If ~/.claude/ is removed while `lite up` holds it open, restoring must recreate the directory rather than crash with FileNotFoundError and strand the backup file, which @@ -216,49 +236,6 @@ class TestBackupRoundTrip: assert not backup_path.exists() -class TestResolveApiKeyHelper: - def test_returns_helper_command_bound_to_the_selected_proxy(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") - helper = resolve_api_key_helper("http://localhost:4000") - assert helper == "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token" - - def test_quotes_a_base_url_containing_shell_metacharacters(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") - helper = resolve_api_key_helper("http://example.com/path; rm -rf /") - assert helper == "/usr/local/bin/lite --base-url 'http://example.com/path; rm -rf /' auth print-token" - - def test_raises_when_lite_not_on_path(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: None) - with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): - resolve_api_key_helper("http://localhost:4000") - - def test_windows_quotes_for_cmd_exe_instead_of_posix_sh(self, monkeypatch): - """cmd.exe takes a single quote literally, so a POSIX-quoted backslashed path is unrunnable.""" - lite_exe = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" - monkeypatch.setattr(shutil, "which", lambda name: lite_exe) - - helper = resolve_api_key_helper("https://gateway.example.com", platform="win32") - - assert helper == f'"{lite_exe}" "--base-url" "https://gateway.example.com" "auth" "print-token"' - - def test_windows_keeps_a_spaced_path_and_a_metacharacter_url_as_single_tokens(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: "C:\\Program Files\\LiteLLM\\lite.EXE") - - helper = resolve_api_key_helper("https://gateway.example.com/?a=1&b=2", platform="win32") - - assert helper == ( - '"C:\\Program Files\\LiteLLM\\lite.EXE" "--base-url" "https://gateway.example.com/?a=1&b=2" ' - '"auth" "print-token"' - ) - - def test_non_windows_platforms_keep_posix_quoting(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") - - helper = resolve_api_key_helper("http://example.com/path; rm -rf /", platform="darwin") - - assert helper == "/usr/local/bin/lite --base-url 'http://example.com/path; rm -rf /' auth print-token" - - def _make_ctx(base_url): return click.Context(click.Command("test"), obj={"base_url": base_url}) @@ -307,7 +284,8 @@ def _capture_login(monkeypatch, on_login=lambda: None): login_calls = [] @click.pass_context - def fake_login(ctx, pkce=False): + def fake_login(ctx, config_claude=False, pkce=False): + assert config_claude is False, "`lite up` patches settings itself; the login it starts must not also configure" login_calls.append((ctx.obj["base_url"], pkce)) on_login() @@ -317,8 +295,8 @@ def _capture_login(monkeypatch, on_login=lambda: None): class TestEnsureFreshLogin: """A token that is fresh but was issued for a *different* proxy must not be trusted: without - this check, a user logged into proxy A who runs `up --base-url proxy-b` would silently get an - apiKeyHelper wired up around proxy A's real token, which print-token would then hand to proxy B.""" + this check, a user logged into proxy A who runs `up --base-url proxy-b` would silently get proxy A's + real token written into settings pointed at proxy B.""" def test_reuses_a_fresh_token_issued_for_the_same_proxy(self, monkeypatch): _FakeTokenStore( @@ -327,7 +305,7 @@ class TestEnsureFreshLogin: monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) login_calls = _capture_login(monkeypatch) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [] @@ -339,7 +317,7 @@ class TestEnsureFreshLogin: monkeypatch, on_login=lambda: store.log_in({"key": "sk-b", "base_url": "http://proxy-b:4000"}, "sk-b") ) - _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + ensure_fresh_login(_make_ctx("http://proxy-b:4000")) assert login_calls == [("http://proxy-b:4000", False)] assert store.key_requests == ["http://proxy-b:4000", "http://proxy-b:4000"] @@ -353,7 +331,7 @@ class TestEnsureFreshLogin: on_login=lambda: store.log_in({"key": "sk-a", "base_url": "http://proxy-a:4000"}, "sk-a"), ) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [("http://proxy-a:4000", False)] @@ -363,7 +341,7 @@ class TestEnsureFreshLogin: monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) with pytest.raises(UpError, match="Run `lite login` first"): - _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + ensure_fresh_login(_make_ctx("http://proxy-b:4000")) def test_trusts_a_pkce_credential_that_was_renewed_on_the_way_in(self, monkeypatch): """A --pkce key inside its freshness buffer is renewed by `get_stored_api_key`, so `lite up` @@ -377,7 +355,7 @@ class TestEnsureFreshLogin: ) login_calls = _capture_login(monkeypatch) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [] assert store.key_requests == ["http://proxy-a:4000"] @@ -390,7 +368,7 @@ class TestEnsureFreshLogin: on_login=lambda: store.log_in(_pkce_record("http://proxy-a:4000", seconds_left=86_400), "sk-pkce-fresh"), ) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [("http://proxy-a:4000", True)] @@ -399,7 +377,7 @@ class TestEnsureFreshLogin: _FakeTokenStore(monkeypatch, _pkce_record("http://proxy-a:4000", seconds_left=-10), {}) with pytest.raises(UpError, match="Run `lite login --pkce` first"): - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) def test_trusts_the_key_the_cli_group_already_resolved_instead_of_reading_the_token_file_again( self, monkeypatch @@ -409,7 +387,7 @@ class TestEnsureFreshLogin: store = _FakeTokenStore(monkeypatch, _pkce_record("http://proxy-a:4000", seconds_left=86_400), {}) login_calls = _capture_login(monkeypatch) - _ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key="sk-pkce-renewed-by-the-group")) + ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key="sk-pkce-renewed-by-the-group")) assert login_calls == [] assert store.key_requests == [] @@ -423,7 +401,7 @@ class TestEnsureFreshLogin: ) with pytest.raises(UpError, match="Run `lite login --pkce` first"): - _ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) + ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) assert store.key_requests == [] @@ -435,7 +413,7 @@ class TestEnsureFreshLogin: on_login=lambda: store.log_in(_pkce_record("http://proxy-a:4000", seconds_left=86_400), "sk-pkce-fresh"), ) - _ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) + ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) assert login_calls == [("http://proxy-a:4000", True)] assert store.key_requests == ["http://proxy-a:4000"] @@ -500,11 +478,13 @@ class TestUpCommand: settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) original = {"theme": "dark"} settings_path.write_text(json.dumps(original)) + settings_path.chmod(0o644) captured = {} def fake_wait(self, timeout=None): captured["settings"] = json.loads(settings_path.read_text()) + captured["settings_mode"] = stat.S_IMODE(settings_path.stat().st_mode) captured["backup_existed"] = backup_path.exists() return True @@ -514,10 +494,6 @@ class TestUpCommand: patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), patch(f"{UP_MODULE}.verify_proxy_key"), - patch( - f"{UP_MODULE}.resolve_api_key_helper", - return_value="/usr/local/bin/lite auth print-token", - ), patch(f"{UP_MODULE}.signal.signal"), patch(f"{UP_MODULE}.atexit.register"), patch("threading.Event.wait", new=fake_wait), @@ -529,7 +505,11 @@ class TestUpCommand: assert captured["settings"]["theme"] == "dark" assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true" - assert captured["settings"]["apiKeyHelper"] == "/usr/local/bin/lite auth print-token" + assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-fresh" + # The file now carries the key, so the umask (and the file's earlier 0644) must not decide who reads it. + assert captured["settings_mode"] == 0o600 + assert "apiKeyHelper" not in captured["settings"] + assert captured["settings"]["statusLine"]["command"].endswith("statusline.py") assert json.loads(settings_path.read_text()) == original assert not backup_path.exists() @@ -598,7 +578,7 @@ class TestUpCanInvokeTheRealLoginCommand: @click.pass_context def driver(ctx): ctx.obj = {"base_url": "http://127.0.0.1:9"} - ctx.invoke(real_login, pkce=False) + ctx.invoke(real_login, config_claude=False, pkce=False) with patch( f"{AUTH_MODULE}._start_cli_sso_flow", @@ -618,12 +598,9 @@ class TestUpCanInvokeTheRealLoginCommand: @click.pass_context def driver(ctx): ctx.obj = {"base_url": "http://127.0.0.1:9"} - ctx.invoke(real_login, pkce=False) + ctx.invoke(real_login, config_claude=False, pkce=False) - with ( - patch(f"{AUTH_MODULE}.CLAUDE_SETTINGS_PATH", settings_path), - patch(f"{AUTH_MODULE}._start_cli_sso_flow", side_effect=RuntimeError("stop")), - ): - CliRunner().invoke(driver, [], standalone_mode=False) + with patch(f"{AUTH_MODULE}._start_cli_sso_flow", side_effect=RuntimeError("stop")): + CliRunner().invoke(driver, [], standalone_mode=False, env={"CLAUDE_CONFIG_DIR": str(tmp_path)}) assert not settings_path.exists() diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 7d5fc1a3544..4e2059ac30b 100644 --- a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -1,4 +1,5 @@ import asyncio +import hashlib import json from typing import Iterable, List, Optional, Tuple from unittest.mock import patch @@ -7,6 +8,7 @@ import pytest from redis.asyncio import Redis from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AUTH_CACHE_INVALIDATION_CHANNEL, AuthCacheInvalidationSubscriber, @@ -145,6 +147,35 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] +@pytest.mark.asyncio +async def test_subscriber_deletes_key_object_partition_entry_on_message() -> None: + """ + LIT-7563 moved user-key objects into their own in-memory partition; a key + invalidation broadcast must still evict the hashed-token entry there, or a + deleted key keeps authenticating on other workers until its TTL expires. + """ + hashed_token = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest() + cache = UserApiKeyCache() + cache.set_cache(hashed_token, UserAPIKeyAuth(token=hashed_token), model_type=UserAPIKeyAuth) + assert cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is not None + + pubsub = _QueuePubSub(initial_messages=[_invalidation_message(hashed_token)]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + ) + subscriber.start() + try: + for _ in range(200): + if cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is None: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is None + + @pytest.mark.asyncio async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None: """ diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 66f77db6da9..ecb2375d495 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -1,30 +1,33 @@ import copy +import json import sys from types import ModuleType, SimpleNamespace +from typing import Final +from unittest.mock import patch import pytest - +import litellm +from litellm.caching.caching import DualCache +from litellm.constants import MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( + _serialize_scan_metadata_header, add_guardrail_scan_id, add_policy_to_applied_policies_header, decrypt_callback_vars, encrypt_callback_vars, get_logging_caching_headers, - initialize_callbacks_on_proxy, get_remaining_tokens_and_requests_from_request_data, + initialize_callbacks_on_proxy, normalize_callback_names, + process_callback, sanitize_openai_provider_metadata, strip_callback_config, ) -import litellm -from litellm.caching.caching import DualCache -from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging - -from unittest.mock import patch -from litellm.proxy.common_utils.callback_utils import process_callback +from litellm.types.guardrails import GuardrailEventHooks def test_get_remaining_tokens_and_requests_from_request_data(): @@ -189,20 +192,109 @@ def test_get_logging_caching_headers_merges_metadata_and_litellm_metadata(): assert headers["x-litellm-policy-sources"] == "global-baseline=team_default" +def _record( + request_data: dict[str, object], + scan_id: str | None, + guardrail_name: str = "airs", + provider: str = "panw_prisma_airs", + stage: GuardrailEventHooks = GuardrailEventHooks.pre_call, +) -> None: + add_guardrail_scan_id( + request_data=request_data, scan_id=scan_id, guardrail_name=guardrail_name, provider=provider, stage=stage + ) + + def test_add_guardrail_scan_id_dedupes_and_becomes_response_header(): request_data = {"litellm_metadata": {}} - add_guardrail_scan_id(request_data=request_data, scan_id="scan-1") - add_guardrail_scan_id(request_data=request_data, scan_id="scan-1") - add_guardrail_scan_id(request_data=request_data, scan_id="scan-2") - add_guardrail_scan_id(request_data=request_data, scan_id=None) + _record(request_data, "scan-1") + _record(request_data, "scan-1") + _record(request_data, "scan-2") + _record(request_data, None) assert request_data["litellm_metadata"]["guardrail_scan_ids"] == ("scan-1", "scan-2") assert get_logging_caching_headers(request_data)["x-litellm-guardrail-scan-id"] == "scan-1,scan-2" -def test_get_logging_caching_headers_omits_scan_id_header_without_scans(): - assert "x-litellm-guardrail-scan-id" not in get_logging_caching_headers({"litellm_metadata": {}}) +def test_scan_metadata_header_maps_each_id_to_its_guardrail_stage_and_provider(): + request_data: Final[dict[str, object]] = {"litellm_metadata": {}} + + _record( + request_data, "scan-1", guardrail_name="airs", provider="panw_prisma_airs", stage=GuardrailEventHooks.pre_call + ) + _record( + request_data, "mod-1", guardrail_name="mod", provider="openai_moderation", stage=GuardrailEventHooks.pre_call + ) + _record( + request_data, "scan-2", guardrail_name="airs", provider="panw_prisma_airs", stage=GuardrailEventHooks.post_call + ) + _record( + request_data, "scan-2", guardrail_name="airs", provider="panw_prisma_airs", stage=GuardrailEventHooks.post_call + ) + _record(request_data, None, guardrail_name="mod", provider="openai_moderation", stage=GuardrailEventHooks.post_call) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == "scan-1,mod-1,scan-2" + assert json.loads(headers["x-litellm-guardrail-scan-metadata"]) == [ + {"guardrail": "airs", "stage": "pre_call", "provider": "panw_prisma_airs", "scan_id": "scan-1"}, + {"guardrail": "mod", "stage": "pre_call", "provider": "openai_moderation", "scan_id": "mod-1"}, + {"guardrail": "airs", "stage": "post_call", "provider": "panw_prisma_airs", "scan_id": "scan-2"}, + ] + + +def test_scan_metadata_keeps_same_id_reused_across_stages(): + request_data: Final[dict[str, object]] = {"metadata": {}} + + _record(request_data, "scan-1", stage=GuardrailEventHooks.pre_call) + _record(request_data, "scan-1", stage=GuardrailEventHooks.post_call) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == "scan-1" + assert [entry["stage"] for entry in json.loads(headers["x-litellm-guardrail-scan-metadata"])] == [ + "pre_call", + "post_call", + ] + + +def test_scan_metadata_header_drops_trailing_entries_to_stay_within_length_limit(): + request_data: Final[dict[str, object]] = {"litellm_metadata": {}} + scan_ids: Final = tuple(f"0f9c4b7e-3d2a-4c1b-9e8f-{index:012d}" for index in range(40)) + for scan_id in scan_ids: + _record(request_data, scan_id, stage=GuardrailEventHooks.post_call) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == ",".join(scan_ids) + header: Final = headers["x-litellm-guardrail-scan-metadata"] + assert len(header) <= MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH + kept: Final = json.loads(header) + assert 1 < len(kept) < len(scan_ids) + assert [entry["scan_id"] for entry in kept] == list(scan_ids[: len(kept)]) + + +def test_serialize_scan_metadata_header_keeps_exactly_the_entries_that_fit(): + entries: Final = ({"scan_id": "a"}, {"scan_id": "b"}, {"scan_id": "c"}) + two_entries: Final = '[{"scan_id":"a"},{"scan_id":"b"}]' + + assert _serialize_scan_metadata_header(entries, max_length=len(two_entries)) == two_entries + assert _serialize_scan_metadata_header(entries, max_length=len(two_entries) - 1) == '[{"scan_id":"a"}]' + assert _serialize_scan_metadata_header(entries, max_length=len(two_entries) + 1) == two_entries + assert _serialize_scan_metadata_header(entries, max_length=1000) == json.dumps(entries, separators=(",", ":")) + assert _serialize_scan_metadata_header(entries, max_length=5) is None + assert _serialize_scan_metadata_header((), max_length=1000) is None + + +def test_scan_metadata_is_an_internal_metadata_key(): + assert sanitize_openai_provider_metadata({"guardrail_scan_metadata": "x", "keep": "y"}) == {"keep": "y"} + + +def test_get_logging_caching_headers_omits_scan_headers_without_scans(): + headers: Final = get_logging_caching_headers({"litellm_metadata": {}}) + assert headers is not None + assert "x-litellm-guardrail-scan-id" not in headers + assert "x-litellm-guardrail-scan-metadata" not in headers def test_initialize_callbacks_on_proxy_instantiates_compression_interception( diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py index 3bb3f75b9b8..2a0ebf9492a 100644 --- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -824,7 +824,7 @@ class _StopFailingSubscriber(ConfigSyncSubscriber): raise RuntimeError("stop failed") -async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> None: +async def test_proxy_config_subscriber_resyncs_deployments_only() -> None: from litellm.proxy.proxy_server import ProxyConfig cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()])) @@ -852,10 +852,7 @@ async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> await callback() await config.stop_config_sync_subscriber() - assert calls == [ - ("add_deployment", prisma_client, proxy_logging_obj), - ("get_credentials", prisma_client, None), - ] + assert calls == [("add_deployment", prisma_client, proxy_logging_obj)] assert config.config_sync_subscriber is None assert subscriber._task is None diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index 011571a37e0..bc4e756eb65 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -26,9 +26,58 @@ from litellm.proxy.common_utils.http_parsing_utils import ( get_tags_from_request_body, numeric_form_fields, populate_request_with_path_params, + read_raw_json_body, ) +def _starlette_request(body: bytes, content_type: str) -> Request: + scope = { + "type": "http", + "method": "POST", + "path": "/v1/messages", + "headers": [(b"content-type", content_type.encode())], + "query_string": b"", + } + chunks = iter((body,)) + + async def receive(): + return {"type": "http.request", "body": next(chunks, b""), "more_body": False} + + return Request(scope, receive) + + +@pytest.mark.asyncio +async def test_read_raw_json_body_returns_the_bytes_the_parsed_body_came_from(): + body = b'{"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}]}' + request = _starlette_request(body, "application/json") + + assert await _read_request_body(request) == orjson.loads(body) + assert await read_raw_json_body(request) == body + + +@pytest.mark.asyncio +async def test_read_raw_json_body_is_none_until_the_body_has_been_parsed(): + request = _starlette_request(b'{"model": "claude-sonnet-4-5"}', "application/json") + + assert await read_raw_json_body(request) is None + assert await read_raw_json_body(None) is None + + +@pytest.mark.asyncio +async def test_read_raw_json_body_is_none_for_form_bodies(): + request = _starlette_request(b"model=claude-sonnet-4-5", "application/x-www-form-urlencoded") + + assert await _read_request_body(request) == {"model": "claude-sonnet-4-5"} + assert await read_raw_json_body(request) is None + + +@pytest.mark.asyncio +async def test_read_raw_json_body_is_none_for_a_request_that_only_mocks_the_parsed_body_path(): + mock_request = MagicMock() + + assert await read_raw_json_body(mock_request) is None + + @pytest.mark.asyncio async def test_request_body_caching(): """ diff --git a/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py new file mode 100644 index 00000000000..7ef03140093 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py @@ -0,0 +1,109 @@ +"""Model identity survives Claude Code presentation, filtering and configured alias precedence.""" + +from itertools import combinations + +import pytest + +from litellm import Router +from litellm.proxy.common_utils.model_listing_utils import ( + ClaudeCodeRoutingNames, + claude_code_group_name, + claude_code_model_id, + claude_code_requested_group, + claude_code_view_ids, +) + + +def _encoded(name): + return "claude-router-" + name.encode().hex() + + +def _marked(name): + return f"{_encoded(name)}[1m]" + + +def _row(name, limit=1000000): + return {"id": name, "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": limit} + + +def _router(*names, aliases=None): + return Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} + for name in names + ], + model_group_alias=aliases, + ) + + +@pytest.mark.parametrize("limit", [None, 999999, 1000000]) +@pytest.mark.parametrize("name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"]) +def test_listing_round_trips_entire_source_name(name, limit): + names = frozenset({name}) + view = claude_code_model_id(name, limit, names) + assert (claude_code_group_name(view, names) or view) == name + if "claude" not in name: + assert view.startswith(_encoded(name)) + assert ("[1m]" in view.lower()) == (limit == 1000000 or "[1m]" in name.lower() and "claude" in name) + + +def test_collision_matrix_round_trips_without_duplicate_ids(): + universe = ("foo", "foo[1m]", "claude-router-foo", _encoded("foo"), _encoded("foo") + "[1m]", "claude-opus-5", "claude-opus-5[1m]") + for pair in combinations(universe, 2): + for visible in (pair, pair[:1], pair[1:]): + names = frozenset(pair) + view = claude_code_view_ids(tuple(_row(n) for n in visible), {"user-agent": "claude-code/2.1.267"}, names) + assert len(set(view.values())) == len(visible) + assert all((claude_code_group_name(shown, names) or shown) == source for source, shown in view.items()) + + +@pytest.mark.parametrize("spelling", ["claude-router-foo", "claude-router-ff", "claude-router-66 6f6f", "claude-router-666F6F", "claude-router-", _encoded("missing")]) +def test_unknown_or_noncanonical_ids_are_never_guessed(spelling): + assert claude_code_group_name(spelling, frozenset({"foo"})) is None + + +@pytest.mark.parametrize("headers,enabled", [ + ({"user-agent": "claude-code/2.1.267"}, True), + ({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True), + ({"x-gateway-client": "Claude-Code"}, True), + ({"user-agent": "anthropic-sdk-python/0.40"}, False), + ({}, False), +]) +def test_only_claude_code_gets_the_view(headers, enabled): + rows = (_row("foo"), _row("claude-opus-5")) + view = claude_code_view_ids(rows, headers, frozenset(row["id"] for row in rows)) + assert dict(view) == ({"foo": _encoded("foo") + "[1m]", "claude-opus-5": "claude-opus-5[1m]"} if enabled else {}) + + +@pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "team", "wildcard"]) +def test_configured_names_outrank_generated_ids_even_when_hidden_from_listing(monkeypatch, layer): + import litellm + + encoded = _encoded("foo") + alias = {encoded: "other"} + monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) + router = _router("foo", "other", *( (encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), aliases=alias if layer == "router" else None) + maps = (alias,) if layer in ("key", "team") else () + names = ClaudeCodeRoutingNames(router, None, maps) + assert claude_code_requested_group(encoded, router, None, maps) is None + assert claude_code_view_ids((_row("foo", None),), {"x-gateway-client": "claude-code"}, names)["foo"] == "foo" + + +@pytest.mark.parametrize("source", ["foo", "foo[1m]", "世界"]) +def test_mutation_breaking_the_hex_name_cannot_route_to_the_source(source): + router = _router(source) + encoded = _encoded(source) + malformed = encoded[:-1] + ("0" if encoded[-1] != "0" else "1") + assert claude_code_requested_group(malformed, router, None) is None + assert claude_code_requested_group(_marked(source), router, None) == source + + +def test_team_public_name_uses_the_same_scope_at_list_and_request(): + router = Router(model_list=[{ + "model_name": "model_name_team-a_id", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + "model_info": {"team_id": "team-a", "team_public_model_name": "shared"}, + }]) + shown = claude_code_view_ids((_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a"))["shared"] + assert claude_code_requested_group(shown, router, "team-a") == "shared" + assert claude_code_requested_group(shown, router, "team-b") is None diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py new file mode 100644 index 00000000000..8b653ddfb71 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -0,0 +1,145 @@ +import json + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) + + +@pytest.mark.parametrize( + "status_code, expected_type", + [ + (400, "invalid_request_error"), + (401, "authentication_error"), + (403, "permission_error"), + (404, "invalid_request_error"), + (408, "invalid_request_error"), + (422, "invalid_request_error"), + (429, "rate_limit_error"), + (499, "invalid_request_error"), + (500, "internal_server_error"), + (502, "internal_server_error"), + (503, "internal_server_error"), + ], +) +def test_status_code_decides_the_type_when_the_exception_carries_none(status_code: int, expected_type: str): + """A route that raises a bare HTTPException carries no error type, so the status it + answered with is the only thing left to name the OpenAI type from.""" + assert openai_error_type(HTTPException(status_code=status_code, detail="boom"), status_code) == expected_type + + +def test_a_carried_type_wins_over_the_one_the_status_would_imply(): + """A ProxyException raised mid-request already names its own type, and relabelling a + 402 budget_exceeded as the status map's guess would lose what the client branches on.""" + carried = ProxyException( + message="Budget has been exceeded", + type=ProxyErrorTypes.budget_exceeded.value, + param=None, + code=400, + ) + + assert openai_error_type(carried, 400) == ProxyErrorTypes.budget_exceeded.value + + +@pytest.mark.parametrize("carried_type", [None, 400, {"type": "invalid_request_error"}, ["invalid_request_error"]]) +def test_a_non_string_carried_type_falls_back_to_the_status(carried_type: object): + """OpenAI types error.type as a string, so anything else on the exception is not one and + must not reach the wire the way the literal "None" used to.""" + + class _Carrier(Exception): + type = carried_type + + assert openai_error_type(_Carrier("boom"), 401) == "authentication_error" + + +def test_the_type_is_never_the_string_none_after_a_json_round_trip(): + """The bug this module exists for: json.dumps of a "None" default is indistinguishable + from a real type to a client's error handler.""" + payload = json.loads( + json.dumps( + { + "type": openai_error_type(HTTPException(status_code=400, detail="boom"), 400), + "param": openai_error_param(HTTPException(status_code=400, detail="boom")), + } + ) + ) + + assert payload == {"type": "invalid_request_error", "param": None} + + +def test_a_carried_param_names_the_offending_field(): + carried = ProxyException(message="Invalid purpose", type="invalid_request_error", param="purpose", code=400) + + assert openai_error_param(carried) == "purpose" + + +@pytest.mark.parametrize("exc", [HTTPException(status_code=400, detail="boom"), ValueError("boom"), None]) +def test_param_is_json_null_when_the_exception_names_no_field(exc: Exception | None): + assert openai_error_param(exc) is None + + +def test_a_non_string_carried_param_is_json_null(): + class _Carrier(Exception): + param = 42 + + assert openai_error_param(_Carrier("boom")) is None + + +def test_a_carried_status_code_wins_over_the_default(): + assert error_status_code(HTTPException(status_code=429, detail="slow down"), 400) == 429 + + +@pytest.mark.parametrize("default", [400, 500]) +def test_the_default_status_stands_when_the_exception_carries_none(default: int): + assert error_status_code(ValueError("boom"), default) == default + + +@pytest.mark.parametrize("carried_status", [True, False, "429", None, 429.0]) +def test_a_non_int_carried_status_falls_back_to_the_default(carried_status: object): + """True is an int in Python but not an HTTP status, and a stringified one would break + every caller that compares the code numerically.""" + + class _Carrier(Exception): + status_code = carried_status + + assert error_status_code(_Carrier("boom"), 500) == 500 + + +def test_a_proxy_exception_keeps_the_status_it_was_raised_with(): + """ProxyException stores its status as the string ``code`` rather than ``status_code``, + so a route tail that rewraps one used to answer a 4xx rejection as a 500.""" + rejection = ProxyException(message="session_id is required", type="bad_request_error", param="session_id", code=400) + + assert error_status_code(rejection, 500) == 400 + + +@pytest.mark.parametrize("carried_code", [None, "None", "", "rate_limited", "4xx", 404]) +def test_a_code_that_is_not_a_decimal_string_falls_back_to_the_default(carried_code: object): + """Only ProxyException's stringified status is a status; ``code`` on anything else + (OpenAI's ``invalid_api_key``, a stray int) says nothing about the HTTP answer.""" + + class _Carrier(Exception): + code = carried_code + + assert error_status_code(_Carrier("boom"), 500) == 500 + + +def test_a_status_code_wins_over_a_stringified_code(): + class _Carrier(Exception): + status_code = 429 + code = "400" + + assert error_status_code(_Carrier("boom"), 500) == 429 + + +def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): + """The two helpers compose at every call site: the status the exception carries is what + names its type, not the default the route would have used.""" + exc = HTTPException(status_code=403, detail="blocked by policy") + + assert openai_error_type(exc, error_status_code(exc, 400)) == "permission_error" diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index f0fbdea4e85..6f7c20166c5 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -393,6 +393,51 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc assert not proxy_server.MODEL_RECONCILE_LOCK.locked() +@pytest.mark.asyncio +async def test_resync_model_deployments_loads_db_credentials_before_reconciling_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unittest.mock import AsyncMock, MagicMock + + import litellm + import litellm.proxy.proxy_server as proxy_server + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments + from litellm.types.utils import CredentialItem + + rows: Final = [MagicMock()] + prisma_client: Final = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows) + router: Final = MagicMock() + router.get_model_list.return_value = [] + installed: Final = MagicMock() + + async def load_credentials_from_db(prisma_client: object) -> None: + CredentialAccessor.upsert_credentials( + [ + CredentialItem( + credential_name="openai-cred", + credential_values={"api_key": "sk-from-db"}, + credential_info={}, + ) + ] + ) + + def install_models(db_models: object) -> None: + installed(db_models=db_models, credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(litellm, "credential_list", []) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", load_credentials_from_db) + monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", install_models) + + assert await _resync_model_deployments("model-created-on-a-sibling-replica") is True + installed.assert_called_once_with(db_models=rows, credential={"api_key": "sk-from-db"}) + + @pytest.mark.asyncio async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch): from unittest.mock import AsyncMock, MagicMock diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 56c0efb41d2..560953f0b51 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -403,7 +403,7 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): { "table": "enduser", "op": "update_many", - "where": {"user_id": {"in": ["test-enduser-1"]}}, + "where": {"budget_id": {"in": ["test-budget-1"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, } ] @@ -504,7 +504,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): { "table": "enduser", "op": "update_many", - "where": {"user_id": {"in": ["test-enduser-1"]}}, + "where": {"budget_id": {"in": ["test-budget-1"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, } ] @@ -524,6 +524,7 @@ _LINKED_TABLE_CASES = [ ("org", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("tag", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("model_access_group", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), + ("enduser", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ] @@ -553,6 +554,48 @@ def test_budget_table_reset_zeroes_spend_on_every_linked_table( assert writes[0]["data"] == {"spend": 0} +_POSTGRES_MAX_BIND_VARIABLES: Final = 32767 + + +def _bind_count(where: Dict[str, Any]) -> int: + """Bind variables one prisma where-clause compiles to: each scalar is one + placeholder and an ``in`` list contributes one per element.""" + return sum(len(value["in"]) if isinstance(value, dict) and "in" in value else 1 for value in where.values()) + + +@pytest.mark.parametrize("population", [3, 40_000], ids=["small", "over-pg-bind-ceiling"]) +def test_enduser_reset_bind_count_does_not_scale_with_population(reset_budget_job, mock_prisma_client, population): + """Regression for #40564. + + Enumerating every dependent user id put one bind variable per customer into + a single prepared statement. Past PostgreSQL's ceiling the statement could + not be parsed at all, so the whole atomic cascade rolled back, + budget_reset_at never advanced, and the tier stayed due on every later tick + forever. Matching on the budget link keeps the statement the same size no + matter how many customers share a tier. + """ + budget = _budget_row(budget_id="shared-tier", budget_duration="1d") + mock_prisma_client.data["budget"] = [budget] + mock_prisma_client.data["enduser"] = [ + types.SimpleNamespace( + spend=1.0, + litellm_budget_table=budget, + user_id=f"cust-{index:08d}", + budget_id="shared-tier", + ) + for index in range(population) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + writes = _batch_writes(mock_prisma_client, "enduser") + assert [_bind_count(write["where"]) for write in writes] == [2], ( + f"the cascade must not enumerate {population} user ids: past " + f"{_POSTGRES_MAX_BIND_VARIABLES} binds PostgreSQL refuses the statement, got {writes[:1]}" + ) + assert _batch_writes(mock_prisma_client, "budget")[0]["data"]["budget_reset_at"] is not None + + def test_budget_table_reset_writes_nothing_when_no_budget_is_due(reset_budget_job, mock_prisma_client): """Nothing due means no transaction is opened at all.""" asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -720,14 +763,22 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - # Both end users are zeroed by the same committed statement. - enduser_writes = _batch_writes(mock_prisma_client, "enduser") - assert len(enduser_writes) == 1, f"Expected a single enduser write, got {enduser_writes}" - assert set(enduser_writes[0]["where"]["user_id"]["in"]) == { - "enduser-explicit", - "enduser-implicit", - } - assert enduser_writes[0]["data"] == {"spend": 0} + # Both end users are zeroed: the linked rows on the tier's budget_id, the + # implicit ones on the NULL branch that stands in for the default tier. + assert _batch_writes(mock_prisma_client, "enduser") == [ + { + "table": "enduser", + "op": "update_many", + "where": {"budget_id": {"in": [default_budget_id]}, "spend": {"gt": 0}}, + "data": {"spend": 0}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"budget_id": None, "spend": {"gt": 0}}, + "data": {"spend": 0}, + }, + ] # Verify find_many was called to fetch NULL-budget-id end users find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls @@ -3043,13 +3094,13 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( assert { "table": "enduser", "op": "update_many", - "where": {"user_id": {"in": ["enduser-roll"]}, "spend": {"gt": 10.0}}, + "where": {"budget_id": "budget-roll", "spend": {"gt": 10.0}}, "data": {"spend": {"decrement": 10.0}}, } in enduser_writes assert { "table": "enduser", "op": "update_many", - "where": {"user_id": {"in": ["enduser-roll"]}, "spend": {"lte": 10.0}}, + "where": {"budget_id": "budget-roll", "spend": {"gt": 0, "lte": 10.0}}, "data": {"spend": 0}, } in enduser_writes diff --git a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py index ca4d62737b6..54e0aa74a25 100644 --- a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py +++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py @@ -23,7 +23,7 @@ from litellm.proxy.common_utils.scheduled_job_stagger import ( ) OPERATOR_CRON_JOB_ID = "spend_log_cleanup_job" -SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "get_credentials_job", "add_deployment_job") +SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "add_deployment_job") async def _noop() -> None: ... diff --git a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py index 33e0d8bf38e..2d5d76ed542 100644 --- a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py @@ -1,3 +1,4 @@ +import hashlib import json from typing import Any @@ -10,10 +11,14 @@ from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + end_user_cache_key, get_management_object_ttl, + is_user_key_cache_key, ) from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum +HASHED_TOKEN = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest() + class CapturingInMemoryCache(InMemoryCache): """Records ``ttl`` passed into ``set_cache`` (what DualCache injects).""" @@ -204,9 +209,7 @@ class TestUserApiKeyCache: # Bypass UserApiKeyCache.serialize: CacheCodec rejects non-dict cached values # for dict-based models (deserialize returns None). - await cache.in_memory_cache.async_set_cache( - key="k", value="invalid-payload-not-a-dict" - ) + await cache.in_memory_cache.async_set_cache(key="k", value="invalid-payload-not-a-dict") value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth) assert value is None @@ -224,6 +227,141 @@ class TestUserApiKeyCache: fake.set_cache("k2", {"ok": NotSerializable()}) +class TestUserKeyObjectPartition: + """ + Regression for LIT-7563: user-key objects share one 200-entry ``InMemoryCache`` with + every other management object, so end-user / team / tag churn evicts hot keys and + forces a ``LiteLLM_VerificationToken`` lookup on the next request. + """ + + @pytest.mark.parametrize( + ("key", "expected"), + [ + (HASHED_TOKEN, True), + (HASHED_TOKEN.upper(), False), + (f"team_id:{HASHED_TOKEN}", False), + (end_user_cache_key("u1"), False), + ("sk-lit7563-hot-key", False), + ], + ) + def test_is_user_key_cache_key(self, key: str, expected: bool): + assert is_user_key_cache_key(key) is expected + + @pytest.mark.asyncio + async def test_management_object_churn_does_not_evict_key_object(self): + cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2)) + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth, ttl=100) + for i in range(2): + await cache.async_set_cache(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}, ttl=200) + + key_obj = await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + assert cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"} + assert HASHED_TOKEN not in cache.in_memory_cache.cache_dict + + def test_sync_write_and_read_route_to_key_object_partition(self): + cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2)) + cache.set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth, ttl=100) + for i in range(2): + cache.set_cache(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}, ttl=200) + + key_obj = cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + + @pytest.mark.asyncio + async def test_redis_hit_backfills_key_object_partition_with_configured_ttl(self): + redis = FakeRedisCache() + writer = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=30) + await writer.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + + key_partition = CapturingInMemoryCache() + reader = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=30, key_object_in_memory_cache=key_partition) + key_obj = await reader.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + assert key_partition.last_ttl == 30 + assert HASHED_TOKEN not in reader.in_memory_cache.cache_dict + + @pytest.mark.asyncio + async def test_update_cache_ttl_applies_to_key_object_partition(self): + key_partition = CapturingInMemoryCache() + cache = UserApiKeyCache(default_in_memory_ttl=60, key_object_in_memory_cache=key_partition) + cache.update_cache_ttl(default_in_memory_ttl=7, default_redis_ttl=7) + + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + + assert key_partition.last_ttl == 7 + + @pytest.mark.asyncio + async def test_attach_redis_cache_applies_to_key_object_partition(self): + redis = FakeRedisCache() + cache = UserApiKeyCache() + cache.attach_redis_cache(redis) + + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + + other_worker = UserApiKeyCache(redis_cache=redis) + key_obj = await other_worker.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + + @pytest.mark.asyncio + async def test_delete_removes_key_object_from_partition_and_redis(self): + redis = FakeRedisCache() + cache = UserApiKeyCache(redis_cache=redis) + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is not None + + cache.delete_cache(HASHED_TOKEN) + + assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None + assert await redis.async_get_cache(HASHED_TOKEN) is None + + @pytest.mark.asyncio + async def test_async_delete_removes_key_object_from_partition_and_redis(self): + redis = FakeRedisCache() + cache = UserApiKeyCache(redis_cache=redis) + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + + await cache.async_delete_cache(HASHED_TOKEN) + + assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None + assert await redis.async_get_cache(HASHED_TOKEN) is None + + @pytest.mark.asyncio + async def test_pipeline_write_routes_each_entry_to_its_partition(self): + cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2)) + await cache.async_set_cache_pipeline( + [(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN))] + + [(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}) for i in range(2)], + ttl=100, + ) + + key_obj = await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + assert HASHED_TOKEN not in cache.in_memory_cache.cache_dict + assert cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"} + + def test_flush_clears_key_object_partition(self): + cache = UserApiKeyCache() + cache.set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + cache.set_cache(end_user_cache_key("u1"), {"user_id": "u1"}) + + cache.flush_cache() + + assert cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None + assert cache.get_cache(end_user_cache_key("u1")) is None + + def test_in_memory_cache_for_routes_by_key(self): + cache = UserApiKeyCache() + assert cache.in_memory_cache_for(HASHED_TOKEN) is cache.key_object_cache.in_memory_cache + assert cache.in_memory_cache_for(end_user_cache_key("u1")) is cache.in_memory_cache + + class TestManagementObjectTTL: """ Regression for LIT-3338: ``general_settings.user_api_key_cache_ttl`` (which the @@ -238,19 +376,13 @@ class TestManagementObjectTTL: def test_falls_back_to_constant_when_no_default_configured(self): cache = UserApiKeyCache() assert cache.default_in_memory_ttl is None - assert ( - get_management_object_ttl(cache) - == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL - ) + assert get_management_object_ttl(cache) == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL def test_resolves_on_a_plain_dual_cache(self): # Many call sites are typed UserApiKeyCache but exercised in tests with a # bare DualCache; the resolver must work on the base type, not just the subclass. assert get_management_object_ttl(DualCache(default_in_memory_ttl=300)) == 300 - assert ( - get_management_object_ttl(DualCache()) - == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL - ) + assert get_management_object_ttl(DualCache()) == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL @pytest.mark.asyncio async def test_management_write_uses_configured_ttl_over_constant(self): @@ -260,9 +392,7 @@ class TestManagementObjectTTL: redis_cache=FakeRedisCache(), default_in_memory_ttl=300, ) - assert get_management_object_ttl(cache) != ( - DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL - ) + assert get_management_object_ttl(cache) != (DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL) await cache.async_set_cache( "team_id:abc", diff --git a/tests/test_litellm/proxy/db/conftest.py b/tests/test_litellm/proxy/db/conftest.py index 6f67b91ac1d..d3226b0ec50 100644 --- a/tests/test_litellm/proxy/db/conftest.py +++ b/tests/test_litellm/proxy/db/conftest.py @@ -1,6 +1,12 @@ +import json import os +import signal +import sys +import time from collections.abc import Generator -from typing import Optional +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Optional import pytest @@ -75,3 +81,85 @@ def reset_entra_token_provider_cache() -> Generator[None, None, None]: def unset_database_url(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DATABASE_URL", "about-to-be-unset") monkeypatch.delenv("DATABASE_URL") + + +FAKE_PRISMA_CLI = """#!{python} +import json +import os +import pathlib +import subprocess +import sys +import time + +calls_file = pathlib.Path(os.environ["FAKE_PRISMA_CALLS"]) +earlier_calls = calls_file.read_text().splitlines() if calls_file.exists() else [] +with calls_file.open("a") as log: + print(json.dumps(sys.argv[1:]), file=log) +if not earlier_calls and os.environ.get("FAKE_PRISMA_HANG_FIRST"): + grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"]) + pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid)) + time.sleep(600) +sys.exit(0) +""" + + +@dataclass(frozen=True, slots=True) +class FakePrismaCli: + """A stand-in `prisma` on PATH, recording every invocation. + + With FAKE_PRISMA_HANG_FIRST set it hangs on its first call from a process tree + of its own, the way the real CLI wraps Node around a Rust schema engine, so a + timeout that kills only the direct child leaves the rest of that tree running. + """ + + calls_file: Path + grandchild_pidfile: Path + + @property + def calls(self) -> list[list[str]]: + if not self.calls_file.exists(): + return [] + return [json.loads(line) for line in self.calls_file.read_text().splitlines()] + + def grandchild_is_gone(self, within_seconds: float) -> bool: + pid: Final = int(self.grandchild_pidfile.read_text()) + deadline: Final = time.monotonic() + within_seconds + while time.monotonic() < deadline: + if os.name != "nt": + try: + reaped_pid, _ = os.waitpid(pid, os.WNOHANG) + if reaped_pid == pid: + return True + except ChildProcessError: + pass + try: + os.kill(pid, 0) + except ProcessLookupError: + return True + time.sleep(0.05) + return False + + +@pytest.fixture +def fake_prisma_cli(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[FakePrismaCli, None, None]: + bin_dir = tmp_path / "fakebin" + bin_dir.mkdir() + script = bin_dir / "prisma" + script.write_text(FAKE_PRISMA_CLI.format(python=sys.executable)) + script.chmod(0o755) + cli = FakePrismaCli( + calls_file=tmp_path / "calls.jsonl", + grandchild_pidfile=tmp_path / "grandchild.pid", + ) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}") + monkeypatch.setenv("FAKE_PRISMA_CALLS", str(cli.calls_file)) + monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(cli.grandchild_pidfile)) + monkeypatch.setenv("LITELLM_PRISMA_COMMAND_TIMEOUT", "1") + monkeypatch.delenv("FAKE_PRISMA_HANG_FIRST", raising=False) + yield cli + if cli.grandchild_pidfile.exists(): + try: + os.kill(int(cli.grandchild_pidfile.read_text()), signal.SIGKILL) + except ProcessLookupError: + pass + assert cli.grandchild_is_gone(within_seconds=5) diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py index ecd5c5f50c0..4684c3213d6 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -1,4 +1,5 @@ import json +import logging from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch @@ -6,6 +7,7 @@ import pytest from fastapi.testclient import TestClient +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager @@ -215,6 +217,22 @@ async def test_redis_error_handling(pod_lock_manager, mock_redis): ) +@pytest.mark.asyncio +async def test_lock_refused_by_the_open_circuit_breaker_is_not_logged_as_an_error(pod_lock_manager, mock_redis, caplog): + """Every cron job retries its lock on a timer, so an open breaker must not add an error line per cycle.""" + refused = RedisCircuitBreakerOpenError("Redis circuit breaker is open - skipping async_set_cache") + mock_redis.async_set_cache.side_effect = refused + mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id + mock_redis.async_delete_cache.side_effect = refused + + with caplog.at_level(logging.ERROR): + acquired = await pod_lock_manager.acquire_lock(cronjob_id="test_job") + await pod_lock_manager.release_lock(cronjob_id="test_job") + + assert acquired is False + assert caplog.records == [] + + @pytest.mark.asyncio async def test_bytes_handling(pod_lock_manager, mock_redis): """ diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py index 4ea655b8871..609dd13afc2 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py @@ -4,7 +4,7 @@ selection, the non-partitioned no-op safety path, and the drop/ensure SQL flow. """ from contextlib import asynccontextmanager -from datetime import date, datetime, timezone +from datetime import date, datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock import pytest @@ -19,7 +19,6 @@ from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import ( upcoming_partitions, ) - DDL_TIMEOUT_MS = 30000 @@ -28,19 +27,23 @@ def _budget(ms: "int | None" = DDL_TIMEOUT_MS): return lambda: ms -def _wire_tx(db) -> list[str]: +def _wire_tx(db) -> "tuple[list[str], list[int | timedelta | None]]": """ Model the prisma seam the partition DDL uses. Every statement this manager issues, DDL and catalog query alike, runs inside db.tx() so it can carry SET LOCAL timeouts. Those SET LOCAL statements are - collected in the returned list rather than forwarded, so assertions on - db.execute_raw and db.query_raw still see only the real statements. + collected in the first returned list rather than forwarded, so assertions on + db.execute_raw and db.query_raw still see only the real statements. The + second list records the interactive-transaction timeout each tx was opened + with. """ session_settings: list[str] = [] + tx_timeouts: list[int | timedelta | None] = [] @asynccontextmanager - async def _tx(): + async def _tx(*, max_wait: "int | timedelta | None" = None, timeout: "int | timedelta | None" = None): + tx_timeouts.append(timeout) tx = MagicMock() async def _execute_raw(sql, *args): @@ -57,7 +60,7 @@ def _wire_tx(db) -> list[str]: yield tx db.tx = _tx - return session_settings + return session_settings, tx_timeouts def test_period_start_per_interval(): @@ -227,7 +230,7 @@ async def test_partition_ddl_carries_a_statement_and_lock_timeout(): } ] ) - session_settings = _wire_tx(client.db) + session_settings, _ = _wire_tx(client.db) await mgr.ensure_partitions(client, _budget(7000)) await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), _budget(7000)) @@ -249,7 +252,7 @@ async def test_catalog_queries_carry_a_statement_timeout(): mgr = SpendLogsPartitionManager() client = MagicMock() client.db.query_raw = AsyncMock(return_value=[]) - session_settings = _wire_tx(client.db) + session_settings, _ = _wire_tx(client.db) await mgr.is_partitioned(client, _budget(4000)) assert session_settings == ["SET LOCAL statement_timeout = 4000"], ( @@ -263,6 +266,34 @@ async def test_catalog_queries_carry_a_statement_timeout(): ) +@pytest.mark.asyncio +async def test_partition_transactions_outlive_their_statement_bound(): + """ + prisma's interactive transaction has its own timeout, 5s by default, which + keeps ticking while a statement waits on the partition lock. A tx shorter + than the SET LOCAL bound it carries is closed mid-lock-wait and the engine + then answers the next call with a 422, so the partition is silently not + created. A tx equal to the bound is closed too: the statement can consume + its whole bound waiting on the lock, then still needs to commit. + """ + mgr = SpendLogsPartitionManager() + client = MagicMock() + client.db.execute_raw = AsyncMock(return_value=0) + client.db.query_raw = AsyncMock(return_value=[]) + _, tx_timeouts = _wire_tx(client.db) + + await mgr.is_partitioned(client, _budget()) + await mgr.ensure_partitions(client, _budget()) + await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), _budget()) + + assert len(tx_timeouts) > 0 + for tx_timeout in tx_timeouts: + assert isinstance(tx_timeout, timedelta) + assert tx_timeout > timedelta(milliseconds=DDL_TIMEOUT_MS), ( + f"tx timeout {tx_timeout} does not outlive its {DDL_TIMEOUT_MS}ms statement bound" + ) + + @pytest.mark.asyncio async def test_partition_loops_stop_when_the_budget_runs_out_mid_way(): """ @@ -308,17 +339,15 @@ async def test_partition_maintenance_issues_nothing_when_the_budget_is_already_s def test_unsupported_interval_raises(): - with pytest.raises(ValueError, match='Unsupported partition interval: year'): + with pytest.raises(ValueError, match="Unsupported partition interval: year"): period_start(date(2026, 6, 1), "year") - with pytest.raises(ValueError, match='Unsupported partition interval: year'): + with pytest.raises(ValueError, match="Unsupported partition interval: year"): next_period_start(date(2026, 6, 1), "year") def test_parse_partition_upper_bound_unparseable_to_value_is_none(): """A TO(...) value that is not a valid timestamp must not raise; return None.""" - assert ( - parse_partition_upper_bound("FOR VALUES FROM ('x') TO ('not-a-date')") is None - ) + assert parse_partition_upper_bound("FOR VALUES FROM ('x') TO ('not-a-date')") is None @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index 2ed4f843711..271751a3ff8 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -9,13 +9,15 @@ request-time transaction builder and the flush contract with an injected fake cl import asyncio import json from datetime import datetime +from types import SimpleNamespace +from typing import Final import httpx import pytest from litellm.proxy.db.autorouter_session_rollup import ( - AutoRouterTurnTransaction, UPSERT_AUTOROUTER_SESSION_SQL, + AutoRouterTurnTransaction, build_autorouter_turn_transaction, flush_autorouter_turn_transactions, ) @@ -57,7 +59,8 @@ class TestBuildTransaction: def test_successful_auto_routed_turn_builds_every_field(self): transaction = _build( metadata=_metadata( - usage_object={"prompt_tokens": 90, "cache_read_input_tokens": 5, "cache_creation_input_tokens": 7} + routing_decision={**ROUTING_DECISION, "savings_baseline_model": "anthropic/claude-opus-5"}, + usage_object={"prompt_tokens": 90, "cache_read_input_tokens": 5, "cache_creation_input_tokens": 7}, ) ) assert transaction == AutoRouterTurnTransaction( @@ -70,10 +73,12 @@ class TestBuildTransaction: total_tokens=100, spend=0.01, saved_spend=0.02, + classifier_cost=0.0, covered=True, cache_hit=True, cache_ttl_seconds=300, cache_touched=True, + baseline_model="anthropic/claude-opus-5", ) @pytest.mark.parametrize( @@ -106,11 +111,25 @@ class TestBuildTransaction: transaction = _build() assert transaction is not None and transaction.tier is None + def test_the_baseline_the_turn_was_priced_against_travels_with_the_turn(self): + decision = {**ROUTING_DECISION, "savings_baseline_model": "anthropic/claude-opus-5"} + transaction = _build(metadata=_metadata(routing_decision=decision)) + assert transaction is not None and transaction.baseline_model == "anthropic/claude-opus-5" + + @pytest.mark.parametrize("baseline", [None, "", 3]) + def test_a_decision_without_a_usable_baseline_records_none(self, baseline: object): + decision = {**ROUTING_DECISION, "savings_baseline_model": baseline} + transaction = _build(metadata=_metadata(routing_decision=decision)) + assert transaction is not None and transaction.baseline_model is None + def test_a_priced_classifier_rides_the_turns_spend(self): """The classifier row is excluded from the rollup, so its charge lands here, folded once into the turn that paid for it (GH #38816).""" transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "classifier_cost": 0.005})) assert transaction is not None and transaction.spend == pytest.approx(0.015) + assert transaction.classifier_cost == 0.005 + assert transaction.spend - transaction.classifier_cost == pytest.approx(0.01) + assert transaction.saved_spend == 0.02 @pytest.mark.parametrize( "decision_extra", [{}, {"classifier_cost": 0.0}, {"classifier_cost": "bogus"}, {"classifier_cost": True}] @@ -118,6 +137,8 @@ class TestBuildTransaction: def test_an_unpriced_classifier_leaves_the_spend_alone(self, decision_extra: dict): transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, **decision_extra})) assert transaction is not None and transaction.spend == pytest.approx(0.01) + assert transaction.classifier_cost == 0.0 + assert transaction.saved_spend == 0.02 def test_every_turn_carries_its_own_classifier_charge(self): first = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "classifier_cost": 0.005})) @@ -127,6 +148,7 @@ class TestBuildTransaction: ) assert first is not None and first.spend == pytest.approx(0.015) assert second is not None and second.spend == pytest.approx(0.027) + assert (first.classifier_cost, second.classifier_cost) == (0.005, 0.007) def test_router_name_falls_back_to_the_payload_model_group(self): transaction = _build(metadata=_metadata(routing_decision={"router_type": "complexity"})) @@ -206,6 +228,7 @@ def _transaction( session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, 0, 0), tier: str | None = "medium", + baseline_model: str | None = "anthropic/claude-opus-5", ) -> AutoRouterTurnTransaction: return AutoRouterTurnTransaction( api_key="k1", @@ -217,11 +240,13 @@ def _transaction( total_tokens=100, spend=0.01, saved_spend=0.02, + classifier_cost=0.005, covered=True, cache_hit=False, cache_ttl_seconds=None, cache_touched=False, tier=tier, + baseline_model=baseline_model, ) @@ -249,11 +274,13 @@ class TestFlush: 100, 0.01, 0.02, + 0.005, 1, 0, None, 0, "medium", + "anthropic/claude-opus-5", ) def test_a_connect_error_retries_the_same_statement(self): @@ -279,26 +306,31 @@ class TestFlush: class TestEnqueueSeam: @pytest.mark.asyncio - async def test_update_database_seam_enqueues_only_auto_routed_success(self, monkeypatch: pytest.MonkeyPatch): + @pytest.mark.parametrize("classifier_cost", [0.005, 0.0, None]) + async def test_update_database_seam_enqueues_only_auto_routed_success(self, classifier_cost: float | None): from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter - from litellm.proxy.utils import PrismaClient - monkeypatch.setattr(PrismaClient, "autorouter_turn_transactions", []) - writer = DBSpendUpdateWriter() - fake_prisma = type("P", (), {})() - fake_prisma._autorouter_turn_transactions_lock = asyncio.Lock() - fake_prisma.autorouter_turn_transactions = [] + writer: Final = DBSpendUpdateWriter() + fake_prisma: Final = SimpleNamespace( + _autorouter_turn_transactions_lock=asyncio.Lock(), autorouter_turn_transactions=[] + ) + metadata: Final = _metadata( + routing_decision={**ROUTING_DECISION, "classifier_cost": classifier_cost}, autorouter_savings=-0.003 + ) + for payload in ( + _payload(metadata=json.dumps(metadata)), + _payload(metadata=json.dumps({"usage_object": {"prompt_tokens": 9}})), + _payload(status="failure", metadata=json.dumps(metadata)), + _payload(metadata=json.dumps({**metadata, "internal_call_origin": "autorouter_classifier"})), + ): + await writer._enqueue_autorouter_turn_transaction(payload=payload, prisma_client=fake_prisma) - routed = _payload() - routed["metadata"] = json.dumps(_metadata()) - await writer._enqueue_autorouter_turn_transaction(payload=routed, prisma_client=fake_prisma) - - plain = _payload() - plain["metadata"] = json.dumps({"usage_object": {"prompt_tokens": 9}}) - await writer._enqueue_autorouter_turn_transaction(payload=plain, prisma_client=fake_prisma) - - assert [t.router_name for t in fake_prisma.autorouter_turn_transactions] == ["live-auto"] - assert fake_prisma.autorouter_turn_transactions[0].saved_spend == 0.0 + assert len(fake_prisma.autorouter_turn_transactions) == 1 + transaction: Final = fake_prisma.autorouter_turn_transactions[0] + assert transaction.router_name == "live-auto" + assert transaction.spend == pytest.approx(0.01 + (classifier_cost or 0.0)) + assert transaction.classifier_cost == (classifier_cost or 0.0) + assert transaction.saved_spend == -0.003 def test_every_drain_trigger_reads_the_one_queue_census_owner(): diff --git a/tests/test_litellm/proxy/db/test_check_migration.py b/tests/test_litellm/proxy/db/test_check_migration.py index 9e2f6a1089c..74a6841cb23 100644 --- a/tests/test_litellm/proxy/db/test_check_migration.py +++ b/tests/test_litellm/proxy/db/test_check_migration.py @@ -37,3 +37,35 @@ def test_check_migration_out_of_sync(mocker): check_migration.verbose_logger.exception.assert_called_once() actual_message = check_migration.verbose_logger.exception.call_args[0][0] assert "prisma schema out of sync with db" in actual_message + + +@pytest.mark.timeout(30) +def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fake_prisma_cli, monkeypatch): + """ + `prisma migrate diff` ran unbounded, so a database that never answers hung boot + before uvicorn ever started, and interrupting the proxy orphaned the schema engine. + """ + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1") + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) + assert fake_prisma_cli.calls == [ + ["migrate", "diff", "--from-url", "postgresql://u:p@localhost:9/x", + "--to-schema-datamodel", "./schema.prisma", "--script"] + ] + assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(monkeypatch): + """ + Boot calls this helper directly, so an ImportError here takes the proxy down before + uvicorn starts. An install without the runner must lose the diagnostic, not the proxy. + """ + import sys + + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 11ef911de3e..5e977712a1e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -7,6 +7,7 @@ import re from collections.abc import Callable from contextlib import asynccontextmanager from datetime import datetime, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -2936,16 +2937,16 @@ 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. """ - from litellm.proxy.utils import PrismaClient - db_writer = DBSpendUpdateWriter() prisma = _tool_usage_prisma() - PrismaClient.spend_log_flush_requested.clear() + prisma.spend_log_flush_requested = asyncio.Event() await db_writer._insert_spend_log_to_db( payload={"request_id": "req-1", "call_type": call_type}, @@ -2953,8 +2954,304 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(c ) assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}] - assert PrismaClient.spend_log_flush_requested.is_set() is expects_flush - PrismaClient.spend_log_flush_requested.clear() + assert prisma.spend_log_flush_requested.is_set() is expects_flush + + +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 diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index db524625a93..875dca4bee3 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -11,15 +11,30 @@ clobber a pre-existing ``DATABASE_URL_READ_REPLICA``. A pre-existing ``DATABASE_URL`` (password auth) is likewise left untouched. """ +import datetime +import hashlib import os +import socket +import ssl +import tempfile +import threading import urllib.parse +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final from unittest.mock import patch import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec from pydantic import ValidationError from litellm.proxy.db.db_url_settings import ( + PG_SSL_REQUEST, DatabaseURLSettings, + translate_libpq_ssl_params, unsupported_db_scheme, unsupported_db_scheme_message, ) @@ -35,6 +50,7 @@ _MANAGED_DB_ENV_VARS = ( "IAM_TOKEN_DB_AUTH", "AZURE_POSTGRESQL_AUTH", "DATABASE_DISABLE_PREPARED_STATEMENTS", + "DATABASE_MAX_IDLE_CONNECTION_LIFETIME", "DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA", @@ -111,7 +127,7 @@ def test_assembles_writer_url_when_iam_enabled(monkeypatch): assert ( os.environ["DATABASE_URL"] - == "postgresql://litellm:WRITER_TOKEN@writer.example.com:5432/litellm_db" + == "postgresql://litellm:WRITER_TOKEN@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) # Reader was never configured, so it must not have been set. assert "DATABASE_URL_READ_REPLICA" not in os.environ @@ -130,7 +146,9 @@ def test_a_pre_encoded_iam_user_survives_url_assembly(monkeypatch): with _stub_iam_token("WRITER_TOKEN"): assert _apply() is True - assert os.environ["DATABASE_URL"] == "postgresql://svc%40corp:WRITER_TOKEN@writer.example.com:5432/litellm_db" + assert os.environ["DATABASE_URL"] == ( + "postgresql://svc%40corp:WRITER_TOKEN@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60" + ) def test_an_unreadable_toggle_fails_the_settings_model(monkeypatch): @@ -168,7 +186,7 @@ def test_reader_url_assembled_when_host_set_and_url_unset(monkeypatch): assert ( os.environ["DATABASE_URL_READ_REPLICA"] - == "postgresql://litellm:READER_TOKEN@reader.example.com:5432/litellm_db" + == "postgresql://litellm:READER_TOKEN@reader.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) @@ -191,7 +209,7 @@ def test_reader_url_not_clobbered_when_already_set(monkeypatch): assert ( os.environ["DATABASE_URL_READ_REPLICA"] - == "postgresql://app:secret@reader.example.com:5432/litellm_db" + == "postgresql://app:secret@reader.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) @@ -222,7 +240,8 @@ def test_reader_field_fallbacks_default_to_writer_values(monkeypatch): assert ( os.environ["DATABASE_URL_READ_REPLICA"] - == "postgresql://litellm:READER_TOKEN@reader.example.com:5432/litellm_db?schema=public" + == "postgresql://litellm:READER_TOKEN@reader.example.com:5432/litellm_db" + "?schema=public&max_idle_connection_lifetime=60" ) @@ -242,7 +261,7 @@ def test_assembles_writer_url_when_azure_entra_enabled(monkeypatch): assert os.environ["DATABASE_URL"] == ( "postgresql://litellm%40contoso.onmicrosoft.com:ENTRA_TOKEN" - "@writer.postgres.database.azure.com:5432/litellm_db" + "@writer.postgres.database.azure.com:5432/litellm_db?max_idle_connection_lifetime=60" ) assert os.environ["AZURE_POSTGRESQL_AUTH"] == "True" assert "IAM_TOKEN_DB_AUTH" not in os.environ @@ -261,7 +280,7 @@ def test_azure_reader_url_assembled_from_writer_fallbacks(monkeypatch): assert os.environ["DATABASE_URL_READ_REPLICA"] == ( "postgresql://litellm%40contoso.onmicrosoft.com:ENTRA_TOKEN" - "@reader.postgres.database.azure.com:5432/litellm_db?schema=public" + "@reader.postgres.database.azure.com:5432/litellm_db?schema=public&max_idle_connection_lifetime=60" ) @@ -357,7 +376,7 @@ def test_assembles_writer_url_from_password(monkeypatch): assert _apply() is True assert ( os.environ["DATABASE_URL"] - == "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db" + == "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) @@ -370,16 +389,14 @@ def test_writer_password_is_percent_encoded(monkeypatch): assert _apply() is True assert ( os.environ["DATABASE_URL"] - == "postgresql://litellm:p%40ss%2Fw%3Ard@writer.example.com:5432/litellm_db" + == "postgresql://litellm:p%40ss%2Fw%3Ard@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) def test_writer_url_not_clobbered_when_already_set(monkeypatch): """An operator-pinned DATABASE_URL (e.g. helm's $(VAR) assembly) always wins over the discrete fields.""" - monkeypatch.setenv( - "DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db" - ) + monkeypatch.setenv("DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db") monkeypatch.setenv("DATABASE_HOST", "writer.example.com") monkeypatch.setenv("DATABASE_USER", "litellm") monkeypatch.setenv("DATABASE_NAME", "litellm_db") @@ -388,7 +405,7 @@ def test_writer_url_not_clobbered_when_already_set(monkeypatch): assert _apply() is False assert ( os.environ["DATABASE_URL"] - == "postgresql://pinned:url@db.example.com:5432/litellm_db" + == "postgresql://pinned:url@db.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) @@ -400,7 +417,7 @@ def test_writer_url_passwordless(monkeypatch): assert _apply() is True assert ( os.environ["DATABASE_URL"] - == "postgresql://litellm@writer.example.com:5432/litellm_db" + == "postgresql://litellm@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) @@ -415,7 +432,7 @@ def test_database_username_alias(monkeypatch): assert _apply() is True assert ( os.environ["DATABASE_URL"] - == "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db" + == "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) @@ -429,7 +446,7 @@ def test_password_reader_falls_back_to_writer_password(monkeypatch): assert _apply() is True assert ( os.environ["DATABASE_URL_READ_REPLICA"] - == "postgresql://litellm:s3cr3t@reader.example.com:5432/litellm_db" + == "postgresql://litellm:s3cr3t@reader.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) @@ -445,7 +462,7 @@ def test_password_reader_uses_own_credentials(monkeypatch): assert _apply() is True assert ( os.environ["DATABASE_URL_READ_REPLICA"] - == "postgresql://litellm_ro:ro_pw@reader.example.com:5432/litellm_db" + == "postgresql://litellm_ro:ro_pw@reader.example.com:5432/litellm_db?max_idle_connection_lifetime=60" ) @@ -511,9 +528,7 @@ def test_apply_to_env_rejects_pinned_sqlite_direct_url(monkeypatch): def test_apply_to_env_rejects_pinned_non_postgres_reader(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") - monkeypatch.setenv( - "DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db" - ) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db") with pytest.raises(RuntimeError, match=r"DATABASE_URL_READ_REPLICA.*mysql"): _apply() @@ -538,15 +553,11 @@ def test_reader_inherits_writer_connection_params(monkeypatch): "DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db?connection_limit=3&pool_timeout=20&pgbouncer=true", ) - monkeypatch.setenv( - "DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db" - ) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db") _apply() - query = urllib.parse.parse_qs( - urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query - ) + query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query) assert query["connection_limit"] == ["3"] assert query["pool_timeout"] == ["20"] assert query["pgbouncer"] == ["true"] @@ -564,9 +575,7 @@ def test_reader_keeps_its_own_pinned_connection_params(monkeypatch): _apply() - query = urllib.parse.parse_qs( - urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query - ) + query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query) assert query["connection_limit"] == ["50"] assert query["pool_timeout"] == ["20"] @@ -641,19 +650,19 @@ def test_reader_keeps_its_own_options_when_writer_params_are_appended(monkeypatc assert query["connection_limit"] == ["3"] -def test_reader_url_left_alone_when_writer_has_no_params(monkeypatch): +def test_reader_url_left_alone_when_nothing_is_missing(monkeypatch): """No params to inherit must mean the reader URL is not rewritten at all.""" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") monkeypatch.setenv( "DATABASE_URL_READ_REPLICA", - "postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp", + "postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp&max_idle_connection_lifetime=45", ) _apply() assert ( os.environ["DATABASE_URL_READ_REPLICA"] - == "postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp" + == "postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp&max_idle_connection_lifetime=45" ) @@ -671,7 +680,7 @@ def test_disable_prepared_statements_appends_pgbouncer_to_assembled_writer(monke assert _apply() is True assert os.environ["DATABASE_URL"] == ( - "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?pgbouncer=true" + "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?pgbouncer=true&max_idle_connection_lifetime=60" ) assert "DIRECT_URL" not in os.environ @@ -685,7 +694,9 @@ def test_disable_prepared_statements_appends_pgbouncer_to_pinned_writer(monkeypa monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db") assert _apply() is False - assert os.environ["DATABASE_URL"] == "postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=true" + assert os.environ["DATABASE_URL"] == ( + "postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=true&max_idle_connection_lifetime=60" + ) def test_disable_prepared_statements_respects_a_pinned_pgbouncer_value(monkeypatch): @@ -694,7 +705,9 @@ def test_disable_prepared_statements_respects_a_pinned_pgbouncer_value(monkeypat _apply() - assert os.environ["DATABASE_URL"] == "postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=false" + assert os.environ["DATABASE_URL"] == ( + "postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=false&max_idle_connection_lifetime=60" + ) def test_disable_prepared_statements_applies_to_direct_url(monkeypatch): @@ -704,7 +717,9 @@ def test_disable_prepared_statements_applies_to_direct_url(monkeypatch): _apply() - assert os.environ["DIRECT_URL"] == "postgresql://u:p@direct.example.com:5432/litellm_db?pgbouncer=true" + assert os.environ["DIRECT_URL"] == ( + "postgresql://u:p@direct.example.com:5432/litellm_db?pgbouncer=true&max_idle_connection_lifetime=60" + ) def test_reader_inherits_pgbouncer_from_disable_prepared_statements(monkeypatch): @@ -724,7 +739,9 @@ def test_disable_prepared_statements_off_leaves_urls_alone(monkeypatch): _apply() - assert os.environ["DATABASE_URL"] == "postgresql://u:p@db.example.com:5432/litellm_db" + assert os.environ["DATABASE_URL"] == ( + "postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=60" + ) def test_disable_prepared_statements_rejects_an_unreadable_value(monkeypatch): @@ -760,15 +777,184 @@ def test_libpq_verify_full_and_sslrootcert_become_prisma_strict_sslcert(monkeypa "sslmode": ["require"], "sslcert": ["/certs/rds-bundle.pem"], "sslaccept": ["strict"], + "max_idle_connection_lifetime": ["60"], } +def _issue_cert( + subject: str, issuer: x509.Certificate | None, issuer_key: ec.EllipticCurvePrivateKey | None, ca: bool +) -> tuple[x509.Certificate, ec.EllipticCurvePrivateKey]: + key: Final = ec.generate_private_key(ec.SECP256R1()) + name: Final = x509.Name((x509.NameAttribute(x509.NameOID.COMMON_NAME, subject),)) + now: Final = datetime.datetime.now(datetime.timezone.utc) + builder: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(issuer.subject if issuer else name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(minutes=5)) + .not_valid_after(now + datetime.timedelta(days=1)) + .add_extension(x509.BasicConstraints(ca=ca, path_length=None), critical=True) + .add_extension(x509.SubjectAlternativeName((x509.DNSName("localhost"),)), critical=False) + ) + return builder.sign(issuer_key or key, hashes.SHA256()), key + + +def _pem(cert: x509.Certificate) -> bytes: + return cert.public_bytes(serialization.Encoding.PEM) + + +class _TlsPostgresStub: + """Answers one libpq ``SSLRequest`` with ``S`` and serves ``leaf + intermediate``.""" + + def __init__(self, chain_pem: Path, key_pem: Path) -> None: + self.context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + self.context.load_cert_chain(str(chain_pem), str(key_pem)) + self.listener: Final = socket.create_server(("127.0.0.1", 0)) + self.port: Final[int] = self.listener.getsockname()[1] + self.thread: Final = threading.Thread(target=self._serve, daemon=True) + self.thread.start() + + def _serve(self) -> None: + with self.listener: + while True: + try: + conn: socket.socket = self.listener.accept()[0] + except OSError: + return + with conn: + try: + if conn.recv(8) == PG_SSL_REQUEST: + conn.sendall(b"S") + with self.context.wrap_socket(conn, server_side=True) as tls: + tls.recv(1) + except OSError: + continue + + +@dataclass(frozen=True, slots=True) +class _RdsLikePki: + bundle: Path + wrong_bundle: Path + root: Path + port: int + + +@pytest.fixture +def rds_like_pki(tmp_path: Path) -> Iterator[_RdsLikePki]: + """An RDS-shaped trust setup: the server sends leaf + intermediate, the + bundle holds only self-signed roots, and the right root is not first.""" + root, root_key = _issue_cert("Real Root CA", None, None, ca=True) + decoys: Final = tuple(_issue_cert(f"Decoy Root CA {i}", None, None, ca=True)[0] for i in range(3)) + intermediate, intermediate_key = _issue_cert("Intermediate CA", root, root_key, ca=True) + leaf, leaf_key = _issue_cert("localhost", intermediate, intermediate_key, ca=False) + chain_pem: Final = tmp_path / "server-chain.pem" + chain_pem.write_bytes(_pem(leaf) + _pem(intermediate)) + key_pem: Final = tmp_path / "server.key" + key_pem.write_bytes( + leaf_key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() + ) + ) + bundle: Final = tmp_path / "global-bundle.pem" + bundle.write_bytes(b"".join(_pem(decoy) for decoy in decoys) + _pem(root)) + wrong_bundle: Final = tmp_path / "wrong-bundle.pem" + wrong_bundle.write_bytes(b"".join(_pem(decoy) for decoy in decoys)) + root_pem: Final = tmp_path / "root.pem" + root_pem.write_bytes(_pem(root)) + stub: Final = _TlsPostgresStub(chain_pem, key_pem) + yield _RdsLikePki(bundle=bundle, wrong_bundle=wrong_bundle, root=root_pem, port=stub.port) + stub.listener.close() + + +def _params(url: str) -> tuple[tuple[str, str], ...]: + return tuple(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query, keep_blank_values=True)) + + +def test_multi_root_bundle_is_pinned_to_the_root_the_server_chains_to( + monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki +): + """Prisma's ``sslcert`` loads only the first certificate of the file, so + handing it the whole RDS bundle trusts one region's root and fails with + "unable to get local issuer certificate" everywhere else. The URL Prisma + receives must point at a single-certificate file holding the server's root.""" + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db?sslmode=verify-full&sslrootcert={rds_like_pki.bundle}", + ) + + _apply() + + (sslmode, sslcert, sslaccept, _) = _params(os.environ["DATABASE_URL"]) + assert (sslmode, sslaccept) == (("sslmode", "require"), ("sslaccept", "strict")) + assert sslcert[0] == "sslcert" and sslcert[1] != str(rds_like_pki.bundle) + assert Path(sslcert[1]).read_bytes() == rds_like_pki.root.read_bytes() + + +def test_pinned_root_replaces_a_planted_symlink_instead_of_writing_through_it( + monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki, tmp_path: Path +): + """The pinned file has a predictable name in a shared temp dir, so a symlink + planted there must not redirect the write onto its target.""" + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + root_der: Final = x509.load_pem_x509_certificate(rds_like_pki.root.read_bytes()).public_bytes( + serialization.Encoding.DER + ) + pinned: Final = tmp_path / f"litellm-sslcert-{hashlib.sha256(root_der).hexdigest()[:16]}.pem" + victim: Final = tmp_path / "victim.txt" + victim.write_text("untouched") + pinned.symlink_to(victim) + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db?sslmode=verify-full&sslrootcert={rds_like_pki.bundle}", + ) + + _apply() + + assert ("sslcert", str(pinned)) in _params(os.environ["DATABASE_URL"]) + assert victim.read_text() == "untouched" + assert not pinned.is_symlink() and pinned.read_bytes() == rds_like_pki.root.read_bytes() + + +def test_bundle_without_the_servers_root_is_passed_through_unchanged( + monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki +): + """Nothing in the bundle verifies the server, so no root is pinned and + Prisma keeps rejecting the connection instead of trusting a root the + operator never shipped.""" + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db" + f"?sslmode=verify-full&sslrootcert={rds_like_pki.wrong_bundle}", + ) + + _apply() + + assert ("sslcert", str(rds_like_pki.wrong_bundle)) in _params(os.environ["DATABASE_URL"]) + + +def test_root_cert_resolver_receives_the_urls_host_and_default_port(): + def resolver(cert_path: str, host: str, port: int) -> str: + return f"/pinned/{host}/{port}{cert_path}" + + url: Final = translate_libpq_ssl_params( + "postgresql://u:p@db.example.com/litellm_db?sslmode=verify-full&sslrootcert=/certs/bundle.pem", resolver + ) + + assert ("sslcert", "/pinned/db.example.com/5432/certs/bundle.pem") in _params(url) + + def test_libpq_verify_ca_becomes_prisma_strict(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=verify-ca") _apply() - assert _query(os.environ["DATABASE_URL"]) == {"sslmode": ["require"], "sslaccept": ["strict"]} + assert _query(os.environ["DATABASE_URL"]) == { + "sslmode": ["require"], + "sslaccept": ["strict"], + "max_idle_connection_lifetime": ["60"], + } def test_sslrootcert_alone_turns_on_strict_verification(monkeypatch): @@ -784,6 +970,7 @@ def test_sslrootcert_alone_turns_on_strict_verification(monkeypatch): "sslmode": ["require"], "sslcert": ["/certs/ca.pem"], "sslaccept": ["strict"], + "max_idle_connection_lifetime": ["60"], } @@ -800,11 +987,15 @@ def test_pinned_prisma_ssl_params_win_over_libpq_translation(monkeypatch): "sslmode": ["require"], "sslcert": ["/pinned.pem"], "sslaccept": ["accept_invalid_certs"], + "max_idle_connection_lifetime": ["60"], } def test_prisma_native_ssl_url_is_left_untouched(monkeypatch): - url = "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=require&sslcert=/certs/ca.pem&sslaccept=strict" + url = ( + "postgresql://u:p@db.example.com:5432/litellm_db" + "?sslmode=require&sslcert=/certs/ca.pem&sslaccept=strict&max_idle_connection_lifetime=60" + ) monkeypatch.setenv("DATABASE_URL", url) _apply() @@ -820,4 +1011,84 @@ def test_libpq_ssl_translation_covers_direct_url_and_read_replica(monkeypatch): _apply() for env_var in ("DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA"): - assert _query(os.environ[env_var]) == {"sslmode": ["require"], "sslaccept": ["strict"]}, env_var + assert _query(os.environ[env_var]) == { + "sslmode": ["require"], + "sslaccept": ["strict"], + "max_idle_connection_lifetime": ["60"], + }, env_var + + +def test_default_idle_lifetime_applied_to_pinned_writer_and_direct_url(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db") + monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@direct.example.com:5432/litellm_db") + + assert _apply() is False + + assert os.environ["DATABASE_URL"] == ( + "postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=60" + ) + assert os.environ["DIRECT_URL"] == ( + "postgresql://u:p@direct.example.com:5432/litellm_db?max_idle_connection_lifetime=60" + ) + + +def test_url_pinned_idle_lifetime_wins_over_default_and_env_knob(monkeypatch): + monkeypatch.setenv("DATABASE_MAX_IDLE_CONNECTION_LIFETIME", "45") + monkeypatch.setenv( + "DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=300" + ) + + _apply() + + assert os.environ["DATABASE_URL"] == ( + "postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=300" + ) + + +def test_env_knob_overrides_default_idle_lifetime(monkeypatch): + monkeypatch.setenv("DATABASE_MAX_IDLE_CONNECTION_LIFETIME", "45") + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db") + monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@direct.example.com:5432/litellm_db") + + _apply() + + assert os.environ["DATABASE_URL"] == ( + "postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=45" + ) + assert os.environ["DIRECT_URL"] == ( + "postgresql://u:p@direct.example.com:5432/litellm_db?max_idle_connection_lifetime=45" + ) + + +def test_env_knob_rejects_a_non_integer_value(monkeypatch): + monkeypatch.setenv("DATABASE_MAX_IDLE_CONNECTION_LIFETIME", "soon") + + with pytest.raises(ValidationError, match="DATABASE_MAX_IDLE_CONNECTION_LIFETIME"): + DatabaseURLSettings.from_env() + + +@pytest.mark.parametrize(("knob", "expected"), [(None, "60"), ("45", "45")]) +def test_reader_inherits_the_writer_idle_lifetime(monkeypatch, knob, expected): + if knob is not None: + monkeypatch.setenv("DATABASE_MAX_IDLE_CONNECTION_LIFETIME", knob) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db") + + _apply() + + assert os.environ["DATABASE_URL_READ_REPLICA"] == ( + f"postgresql://u:p@reader.example.com:5432/db?max_idle_connection_lifetime={expected}" + ) + + +def test_reader_keeps_its_own_pinned_idle_lifetime(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") + monkeypatch.setenv( + "DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db?max_idle_connection_lifetime=120" + ) + + _apply() + + assert os.environ["DATABASE_URL_READ_REPLICA"] == ( + "postgresql://u:p@reader.example.com:5432/db?max_idle_connection_lifetime=120" + ) diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index c685d778c0e..3f009137a1c 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -665,10 +665,47 @@ def test_is_deadlock_error_excludes_non_deadlocks(error): assert PrismaDBExceptionHandler.is_deadlock_error(error) is False +READ_ONLY_CONNECTOR_ERROR: Final = ( + "Error occurred during query execution:\nConnectorError(ConnectorError { user_facing_error: None, " + 'kind: QueryError(PostgresError { code: "25006", message: "cannot execute UPDATE in a read-only transaction", ' + 'severity: "ERROR", detail: None, column: None, hint: None }), transient: false })' +) + + +@pytest.mark.parametrize( + "error", + [ + DataError(data={"user_facing_error": {"message": READ_ONLY_CONNECTOR_ERROR}}), + RawQueryError(data={"user_facing_error": {"message": "cannot execute INSERT in a read-only transaction"}}), + PrismaError( + 'PostgresError { code: "25006", message: "kann DELETE in einer Read-Only-Transaktion nicht ausführen" }' + ), + ], +) +def test_is_read_only_transaction_error_matches_sqlstate_25006(error): + assert PrismaDBExceptionHandler.is_read_only_transaction_error(error) is True + + +@pytest.mark.parametrize( + "error", + [ + UniqueViolationError(data={"user_facing_error": {"error_code": "P2002", "meta": {"table": "t"}}}), + PrismaError("can't reach database server"), + RawQueryError(data={"user_facing_error": {"message": "deadlock detected", "meta": {"table": "t"}}}), + httpx.ConnectError("connection refused"), + RuntimeError("cannot execute UPDATE in a read-only transaction"), + ValueError('"25006"'), + ], +) +def test_is_read_only_transaction_error_excludes_other_failures(error): + assert PrismaDBExceptionHandler.is_read_only_transaction_error(error) is False + + MOCKED_PRISMA_PREDICATES: Final = ( PrismaDBExceptionHandler.is_database_infrastructure_error, PrismaDBExceptionHandler.is_database_transport_error, PrismaDBExceptionHandler.is_deadlock_error, + PrismaDBExceptionHandler.is_read_only_transaction_error, PrismaDBExceptionHandler.is_prisma_engine_internal_error, PrismaDBExceptionHandler.is_database_service_unavailable_error, ) diff --git a/tests/test_litellm/proxy/db/test_gateway_request_tracking.py b/tests/test_litellm/proxy/db/test_gateway_request_tracking.py index 93a11a914cb..045261e2d53 100644 --- a/tests/test_litellm/proxy/db/test_gateway_request_tracking.py +++ b/tests/test_litellm/proxy/db/test_gateway_request_tracking.py @@ -8,8 +8,11 @@ from datetime import datetime, timezone import pytest +from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY from litellm.proxy.db.gateway_request_tracking import ( + GATEWAY_REQUESTS_JOB_NAME, GatewayRequestAccumulator, + GatewayRequestRedisBuffer, commit_gateway_requests_to_db, flush_gateway_requests, ) @@ -83,40 +86,47 @@ def test_drain_snapshot_is_not_mutated_by_later_records(): # ── commit ──────────────────────────────────────────────────────────────────── -class FakeTable: - def __init__(self) -> None: - self.upserts: list[dict] = [] - - def upsert(self, *, where: dict, data: dict) -> None: - self.upserts.append({"where": where, "data": data}) - - -class FakeBatcher: - def __init__(self, table: FakeTable) -> None: - self.litellm_dailygatewayrequests = table - - async def __aenter__(self) -> "FakeBatcher": - return self - - async def __aexit__(self, *args: object) -> bool: - return False - - class FakeDB: - def __init__(self, table: FakeTable) -> None: - self._table = table + def __init__(self) -> None: + self.statements: list[tuple[str, tuple[object, ...]]] = [] - def batch_(self) -> FakeBatcher: - return FakeBatcher(self._table) + async def execute_raw(self, query: str, *args: object) -> int: + self.statements.append((query, args)) + return len(args) // 5 class FakePrismaClient: def __init__(self) -> None: - self.table = FakeTable() - self.db = FakeDB(self.table) + self.db = FakeDB() -def test_commit_upserts_one_incrementing_row_per_key(): +def _rows_written(client: FakePrismaClient) -> list[tuple[object, ...]]: + """Every (date, category, route, successful, failed) tuple the database received, in statement order.""" + return [params[i : i + 5] for _, params in client.db.statements for i in range(0, len(params), 5)] + + +def test_commit_increments_with_a_single_statement_for_the_whole_snapshot(): + """One statement per flush is the whole point: the previous per-key upsert cost + the primary (workers x routes) statements per interval.""" + client = FakePrismaClient() + snapshot = { + GatewayRequestKey(date="2026-08-01", category="llm", route=route): ( + GatewayRequestCounts(successful_requests=7, failed_requests=2) + ) + for route in ("/chat/completions", "/embeddings", "/responses", "/v1/messages", "/mcp") + } + + asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) + + assert len(client.db.statements) == 1 + sql, params = client.db.statements[0] + assert sql.count("ON CONFLICT") == 1 + assert sql.count("(NOW() AT TIME ZONE 'UTC'))") == 5 + assert len(params) == 25 + + +def test_commit_sql_adds_to_the_existing_row_instead_of_replacing_it(): + """A worker only knows its own share; the SQL must add EXCLUDED onto the stored count.""" client = FakePrismaClient() snapshot = { GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): ( @@ -126,20 +136,36 @@ def test_commit_upserts_one_incrementing_row_per_key(): asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) - assert len(client.table.upserts) == 1 - written = client.table.upserts[0] - assert written["where"] == { - "date_category_route": { - "date": "2026-08-01", - "category": "llm", - "route": "/chat/completions", - } + sql, params = client.db.statements[0] + assert 'INSERT INTO "LiteLLM_DailyGatewayRequests"' in sql + assert 'ON CONFLICT ("date", "category", "route") DO UPDATE SET' in sql + assert ( + '"successful_requests" = "LiteLLM_DailyGatewayRequests"."successful_requests" + EXCLUDED."successful_requests"' + in sql + ) + assert '"failed_requests" = "LiteLLM_DailyGatewayRequests"."failed_requests" + EXCLUDED."failed_requests"' in sql + assert params == ("2026-08-01", "llm", "/chat/completions", 7, 2) + + +def test_commit_placeholders_line_up_with_params(): + """$n positions are generated per row; a drift here silently swaps a route for a count.""" + client = FakePrismaClient() + snapshot = { + GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): ( + GatewayRequestCounts(successful_requests=1, failed_requests=0) + ), + GatewayRequestKey(date="2026-08-01", category="mcp", route="/mcp"): ( + GatewayRequestCounts(successful_requests=0, failed_requests=3) + ), } - assert written["data"]["update"] == { - "successful_requests": {"increment": 7}, - "failed_requests": {"increment": 2}, - } - assert written["data"]["create"]["successful_requests"] == 7 + + asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) + + sql, params = client.db.statements[0] + assert "($1::text, $2::text, $3::text, $4::bigint, $5::bigint," in sql + assert "($6::text, $7::text, $8::text, $9::bigint, $10::bigint," in sql + assert "$11" not in sql + assert params == ("2026-08-01", "llm", "/chat/completions", 1, 0, "2026-08-01", "mcp", "/mcp", 0, 3) def test_commit_is_deterministically_ordered(): @@ -154,17 +180,14 @@ def test_commit_is_deterministically_ordered(): asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) - written_order = [ - (row["where"]["date_category_route"]["date"], row["where"]["date_category_route"]["category"]) - for row in client.table.upserts - ] + written_order = [(row[0], row[1]) for row in _rows_written(client)] assert written_order == [("2026-08-01", "llm"), ("2026-08-01", "mcp"), ("2026-08-02", "llm")] def test_commit_skips_the_database_entirely_when_nothing_accumulated(): client = FakePrismaClient() asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot={})) - assert client.table.upserts == [] + assert client.db.statements == [] # ── flush ───────────────────────────────────────────────────────────────────── @@ -177,12 +200,12 @@ def test_flush_drains_and_commits(): asyncio.run(flush_gateway_requests(client, acc)) - assert len(client.table.upserts) == 1 + assert len(client.db.statements) == 1 assert acc.drain() == {} class ExplodingDB: - def batch_(self): + async def execute_raw(self, query: str, *args: object) -> int: raise RuntimeError("db gone") @@ -208,10 +231,7 @@ def test_failed_flush_keeps_counts_for_the_next_attempt(): client = FakePrismaClient() asyncio.run(flush_gateway_requests(client, acc)) - assert client.table.upserts[0]["data"]["update"] == { - "successful_requests": {"increment": 1}, - "failed_requests": {"increment": 1}, - } + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)] def test_restored_counts_merge_with_requests_recorded_meanwhile(): @@ -223,5 +243,272 @@ def test_restored_counts_merge_with_requests_recorded_meanwhile(): client = FakePrismaClient() asyncio.run(flush_gateway_requests(client, acc)) - assert len(client.table.upserts) == 1 - assert client.table.upserts[0]["data"]["update"]["successful_requests"] == {"increment": 2} + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +class ExplodingDBWithInFlightRequest: + """Fails the write after a request has been recorded while it was in flight.""" + + def __init__(self, accumulator: GatewayRequestAccumulator) -> None: + self.accumulator = accumulator + + async def execute_raw(self, query: str, *args: object) -> int: + _record(self.accumulator, 500) + raise RuntimeError("db gone") + + +class ExplodingClientWithInFlightRequest: + def __init__(self, accumulator: GatewayRequestAccumulator) -> None: + self.db = ExplodingDBWithInFlightRequest(accumulator) + + +def test_restore_keeps_requests_recorded_while_the_failed_write_was_in_flight(): + acc = GatewayRequestAccumulator() + _record(acc, 200) + asyncio.run(flush_gateway_requests(ExplodingClientWithInFlightRequest(acc), acc)) + + client = FakePrismaClient() + asyncio.run(flush_gateway_requests(client, acc)) + + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)] + + +# ── redis buffer ────────────────────────────────────────────────────────────── + + +class FakeRedis: + def __init__(self) -> None: + self.lists: dict[str, list[str]] = {} + + async def async_rpush(self, key: str, values: list[str]) -> int: + self.lists.setdefault(key, []).extend(values) + return len(self.lists[key]) + + async def async_lpop(self, key: str, count: int) -> list[str] | None: + queue = self.lists.get(key, []) + if not queue: + return None + popped, self.lists[key] = queue[:count], queue[count:] + return popped + + +class FakePodLock: + def __init__(self, *, leader: bool) -> None: + self.leader = leader + self.held: list[str] = [] + self.released: list[str] = [] + + async def acquire_lock(self, cronjob_id: str) -> bool: + self.held.append(cronjob_id) + return self.leader + + async def release_lock(self, cronjob_id: str) -> None: + self.released.append(cronjob_id) + + +class FakeLease: + """Redis-side view of the job lock: SET NX by pod id, re-entrant for the holder, freed only by release or TTL.""" + + def __init__(self) -> None: + self.holder: str | None = None + + +class FakeLeasePodLock: + def __init__(self, lease: FakeLease, pod_id: str) -> None: + self.lease = lease + self.pod_id = pod_id + + async def acquire_lock(self, cronjob_id: str) -> bool: + if self.lease.holder is None: + self.lease.holder = self.pod_id + return self.lease.holder == self.pod_id + + async def release_lock(self, cronjob_id: str) -> None: + if self.lease.holder == self.pod_id: + self.lease.holder = None + + +def _buffer(redis: FakeRedis, *, leader: bool) -> tuple[GatewayRequestRedisBuffer, FakePodLock]: + lock = FakePodLock(leader=leader) + return GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=lock), lock # pyright: ignore[reportArgumentType] # duck-typed fakes + + +def test_non_leader_workers_push_to_redis_and_never_touch_the_database(): + redis = FakeRedis() + client = FakePrismaClient() + for _ in range(3): + acc = GatewayRequestAccumulator() + _record(acc, 200) + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + assert client.db.statements == [] + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3 + + +def test_leader_folds_every_workers_snapshot_into_one_statement(): + """Fifty workers each flushing the same routes must cost the primary one statement, not fifty.""" + redis = FakeRedis() + client = FakePrismaClient() + for _ in range(50): + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 500, route="/responses") + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + leader_acc = GatewayRequestAccumulator() + _record(leader_acc, 200) + leader, lock = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, leader_acc, leader)) + + assert len(client.db.statements) == 1 + assert _rows_written(client) == [ + (_today(), "llm", "/chat/completions", 51, 0), + (_today(), "llm", "/responses", 0, 50), + ] + assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + assert lock.held == [GATEWAY_REQUESTS_JOB_NAME] + assert lock.released == [] + + +def test_leader_keeps_the_lease_so_staggered_pods_cost_one_statement_per_interval(): + """Pods flush on their own clocks; without the lease each one would win the lock in turn and commit alone.""" + redis = FakeRedis() + client = FakePrismaClient() + lease = FakeLease() + pods = tuple( + GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=FakeLeasePodLock(lease, f"pod-{i}")) # pyright: ignore[reportArgumentType] # duck-typed fakes + for i in range(4) + ) + + for _interval in range(3): + for pod in pods: + acc = GatewayRequestAccumulator() + _record(acc, 200) + asyncio.run(flush_gateway_requests(client, acc, pod)) + + assert lease.holder == "pod-0" + assert len(client.db.statements) == 3 + assert [row[3] for row in _rows_written(client)] == [1, 4, 4] + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3 + + +def test_leader_drains_a_backlog_deeper_than_one_capped_pop(): + """More workers than MAX_REDIS_BUFFER_DEQUEUE_COUNT must not leave a growing tail queued behind the cap.""" + redis = FakeRedis() + client = FakePrismaClient() + workers = MAX_REDIS_BUFFER_DEQUEUE_COUNT * 2 + 1 + for _ in range(workers): + acc = GatewayRequestAccumulator() + _record(acc, 200) + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + leader, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader)) + + assert len(client.db.statements) == 1 + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", workers, 0)] + assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + + +def test_leader_with_nothing_buffered_writes_nothing(): + redis = FakeRedis() + client = FakePrismaClient() + leader, lock = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader)) + + assert client.db.statements == [] + assert lock.released == [] + + +def test_leader_requeues_to_redis_when_the_database_commit_fails(): + """Counts popped from Redis are gone from every worker; a failed commit must put them back.""" + redis = FakeRedis() + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 200) + leader, lock = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(ExplodingClient(), acc, leader)) + + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1 + assert lock.released == [] + assert acc.drain() == {} + + client = FakePrismaClient() + retry, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), retry)) + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +class ExplodingRedis(FakeRedis): + async def async_rpush(self, key: str, values: list[str]) -> int: + raise RuntimeError("redis gone") + + +class UnreadableRedis(FakeRedis): + async def async_lpop(self, key: str, count: int) -> list[str] | None: + raise RuntimeError("redis gone mid-flush") + + +class UnwritableRedis(FakeRedis): + """Pops succeed, pushes fail: a Redis that went read-only between the leader's pop and its re-queue.""" + + async def async_rpush(self, key: str, values: list[str]) -> int: + raise RuntimeError("redis read-only") + + +def test_leader_keeps_popped_counts_in_memory_when_both_the_database_and_the_requeue_fail(): + """The pop removed the only copy; if Redis will not take it back the leader itself must carry it.""" + redis = FakeRedis() + worker_acc = GatewayRequestAccumulator() + _record(worker_acc, 200) + _record(worker_acc, 200) + worker, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(FakePrismaClient(), worker_acc, worker)) + + degraded = UnwritableRedis() + degraded.lists = redis.lists + leader_acc = GatewayRequestAccumulator() + leader, _ = _buffer(degraded, leader=True) + asyncio.run(flush_gateway_requests(ExplodingClient(), leader_acc, leader)) + assert degraded.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + + client = FakePrismaClient() + retry, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, leader_acc, retry)) + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +def test_leader_whose_redis_read_fails_leaves_the_pushed_rows_for_the_next_flush(): + """The scheduler job must not raise, and nothing is popped so nothing needs restoring anywhere.""" + redis = UnreadableRedis() + acc = GatewayRequestAccumulator() + _record(acc, 200) + client = FakePrismaClient() + leader, _ = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(client, acc, leader)) + + assert client.db.statements == [] + assert acc.drain() == {} + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1 + + +def test_failed_redis_push_keeps_counts_locally_for_the_next_flush(): + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 500) + buffer, lock = _buffer(ExplodingRedis(), leader=True) + + asyncio.run(flush_gateway_requests(FakePrismaClient(), acc, buffer)) + + assert lock.held == [] + assert acc.drain() == { + GatewayRequestKey(date=_today(), category="llm", route="/chat/completions"): ( + GatewayRequestCounts(successful_requests=1, failed_requests=1) + ) + } diff --git a/tests/test_litellm/proxy/db/test_health_check_latest.py b/tests/test_litellm/proxy/db/test_health_check_latest.py new file mode 100644 index 00000000000..529e4d8f2e9 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_health_check_latest.py @@ -0,0 +1,111 @@ +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.health_check_latest import ( + LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, + LATEST_HEALTH_CHECKS_SQL, + fetch_latest_health_checks, + fetch_latest_health_checks_for_models, +) + + +def _prisma(rows): + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=rows) + return prisma + + +def _raw_row(**overrides): + row = { + "health_check_id": "hc-1", + "model_name": "gpt-4", + "model_id": "deployment-abc", + "status": "healthy", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + "response_time_ms": 12.5, + "details": None, + "checked_by": "pod-1", + "checked_at": "2026-08-25T00:00:00+00:00", + "created_at": "2026-08-25T00:00:00+00:00", + "updated_at": "2026-08-25T00:00:00+00:00", + } + return {**row, **overrides} + + +@pytest.mark.asyncio +async def test_fetch_all_runs_one_distinct_on_query_with_no_parameters(): + """The dedup must be in the SQL: prisma find_many(distinct=...) streams the whole history table.""" + prisma = _prisma([]) + assert await fetch_latest_health_checks(prisma) == () + assert prisma.db.query_raw.await_args.args == (LATEST_HEALTH_CHECKS_SQL,) + assert 'DISTINCT ON ("model_id", "model_name")' in LATEST_HEALTH_CHECKS_SQL + assert '"checked_at" DESC' in LATEST_HEALTH_CHECKS_SQL + + +@pytest.mark.asyncio +async def test_raw_datetimes_come_back_tz_aware_with_or_without_an_offset(): + """The save path subtracts checked_at from datetime.now(timezone.utc); a naive value would TypeError.""" + naive = _raw_row(health_check_id="hc-naive", model_id=None, checked_at="2026-08-25T00:00:00") + aware = _raw_row(health_check_id="hc-aware", checked_at="2026-08-25T01:00:00+02:00") + rows = await fetch_latest_health_checks(_prisma([naive, aware])) + assert {row.health_check_id: (row.model_id, row.checked_at) for row in rows} == { + "hc-naive": (None, datetime(2026, 8, 25, 0, 0, tzinfo=timezone.utc)), + "hc-aware": ("deployment-abc", datetime(2026, 8, 24, 23, 0, tzinfo=timezone.utc)), + } + + +@pytest.mark.asyncio +async def test_json_details_decode_from_text_and_pass_through_as_dict(): + rows = await fetch_latest_health_checks( + _prisma( + [ + _raw_row(health_check_id="text", details='{"region": "eu"}'), + _raw_row(health_check_id="dict", details={"region": "us"}), + _raw_row(health_check_id="none", details=None), + ] + ) + ) + assert {row.health_check_id: row.details for row in rows} == { + "text": {"region": "eu"}, + "dict": {"region": "us"}, + "none": None, + } + + +@pytest.mark.asyncio +async def test_fetch_all_degrades_to_no_rows_when_the_query_fails(): + prisma = _prisma([]) + prisma.db.query_raw.side_effect = RuntimeError("db down") + assert await fetch_latest_health_checks(prisma) == () + + +@pytest.mark.asyncio +async def test_fetch_all_degrades_to_no_rows_for_a_malformed_row(): + assert await fetch_latest_health_checks(_prisma([{"unexpected": "shape"}])) == () + + +@pytest.mark.asyncio +async def test_fetch_for_models_binds_the_page_as_the_only_parameter(): + prisma = _prisma([_raw_row()]) + rows = await fetch_latest_health_checks_for_models(prisma, ("gpt-4", "claude-opus")) + assert prisma.db.query_raw.await_args.args == (LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, ["gpt-4", "claude-opus"]) + assert [row.model_name for row in rows] == ["gpt-4"] + assert 'WHERE "model_name" = ANY($1)' in LATEST_HEALTH_CHECKS_FOR_MODELS_SQL + + +@pytest.mark.asyncio +async def test_fetch_for_models_skips_the_database_for_an_empty_page(): + prisma = _prisma([]) + assert await fetch_latest_health_checks_for_models(prisma, ()) == () + prisma.db.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_fetch_for_models_degrades_to_no_rows_when_the_query_fails(): + prisma = _prisma([]) + prisma.db.query_raw.side_effect = RuntimeError("db down") + assert await fetch_latest_health_checks_for_models(prisma, ("gpt-4",)) == () diff --git a/tests/test_litellm/proxy/db/test_pgbouncer.py b/tests/test_litellm/proxy/db/test_pgbouncer.py new file mode 100644 index 00000000000..7b5a0bf10f9 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_pgbouncer.py @@ -0,0 +1,868 @@ +import base64 +import configparser +import json +import logging +import os +import signal +import socket +import stat +import sys +import tempfile +import textwrap +import time +import urllib.parse +from collections import deque +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Final, cast + +import pytest + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.db.pgbouncer import ( + PGBOUNCER_POOLED_ENV_VAR, + PgBouncerError, + PgBouncerPlan, + PgBouncerProcess, + PgBouncerSettings, + PgBouncerTokenRefresher, + PgBouncerTokenSource, + database_url_is_pooled, + export_pooled_database_url, + install_pgbouncer_token, + pgbouncer_version, + plan_pgbouncer, + start_in_container_pgbouncer, + unix_socket_path, + write_pgbouncer_ini, + write_userlist, +) +from litellm.proxy.db.token_auth import AzureEntraTokenAuth, IAMEndpoint + +UPSTREAM: Final = ( + "postgresql://app:p%40ss%27w@db.internal:5433/litellm" + "?schema=public&connection_limit=10&pool_timeout=20" + "&sslmode=require&sslaccept=strict&sslcert=/certs/ca.pem" + "&options=-c%20statement_timeout%3D7000%20-c%20lock_timeout%3D3000" +) +SETTINGS: Final = PgBouncerSettings(enabled=True, port=6543, max_db_connections=8, max_client_conn=400) + + +def _plan(url: str = UPSTREAM, run_as_user: str | None = None) -> PgBouncerPlan: + plan: Final = plan_pgbouncer(url, SETTINGS, Path("/run/pgb"), run_as_user) + assert isinstance(plan, PgBouncerPlan), plan + return plan + + +def _ini(plan: PgBouncerPlan) -> configparser.ConfigParser: + parser: Final = configparser.ConfigParser(interpolation=None) + parser.read_string(plan.ini) + return parser + + +def _query(url: str) -> dict[str, str]: + return dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query, keep_blank_values=True)) + + +class TestPlanPgBouncer: + def test_upstream_route_and_timeouts_move_into_the_pgbouncer_config_without_the_password(self): + ini: Final = _ini(_plan()) + assert ini["databases"]["litellm"] == ( + "host='db.internal' port=5433 dbname='litellm' user='app' " + "connect_query='SET statement_timeout TO ''7000''; SET lock_timeout TO ''3000'''" + ) + + def test_the_auth_file_holds_the_upstream_password_and_the_pool_users_own(self): + plan: Final = _plan() + assert plan.upstream_password == "p@ss'w" + assert plan.userlist("p@ss'w") == f'"app" "p@ss\'w"\n"litellm_pgbouncer" "{plan.pool_password}"\n' + + def test_a_token_with_quotes_is_escaped_the_way_pgbouncer_reads_it(self): + assert _plan().userlist('to"ken').startswith('"app" "to""ken"\n') + + def test_an_upstream_without_a_port_is_reached_on_the_postgres_default(self): + ini: Final = _ini(_plan("postgresql://app:pw@db/litellm")) + assert ini["databases"]["litellm"] == "host='db' port=5432 dbname='litellm' user='app'" + + def test_pool_is_sized_from_settings_in_transaction_mode(self): + pgb: Final = _ini(_plan())["pgbouncer"] + assert pgb["pool_mode"] == "transaction" + assert pgb["max_db_connections"] == "8" + assert pgb["default_pool_size"] == "8" + assert pgb["max_client_conn"] == "400" + assert pgb["auth_type"] == "scram-sha-256" + assert pgb["listen_addr"] == "127.0.0.1" + assert pgb["listen_port"] == "6543" + assert pgb["auth_file"] == "/run/pgb/userlist.txt" + assert pgb["unix_socket_dir"] == "/run/pgb" + + def test_the_pool_user_can_read_the_pgbouncer_console(self): + assert _ini(_plan())["pgbouncer"]["stats_users"] == "litellm_pgbouncer" + + def test_a_database_user_named_like_the_pool_user_is_refused(self): + outcome: Final = plan_pgbouncer( + "postgresql://litellm_pgbouncer:pw@db/litellm", SETTINGS, Path("/run/pgb"), None + ) + assert isinstance(outcome, PgBouncerError) + assert "litellm_pgbouncer" in outcome.reason + + def test_pooled_url_points_prisma_at_loopback_as_the_pool_user_without_prepared_statements(self): + plan: Final = _plan() + pooled: Final = urllib.parse.urlsplit(plan.pooled_url) + assert (pooled.hostname, pooled.port, pooled.path) == ("127.0.0.1", 6543, "/litellm") + assert (pooled.username, pooled.password) == ("litellm_pgbouncer", plan.pool_password) + assert len(plan.pool_password) >= 32 + assert "p%40ss" not in plan.pooled_url + assert _query(plan.pooled_url) == { + "schema": "public", + "connection_limit": "10", + "pool_timeout": "20", + "pgbouncer": "true", + } + + @pytest.mark.parametrize("hop_param", ["channel_binding=require", "gssencmode=require"]) + def test_transport_params_for_the_postgres_hop_stay_off_the_plain_tcp_loopback_url(self, hop_param: str): + pooled: Final = _plan(f"postgresql://app:pw@db/litellm?connection_limit=5&{hop_param}").pooled_url + assert _query(pooled) == {"connection_limit": "5", "pgbouncer": "true"} + + def test_verified_tls_becomes_server_side_verify_full_with_a_ca_copy_in_the_runtime_dir(self): + plan: Final = _plan() + pgb: Final = _ini(plan)["pgbouncer"] + assert pgb["server_tls_sslmode"] == "verify-full" + assert pgb["server_tls_ca_file"] == "/run/pgb/server-ca.pem" + assert plan.ca_source == "/certs/ca.pem" + + def test_unverified_require_stays_require_without_a_ca_file(self): + plan: Final = _plan("postgresql://app:pw@db/litellm?sslmode=require") + pgb: Final = _ini(plan)["pgbouncer"] + assert pgb["server_tls_sslmode"] == "require" + assert "server_tls_ca_file" not in pgb + assert plan.ca_source is None + + def test_every_plan_gets_its_own_pool_password(self): + assert _plan().pool_password != _plan().pool_password + + def test_no_tls_params_default_to_prefer(self): + assert _ini(_plan("postgresql://app:pw@db/litellm"))["pgbouncer"]["server_tls_sslmode"] == "prefer" + + def test_verification_without_a_ca_bundle_is_refused(self): + outcome: Final = plan_pgbouncer( + "postgresql://app:pw@db/litellm?sslmode=require&sslaccept=strict", SETTINGS, Path("/run/pgb"), None + ) + assert isinstance(outcome, PgBouncerError) + assert "sslcert" in outcome.reason + + def test_client_certificates_are_refused(self): + outcome: Final = plan_pgbouncer( + "postgresql://app:pw@db/litellm?sslidentity=/certs/client.p12", SETTINGS, Path("/run/pgb"), None + ) + assert isinstance(outcome, PgBouncerError) + assert "sslidentity" in outcome.reason + + @pytest.mark.parametrize( + "url", + [ + "postgresql://app:pw@db", + "postgresql://:pw@db/litellm", + "postgresql://app:pw@/litellm", + ], + ) + def test_urls_missing_a_route_are_refused(self, url: str): + outcome: Final = plan_pgbouncer(url, SETTINGS, Path("/run/pgb"), None) + assert isinstance(outcome, PgBouncerError) + + def test_a_url_without_a_password_plans_for_a_token_to_be_installed_later(self): + plan: Final = _plan("postgresql://app@db/litellm") + assert plan.upstream_password is None + assert plan.userlist("minted").startswith('"app" "minted"\n') + + def test_every_options_spelling_becomes_a_set_statement(self): + options: Final = urllib.parse.quote("-c a=1 -cb=2 --c=3") + ini: Final = _ini(_plan(f"postgresql://app:pw@db/litellm?options={options}")) + assert ini["databases"]["litellm"].endswith("connect_query='SET a TO ''1''; SET b TO ''2''; SET c TO ''3'''") + + def test_options_that_are_not_settings_are_refused(self): + outcome: Final = plan_pgbouncer( + "postgresql://app:pw@db/litellm?options=-c%20search_path", SETTINGS, Path("/run/pgb"), None + ) + assert isinstance(outcome, PgBouncerError) + assert "options" in outcome.reason + + def test_run_as_user_is_only_written_when_given(self): + assert _ini(_plan(run_as_user="nobody"))["pgbouncer"]["user"] == "nobody" + assert "user" not in _ini(_plan())["pgbouncer"] + + +class TestWritePgBouncerFiles: + def test_files_hold_the_plan_and_are_private_to_the_owner(self, tmp_path: Path): + plan: Final = _plan("postgresql://app:pw@db/litellm") + ini_path: Final = write_pgbouncer_ini(plan, tmp_path, None) + assert isinstance(ini_path, Path), ini_path + userlist_path: Final = write_userlist(plan.userlist("pw"), tmp_path, None) + assert ini_path == tmp_path / "pgbouncer.ini" + assert userlist_path == tmp_path / "userlist.txt" + assert ini_path.read_text() == plan.ini + assert userlist_path.read_text() == plan.userlist("pw") + for path in (ini_path, userlist_path): + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + assert not (tmp_path / "server-ca.pem").exists() + + def test_the_ca_bundle_is_copied_next_to_the_ini_pgbouncer_reads(self, tmp_path: Path): + bundle: Final = tmp_path / "rds-root.pem" + bundle.write_text("-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----\n") + runtime_dir: Final = tmp_path / "run" + runtime_dir.mkdir() + plan: Final = plan_pgbouncer( + f"postgresql://app:pw@db/litellm?sslmode=verify-full&sslcert={bundle}", SETTINGS, runtime_dir, None + ) + assert isinstance(plan, PgBouncerPlan), plan + ini_path: Final = write_pgbouncer_ini(plan, runtime_dir, None) + assert isinstance(ini_path, Path), ini_path + ca_file: Final = Path(_ini(plan)["pgbouncer"]["server_tls_ca_file"]) + assert ca_file.parent == runtime_dir + assert ca_file.read_text() == bundle.read_text() + + def test_an_unreadable_ca_bundle_is_reported(self, tmp_path: Path): + plan: Final = plan_pgbouncer( + f"postgresql://app:pw@db/litellm?sslmode=verify-full&sslcert={tmp_path / 'missing.pem'}", + SETTINGS, + tmp_path, + None, + ) + assert isinstance(plan, PgBouncerPlan), plan + outcome: Final = write_pgbouncer_ini(plan, tmp_path, None) + assert isinstance(outcome, PgBouncerError) + assert "missing.pem" in outcome.reason + assert not (tmp_path / "pgbouncer.ini").exists() + + def test_rewriting_the_userlist_replaces_it_whole_and_leaves_nothing_else_behind(self, tmp_path: Path): + write_userlist('"app" "first"\n', tmp_path, None) + with open(tmp_path / "userlist.txt", encoding="utf-8") as before_rewrite: + write_userlist('"app" "second"\n', tmp_path, None) + assert before_rewrite.read() == '"app" "first"\n' + assert (tmp_path / "userlist.txt").read_text() == '"app" "second"\n' + assert stat.S_IMODE((tmp_path / "userlist.txt").stat().st_mode) == 0o600 + assert sorted(path.name for path in tmp_path.iterdir()) == ["userlist.txt"] + + +class TestPooledUrlMarker: + def test_exporting_the_pooled_url_marks_it_for_the_workers(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv(PGBOUNCER_POOLED_ENV_VAR, "") + monkeypatch.delenv(PGBOUNCER_POOLED_ENV_VAR) + monkeypatch.setenv("DATABASE_URL", "postgresql://app:token@db/litellm") + assert not database_url_is_pooled() + export_pooled_database_url("postgresql://litellm_pgbouncer:pw@127.0.0.1:6432/litellm?pgbouncer=true") + assert os.environ["DATABASE_URL"] == "postgresql://litellm_pgbouncer:pw@127.0.0.1:6432/litellm?pgbouncer=true" + assert database_url_is_pooled() + assert PGBOUNCER_POOLED_ENV_VAR == "LITELLM_PGBOUNCER_POOLED_DATABASE_URL" + + +def _bound_port(sock: socket.socket) -> int: + return cast(tuple[str, int], sock.getsockname())[1] + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return _bound_port(probe) + + +def _fake_pooler( + tmp_path: Path, + port: int, + exit_immediately: bool = False, + port_file: Path | None = None, + bind_delay_seconds: float = 0.0, + version_banner: str = "PgBouncer 1.25.2\nlibevent 2.1.13-stable", + auth_log: Path | None = None, +) -> Path: + """An executable that listens like PgBouncer: on the TCP port first, then on ``.s.PGSQL.`` in the socket dir. + + Port and socket dir come from the ini it is given, else from ``port`` and + ``tmp_path``. With ``port_file`` each start reads the port from that file + instead. ``bind_delay_seconds`` holds the bind back, like a slow start. + ``--version`` prints ``version_banner``. With ``auth_log`` it appends the + ``auth_file`` it reads at startup and on every SIGHUP, one line per read, + like PgBouncer loading its credentials. + """ + script: Final = tmp_path / "fake-pgbouncer" + script.write_text( + textwrap.dedent( + f"""\ + #!{sys.executable} + import configparser, os, pathlib, select, signal, socket, sys, time + if sys.argv[1:] == ["--version"]: + print({version_banner!r}) + sys.exit(0) + if {exit_immediately!r}: + sys.exit(3) + ini = configparser.ConfigParser() + ini.read(sys.argv[1:2]) + if not {auth_log is None!r}: + def load_auth_file(*_): + with open({str(auth_log)!r}, "a") as log: + log.write(repr(pathlib.Path(ini.get("pgbouncer", "auth_file")).read_text()) + "\\n") + load_auth_file() + signal.signal(signal.SIGHUP, load_auth_file) + port = ini.getint("pgbouncer", "listen_port", fallback={port}) + if not {port_file is None!r}: + port = int(pathlib.Path({str(port_file)!r}).read_text()) + socket_dir = ini.get("pgbouncer", "unix_socket_dir", fallback={str(tmp_path)!r}) + socket_path = f"{{socket_dir}}/.s.PGSQL.{{port}}" + time.sleep({bind_delay_seconds!r}) + listener = socket.socket() + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind(("127.0.0.1", port)) + listener.listen() + if os.path.exists(socket_path): + os.unlink(socket_path) + unix_listener = socket.socket(socket.AF_UNIX) + unix_listener.bind(socket_path) + unix_listener.listen() + while True: + for ready in select.select([listener, unix_listener], [], [])[0]: + conn, _ = ready.accept() + conn.close() + """ + ) + ) + script.chmod(0o700) + return script + + +def _listening(port: int) -> bool: + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.5): + return True + except OSError: + return False + + +def _wait_until(condition: Callable[[], bool], timeout_seconds: float = 5.0) -> bool: + deadline: Final = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + if condition(): + return True + time.sleep(0.05) + return False + + +class TestPgBouncerProcess: + def test_start_waits_for_the_listener_and_stop_ends_it(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port) + ) + assert pooler.start() is None + assert _listening(port) + pid: Final = pooler.pid + assert pid is not None + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + def test_a_crashed_pooler_is_restarted_with_a_new_pid(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.1, + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port)) + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + + def test_a_failed_restart_is_retried_until_the_pooler_is_back( + self, tmp_path: Path, caplog: pytest.LogCaptureFixture + ): + port: Final = _free_port() + script: Final = _fake_pooler(tmp_path, port) + pooler: Final = PgBouncerProcess( + argv=(str(script),), port=port, socket_path=unix_socket_path(tmp_path, port), restart_delay_seconds=0.1 + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + hidden: Final = script.rename(tmp_path / "hidden") + with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: any("could not be restarted" in record.message for record in caplog.records)) + assert not _listening(port) + hidden.rename(script) + assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port)) + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + + def test_a_replacement_that_never_listens_is_replaced_again(self, tmp_path: Path, caplog: pytest.LogCaptureFixture): + port: Final = _free_port() + port_file: Final = tmp_path / "port" + port_file.write_text(str(port)) + script: Final = _fake_pooler(tmp_path, port, port_file=port_file) + pooler: Final = PgBouncerProcess( + argv=(str(script),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.1, + ready_timeout_seconds=0.3, + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + wrong_port: Final = _free_port() + port_file.write_text(str(wrong_port)) + with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: _listening(wrong_port)) + assert _wait_until(lambda: any("did not start listening" in record.message for record in caplog.records)) + port_file.write_text(str(port)) + assert _wait_until(lambda: _listening(port)) + assert _wait_until(lambda: not _listening(wrong_port)) + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + + def test_stopping_during_the_restart_delay_leaves_no_pooler_behind(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.3, + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: not _listening(port)) + pooler.stop() + time.sleep(1.0) + assert not _listening(port) + assert pooler.pid == first_pid + + def test_a_stopped_pooler_is_not_restarted(self, tmp_path: Path, caplog: pytest.LogCaptureFixture): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.1, + ) + assert pooler.start() is None + with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + pooler.stop() + time.sleep(0.5) + assert not _listening(port) + assert caplog.records == [] + + def test_a_pooler_that_exits_during_startup_is_reported(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port, exit_immediately=True)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + ) + outcome: Final = pooler.start() + assert isinstance(outcome, PgBouncerError) + assert "status 3" in outcome.reason + + def test_a_missing_binary_is_reported(self, tmp_path: Path): + outcome: Final = PgBouncerProcess( + argv=("/nonexistent/pgbouncer",), port=_free_port(), socket_path=tmp_path / "sock" + ).start() + assert isinstance(outcome, PgBouncerError) + assert "/nonexistent/pgbouncer" in outcome.reason + + def test_a_port_owned_by_someone_else_is_refused_before_spawning(self, tmp_path: Path): + with socket.socket() as squatter: + squatter.bind(("127.0.0.1", 0)) + squatter.listen() + port: Final = _bound_port(squatter) + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port) + ) + outcome: Final = pooler.start() + assert isinstance(outcome, PgBouncerError) + assert f"127.0.0.1:{port} is already in use" in outcome.reason + assert pooler.pid is None + + def test_a_replacement_waits_until_a_squatter_leaves_the_port( + self, tmp_path: Path, caplog: pytest.LogCaptureFixture + ): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.5, + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: not _listening(port)) + with socket.socket() as squatter, caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + squatter.bind(("127.0.0.1", port)) + squatter.listen() + assert _wait_until(lambda: any("already in use" in record.message for record in caplog.records)) + assert pooler.pid == first_pid + assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port)) + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + + def test_a_listener_that_grabs_the_port_after_the_spawn_is_not_taken_for_the_pooler(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=0.5)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + ready_timeout_seconds=3.0, + ) + with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter: + squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + starting: Final = starter.submit(pooler.start) + assert _wait_until(lambda: pooler.pid is not None) + squatter.bind(("127.0.0.1", port)) + squatter.listen() + outcome: Final = starting.result() + assert isinstance(outcome, PgBouncerError) + assert "exited with status 1" in outcome.reason + + def test_a_port_served_by_a_stranger_while_the_pooler_is_still_starting_is_reported(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=30.0)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + ready_timeout_seconds=0.5, + ) + with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter: + squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + starting: Final = starter.submit(pooler.start) + assert _wait_until(lambda: pooler.pid is not None) + squatter.bind(("127.0.0.1", port)) + squatter.listen() + outcome: Final = starting.result() + assert isinstance(outcome, PgBouncerError) + assert f"127.0.0.1:{port} is served by another process" in outcome.reason + pid: Final = pooler.pid + assert pid is not None + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + def test_a_pooler_that_never_listens_times_out(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, _free_port())),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + ready_timeout_seconds=0.5, + ) + outcome: Final = pooler.start() + assert isinstance(outcome, PgBouncerError) + assert "did not start listening" in outcome.reason + pid: Final = pooler.pid + assert pid is not None + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + +NOW: Final = datetime(2026, 9, 10, 12, 0, tzinfo=timezone.utc) +ENDPOINT: Final = IAMEndpoint(host="db", port="5432", user="app", name="litellm") + + +def _entra_jwt(expires_at: datetime) -> str: + payload: Final = base64.urlsafe_b64encode(json.dumps({"exp": int(expires_at.timestamp())}).encode()) + return f"aGVhZGVy.{payload.rstrip(b'=').decode()}.c2ln" + + +def _token_source(*tokens: str | Exception) -> PgBouncerTokenSource: + """A token source handing out ``tokens`` in order, raising the exceptions among them, then repeating the last.""" + pending: Final = deque(tokens) + + def provide() -> str: + outcome: Final = pending.popleft() if len(pending) > 1 else pending[0] + if isinstance(outcome, Exception): + raise outcome + return outcome + + return PgBouncerTokenSource(auth=AzureEntraTokenAuth(token_provider=provide), endpoint=ENDPOINT) + + +class TestPgBouncerTokenRefresher: + def _refresher( + self, + source: PgBouncerTokenSource, + installed: list[str], + install: Callable[[str], None] | None = None, + **timing: float, + ) -> PgBouncerTokenRefresher: + return PgBouncerTokenRefresher( + source, + install if install is not None else installed.append, + now=lambda: NOW.replace(tzinfo=None), + **timing, + ) + + def test_the_next_refresh_is_due_a_buffer_before_the_token_expires(self): + installed: Final[list[str]] = [] + token: Final = _entra_jwt(NOW + timedelta(hours=1)) + refresher: Final = self._refresher(_token_source(token), installed, buffer_seconds=180) + assert refresher.refresh() == 3600 - 180 + assert installed == [token] + + def test_a_token_whose_expiry_cannot_be_read_is_refreshed_on_the_fallback_interval(self): + installed: Final[list[str]] = [] + refresher: Final = self._refresher(_token_source("opaque token"), installed, fallback_seconds=600) + assert refresher.refresh() == 600 + assert installed == ["opaque token"] + + def test_a_token_already_inside_the_buffer_is_refreshed_after_the_retry_delay(self): + token: Final = _entra_jwt(NOW + timedelta(seconds=100)) + refresher: Final = self._refresher(_token_source(token), [], buffer_seconds=180, retry_seconds=30) + assert refresher.refresh() == 30 + + def test_the_token_reaches_the_auth_file_in_wire_form_not_url_encoded(self): + installed: Final[list[str]] = [] + self._refresher(_token_source("to ken/with+odd=chars"), installed).refresh() + assert installed == ["to ken/with+odd=chars"] + + def test_a_failed_mint_is_reported_and_installs_nothing(self): + installed: Final[list[str]] = [] + outcome: Final = self._refresher(_token_source(RuntimeError("no credential")), installed).refresh() + assert isinstance(outcome, PgBouncerError) + assert "Azure Entra token" in outcome.reason + assert "no credential" in outcome.reason + assert installed == [] + + def test_a_token_pgbouncer_cannot_hold_is_refused(self): + installed: Final[list[str]] = [] + outcome: Final = self._refresher(_token_source("x" * 2048), installed).refresh() + assert isinstance(outcome, PgBouncerError) + assert "2047" in outcome.reason + assert installed == [] + + def test_an_auth_file_that_cannot_be_written_is_reported_not_raised(self): + def refuse(_: str) -> None: + raise PermissionError("read-only runtime dir") + + outcome: Final = self._refresher(_token_source("token"), [], install=refuse).refresh() + assert isinstance(outcome, PgBouncerError) + assert "read-only runtime dir" in outcome.reason + + def test_start_fails_when_the_first_token_cannot_be_minted_and_schedules_nothing(self): + installed: Final[list[str]] = [] + refresher: Final = self._refresher( + _token_source(RuntimeError("no credential"), "later"), installed, fallback_seconds=0.05 + ) + assert isinstance(refresher.start(), PgBouncerError) + time.sleep(0.3) + assert installed == [] + + def test_a_failed_renewal_keeps_the_previous_token_until_the_retry_succeeds(self, caplog: pytest.LogCaptureFixture): + installed: Final[list[str]] = [] + refresher: Final = self._refresher( + _token_source("first", RuntimeError("blip"), "third"), + installed, + fallback_seconds=0.05, + retry_seconds=0.05, + ) + with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + assert refresher.start() is None + assert installed == ["first"] + assert _wait_until(lambda: "third" in installed) + assert installed[:2] == ["first", "third"] + assert any("keeps its current Azure Entra token" in record.message for record in caplog.records) + refresher.stop() + settled: Final = len(installed) + time.sleep(0.3) + assert len(installed) == settled + + +def _runtime_dir_listening_on(port: int) -> Path: + matches: Final = tuple( + ini.parent + for ini in Path(tempfile.gettempdir()).glob("litellm-pgbouncer-*/pgbouncer.ini") + if f"listen_port = {port}\n" in ini.read_text() + ) + assert len(matches) == 1, matches + return matches[0] + + +class TestStartInContainerPgBouncer: + def test_returns_the_loopback_url_once_the_pooler_listens(self, tmp_path: Path): + port: Final = _free_port() + auth_log: Final = tmp_path / "auth.log" + binary: Final = _fake_pooler(tmp_path, port, auth_log=auth_log) + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary)) + pooled: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm?connection_limit=5") + assert isinstance(pooled, str), pooled + parsed: Final = urllib.parse.urlsplit(pooled) + assert (parsed.username, parsed.hostname, parsed.port, parsed.path) == ( + "litellm_pgbouncer", + "127.0.0.1", + port, + "/litellm", + ) + assert _query(pooled) == {"connection_limit": "5", "pgbouncer": "true"} + assert _listening(port) + assert auth_log.read_text() == repr(f'"app" "pw"\n"litellm_pgbouncer" "{parsed.password}"\n') + "\n" + + @pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") + def test_a_forked_worker_exiting_leaves_the_pooler_and_its_files_to_the_parent(self, tmp_path: Path): + port: Final = _free_port() + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port))) + exit_hooks: Final[list[Callable[[], None]]] = [] + pooled: Final = start_in_container_pgbouncer( + settings, "postgresql://app:pw@db/litellm", register_exit_hook=exit_hooks.append + ) + assert isinstance(pooled, str), pooled + runtime_dir: Final = _runtime_dir_listening_on(port) + + worker: Final = os.fork() + if worker == 0: + try: + for hook in exit_hooks: + hook() + finally: + os._exit(0) + if not _wait_until(lambda: os.waitpid(worker, os.WNOHANG) != (0, 0)): + os.kill(worker, signal.SIGKILL) + pytest.fail("the forked worker did not exit: an exit hook blocked on state inherited from the parent") + assert _listening(port) + assert (runtime_dir / "pgbouncer.ini").exists() + + for hook in exit_hooks: + hook() + assert _wait_until(lambda: not _listening(port)) + assert not runtime_dir.exists() + + def test_a_bad_upstream_url_is_reported_without_starting_anything(self, tmp_path: Path): + port: Final = _free_port() + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port))) + outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db") + assert isinstance(outcome, PgBouncerError) + assert not _listening(port) + + def test_a_passwordless_url_without_token_auth_is_refused_without_starting_anything(self, tmp_path: Path): + port: Final = _free_port() + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port))) + outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app@db/litellm") + assert isinstance(outcome, PgBouncerError) + assert "IAM_TOKEN_DB_AUTH" in outcome.reason + assert not _listening(port) + + def test_token_auth_mints_the_first_token_into_the_auth_file_before_the_pooler_starts(self, tmp_path: Path): + port: Final = _free_port() + auth_log: Final = tmp_path / "auth.log" + binary: Final = _fake_pooler(tmp_path, port, auth_log=auth_log) + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary)) + token: Final = _entra_jwt(datetime.now(tz=timezone.utc) + timedelta(hours=1)) + pooled: Final = start_in_container_pgbouncer( + settings, + "postgresql://app:stale-token@db/litellm", + token_auth=AzureEntraTokenAuth(token_provider=lambda: token), + ) + assert isinstance(pooled, str), pooled + parsed: Final = urllib.parse.urlsplit(pooled) + assert parsed.username == "litellm_pgbouncer" + assert token not in pooled + assert _listening(port) + assert auth_log.read_text() == repr(f'"app" "{token}"\n"litellm_pgbouncer" "{parsed.password}"\n') + "\n" + + def test_a_first_token_that_cannot_be_minted_is_reported_without_starting_anything(self, tmp_path: Path): + port: Final = _free_port() + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port))) + + def fail() -> str: + raise RuntimeError("no Azure credential") + + outcome: Final = start_in_container_pgbouncer( + settings, "postgresql://app@db/litellm", token_auth=AzureEntraTokenAuth(token_provider=fail) + ) + assert isinstance(outcome, PgBouncerError) + assert "no Azure credential" in outcome.reason + assert not _listening(port) + + def test_a_renewed_token_is_written_and_picked_up_by_the_running_and_by_a_restarted_pooler(self, tmp_path: Path): + port: Final = _free_port() + auth_log: Final = tmp_path / "auth.log" + plan: Final = plan_pgbouncer( + "postgresql://app@db/litellm", PgBouncerSettings(enabled=True, port=port), tmp_path, None + ) + assert isinstance(plan, PgBouncerPlan), plan + ini_path: Final = write_pgbouncer_ini(plan, tmp_path, None) + write_userlist(plan.userlist("first"), tmp_path, None) + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port, auth_log=auth_log)), str(ini_path)), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.1, + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + install_pgbouncer_token(plan, tmp_path, None, pooler, "second") + assert _wait_until(lambda: auth_log.read_text().count("\n") == 2) + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port)) + pooler.stop() + assert auth_log.read_text().splitlines() == [ + repr(plan.userlist("first")), + repr(plan.userlist("second")), + repr(plan.userlist("second")), + ] + + def test_a_pgbouncer_that_survives_a_failed_tcp_bind_is_refused_without_starting(self, tmp_path: Path): + port: Final = _free_port() + binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.18.1\nlibevent 2.1.12-stable") + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary)) + outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm") + assert isinstance(outcome, PgBouncerError) + assert "PgBouncer 1.18" in outcome.reason + assert "1.19" in outcome.reason + assert not _listening(port) + + def test_the_first_version_that_dies_on_a_failed_tcp_bind_is_accepted(self, tmp_path: Path): + port: Final = _free_port() + binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.19.0") + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary)) + pooled: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm") + assert isinstance(pooled, str), pooled + assert urllib.parse.urlsplit(pooled).port == port + assert _listening(port) + + +class TestPgBouncerVersion: + def test_reads_major_and_minor_from_the_banner(self, tmp_path: Path): + assert pgbouncer_version(str(_fake_pooler(tmp_path, _free_port()))) == (1, 25) + + def test_a_binary_that_cannot_run_is_reported(self, tmp_path: Path): + outcome: Final = pgbouncer_version(str(tmp_path / "missing-pgbouncer")) + assert isinstance(outcome, PgBouncerError) + assert "missing-pgbouncer" in outcome.reason + + def test_a_banner_without_a_version_is_reported(self, tmp_path: Path): + outcome: Final = pgbouncer_version(str(_fake_pooler(tmp_path, _free_port(), version_banner="something else"))) + assert isinstance(outcome, PgBouncerError) + assert "something else" in outcome.reason + + +class TestPgBouncerSettings: + def test_reads_the_litellm_pgbouncer_env_vars(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") + monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", "7000") + monkeypatch.setenv("LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", "12") + settings: Final = PgBouncerSettings() + assert (settings.enabled, settings.port, settings.max_db_connections) == (True, 7000, 12) + + def test_defaults_are_off(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False) + assert PgBouncerSettings().enabled is False diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py index f0983d6bf62..963f6a5640f 100644 --- a/tests/test_litellm/proxy/db/test_prisma_client.py +++ b/tests/test_litellm/proxy/db/test_prisma_client.py @@ -10,7 +10,7 @@ from fastapi.testclient import TestClient -from litellm.proxy.db.prisma_client import PrismaWrapper, should_update_prisma_schema +from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema @pytest.fixture(autouse=True) @@ -193,7 +193,10 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client( mock_new_prisma.connect.assert_awaited_once() -def test_db_push_applies_replica_identity_full_when_requested(monkeypatch): +DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"] + + +def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url): """`prisma db push` bypasses litellm-proxy-extras, so it needs its own call into the opt-in REPLICA IDENTITY FULL step.""" from litellm.proxy.db.prisma_client import PrismaManager @@ -208,14 +211,13 @@ def test_db_push_applies_replica_identity_full_when_requested(monkeypatch): staticmethod(lambda: applied.append(True)), ) - with patch("litellm.proxy.db.prisma_client.subprocess.run") as mock_run: - assert PrismaManager.setup_database(use_migrate=False) is True + assert PrismaManager.setup_database(use_migrate=False) is True - assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV] assert applied == [True] -def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch): +def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch, fake_prisma_cli, unset_database_url): """A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the primary key back to ("request_id"), which Postgres rejects; the guard must fail fast with guidance instead of running the push.""" @@ -228,29 +230,23 @@ def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch): monkeypatch.setattr( ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True) ) - with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, asserted never reached - "litellm.proxy.db.prisma_client.subprocess.run" - ) as mock_run: - with pytest.raises(RuntimeError) as err: - PrismaManager.setup_database(use_migrate=False) + with pytest.raises(RuntimeError) as err: + PrismaManager.setup_database(use_migrate=False) assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR - mock_run.assert_not_called() + assert fake_prisma_cli.calls == [] -def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch): +def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_prisma_cli, unset_database_url): from litellm.proxy.db.prisma_client import PrismaManager from litellm_proxy_extras.utils import ProxyExtrasDBManager monkeypatch.setattr( ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False) ) - with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, not SDK logic - "litellm.proxy.db.prisma_client.subprocess.run" - ) as mock_run: - assert PrismaManager.setup_database(use_migrate=False) is True + assert PrismaManager.setup_database(use_migrate=False) is True - assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV] def _entra_jwt(expires_in_seconds: int) -> str: @@ -295,6 +291,56 @@ def test_azure_entra_mint_writes_an_encoded_url_into_the_db_url_env_var(azure_en assert os.environ["DATABASE_URL"] == db_url +@pytest.mark.parametrize( + ("previous_query", "expected_query"), + [ + ("max_idle_connection_lifetime=60", {"max_idle_connection_lifetime": ["60"]}), + ( + "connection_limit=20&pgbouncer=true&max_idle_connection_lifetime=45", + {"connection_limit": ["20"], "pgbouncer": ["true"], "max_idle_connection_lifetime": ["45"]}, + ), + ], +) +def test_token_refresh_keeps_the_connection_params_of_the_url_it_replaces( + azure_env, monkeypatch, previous_query, expected_query +): + old_token = _entra_jwt(60) + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://litellm%40contoso.onmicrosoft.com:{urllib.parse.quote(old_token, safe='')}" + f"@pg.postgres.database.azure.com:5432/litellm_db?{previous_query}", + ) + new_token = _entra_jwt(3600) + + db_url = _azure_wrapper(new_token).get_rds_iam_token() + + assert db_url is not None + assert os.environ["DATABASE_URL"] == db_url + assert urllib.parse.quote(new_token, safe="") in db_url + assert urllib.parse.parse_qs(urllib.parse.urlsplit(db_url).query) == expected_query + + +def test_token_refresh_keeps_the_reader_url_params_separate_from_the_writer(azure_env, monkeypatch): + from litellm.proxy.db.token_auth import IAMEndpoint + + monkeypatch.setenv("DATABASE_URL", "postgresql://w:t@pg:5432/litellm_db?max_idle_connection_lifetime=45") + monkeypatch.setenv( + "DATABASE_URL_READ_REPLICA", "postgresql://r:t@replica:5432/litellm_db?max_idle_connection_lifetime=60" + ) + reader = _azure_wrapper( + _entra_jwt(3600), + db_url_env_var="DATABASE_URL_READ_REPLICA", + iam_endpoint=IAMEndpoint(host="replica", port="5432", user="r", name="litellm_db", schema=None), + ) + + reader_url = reader.get_rds_iam_token() + + assert reader_url is not None + assert reader_url.startswith("postgresql://r:") + assert urllib.parse.parse_qs(urllib.parse.urlsplit(reader_url).query) == {"max_idle_connection_lifetime": ["60"]} + assert os.environ["DATABASE_URL"].endswith("?max_idle_connection_lifetime=45") + + def test_azure_entra_refresh_is_scheduled_off_the_jwt_expiry(azure_env): """Without reading `exp` this falls back to a fixed 600s interval, which silently outlives a token and breaks every reconnect after it lapses (issue #29661).""" @@ -377,3 +423,30 @@ def test_minting_without_the_database_env_vars_names_them(azure_env, monkeypatch with pytest.raises(RuntimeError, match="DATABASE_HOST"): wrapper.get_rds_iam_token() + + +@pytest.mark.timeout(45) +def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_database_url, monkeypatch): + """ + A timed-out `db push` used to leave Node and the schema engine writing the schema, + so the next attempt pushed into a database the abandoned one was still mutating. + """ + monkeypatch.delenv("LITELLM_SET_REPLICA_IDENTITY_FULL", raising=False) + monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1") + + assert PrismaManager.setup_database(use_migrate=False) is True + assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV] + assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot( + fake_prisma_cli, unset_database_url, monkeypatch +): + """ + An ImportError out of setup_database escapes the caller's RuntimeError handler and + kills boot, bypassing the operator's enforce_prisma_migration_check choice. + """ + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert PrismaManager.setup_database(use_migrate=False) is False + assert fake_prisma_cli.calls == [] diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index 10a48941693..e6914213939 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -626,7 +626,7 @@ async def test_direct_reconnect_probe_success_clears_writer_unavailable( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) writer = MagicMock() - writer.query_raw = AsyncMock(return_value=[{"result": 1}]) + writer.query_raw = AsyncMock(return_value=[{"transaction_read_only": "off"}]) reader = MagicMock() routing = RoutingPrismaWrapper(writer=writer, reader=reader) routing._writer_unavailable = True @@ -636,5 +636,5 @@ async def test_direct_reconnect_probe_success_clears_writer_unavailable( with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): await client._run_reconnect_cycle(timeout_seconds=5.0) - writer.query_raw.assert_awaited_once_with("SELECT 1") + writer.query_raw.assert_awaited_once_with("SELECT current_setting('transaction_read_only') AS transaction_read_only") assert routing.writer_unavailable is False diff --git a/tests/test_litellm/proxy/db/test_query_engine_reaper.py b/tests/test_litellm/proxy/db/test_query_engine_reaper.py index efcecb4bc08..5018176854e 100644 --- a/tests/test_litellm/proxy/db/test_query_engine_reaper.py +++ b/tests/test_litellm/proxy/db/test_query_engine_reaper.py @@ -3,6 +3,7 @@ import signal import subprocess import sys import time +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -15,7 +16,6 @@ from litellm.proxy.db.query_engine_reaper import ( _try_reap, list_orphaned_engine_pids, reap_orphaned_engines, - set_child_subreaper, start_query_engine_reaper, terminate_and_reap, terminate_and_reap_all, @@ -79,11 +79,19 @@ class TestListOrphanedEnginePids: class TestSetChildSubreaper: def test_matches_platform_capability(self): - result = set_child_subreaper() - if sys.platform.startswith("linux"): - assert result is True - else: - assert result is False + result: Final = subprocess.run( + [ + sys.executable, + "-c", + "import sys; " + "from litellm.proxy.db.query_engine_reaper import set_child_subreaper; " + "assert set_child_subreaper() is sys.platform.startswith('linux')", + ], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr @pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals and waitpid") diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 8cb3fc665eb..0361d5cfe8f 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -7,6 +7,7 @@ allowed to run: only when the row is missing or belongs to an older window. from __future__ import annotations +import asyncio from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final @@ -14,6 +15,7 @@ from typing import Final import pytest from litellm.caching.dual_cache import DualCache +from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) @@ -42,6 +44,19 @@ class _FakeSpendLogsTable: return [{by[0]: where.get(by[0]), "_sum": {"spend": self._total}}] +class _InFlightCountingTable: + def __init__(self) -> None: + self.in_flight = 0 + self.max_in_flight = 0 + + async def find_unique(self, where: dict[str, str]) -> SimpleNamespace: + self.in_flight += 1 + self.max_in_flight = max(self.max_in_flight, self.in_flight) + await asyncio.sleep(0.001) + self.in_flight -= 1 + return SimpleNamespace(token=where["token"], spend=1.0) + + class _FakePrismaClient: def __init__( self, @@ -55,6 +70,7 @@ class _FakePrismaClient: litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error), litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error), + litellm_verificationtoken=_InFlightCountingTable(), ) @@ -306,6 +322,21 @@ async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_err ) +@pytest.mark.asyncio +async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): + """Per-counter singleflight only collapses duplicates of one key. A cold-cache burst + over many distinct keys must still not flood the prisma engine HTTP pool (LIT-6435).""" + prisma: Final = _FakePrismaClient() + burst: Final = PROXY_DB_LOOKUP_MAX_CONCURRENCY * 5 + + results: Final = await asyncio.gather( + *(SpendCounterReseed.from_db(prisma_client=prisma, counter_key=f"spend:key:hashed-{i}") for i in range(burst)) + ) + + assert results == [1.0] * burst + assert prisma.db.litellm_verificationtoken.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY + + @pytest.mark.asyncio async def test_from_db_still_never_reads_the_end_user_row(): """A cold end-user counter keeps seeding from the cached end-user object the auth diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index be55ac47bde..130b0da000b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -626,6 +626,64 @@ class TestContentFilterGuardrail: assert entry["guardrail_status"] == "success" assert entry["guardrail_response"] == [] + @pytest.mark.asyncio + async def test_streaming_hook_duration_excludes_provider_wait(self): + """ + Streaming post-call: the logged guardrail duration must only cover the + per-chunk scans, not the time spent waiting on the provider between + chunks. PrometheusLogger adds post_call guardrail duration to + litellm_overhead_with_guardrails_latency_metric, so a duration spanning + the whole stream reports LLM generation time as guardrail overhead. + """ + import asyncio + + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = ContentFilterGuardrail( + guardrail_name="test-streaming-duration", + patterns=[ + ContentFilterPattern( + pattern_type="prebuilt", + pattern_name="email", + action=ContentFilterAction.MASK, + ), + ], + event_hook=GuardrailEventHooks.post_call, + ) + + provider_wait_per_chunk = 0.15 + chunks = ("Hello ", "world, reach me at ", "test@example.com ") + + async def slow_stream(): + for i, text in enumerate(chunks): + await asyncio.sleep(provider_wait_per_chunk) + yield ModelResponseStream( + id=f"chunk{i}", + choices=[ + StreamingChoices( + delta=Delta(content=text), + index=0, + finish_reason="stop" if i == len(chunks) - 1 else None, + ) + ], + model="gpt-4", + ) + + request_data = {"messages": [{"role": "user", "content": "Hi"}], "model": "gpt-4o", "metadata": {}} + + async for _ in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), + response=slow_stream(), + request_data=request_data, + ): + pass + + entry = request_data["metadata"]["standard_logging_guardrail_information"][0] + stream_wall_clock = entry["end_time"] - entry["start_time"] + assert stream_wall_clock >= provider_wait_per_chunk * len(chunks) + assert entry["masked_entity_count"].get("email", 0) >= 1 + assert 0 < entry["duration"] < provider_wait_per_chunk, entry["duration"] + @pytest.mark.asyncio async def test_streaming_hook_logs_guardrail_information_mask(self): """ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 2b43720a126..615d06b0f42 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -3,14 +3,19 @@ Test OpenAI Moderation Guardrail """ +import json import os +from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( OpenAIModerationGuardrail, ) @@ -989,3 +994,30 @@ async def test_openai_moderation_initialize_guardrail_forwards_streaming_flags() assert guardrail.streaming_sampling_rate == 2 finally: litellm.logging_callback_manager._reset_all_callbacks() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("input_type", "stage"), [("request", "pre_call"), ("response", "post_call")]) +async def test_openai_moderation_records_moderation_id_as_scan_metadata(input_type: str, stage: str): + """Each moderation call's id is exposed with the guardrail name, stage and provider that produced it.""" + payload: Final = { + "id": f"modr-{stage}", + "model": "omni-moderation-latest", + "results": [{"flagged": False, "categories": {}, "category_scores": {}, "category_applied_input_types": {}}], + } + http_client: Final = AsyncHTTPHandler() + http_client.client = httpx.AsyncClient(transport=httpx.MockTransport(lambda _: httpx.Response(200, json=payload))) + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail: Final = OpenAIModerationGuardrail(guardrail_name="openai-mod") + guardrail.async_handler = http_client + request_data: Final[dict[str, object]] = {"metadata": {}} + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data=request_data, input_type=input_type) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == f"modr-{stage}" + assert json.loads(headers["x-litellm-guardrail-scan-metadata"]) == [ + {"guardrail": "openai-mod", "stage": stage, "provider": "openai_moderation", "scan_id": f"modr-{stage}"} + ] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 7da55f22bda..c0762edec92 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3,6 +3,8 @@ Unit tests for Bedrock Guardrails """ import json +import asyncio +from datetime import datetime, timezone import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -28,6 +30,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockTextContent, ) from litellm.types.utils import CallTypes, ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.mark.asyncio @@ -5842,3 +5845,36 @@ async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): assert response["action"] == "NONE" assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_apply_guardrail_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the ApplyGuardrail request is signed with SigV4, and botocore + refreshes expiring credentials inside that signing with a blocking HTTP call, so it must run + on a worker thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") + probe = EventLoopProbe() + allowed = httpx.Response( + 200, + json={"action": "NONE", "outputs": [], "assessments": []}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=allowed)): + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await guardrail._post_apply_guardrail_content( + content=[{"text": {"text": "hello"}}], + base_request_data={"source": "INPUT"}, + credentials=probe.credentials(), + aws_region_name="us-east-1", + api_key=None, + request_data={}, + event_type=GuardrailEventHooks.pre_call, + start_time=datetime.now(timezone.utc), + completed_chunk_usages=[], + ) + await release + + assert response["action"] == "NONE" + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index f4af77d2e40..bbc8fd539a3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -5,10 +5,12 @@ All Bedrock HTTP calls are mocked; no real AWS calls are made. """ import json +import asyncio import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest +import httpx from fastapi import HTTPException @@ -21,6 +23,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailResponse, ) from litellm.types.utils import Choices, Message, ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe CONTENT_FILTER_CHECKS = {"contentFilter": {"categories": [{"category": "VIOLENCE"}]}} @@ -861,3 +864,33 @@ async def test_checks_bearer_token_never_runs_the_sigv4_credential_chain(monkeyp {"check": "contentFilter", "category": "VIOLENCE", "severityScore": 0.8} ] assert post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_invoke_guardrail_checks_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the checks request is signed with SigV4, and botocore refreshes + expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker + thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + g = BedrockGuardrail(checks=CONTENT_FILTER_CHECKS, content_filter_threshold=0.5) + probe = EventLoopProbe() + allowed = httpx.Response( + 200, + json={"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.1}]}}}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ) + + with ( + patch.object(g, "_load_credentials", return_value=(probe.credentials(), "us-east-1")), + patch.object(g.async_handler, "post", new=AsyncMock(return_value=allowed)), + ): + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await g.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"messages": []}, + ) + await release + + assert response == BedrockGuardrailResponse() + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index a49d7723bcc..d8eeb8d2b8a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -35,7 +35,6 @@ import litellm from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( HeadroomGuardrail, - extract_hashes_from_messages, has_headroom_retrieve_tool, HEADROOM_RETRIEVE_TOOL_NAME, ) @@ -93,7 +92,11 @@ def _make_guardrail(**kwargs) -> HeadroomGuardrail: return HeadroomGuardrail(**defaults) -def _make_compress_response(messages: list, status: int = 200) -> MagicMock: +def _make_compress_response( + messages: list, + status: int = 200, + ccr_hashes: list[str] | None = None, +) -> MagicMock: mock = MagicMock() mock.status_code = status mock.json.return_value = { @@ -102,6 +105,7 @@ def _make_compress_response(messages: list, status: int = 200) -> MagicMock: "tokens_after": 100, "compression_ratio": 0.1, "transforms_applied": ["router:smart_crusher:0.35"], + **({} if ccr_hashes is None else {"ccr_hashes": ccr_hashes}), } mock.text = "" return mock @@ -335,7 +339,7 @@ async def test_apply_guardrail_injects_retrieve_tool_when_hashes_present( texts=["A" * 5000], structured_messages=ORIGINAL_MESSAGES, ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH) + mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]) with patch.object( guardrail.async_handler, @@ -390,7 +394,7 @@ async def test_apply_guardrail_preserves_existing_tools_when_injecting( structured_messages=ORIGINAL_MESSAGES, tools=[existing_tool], ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH) + mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]) with patch.object( guardrail.async_handler, @@ -489,17 +493,18 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( original_content = "This is the full compressed content." mock_retrieve = _make_retrieve_response(original_content) + # Registered hashes are lowercase; a model may echo the marker's hex in uppercase. tool_calls = [ { "id": "call_abc123", "type": "function", "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": {"hash": "b573993006976af767214fac"}, + "arguments": {"hash": "B573993006976AF767214FAC"}, } ] response = _make_openai_response_with_tool_call( tool_name=HEADROOM_RETRIEVE_TOOL_NAME, - arguments={"hash": "b573993006976af767214fac"}, + arguments={"hash": "B573993006976AF767214FAC"}, tool_id="call_abc123", ) messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] @@ -843,33 +848,110 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag assert tool_result_block["content"] == original_content -def test_extract_hashes_from_messages_finds_hashes(): +HASH_SHAPED_HISTORY = [ + {"role": "system", "content": "You are Claude Code."}, + {"role": "user", "content": [{"type": "text", "text": "Run git log."}]}, + { + "role": "assistant", + "content": [{"type": "text", "text": "Done."}], + "tool_calls": [{"id": "tu_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "tu_1", "content": "hash=3f2a9c1d7e5b4a6f8c0d1e2f9a8b7c6d5e4f3a2b"}, + {"role": "user", "content": "Please fetch hash=deadbeef000000000000dead for me."}, +] + + +async def _apply(guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None) -> dict: + request_data = {"model": "claude-sonnet-5"} + + def _echo(**kwargs): + return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"])), ccr_hashes=ccr_hashes) + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + return await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))), + request_data=request_data, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_hash_shaped_text_in_history_never_injects_retrieve_tool(guardrail: HeadroomGuardrail): + """Regression for LIT-7086: a git SHA in a tool result and a hash= string the + caller typed both look like markers, but the service stored nothing, so the + tool must not appear and no hash may be registered as issued. Covers a + service that omits ccr_hashes, returns it empty, or returns a non-list.""" + for ccr_hashes in (None, [], "b573993006976af767214fac"): + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes) + + assert not has_headroom_retrieve_tool(result.get("tools") or []) + assert not guardrail._issued_hashes_by_call_id + + +@pytest.mark.asyncio +async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardrail: HeadroomGuardrail): + """Only the hashes the service reports in ccr_hashes are honored, in the + service's own 12 to 24 hex grammar; anything else is dropped because each + entry is interpolated into the /v1/retrieve URL.""" + result = await _apply( + guardrail, + HASH_SHAPED_HISTORY, + ccr_hashes=["98CA69107318", "b573993006976af767214fac", "../../etc/passwd", "tooshort", 42], + ) + + assert has_headroom_retrieve_tool(result.get("tools") or []) + (issued, _expiry), = guardrail._issued_hashes_by_call_id.values() + assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"}) + + +@pytest.mark.asyncio +async def test_ccr_retrieval_disabled_ignores_service_declared_hashes(monkeypatch: pytest.MonkeyPatch): + """`ccr_retrieval: false` in config.yaml compresses without any retrieval + round trip, so it has to reach the instance through the initializer.""" + from litellm.proxy.guardrails.guardrail_hooks.headroom import initialize_guardrail + from litellm.types.guardrails import LitellmParams + + monkeypatch.setattr(litellm.logging_callback_manager, "add_litellm_callback", lambda callback: None) + params = LitellmParams(guardrail="headroom", mode="pre_call", api_base=FAKE_API_BASE, ccr_retrieval=False) + guardrail = initialize_guardrail(params, {"guardrail_name": "headroom", "litellm_params": params}) + + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=["b573993006976af767214fac"]) + + assert result["structured_messages"][-1] == HASH_SHAPED_HISTORY[-1] + assert not has_headroom_retrieve_tool(result.get("tools") or []) + assert not guardrail._issued_hashes_by_call_id + + +@pytest.mark.asyncio +async def test_anthropic_assistant_history_never_reaches_compression_service(guardrail: HeadroomGuardrail): + """The public Anthropic handler translates assistant content blocks to a + string before Headroom sees them, so model-authored rows must be excluded + from the compression payload rather than protected by their content shape.""" + from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler + + table = "| Guardrail | Model |\n|---|---|\n" + "\n".join(f"| gr-{i} | model-{i} |" for i in range(40)) messages = [ - {"role": "user", "content": "Retrieve more: hash=b573993006976af767214fac"}, - {"role": "assistant", "content": "Also: hash=aabbccdd001122334455aabb"}, + {"role": "user", "content": [{"type": "text", "text": "List the guardrails."}]}, + {"role": "assistant", "content": [{"type": "text", "text": table}]}, + {"role": "user", "content": [{"type": "text", "text": "Earlier follow-up. " + "B" * 5000}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Noted."}]}, + {"role": "user", "content": "Re-print the table."}, ] - hashes = extract_hashes_from_messages(messages) - assert "b573993006976af767214fac" in hashes - assert "aabbccdd001122334455aabb" in hashes + sent: dict = {} + def _echo(**kwargs): + sent["messages"] = kwargs["json"]["messages"] + return _make_compress_response(json.loads(json.dumps(sent["messages"]))) -def test_extract_hashes_from_messages_ignores_short_hashes(): - messages = [{"role": "user", "content": "hash=tooshort"}] - hashes = extract_hashes_from_messages(messages) - assert not hashes + data = {"model": "claude-sonnet-5", "messages": json.loads(json.dumps(messages))} + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + result = await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + assert [row["role"] for row in sent["messages"]] == ["user", "user"] + assert sent["messages"][1]["content"] == "Earlier follow-up. " + "B" * 5000 + assert table not in json.dumps(sent["messages"]) + assert result["messages"][1]["content"] == [{"type": "text", "text": table}] -def test_extract_hashes_from_list_content_blocks(): - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "hash=b573993006976af767214fac found here"}, - ], - } - ] - hashes = extract_hashes_from_messages(messages) - assert "b573993006976af767214fac" in hashes def test_has_headroom_retrieve_tool_recognizes_anthropic_native_shape(): @@ -1001,7 +1083,7 @@ async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstre guardrail.async_handler, "post", new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH), + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]), ): result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) @@ -1793,7 +1875,7 @@ async def test_apply_guardrail_restores_rewritten_all_text_row( ) compressed = _echo_wire_view() compressed[0]["content"] = "compressed history. Retrieve more: hash=b573993006976af767214fac" - mock_response = _make_compress_response(compressed) + mock_response = _make_compress_response(compressed, ccr_hashes=["b573993006976af767214fac"]) with patch.object( guardrail.async_handler, @@ -1819,7 +1901,7 @@ async def test_apply_guardrail_restores_rewritten_all_text_row( assert history_content[0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} # Mixed row passes through byte-identical. assert messages[2]["content"] == PARTS_MESSAGES[2]["content"] - # Hashes inside restored parts still drive retrieve-tool injection. + # The service-declared hash still drives retrieve-tool injection on a restored row. assert has_headroom_retrieve_tool(result.get("tools") or []) @@ -2159,7 +2241,7 @@ async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_c guardrail.async_handler, "post", new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH), + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]), ): result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion) @@ -2345,7 +2427,11 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( AGENTIC_MESSAGES = [ {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, {"role": "user", "content": "H" * 5000}, - {"role": "assistant", "content": "Older answer. " + "O" * 5000}, + { + "role": "assistant", + "content": "Older answer. " + "O" * 5000, + "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, {"role": "tool", "tool_call_id": "old_1", "content": "older tool output " + "T" * 5000}, { "role": "assistant", @@ -2420,23 +2506,21 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): """Negative control: protection must not turn compression into a no-op.""" compressed_history = [ {"role": "user", "content": "hist. hash=b573993006976af767214fac"}, - {"role": "assistant", "content": "older. hash=a73993006976af767214fac1"}, {"role": "tool", "tool_call_id": "old_1", "content": "older tool. hash=c73993006976af767214fac2"}, ] wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES, returned=compressed_history) - # Exactly the three history rows go to the service, in order. - assert [row["role"] for row in wire] == ["user", "assistant", "tool"] + # Older user and tool rows go to the service, in order. Every assistant row + # stays out, but the tool results those turns asked for remain compressible. + assert [row["role"] for row in wire] == ["user", "tool"] assert wire[0]["content"] == "H" * 5000 - assert wire[2]["tool_call_id"] == "old_1" + assert wire[1]["tool_call_id"] == "old_1" messages = result["structured_messages"] assert len(messages) == len(AGENTIC_MESSAGES) assert messages[1] == compressed_history[0] - assert messages[2] == compressed_history[1] - assert messages[3] == compressed_history[2] - # Hashes in the compressed history still drive retrieve-tool injection. - assert has_headroom_retrieve_tool(result.get("tools") or []) + assert messages[2] == AGENTIC_MESSAGES[2] + assert messages[3] == compressed_history[1] # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py index 4444cd693ff..d57a91d45bf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -6,16 +6,19 @@ and allows requests with only registered servers. Covers both /chat/completions and /responses API paths (same pre_call_hook logic, different call_type). """ +from typing import Literal from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException +import litellm from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.mcp_security import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( MCPSecurityGuardrail, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams @pytest.fixture @@ -182,3 +185,23 @@ class TestMCPSecurityGuardrailPreCall: call_type="acompletion", ) assert result == data + + +class TestInitializeGuardrail: + @pytest.mark.parametrize( + "configured,expected", + [("block", "block"), ("alert", "alert"), (None, "alert"), ("warn", "alert"), ("end_session", "alert")], + ) + def test_on_violation_from_litellm_params( + self, + configured: Literal["block", "alert", "warn", "end_session"] | None, + expected: Literal["block", "alert"], + ): + litellm_params = LitellmParams(guardrail="mcp_security", mode="pre_call", on_violation=configured) + guardrail = Guardrail(guardrail_name="mcp-security-block", litellm_params=litellm_params) + + result = initialize_guardrail(litellm_params=litellm_params, guardrail=guardrail) + + assert isinstance(result, MCPSecurityGuardrail) + assert result.on_violation == expected + assert result in litellm.callbacks diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 8f29ba66814..3d7c6e06d94 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -12,6 +12,7 @@ This test file follows LiteLLM's testing patterns and covers: import copy import json from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -5785,7 +5786,14 @@ class TestPanwAirsScanIdExposure: headers = get_logging_caching_headers(data) assert headers["x-litellm-guardrail-scan-id"] == "scan-abc-123" - assert "x-litellm-guardrail-scan-metadata" not in headers + assert json.loads(headers["x-litellm-guardrail-scan-metadata"]) == [ + { + "guardrail": handler.guardrail_name, + "stage": "pre_call", + "provider": "panw_prisma_airs", + "scan_id": "scan-abc-123", + } + ] @pytest.mark.asyncio async def test_request_and_response_scan_ids_are_both_exposed(self, user_api_key_dict): @@ -5809,6 +5817,26 @@ class TestPanwAirsScanIdExposure: headers = get_logging_caching_headers(data) assert headers["x-litellm-guardrail-scan-id"] == "scan-abc-123,scan-response-456" + assert [(e["stage"], e["scan_id"]) for e in json.loads(headers["x-litellm-guardrail-scan-metadata"])] == [ + ("pre_call", "scan-abc-123"), + ("post_call", "scan-response-456"), + ] + + @pytest.mark.asyncio + async def test_apply_guardrail_response_scan_is_tagged_post_call(self): + from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers + + handler: Final = self._handler(self.ALLOW_SCAN_RESULT) + request_data: Final[dict[str, object]] = {"litellm_call_id": "test-call-id", "model": "gpt-4", "metadata": {}} + + await handler.apply_guardrail( + inputs={"texts": ["Hello world"]}, request_data=request_data, input_type="response" + ) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + entries: Final = json.loads(headers["x-litellm-guardrail-scan-metadata"]) + assert [(e["stage"], e["provider"]) for e in entries] == [("post_call", "panw_prisma_airs")] @pytest.mark.asyncio async def test_repeated_scan_id_is_not_duplicated(self, user_api_key_dict): @@ -5850,6 +5878,8 @@ class TestPanwAirsScanIdExposure: assert "guardrail_scan_ids" in _UNTRUSTED_METADATA_CONTROL_FIELDS assert "guardrail_scan_ids" in _UNTRUSTED_ROOT_CONTROL_FIELDS + assert "guardrail_scan_metadata" in _UNTRUSTED_METADATA_CONTROL_FIELDS + assert "guardrail_scan_metadata" in _UNTRUSTED_ROOT_CONTROL_FIELDS class TestPanwAirsBlockedErrorDetailPassthrough: """Regression tests for the full AIRS scan response on blocks. diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 0dbd4591ac9..87a1b84acc5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -3,6 +3,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) """ import json +import logging import re from unittest.mock import patch @@ -520,6 +521,80 @@ class TestToolPermissionGuardrail: ) assert excinfo.value.status_code == 400 + @pytest.mark.asyncio + async def test_async_pre_call_hook_without_tools_logs_skip_at_debug(self, caplog): + data = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + result = await self.guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="completion", + ) + + assert result is data + skip_levels = [r.levelno for r in caplog.records if "No tools or functions in data" in r.getMessage()] + assert skip_levels == [logging.DEBUG] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + @pytest.mark.asyncio + async def test_async_pre_call_hook_denied_tool_logs_at_info(self, caplog): + data = {"tools": [{"type": "function", "function": {"name": "Read"}}]} + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException): + await self.guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="completion", + ) + + denied_levels = [ + r.levelno + for r in caplog.records + if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'" + ] + assert denied_levels == [logging.INFO] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + @pytest.mark.asyncio + async def test_async_post_call_success_hook_denied_tool_logs_at_info(self, caplog): + tool_call = {"function": {"name": "Read", "arguments": "{}"}, "type": "function"} + response = ModelResponse(choices=[Choices(message={"tool_calls": [tool_call]})]) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await self.guardrail.async_post_call_success_hook( + data={"guardrails": ["test-tool-permission"]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + denied_levels = [ + r.levelno + for r in caplog.records + if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'" + ] + assert denied_levels == [logging.INFO] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + def test_parse_tool_call_arguments_malformed_json_logs_warning(self, caplog): + tool_call = ChatCompletionMessageToolCall(function={"name": "Bash", "arguments": "{not json"}, id="call_1") + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + parsed, error = self.guardrail._parse_tool_call_arguments(tool_call) + + assert parsed is None + assert error == "arguments could not be parsed" + warning_messages = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert len(warning_messages) == 1 + assert warning_messages[0].startswith("Tool Permission Guardrail: Failed to decode arguments for tool Bash") + @pytest.mark.asyncio async def test_async_pre_call_hook_blocks_legacy_functions(self): data = { diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index a579370ad3c..3a7ae7aba61 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1034,7 +1034,7 @@ class TestStreamingTransform: ) emitted = [] - async for item in handler._emit_streaming_http_error( + async for item in handler.emit_streaming_http_error( exc, call_type=CallTypes.asend_message.value, responses_so_far=[{"id": "req-1"}], diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index db2f94306fb..0d723d6671f 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -56,13 +56,13 @@ class TestPolicyFromLitellmParams: class _FakeRouter: - """Minimal stand-in for litellm.Router.get_model_list, for policy_for_model.""" + """Minimal stand-in for litellm.Router.deployments_for_request, for policy_for_model.""" def __init__(self, deployments: list[dict[str, Any]]): self._deployments = deployments - def get_model_list(self, model_name, team_id=None): - return [d for d in self._deployments if d.get("model_name") == model_name] + def deployments_for_request(self, model, request_kwargs): + return [d for d in self._deployments if d.get("model_name") == model] def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[str, Any]: @@ -78,23 +78,23 @@ def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[ class TestPolicyForModel: def test_no_router_returns_none(self): - assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=None, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_no_marker_deployment_returns_none(self): router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_marker_deployment_without_policy_returns_none(self): router = _FakeRouter( [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] ) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_marker_deployment_with_policy_is_found(self): router = _FakeRouter( [_marker({"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"})] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) def test_picks_the_marker_whose_tags_the_request_carries(self): @@ -107,8 +107,8 @@ class TestPolicyForModel: ] ) - eu = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) - us = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + eu = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("eu",)) + us = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) assert eu == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) assert us == AutoRouterCompressionPolicy(routing="headroom-us", model=None) @@ -116,7 +116,7 @@ class TestPolicyForModel: def test_untagged_marker_matches_any_request(self): router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) policy = policy_for_model( - llm_router=router, model_alias="smart-router", team_id=None, request_tags=("anything",) + llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("anything",) ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) @@ -128,14 +128,14 @@ class TestPolicyForModel: _marker({"auto_router_routing_compression": "headroom-default"}), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None) def test_no_untagged_fallback_means_no_policy(self): """No matching marker means no policy, not an unrelated slice's compression.""" router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"])]) assert ( - policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) is None + policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) is None ) def test_tag_scoped_marker_takes_precedence_over_untagged(self): @@ -147,7 +147,7 @@ class TestPolicyForModel: _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("eu",)) assert policy == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index fe2cd819717..530f8ffd854 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1361,6 +1361,22 @@ async def test_patch_guardrail_endpoint( assert "Failed to update" in str(mock_logger.warning.call_args) +@pytest.mark.asyncio +async def test_patch_guardrail_rejects_mcp_only_on_violation_with_422(mocker, mock_guardrail_registry): + mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock()) # test-quality-ok: endpoint has no DI seam + mocker.patch( # test-quality-ok: endpoint has no DI seam + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry + ) + request = PatchGuardrailRequest(litellm_params=BaseLitellmParams(on_violation="block")) + + with pytest.raises(HTTPException) as exc_info: + await patch_guardrail("test-guardrail-id", request, user_api_key_dict=MOCK_ADMIN_USER) + + assert exc_info.value.status_code == 422 + assert "only supported by guardrail='mcp_security'" in str(exc_info.value.detail) + mock_guardrail_registry.update_guardrail_in_db.assert_not_called() + + @pytest.mark.parametrize( "scenario,expected_result,expected_exception", [ diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index ab4e15ff423..e650f796f29 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -31,6 +31,7 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): "mode": "during_call", "default_on": True, "file_sanitization_fail_open": False, + "block_on_file_modify": False, }, } ], @@ -43,9 +44,11 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): assert registered[0].default_on is True assert registered[0].event_hook == "during_call" assert registered[0].file_sanitization_fail_open is False + assert registered[0].block_on_file_modify is False config_model = registered[0].get_config_model() assert config_model is not None assert config_model().file_sanitization_fail_open is True + assert config_model().block_on_file_modify is True def test_prompt_security_guard_config_no_api_key(monkeypatch: pytest.MonkeyPatch): @@ -379,6 +382,121 @@ async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch): assert result is not None +@pytest.mark.asyncio +async def test_file_sanitization_modify_blocks_by_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True + ) + csv_data = b"name,email\nAlice,alice@example.com\n" + item = { + "type": "file", + "file": { + "data": base64.b64encode(csv_data).decode(), + "mime_type": "text/csv", + }, + } + upload_response = Response( + json={"jobId": "modify-job"}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"), + ) + poll_response = Response( + json={ + "status": "done", + "content": "name,email\nAlice,[REDACTED]\n", + "metadata": {"action": "modify", "violations": ["Sensitive Data"]}, + }, + status_code=200, + request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"), + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)): + with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)): + with pytest.raises(HTTPException) as exc_info: + await guardrail._process_document_item(item, None) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == "Document blocked by Prompt Security. Violations: Sensitive Data" + + +@pytest.mark.asyncio +async def test_standalone_image_sanitization_modify_blocks_by_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True + ) + guardrail.poll_interval = 0 + image_url = "data:image/png;base64," + base64.b64encode(b"image-content").decode() + upload_response = Response( + json={"jobId": "modify-image-job"}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"), + ) + poll_response = Response( + json={ + "status": "done", + "content": "Email: [REDACTED]", + "metadata": {"action": "modify", "violations": ["Sensitive Data"]}, + }, + status_code=200, + request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"), + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)): + with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)): + with pytest.raises(HTTPException) as exc_info: + await guardrail._process_standalone_images([image_url], None) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == "Image blocked by Prompt Security. Violations: Sensitive Data" + + +@pytest.mark.asyncio +async def test_file_sanitization_modify_can_rewrite_when_blocking_disabled(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + block_on_file_modify=False, + ) + csv_data = b"name,email\nAlice,alice@example.com\n" + item = { + "type": "file", + "file": { + "data": base64.b64encode(csv_data).decode(), + "mime_type": "text/csv", + }, + } + upload_response = Response( + json={"jobId": "modify-job"}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"), + ) + poll_response = Response( + json={ + "status": "done", + "content": "name,email\nAlice,[REDACTED]\n", + "metadata": {"action": "modify", "violations": ["Sensitive Data"]}, + }, + status_code=200, + request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"), + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)): + with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)): + result = await guardrail._process_document_item(item, None) + + assert base64.b64decode(result["file"]["data"]) == b"name,email\nAlice,[REDACTED]\n" + + @pytest.mark.asyncio @pytest.mark.parametrize( "timeout", diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e9e58347337..527d46931fe 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1058,6 +1058,35 @@ async def test_health_services_endpoint_galileo(status, error_message): mock_instance.async_health_check.assert_awaited_once() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status,error_message", + [ + ("healthy", ""), + ("unhealthy", "PointFive authentication failed"), + ], +) +async def test_health_services_endpoint_pointfive(monkeypatch, status, error_message): + import litellm.integrations.pointfive as pointfive_package + + mock_instance = MagicMock() + mock_instance.async_health_check = AsyncMock(return_value={"status": status, "error_message": error_message}) + logger_class = MagicMock(return_value=mock_instance) + monkeypatch.setattr(pointfive_package, "PointFiveLogger", logger_class) + + result = await health_services_endpoint(user_api_key_dict=_pointfive_admin(), service="pointfive") + + if status == "healthy": + assert result["status"] == "healthy" + assert result["message"] == "PointFive is healthy" + else: + assert result["status"] == "unhealthy" + assert result["message"] == error_message + mock_instance.async_health_check.assert_awaited_once() + # A check that left the periodic flush running would leak a flusher per press of the ui test button. + logger_class.assert_called_once_with(start_periodic_flush=False) + + @pytest.mark.asyncio async def test_health_services_endpoint_datadog_llm_observability(): """ @@ -1301,6 +1330,33 @@ def test_health_readiness_details_returns_diagnostic_fields(monkeypatch): assert "cache" in response_data +@pytest.mark.parametrize( + "general_settings, expected_warning", + [ + ({}, True), + ({"disable_env_credential_login": True}, False), + ], +) +def test_health_readiness_details_reports_env_credential_login_warning(monkeypatch, general_settings, expected_warning): + """ + The Admin UI banner is driven by this flag: it must be True while + env-credential login is possible and False once + `disable_env_credential_login` turns that login path off. + """ + app = FastAPI() + app.include_router(_health_endpoints_module.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + client = TestClient(app) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + response = client.get("/health/readiness/details") + + assert response.status_code == 200, response.text + assert response.json()["show_env_credential_login_warning"] is expected_warning + + def test_health_readiness_allows_explicit_legacy_public_details(monkeypatch): """ Operators can explicitly preserve the legacy public readiness payload. @@ -3104,3 +3160,60 @@ def test_test_model_connection_accepts_image_edit_mode(monkeypatch): assert response.status_code == 200, response.text assert response.json()["status"] == "success" + + +def _pointfive_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(token="admin-token", user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) + + +@pytest.mark.asyncio +async def test_health_services_endpoint_pointfive_without_a_key_is_unhealthy_not_a_server_error(monkeypatch): + """ + The logger refuses to start without an api key. + + That refusal is the answer the operator asked for, so it has to come back as an + unhealthy result rather than a 500 from the endpoint. + """ + import litellm.integrations.pointfive as pointfive_package + + def refuse(**_): + raise ValueError("pointfive logging requires an api key. Set POINTFIVE_API_KEY") + + monkeypatch.setattr(pointfive_package, "PointFiveLogger", refuse) + + result = await health_services_endpoint(user_api_key_dict=_pointfive_admin(), service="pointfive") + + assert result["status"] == "unhealthy" + assert "requires an api key" in result["message"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "role", + [ + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + LitellmUserRoles.TEAM, + LitellmUserRoles.CUSTOMER, + ], +) +async def test_health_services_endpoint_pointfive_blocks_non_admin(monkeypatch, role): + """ + The ping travels on the proxy-wide PointFive credential and stamps liveness at PointFive. + + A tenant key must not be able to keep an integration looking alive, or read back + account-level authentication failures through it. + """ + import litellm.integrations.pointfive as pointfive_package + from litellm.proxy._types import ProxyException + + logger_class = MagicMock() + monkeypatch.setattr(pointfive_package, "PointFiveLogger", logger_class) + + with pytest.raises(ProxyException) as raised: + await health_services_endpoint( + user_api_key_dict=UserAPIKeyAuth(token="t", user_id="u", user_role=role), service="pointfive" + ) + + assert str(raised.value.code) == "403" + logger_class.assert_not_called() diff --git a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py index f716a8533d8..dc7ebf6e2ec 100644 --- a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py +++ b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py @@ -1,3 +1,7 @@ +import os +import subprocess +import sys +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -65,3 +69,19 @@ async def test_execute_code_loop_dispatches_litellm_skill_tool(): mock_exec.assert_awaited_once() assert mock_exec.await_args.args[0] == SKILL_TOOL_NAME assert result is final_response + + +def test_importing_proxy_hooks_does_not_mutate_litellm_callbacks(): + script: Final = ( + "import litellm; " + "litellm.callbacks = []; " + "import litellm.proxy.hooks; " + "assert len(litellm.callbacks) == 0, f'mutated: {litellm.callbacks}'" + ) + result: Final = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + ) + assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}" diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py index edc921a40a3..e3e39a87009 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -7,6 +7,7 @@ response-shape helpers the v3 limiter's post-call hooks rely on. """ import base64 +import logging import socket import uuid from collections.abc import Mapping, Sequence @@ -16,6 +17,7 @@ from typing import Final import pytest from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.constants import BATCH_ENQUEUED_TOKEN_TTL_SECONDS from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.batch_enqueued_tokens import ( @@ -439,3 +441,29 @@ async def test_redis_lua_path_full_lifecycle(): refill = await store.reserve(tokens=50, scopes=(key_scope, team_scope)) assert isinstance(refill, BatchEnqueuedTokenReservation) await store.refund(refill) + + +class _OpenBreakerRedis: + def async_register_script(self, script: str): + async def refused(keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> object: + raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") + + return refused + + +@pytest.mark.asyncio +async def test_an_open_circuit_breaker_keeps_reservations_in_memory_without_a_warning(caplog): + scope = _scope(limit=100) + store = BatchEnqueuedTokenStore( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=_OpenBreakerRedis(), default_in_memory_ttl=60)) # pyright: ignore[reportArgumentType] # duck-typed Redis double + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + reservation = await store.reserve(tokens=60, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + assert reservation.backend == "memory" + await store.save_reservation("batch_quiet", reservation) + assert await store.pop_reservation("batch_quiet") == reservation + + assert not [record for record in caplog.records if record.levelno >= logging.WARNING] + assert sum("circuit breaker is open" in record.getMessage() for record in caplog.records) == 3 diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 0ff8b67b1a7..527449bbc48 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1861,3 +1861,87 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(monkeypatch): ) assert capacity_blocked.value.status_code == 429 assert "Model capacity reached" in capacity_blocked.value.detail["error"] + + +@pytest.mark.asyncio +async def test_post_call_success_hook_attaches_priority_headers_to_dict_response(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + RateLimitResponse, + RateLimitStatus, + get_or_create_request_stash, + ) + + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=75, + limit_remaining=74, + rate_limit_type="requests", + descriptor_key="priority_model", + ) + ], + ) + response = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [], + "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}}, + } + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + assert additional_headers["x-litellm-attempted-retries"] == 0 + assert additional_headers["x-ratelimit-priority_model-limit-requests"] == 75 + assert additional_headers["x-ratelimit-priority_model-remaining-requests"] == 74 + assert additional_headers["x-litellm-priority"] == "premium" + assert additional_headers["x-litellm-rate-limiter-version"] == "v3" + + +@pytest.mark.asyncio +async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}), + response=response, + ) + + assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_metadata", "expected_priority_header"), + [ + ({"priority": "优先"}, None), + ({"priority": "high"}, "high"), + ({}, "default"), + ], +) +async def test_post_call_success_hook_priority_header_is_always_http_encodable(team_metadata, expected_priority_header): + from starlette.responses import Response + + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": [], "_hidden_params": {}} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(team_id="team-1", team_metadata=team_metadata), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + http_response = Response(headers={key: str(value) for key, value in additional_headers.items()}) + assert http_response.headers.get("x-litellm-priority") == expected_priority_header + assert http_response.headers["x-litellm-rate-limiter-version"] == "v3" diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py index 879e2d65c7a..a1b3f313814 100644 --- a/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py @@ -10,10 +10,12 @@ Tests that session-scoped budget tracking works correctly: from unittest.mock import patch +import logging import pytest from fastapi import HTTPException from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.max_budget_per_session_limiter import ( _PROXY_MaxBudgetPerSessionHandler, @@ -163,3 +165,37 @@ async def test_no_agent_id_passes(): call_type="", ) assert result is None + + +class _OpenBreakerRedis: + def __init__(self) -> None: + from litellm.caching.redis_cache import RedisCircuitBreaker + + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + for _ in range(3): + self._circuit_breaker.record_failure() + + @_redis_circuit_breaker_guard + async def async_get_cache(self, key, **kwargs): + raise AssertionError("never reached") + + def async_register_script(self, script): + @_redis_circuit_breaker_guard + async def refused(_self, keys, args): + raise AssertionError("never reached") + + return lambda keys, args: refused(self, keys, args) + + +@pytest.mark.asyncio +async def test_an_open_circuit_breaker_reads_session_spend_locally_without_a_warning(caplog): + cache = DualCache(redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + handler = _PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(cache)) + caplog.clear() + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + spend = await handler._get_current_spend("{session_budget:quiet}:spend") + + assert spend == 0.0 + assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] + assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 4003286d887..10c0bb88a82 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -6171,3 +6171,116 @@ async def test_success_hook_leaves_stash_untouched_for_non_batch_responses(): data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5)) ) assert get_request_stash().batch_enqueued_reservation == reservation + + +@pytest.mark.asyncio +async def test_post_call_success_hook_attaches_ratelimit_headers_to_dict_response(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=100, + limit_remaining=99, + rate_limit_type="requests", + descriptor_key="model_saturation_check", + ) + ], + ) + response = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [], + "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}}, + } + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-dict-response")), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + assert additional_headers["x-litellm-attempted-retries"] == 0 + assert additional_headers["x-ratelimit-model_saturation_check-limit-requests"] == 100 + assert additional_headers["x-ratelimit-model_saturation_check-remaining-requests"] == 99 + + +@pytest.mark.asyncio +async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=100, + limit_remaining=99, + rate_limit_type="requests", + descriptor_key="model_saturation_check", + ) + ], + ) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-raw-dict")), + response=response, + ) + + assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + +class _OpenBreakerRedis: + def async_register_script(self, script: str): + async def refused(keys, args): + from litellm.caching.redis_cache import RedisCircuitBreakerOpenError + + raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") + + return refused + + async def async_increment_pipeline(self, increment_list, **kwargs): + from litellm.caching.redis_cache import RedisCircuitBreakerOpenError + + raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") + + +@pytest.mark.asyncio +async def test_an_open_circuit_breaker_falls_back_to_the_pipeline_without_a_warning(caplog): + from litellm.types.caching import RedisPipelineIncrementOperation + + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=_OpenBreakerRedis())) # pyright: ignore[reportArgumentType] # duck-typed Redis double + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + await handler.async_increment_tokens_with_ttl_preservation( + pipeline_operations=[RedisPipelineIncrementOperation(key="quiet_key", increment_value=10.0, ttl=60)] + ) + + assert await handler.internal_usage_cache.dual_cache.async_get_cache("quiet_key") == 10.0 + assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] + + +@pytest.mark.asyncio +async def test_an_open_circuit_breaker_reads_the_sliding_window_locally_without_a_warning(caplog): + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=_OpenBreakerRedis())) # pyright: ignore[reportArgumentType] # duck-typed Redis double + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + values = await handler._execute_redis_batch_rate_limiter_script( + ["{quiet}:window", "{quiet}:counter"], now_int=int(time.time()) + ) + + assert isinstance(values, list) + assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] + assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index f3fd9f41cc2..c96cfc3ee4a 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,18 +1,24 @@ - +import asyncio from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth +from litellm.proxy.collector import SpendEventConsumer +from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( _get_budget_reservation_from_metadata, _ProxyDBLogger, _should_track_cost_callback, _update_database_and_spend_counters, + run_spend_event, ) -from litellm.types.utils import CallTypes, Usage +from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event +from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, Usage @pytest.mark.asyncio @@ -70,9 +76,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 +340,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 +435,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 +457,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 +493,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 +537,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 +757,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 +1181,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 +1768,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 +1847,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 +1945,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 +1991,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: @@ -2036,3 +2102,248 @@ async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_un ) assert charged == ("premium", "tier0") + + +def _offload_kwargs() -> dict: + big_prompt = "x" * 10_000 + reservation = {"reserved_cost": 0.5, "entries": [{"counter_key": "key:hash-1", "reserved_cost": 0.5}]} + return { + "litellm_call_id": "call-1", + "call_type": "acompletion", + "model": "gpt-4o", + "custom_llm_provider": "openai", + "stream": False, + "cache_hit": None, + "response_cost": 0.0125, + "completion_start_time": datetime(2026, 1, 1, 0, 0, 1), + "messages": [{"role": "user", "content": big_prompt}], + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + "litellm_params": { + "api_base": "https://api.openai.com", + "preset_cache_key": None, + "proxy_server_request": {"body": {"messages": [{"role": "user", "content": big_prompt}]}}, + "metadata": { + "user_api_key": "hash-1", + "user_api_key_hash": "hash-1", + "user_api_key_alias": "alias-1", + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-1", + "user_api_key_org_id": "org-1", + "user_api_key_end_user_id": "end-user-1", + "user_api_key_auth": UserAPIKeyAuth(api_key="hash-1", budget_reservation=reservation), + "model_group": "gpt-4o", + "model_info": {"id": "deployment-1"}, + "tags": ["tag-a"], + }, + }, + "standard_logging_object": { + "id": "chatcmpl-1", + "trace_id": "trace-1", + "response_cost": 0.0125, + "model": "gpt-4o-2024-08-06", + "model_id": "deployment-1", + "model_group": "gpt-4o", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "prompt_tokens": 5000, + "completion_tokens": 4000, + "total_tokens": 9000, + "request_tags": ["tag-a"], + "request_model_access_groups": ["premium"], + "messages": [{"role": "user", "content": big_prompt}], + "response": {"choices": [{"message": {"content": "y" * 10_000}}]}, + "model_parameters": {"temperature": 0.1}, + "metadata": { + "user_api_key_hash": "hash-1", + "user_api_key_end_user_id": "end-user-1", + "usage_object": {"prompt_tokens": 5000, "completion_tokens": 4000, "total_tokens": 9000}, + }, + "hidden_params": {"litellm_overhead_time_ms": 3}, + "model_map_information": {}, + "cost_breakdown": {"input_cost": 0.0125, "output_cost": 0.0}, + }, + } + + +def _offload_response() -> ModelResponse: + return ModelResponse( + id="chatcmpl-1", + model="gpt-4o-2024-08-06", + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call-1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + "finish_reason": "tool_calls", + } + ], + usage=Usage(prompt_tokens=5000, completion_tokens=4000, total_tokens=9000), + ) + + +class _RecordingHandler: + def __init__(self) -> None: + self.lines: list[bytes] = [] # mutable-ok: test double records the events the sidecar received + + async def __call__(self, line: bytes) -> None: + self.lines.append(line) + + +async def _no_fallback(line: bytes) -> None: + raise AssertionError("the sidecar was reachable, nothing should fall back") + + +@pytest.mark.asyncio +async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_skips_the_pipeline(tmp_path): + handler = _RecordingHandler() + consumer = SpendEventConsumer(handler) + address = UnixAddress(path=str(tmp_path / "spend.sock")) + server = await consumer.serve(address) + producer = SpendEventProducer( + address=address, on_unavailable="fallback", buffer_size=10, connect_timeout=1.0, fallback=_no_fallback + ) + logger = _ProxyDBLogger(producer) + + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as counters, + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + await logger.async_log_success_event(_offload_kwargs(), _offload_response(), datetime.now(), datetime.now()) + await producer.close(drain_timeout=5.0) + server.close() + assert await consumer.drain(timeout=5.0) == 0 + + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_awaited() + counters.assert_not_awaited() + assert producer.stats().sent == 1 + assert len(handler.lines) == 1 + assert len(handler.lines[0]) < 4_000 + event = decode_spend_event(handler.lines[0]) + assert not isinstance(event, SpendEventDecodeError) + assert event.litellm_params["metadata"]["user_api_key_team_id"] == "team-1" + assert event.response_cost == 0.0125 + + +@pytest.mark.asyncio +async def test_async_log_success_event_keeps_batch_retrieves_in_process(): + producer = SpendEventProducer( + address=UnixAddress(path="/nonexistent/spend.sock"), + on_unavailable="drop", + buffer_size=10, + connect_timeout=1.0, + fallback=_no_fallback, + ) + logger = _ProxyDBLogger(producer) + kwargs = {**_offload_kwargs(), "call_type": CallTypes.aretrieve_batch.value} + completed_batch = LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + output_file_id="file-out", + object="batch", + status="completed", + ) + + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ), + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True) + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + await logger.async_log_success_event(kwargs, completed_batch, datetime.now(), datetime.now()) + + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + assert producer.stats().queued == 0 + + +async def _spend_row_written_by(run) -> tuple[SpendLogsPayload, dict, tuple[str, ...]]: + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as counters, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True) + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + await run() + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + written = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs + counters.assert_awaited_once() + counted = dict(counters.await_args.kwargs) + + row = get_logging_payload( + kwargs=written["kwargs"], + response_obj=written["completion_response"], + start_time=written["start_time"], + end_time=written["end_time"], + ) + return row, counted, response_tool_call_names(written["completion_response"]) + + +@pytest.mark.asyncio +async def test_sidecar_writes_the_same_spend_row_and_counters_as_the_in_process_path(): + start_time = datetime(2026, 1, 1, 0, 0, 0) + end_time = datetime(2026, 1, 1, 0, 0, 2) + + async def in_process() -> None: + await _ProxyDBLogger().async_log_success_event(_offload_kwargs(), _offload_response(), start_time, end_time) + + async def via_sidecar() -> None: + line = build_spend_event(_offload_kwargs(), _offload_response(), start_time, end_time, store_bodies=False) + assert isinstance(line, bytes) + await run_spend_event(line) + + in_process_row, in_process_counters, in_process_tools = await _spend_row_written_by(in_process) + sidecar_row, sidecar_counters, sidecar_tools = await _spend_row_written_by(via_sidecar) + + assert sidecar_row == in_process_row + assert in_process_row["spend"] == 0.0125 + assert in_process_row["team_id"] == "team-1" + assert in_process_row["end_user"] == "end-user-1" + assert in_process_row["total_tokens"] == 9000 + assert in_process_row["model_id"] == "deployment-1" + assert in_process_row["request_tags"] == '["tag-a"]' + assert in_process_row["messages"] == "{}" + assert in_process_row["response"] == "{}" + assert sidecar_counters == in_process_counters + assert in_process_counters["token"] == "hash-1" + assert in_process_counters["response_cost"] == 0.0125 + assert in_process_counters["budget_reservation"]["reserved_cost"] == 0.5 + assert in_process_counters["model_access_groups"] == ("premium",) + assert sidecar_tools == in_process_tools == ("get_weather",) + + +@pytest.mark.asyncio +async def test_sidecar_ignores_an_undecodable_event(): # test-quality-ok: a discarded event has no observable output other than the DB writer never being reached + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + await run_spend_event(b"garbage\n") + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_awaited() diff --git a/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py b/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py index 35c0f8deaf1..463d3c7ef5e 100644 --- a/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py +++ b/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py @@ -6,6 +6,7 @@ This feature allows guardrails to route requests to a different model All subsequent requests in the same session are routed to the same model. """ +import logging import asyncio from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -13,12 +14,14 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm.exceptions import SensitiveDataRouteException from litellm.integrations.custom_guardrail import ( CustomGuardrail, get_session_id_from_request_data, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import InternalUsageCache from litellm.proxy.hooks.sensitive_data_routing import ( _PROXY_SensitiveDataRoutingHandler, SENSITIVE_ROUTING_CACHE_PREFIX, @@ -1034,3 +1037,35 @@ class TestPreCallHookDeferredRouting: metrics_kwargs = prom._record_guardrail_metrics.call_args.kwargs assert metrics_kwargs["status"] == "intervened" assert metrics_kwargs["error_type"] is None + + +class _OpenBreakerRedis: + def __init__(self) -> None: + from litellm.caching.redis_cache import RedisCircuitBreaker + + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + for _ in range(3): + self._circuit_breaker.record_failure() + + @_redis_circuit_breaker_guard + async def async_get_cache(self, key, **kwargs): + raise AssertionError("never reached") + + @_redis_circuit_breaker_guard + async def async_set_cache(self, key, value, **kwargs): + raise AssertionError("never reached") + + +@pytest.mark.asyncio +async def test_an_open_circuit_breaker_keeps_session_routing_in_memory_without_a_warning(caplog): + cache = DualCache(redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + handler = _PROXY_SensitiveDataRoutingHandler(internal_usage_cache=InternalUsageCache(cache)) + caplog.clear() + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + await handler.set_session_routing("quiet-session", "safe-model") + routed = await handler._get_routed_model("quiet-session", None) + + assert routed == "safe-model" + assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] + assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 2839acab6b0..bdaca9ffc2d 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -3670,5 +3670,42 @@ async def test_post_call_success_hook_contains_header_merge_failures( ) +@pytest.mark.asyncio +async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loop(rate_limiter): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + handler, _cache = rate_limiter + stash = get_or_create_request_stash() + warm_tokenizer("claude-fable-5") + data: dict[str, object] = { + "model": "claude-fable-5", + "messages": [{"role": "user", "content": text * 100}], + } + itpm_descriptor = { + "key": PROJECT_ITPM_DESCRIPTOR_KEY, + "value": "proj-loop:claude-fable-5", + "rate_limit": {"tokens_per_unit": 10_000_000, "window_size": 60}, + } + + _, took, lags = await timed_with_loop_lags( + lambda: handler._reserve_project_io_tokens_or_raise( + descriptors=[itpm_descriptor], + data=data, + requested_model="claude-fable-5", + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-itpm-loop"), project_id="proj-loop"), + tpm_reservation_scopes=[], + tpm_reservation_amount=0, + ) + ) + + assert stash.rate_limit_response is not None + assert_loop_stayed_free(took, lags) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index 203391aadad..d8b3eef98bd 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -5,12 +5,12 @@ from typing import Any, Dict import orjson import pytest -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from starlette.requests import Request from starlette.responses import Response -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.image_endpoints import endpoints @@ -167,3 +167,47 @@ def test_image_edit_multipart_n_that_is_not_a_number_is_left_alone(monkeypatch): assert response.status_code == 200 assert captured["n"] == "two" + + +@pytest.mark.asyncio +async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(monkeypatch: pytest.MonkeyPatch): + """A bare HTTPException carries no type or param, so the tail used to ship the + literal string "None" in both fields.""" + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + async def fake_pre_call_hook(*, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> dict[str, object]: + return data + + async def fake_post_call_failure_hook(**_: object) -> None: + return None + + async def failing_route_request(**_: object) -> None: + raise HTTPException( + status_code=404, detail={"error": "image_generation: Invalid model name passed in model=dall-e-3"} + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", + SimpleNamespace(pre_call_hook=fake_pre_call_hook, post_call_failure_hook=fake_post_call_failure_hook), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version") + monkeypatch.setattr("litellm.proxy.image_endpoints.endpoints.route_request", failing_route_request) + + body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk"}) + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + request = Request({"type": "http", "method": "POST", "path": "/v1/images/generations", "headers": []}, receive) + + with pytest.raises(ProxyException) as raised: + await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()) + + assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "404") diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py index 93dc429168f..bb71d67f24e 100644 --- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py +++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py @@ -265,7 +265,7 @@ class TestSuggesterRejectsModelsWithoutToolCalling: def test_a_model_without_forced_tool_choice_support_remains_eligible(self, local_model_cost_map): supported_params = litellm.get_supported_openai_params( - model="amazon.nova-pro-v1:0", + model="meta.llama4-scout-17b-instruct-v1:0", custom_llm_provider="bedrock", ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 21f8d985f22..ef843adad98 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -10,7 +10,6 @@ import pytest from fastapi import HTTPException from pydantic import ValidationError - from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, @@ -21,11 +20,11 @@ from litellm.proxy.management_endpoints.auto_router_endpoints import ( preview_auto_router_routing, ) from litellm.router import Router -from litellm.types.utils import Choices, Message, ModelResponse from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterBenchmarksResponse, AutoRouterRoutingTestRequest, ) +from litellm.types.utils import Choices, Message, ModelResponse ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test", user_id="admin") @@ -489,6 +488,7 @@ class TestAutoRouterBenchmarks: monkeypatch: pytest.MonkeyPatch, rows: Sequence[Mapping[str, object]], model_list: Sequence[object], + api_key: str | None = None, ) -> AutoRouterBenchmarksResponse: from litellm.proxy import proxy_server from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks @@ -503,6 +503,7 @@ class TestAutoRouterBenchmarks: user_api_key_dict=ADMIN, start_date="2026-07-01", end_date="2026-08-01", + api_key=api_key, ) ROW = _SessionAggRow( @@ -527,6 +528,8 @@ class TestAutoRouterBenchmarks: total_tokens=4000, spend=10.0, saved_spend=30.0, + classifier_cost=0.4, + classifier_cost_recorded_turns=40, session_seconds=400.0, ) @@ -565,6 +568,7 @@ class TestAutoRouterBenchmarks: totals = _benchmark_totals(losing) assert totals.baseline_spend == 5.0 assert totals.saved_pct == -100.0 + assert totals.classifier_cost == 0.4 def test_an_empty_window_folds_to_zeros(self): from litellm.proxy.management_endpoints.auto_router_endpoints import ( @@ -577,6 +581,7 @@ class TestAutoRouterBenchmarks: assert totals.turns == 0 assert totals.saved_pct == 0.0 assert totals.cache.hit_rate_pct == 0.0 + assert totals.classifier_cost == 0.0 def test_totals_sum_counters_across_groups_before_deriving_ratios(self): from litellm.proxy.management_endpoints.auto_router_endpoints import ( @@ -652,11 +657,44 @@ class TestAutoRouterBenchmarks: user_api_key_dict=ADMIN, start_date="2026-07-01", end_date="2026-08-01", + api_key="key-hash", ) - assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00") + assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", "key-hash") assert response.routers_in_scope == 1 assert response.groups[0].router_name == "live-auto" assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0 + assert response.groups[0].classifier_cost == response.totals.classifier_cost == 0.4 + assert response.totals.spend - response.totals.classifier_cost == pytest.approx(9.6) + + @pytest.mark.asyncio + @pytest.mark.parametrize("recorded_turns", [0, 3, 10]) + async def test_classifier_subtotals_require_every_included_turn_to_be_recorded( + self, recorded_turns: int, monkeypatch: pytest.MonkeyPatch + ): + other: Final = self.ROW.model_copy( + update={ + "router_name": "other-auto", + "sessions": 1, + "turns": 10, + "spend": 2.0, + "saved_spend": -0.5, + "classifier_cost": recorded_turns * 0.02, + "classifier_cost_recorded_turns": recorded_turns, + } + ) + response: Final = await self._benchmarks( + monkeypatch, rows=[self.ROW.model_dump(), other.model_dump()], model_list=[] + ) + wire: Final = response.model_dump() + assert wire["groups"][0]["classifier_cost"] == 0.4 + assert wire["groups"][1]["classifier_cost"] == (pytest.approx(0.2) if recorded_turns == 10 else None) + assert wire["totals"]["classifier_cost"] == (pytest.approx(0.6) if recorded_turns == 10 else None) + assert response.totals.turns == 50 + assert response.totals.spend == 12.0 + assert response.totals.saved_spend == 29.5 + assert response.totals.baseline_spend == 41.5 + assert response.totals.saved_pct == 71.1 + assert response.totals.saved_per_session == 5.9 @pytest.mark.asyncio @pytest.mark.parametrize( @@ -720,6 +758,7 @@ class TestAutoRouterBenchmarks: assert (idle.cache.hit_rate_pct, idle.cache.coverage_pct) == (0.0, 0.0) assert idle.cache.same_model.turns == idle.cache.return_to_tier.hits == 0 assert idle.tier_turns == {} + assert idle.classifier_cost == 0.0 @pytest.mark.asyncio @pytest.mark.parametrize( @@ -796,7 +835,6 @@ class TestAutoRouterBenchmarks: from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock - from litellm.proxy.management_endpoints.auto_router_endpoints import ( get_shadow_eval_job, list_shadow_eval_jobs, @@ -806,6 +844,143 @@ from litellm.proxy.management_endpoints.auto_router_endpoints import ( from litellm.types.management_endpoints.auto_router_endpoints import SHADOW_EVAL_TURN_VALVE, StartShadowEvalRequest VIEWER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, api_key="sk-view", user_id="viewer") + + +class TestAutoRouterSession: + """GET /auto_router/session: a key reads its own session's routed model and savings, nothing else.""" + + ROW = { + "router_name": "claude-auto", + "router_type": "complexity", + "first_turn_at": datetime(2026, 9, 1, 12, 0, 0), + "last_turn_at": datetime(2026, 9, 1, 12, 5, 0), + "turns": 3, + "last_model": "anthropic/claude-sonnet-5", + "spend": 0.14, + "saved_spend": 0.24, + "classifier_cost": 0.0, + "tier_turns": {"simple": 1, "complex": 2}, + "baseline_models": {"anthropic/claude-opus-5": 3}, + } + + @staticmethod + def _rig(monkeypatch: pytest.MonkeyPatch, rows: Sequence[Mapping[str, object]]): + from litellm.proxy import proxy_server + + lookups: list[tuple[Mapping[str, object], Mapping[str, object]]] = [] + + class _Table: + async def find_first(self, where: Mapping[str, object], order: Mapping[str, object]): + lookups.append((where, order)) + matching = [r for r in rows if (r["api_key"], r["session_id"]) == (where["api_key"], where["session_id"])] + return max(matching, key=lambda r: r["last_turn_at"], default=None) + + monkeypatch.setattr( + proxy_server, "prisma_client", type("P", (), {"db": type("D", (), {"litellm_autoroutersession": _Table()})()})() + ) + return lookups + + @pytest.mark.asyncio + async def test_a_key_reads_its_own_session_with_the_baseline_its_turns_were_priced_against( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + caller = UserAPIKeyAuth(api_key="sk-caller") + self._rig(monkeypatch, [{**self.ROW, "api_key": caller.api_key, "session_id": "sess-1"}]) + response = await get_auto_router_session(user_api_key_dict=caller, session_id="sess-1") + assert response.model_dump() == { + "session_id": "sess-1", + "router_name": "claude-auto", + "router_type": "complexity", + "turns": 3, + "last_model": "anthropic/claude-sonnet-5", + "spend": 0.14, + "saved_spend": 0.24, + "baseline_spend": pytest.approx(0.38), + "baseline_model": "anthropic/claude-opus-5", + "baseline_models": {"anthropic/claude-opus-5": 3}, + } + + @pytest.mark.asyncio + async def test_another_keys_session_is_a_404_even_for_an_admin(self, monkeypatch: pytest.MonkeyPatch): + # The scope is the caller's own key hash, exactly what the spend writer keyed the row under; + # an admin wanting every key's sessions has /auto_router/benchmarks. + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + other = UserAPIKeyAuth(api_key="sk-other") + lookups = self._rig(monkeypatch, [{**self.ROW, "api_key": other.api_key, "session_id": "sess-1"}]) + with pytest.raises(HTTPException) as err: + await get_auto_router_session(user_api_key_dict=ADMIN, session_id="sess-1") + assert err.value.status_code == 404 + assert lookups == [({"api_key": ADMIN.api_key, "session_id": "sess-1"}, {"last_turn_at": "desc"})] + assert ADMIN.api_key != "sk-test" + + @pytest.mark.asyncio + async def test_the_sessions_most_recently_active_router_is_the_one_reported(self, monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + older = {**self.ROW, "api_key": ADMIN.api_key, "session_id": "s", "router_name": "old-auto"} + newer = { + **self.ROW, + "api_key": ADMIN.api_key, + "session_id": "s", + "router_name": "new-auto", + "last_turn_at": datetime(2026, 9, 1, 13, 0, 0), + } + self._rig(monkeypatch, [older, newer]) + response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") + assert response.router_name == "new-auto" + + @pytest.mark.asyncio + async def test_a_reconfigured_router_keeps_the_label_the_money_was_priced_against( + self, monkeypatch: pytest.MonkeyPatch + ): + # The proxy's router now prices against a different baseline, but the row's money was priced + # against opus for two of three turns, and the label says so; the full split is on the response. + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + priced = {"anthropic/claude-opus-5": 2, "anthropic/claude-sonnet-5": 1} + self._rig(monkeypatch, [{**self.ROW, "api_key": ADMIN.api_key, "session_id": "s", "baseline_models": priced}]) + response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") + assert response.baseline_model == "anthropic/claude-opus-5" + assert response.baseline_models == priced + + @pytest.mark.asyncio + async def test_a_session_whose_turns_recorded_no_baseline_reports_the_money_without_a_name( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + self._rig(monkeypatch, [{**self.ROW, "api_key": ADMIN.api_key, "session_id": "s", "baseline_models": {}}]) + response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") + assert response.baseline_model is None + assert response.baseline_spend == pytest.approx(0.38) + + @pytest.mark.asyncio + async def test_an_oversized_client_session_id_is_bounded_like_the_writer_bounded_it( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.db.autorouter_session_rollup import bounded_session_id + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + long_id = "s" * 300 + self._rig(monkeypatch, [{**self.ROW, "api_key": ADMIN.api_key, "session_id": bounded_session_id(long_id)}]) + response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id=long_id) + assert response.session_id == long_id + assert response.turns == 3 + + @pytest.mark.asyncio + async def test_without_a_database_the_endpoint_says_so(self, monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + monkeypatch.setattr(proxy_server, "prisma_client", None) + with pytest.raises(HTTPException) as err: + await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") + assert err.value.status_code == 500 + + NON_ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="user") @@ -1131,7 +1306,9 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp assert ( len( { - frozenset((k, tuple(v) if isinstance(v, list) else v) for k, v in row.items() if k not in ("target_id", "id")) + frozenset( + (k, tuple(v) if isinstance(v, list) else v) for k, v in row.items() if k not in ("target_id", "id") + ) for row in rows } ) @@ -1258,8 +1435,8 @@ async def test_start_shadow_eval_accepts_an_sdk_judge_when_anthropic_secret_look monkeypatch: pytest.MonkeyPatch, ) -> None: import litellm - from litellm.integrations.custom_secret_manager import CustomSecretManager import litellm.proxy.proxy_server as proxy_server + from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem class AnthropicSecretManager(CustomSecretManager): @@ -1821,9 +1998,10 @@ async def test_list_shadow_eval_jobs_rejects_a_lone_filter_half(monkeypatch: pyt @pytest.mark.asyncio async def test_start_shadow_eval_concurrent_unique_violation_is_a_409(monkeypatch: pytest.MonkeyPatch): - import litellm.proxy.proxy_server as proxy_server from prisma.errors import UniqueViolationError + import litellm.proxy.proxy_server as proxy_server + _configure_anthropic_sdk_judge(monkeypatch) prisma = _shadow_prisma() prisma.db.litellm_shadowevaljob.create_many = AsyncMock( diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 37a54c4901a..6cd900cb041 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -492,7 +492,7 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash( @pytest.mark.asyncio async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_spend_logs(): - """A dirty key no table can explain costs two digest lookups, never a token page walk or a SpendLogs scan.""" + """Without a spend-log window a dirty key no table can explain costs two digest lookups and never a token page walk.""" from litellm.proxy.utils import hash_token double_hashed = hash_token("b" * 64) @@ -518,6 +518,93 @@ async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_s assert all("take" not in call.kwargs and "skip" not in call.kwargs for call in token_lookups) +def _spend_log_transaction(mock_prisma: MagicMock, rows: list[dict[str, str | None]]) -> AsyncMock: + transaction = MagicMock() + transaction.execute_raw = AsyncMock(return_value=0) + transaction.query_raw = AsyncMock(return_value=rows) + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + return transaction.query_raw + + +def _spend_log_row(digest: str, key_alias: str, user_id: str) -> dict[str, str | None]: + return { + "digest": digest, + "first_alias": key_alias, + "last_alias": key_alias, + "first_team": None, + "last_team": None, + "first_owner": user_id, + "last_owner": user_id, + } + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_permanent_miss_with_a_window_reads_spend_logs_once_within_it(): + from litellm.proxy.utils import hash_token + + double_hashed = hash_token("permanent-miss-with-window-6852") + window = (datetime(2024, 1, 1), datetime(2024, 1, 4)) + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + spend_log_query_raw = _spend_log_transaction(mock_prisma, []) + + result = await get_api_key_metadata(prisma_client=mock_prisma, api_keys={double_hashed}, spend_logs_window=window) + + assert double_hashed not in result + assert mock_prisma.db.query_raw.await_count == 2 + ((_, digests, start, end),) = [call.args for call in spend_log_query_raw.call_args_list] + assert digests == [double_hashed] + assert (start, end) == window + + +@pytest.mark.asyncio +async def test_get_daily_activity_recovers_a_session_key_alias_from_spend_logs_around_the_page_dates(): + from litellm.proxy.utils import hash_token + + session_digest = hash_token("cli-session-daily-activity-6852") + records = [_daily_user_spend_record(user_id="session-user", api_key=session_digest, spend=1.5)] + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_table = MagicMock() + mock_table.count = AsyncMock(return_value=len(records)) + mock_table.find_many = AsyncMock(return_value=records) + mock_prisma.db.litellm_dailyuserspend = mock_table + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="session-user", user_email="session@example.com")] + ) + + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + spend_log_query_raw = _spend_log_transaction( + mock_prisma, [_spend_log_row(session_digest, "cli-session-alias", "session-user")] + ) + + result = await get_daily_activity( + prisma_client=mock_prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + page=1, + page_size=1000, + ) + + key_metadata = result.results[0].breakdown.api_keys[session_digest].metadata + assert key_metadata.key_alias == "cli-session-alias" + assert key_metadata.user_email == "session@example.com" + ((_, digests, start, end),) = [call.args for call in spend_log_query_raw.call_args_list] + assert digests == [session_digest] + assert (start, end) == (datetime(2023, 12, 31), datetime(2024, 1, 3)) + + def test_key_metadata_includes_recovered_user_email(): from litellm.proxy.management_endpoints.common_daily_activity import _key_metadata @@ -2105,3 +2192,48 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown(): # Rollups with the entity bit set must still land in their usual buckets assert daily.breakdown.models["gpt-4o"].metrics.spend == 18.0 assert daily.breakdown.api_keys["key-1"].metrics.spend == 12.0 + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_resolves_session_key_via_spend_log_window(): + from litellm.proxy.utils import hash_token + + session_digest = hash_token("cli-session-user-42") + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="user-42", user_email="user42@example.com")] + ) + + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + spend_log_query_raw = _spend_log_transaction( + mock_prisma, [_spend_log_row(session_digest, "cli-session-user-42", "user-42")] + ) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={session_digest}, + spend_logs_window=(datetime(2026, 9, 7), datetime(2026, 9, 10)), + ) + + assert result[session_digest]["key_alias"] == "cli-session-user-42" + assert result[session_digest]["user_id"] == "user-42" + assert result[session_digest]["user_email"] == "user42@example.com" + ((_, digests, start, end),) = [call.args for call in spend_log_query_raw.call_args_list] + assert digests == [session_digest] + assert (start, end) == (datetime(2026, 9, 7), datetime(2026, 9, 10)) + + +def test_spend_logs_window_pads_min_minus_one_day_and_max_plus_two_days(): + from litellm.proxy.management_endpoints.common_daily_activity import _spend_logs_window + + window = _spend_logs_window({"2026-09-08", "2026-09-05", "not-a-date"}) + + assert window == (datetime(2026, 9, 4), datetime(2026, 9, 10)) + + +def test_spend_logs_window_is_none_when_no_date_parses(): + from litellm.proxy.management_endpoints.common_daily_activity import _spend_logs_window + + assert _spend_logs_window({"garbage", ""}) is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index ec62cc47018..7ece35ceedf 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -4,13 +4,17 @@ Tests for cost tracking settings management endpoints. Tests the GET and PATCH endpoints for managing cost discount configuration. """ +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from pydantic import ValidationError import litellm +from litellm._internal_context import pinned_billing_time +from litellm.proxy._types import CostEstimateRequest from litellm.proxy.management_endpoints.cost_tracking_settings import router from litellm.proxy.proxy_server import app @@ -789,13 +793,13 @@ INPUT_TOKENS = 1000 OUTPUT_TOKENS = 500 -def _router_pricing(**pricing: float) -> MagicMock: +def _router_pricing(model: str = AN_UNDERLYING_MODEL, **pricing: float) -> MagicMock: mock_router = MagicMock() mock_router.get_model_list.return_value = [ { "model_name": AN_ALIAS, "litellm_params": { - "model": AN_UNDERLYING_MODEL, + "model": model, "custom_llm_provider": "openai", **pricing, }, @@ -811,9 +815,7 @@ async def _estimate(mock_router: MagicMock | None, model: str = AN_ALIAS, **over request = CostEstimateRequest( model=model, - input_tokens=INPUT_TOKENS, - output_tokens=OUTPUT_TOKENS, - **overrides, + **{"input_tokens": INPUT_TOKENS, "output_tokens": OUTPUT_TOKENS, **overrides}, ) with patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point "litellm.proxy.proxy_server.llm_router", mock_router @@ -909,3 +911,299 @@ class TestEstimateCostPeriodTotals: assert response.cost_per_request == pytest.approx(0.0022) assert response.daily_margin_cost == pytest.approx(0.02) assert response.daily_cost == pytest.approx(0.22) + + +CACHE_READ_TOKENS = 800 +CACHE_CREATION_TOKENS = 100 +REASONING_TOKENS = 200 +TEXT_INPUT_TOKENS = INPUT_TOKENS - CACHE_READ_TOKENS - CACHE_CREATION_TOKENS +TEXT_OUTPUT_TOKENS = OUTPUT_TOKENS - REASONING_TOKENS + + +async def _estimate_with_cache_and_reasoning(mock_router: MagicMock | None, model: str = AN_ALIAS, **overrides: int): + return await _estimate( + mock_router, + model=model, + cache_read_input_tokens=CACHE_READ_TOKENS, + cache_creation_input_tokens=CACHE_CREATION_TOKENS, + reasoning_tokens=REASONING_TOKENS, + **overrides, + ) + + +class TestEstimateCostCacheAndReasoningTokens: + @pytest.mark.asyncio + async def test_a_mapped_model_bills_cache_and_reasoning_tokens_at_their_own_rates(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "output_cost_per_reasoning_token": 1e-5, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate_with_cache_and_reasoning(None, model=A_MAPPED_MODEL, num_requests_per_day=10) + + assert response.cache_read_cost_per_request == pytest.approx(CACHE_READ_TOKENS * 3e-7) + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 3.75e-6) + assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 1e-5) + assert response.input_cost_per_request == pytest.approx( + TEXT_INPUT_TOKENS * 3e-6 + CACHE_READ_TOKENS * 3e-7 + CACHE_CREATION_TOKENS * 3.75e-6 + ) + assert response.output_cost_per_request == pytest.approx(TEXT_OUTPUT_TOKENS * 15e-6 + REASONING_TOKENS * 1e-5) + assert response.cost_per_request == pytest.approx( + response.input_cost_per_request + response.output_cost_per_request + ) + assert response.daily_cache_read_cost == pytest.approx(10 * CACHE_READ_TOKENS * 3e-7) + assert response.daily_cache_creation_cost == pytest.approx(10 * CACHE_CREATION_TOKENS * 3.75e-6) + assert response.daily_reasoning_cost == pytest.approx(10 * REASONING_TOKENS * 1e-5) + assert response.monthly_cache_read_cost is None + assert response.cache_read_input_token_cost == pytest.approx(3e-7) + assert response.cache_creation_input_token_cost == pytest.approx(3.75e-6) + assert response.output_cost_per_reasoning_token == pytest.approx(1e-5) + assert ( + response.cache_read_input_tokens, + response.cache_creation_input_tokens, + response.reasoning_tokens, + ) == (CACHE_READ_TOKENS, CACHE_CREATION_TOKENS, REASONING_TOKENS) + + @pytest.mark.asyncio + async def test_a_model_without_cache_or_reasoning_prices_estimates_what_the_proxy_bills(self, monkeypatch): + """The cost calculator bills cache tokens of a cost-map model without cache prices at zero + and its reasoning tokens at the output rate. The estimate reports those effective rates.""" + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + {"input_cost_per_token": 5e-6, "output_cost_per_token": 6e-6, "litellm_provider": "openai", "mode": "chat"}, + ) + + response = await _estimate_with_cache_and_reasoning(None, model=A_MAPPED_MODEL) + + assert response.cache_read_cost_per_request == 0.0 + assert response.cache_creation_cost_per_request == 0.0 + assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 6e-6) + assert response.input_cost_per_request == pytest.approx(TEXT_INPUT_TOKENS * 5e-6) + assert response.cost_per_request == pytest.approx(TEXT_INPUT_TOKENS * 5e-6 + OUTPUT_TOKENS * 6e-6) + assert response.cache_read_input_token_cost == 0.0 + assert response.cache_creation_input_token_cost == 0.0 + assert response.output_cost_per_reasoning_token == pytest.approx(6e-6) + + @pytest.mark.asyncio + async def test_a_request_without_cache_or_reasoning_tokens_estimates_as_before(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "output_cost_per_reasoning_token": 1e-5, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate(None, model=A_MAPPED_MODEL, num_requests_per_day=10) + + assert response.cost_per_request == pytest.approx(INPUT_TOKENS * 3e-6 + OUTPUT_TOKENS * 15e-6) + assert response.cache_read_cost_per_request == 0.0 + assert response.cache_creation_cost_per_request == 0.0 + assert response.reasoning_cost_per_request == 0.0 + assert response.daily_cache_read_cost == 0.0 + assert response.daily_reasoning_cost == 0.0 + + @pytest.mark.asyncio + async def test_a_custom_priced_deployment_bills_cache_and_reasoning_tokens_from_its_flat_rates(self): + response = await _estimate_with_cache_and_reasoning( + _router_pricing(input_cost_per_token=1e-6, output_cost_per_token=2e-6, cache_read_input_token_cost=1e-7) + ) + + assert response.cache_read_cost_per_request == pytest.approx(CACHE_READ_TOKENS * 1e-7) + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 1e-6) + assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 2e-6) + assert response.cost_per_request == pytest.approx( + TEXT_INPUT_TOKENS * 1e-6 + CACHE_READ_TOKENS * 1e-7 + CACHE_CREATION_TOKENS * 1e-6 + OUTPUT_TOKENS * 2e-6 + ) + assert response.cache_read_input_token_cost == pytest.approx(1e-7) + assert response.cache_creation_input_token_cost == pytest.approx(1e-6) + assert response.output_cost_per_reasoning_token == pytest.approx(2e-6) + + @pytest.mark.asyncio + async def test_a_custom_priced_deployment_of_a_mapped_model_inherits_its_built_in_cache_rates(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 5e-6, + "output_cost_per_token": 6e-6, + "cache_read_input_token_cost": 5e-7, + "cache_creation_input_token_cost": 6.25e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate_with_cache_and_reasoning( + _router_pricing(model=A_MAPPED_MODEL, input_cost_per_token=1e-6, output_cost_per_token=2e-6) + ) + + assert response.cache_read_cost_per_request == pytest.approx(CACHE_READ_TOKENS * 5e-7) + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 6.25e-6) + assert response.input_cost_per_request == pytest.approx( + TEXT_INPUT_TOKENS * 1e-6 + CACHE_READ_TOKENS * 5e-7 + CACHE_CREATION_TOKENS * 6.25e-6 + ) + assert response.cache_read_input_token_cost == pytest.approx(5e-7) + assert response.cache_creation_input_token_cost == pytest.approx(6.25e-6) + + @pytest.mark.asyncio + async def test_a_tiered_model_reports_the_rates_its_lines_were_billed_at(self, monkeypatch): + """Above a token tier the calculator bills every line at the tier's rate, so the reported + rates must be the tier's too: each line equals its token count times the rate next to it.""" + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate( + None, + model=A_MAPPED_MODEL, + input_tokens=250_000, + cache_read_input_tokens=200_000, + cache_creation_input_tokens=10_000, + output_tokens=1_000, + reasoning_tokens=200, + ) + + assert response.input_cost_per_token == pytest.approx(6e-6) + assert response.output_cost_per_token == pytest.approx(3e-5) + assert response.cache_read_input_token_cost == pytest.approx(6e-7) + assert response.cache_creation_input_token_cost == pytest.approx(7.5e-6) + assert response.output_cost_per_reasoning_token == pytest.approx(3e-5) + assert response.cache_read_cost_per_request == pytest.approx(200_000 * response.cache_read_input_token_cost) + assert response.cache_creation_cost_per_request == pytest.approx( + 10_000 * response.cache_creation_input_token_cost + ) + assert response.reasoning_cost_per_request == pytest.approx(200 * response.output_cost_per_reasoning_token) + assert response.input_cost_per_request == pytest.approx( + 40_000 * response.input_cost_per_token + + response.cache_read_cost_per_request + + response.cache_creation_cost_per_request + ) + assert response.output_cost_per_request == pytest.approx(1_000 * response.output_cost_per_token) + + @pytest.mark.asyncio + async def test_a_quote_prices_its_totals_and_its_rates_at_the_same_moment(self, monkeypatch): + """The totals and the reported rates resolve off-peak pricing on separate paths. A quote + taken as a window opens must not bill on one side of it and report rates from the other.""" + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "off_peak_pricing": { + "hours_utc": "02:00-03:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 5e-6, + }, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + with pinned_billing_time(datetime(2026, 1, 1, 2, 30, tzinfo=timezone.utc)): + response = await _estimate(None, model=A_MAPPED_MODEL) + + assert response.input_cost_per_token == pytest.approx(1e-6) + assert response.output_cost_per_token == pytest.approx(5e-6) + assert response.input_cost_per_request == pytest.approx(INPUT_TOKENS * response.input_cost_per_token) + assert response.output_cost_per_request == pytest.approx(OUTPUT_TOKENS * response.output_cost_per_token) + + + @pytest.mark.asyncio + async def test_an_unrouted_model_reports_the_rates_of_the_provider_the_calculator_inferred(self, monkeypatch): + """The cost calculator infers a provider this endpoint never resolved, and the provider decides + whether a tier threshold is inclusive. xai bills a request sitting exactly on the 200k threshold + at the tier rate, so the reported rates have to be the tier's rather than the sub-tier base.""" + an_xai_model = "xai/tiered-model" + monkeypatch.setitem( + litellm.model_cost, + an_xai_model, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "litellm_provider": "xai", + "mode": "chat", + }, + ) + + response = await _estimate( + None, + model=an_xai_model, + input_tokens=200_000, + cache_read_input_tokens=100_000, + output_tokens=1_000, + ) + + assert response.input_cost_per_token == pytest.approx(6e-6) + assert response.output_cost_per_token == pytest.approx(3e-5) + assert response.cache_read_input_token_cost == pytest.approx(6e-7) + assert response.cache_read_cost_per_request == pytest.approx(100_000 * response.cache_read_input_token_cost) + assert response.input_cost_per_request == pytest.approx( + 100_000 * response.input_cost_per_token + response.cache_read_cost_per_request + ) + assert response.output_cost_per_request == pytest.approx(1_000 * response.output_cost_per_token) + + +class TestCostEstimateRequestTokenSubsets: + def test_cache_tokens_beyond_the_input_tokens_are_rejected(self): + with pytest.raises(ValidationError, match="cannot exceed input_tokens"): + CostEstimateRequest( + model=AN_ALIAS, + input_tokens=INPUT_TOKENS, + output_tokens=OUTPUT_TOKENS, + cache_read_input_tokens=INPUT_TOKENS, + cache_creation_input_tokens=1, + ) + + def test_reasoning_tokens_beyond_the_output_tokens_are_rejected(self): + with pytest.raises(ValidationError, match="cannot exceed output_tokens"): + CostEstimateRequest( + model=AN_ALIAS, + input_tokens=INPUT_TOKENS, + output_tokens=OUTPUT_TOKENS, + reasoning_tokens=OUTPUT_TOKENS + 1, + ) + + def test_the_endpoint_answers_422_when_cache_tokens_exceed_input_tokens(self): + response = client.post( + "/cost/estimate", + headers={"Authorization": "Bearer sk-1234"}, + json={"model": AN_ALIAS, "input_tokens": 1000, "output_tokens": 100, "cache_read_input_tokens": 8000}, + ) + + assert response.status_code == 422 + assert "cannot exceed input_tokens" in response.text diff --git a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py index 81226981089..0ecc4f8d7cb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py +++ b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py @@ -8,6 +8,7 @@ proof-of-fix (real proxy + DB) is performed separately on the repro server. import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -457,6 +458,65 @@ async def test_scan_covered_tables_classifies_legacy_and_v2(salt_key, monkeypatc assert by_loc["credentials"].legacy == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("column", ("static_headers", "env")) +@pytest.mark.parametrize("algorithm", ("xsalsa20-poly1305", "aes-256-gcm")) +@pytest.mark.parametrize("as_json", (False, True)) +@pytest.mark.parametrize( + "case", ("legacy", "encrypted", "wrong-key", "corrupt", "invalid-shape", "invalid-scalar", "empty", "null") +) +async def test_check_classifies_mcp_secret_maps( + salt_key: str, + monkeypatch: pytest.MonkeyPatch, + column: str, + algorithm: str, + as_json: bool, + case: str, +) -> None: + monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": algorithm}) + plaintext: Final = {"Authorization": "v2:gcm:operator-text", "CUSTOM": "litellm_enc::literal\n café "} + ciphertext: Final = encrypt_value_helper(json.dumps(plaintext)) + cases: Final[dict[str, object]] = { + "legacy": plaintext, + "encrypted": ciphertext, + "wrong-key": encrypt_value_helper(json.dumps(plaintext), new_encryption_key="different-map-salt"), + "corrupt": ciphertext[:-4] + "AAAA", + "invalid-shape": encrypt_value_helper(json.dumps({"Authorization": 42})), + "invalid-scalar": "null", + "empty": {}, + "null": None, + } + value: Final = json.dumps(cases[case]) if as_json and case != "null" else cases[case] + row: Final = SimpleNamespace(**{column: value}) + client: Final = MagicMock() + _empty_covered_tables(client) + client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + client.db.litellm_mcpservertable.update = AsyncMock() + client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + client.db.litellm_config.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + + report: Final = await cm.check_encryption(client) + expected_legacy: Final = int(case == "legacy" or (case == "encrypted" and algorithm == "xsalsa20-poly1305")) + expected_v2: Final = int(case == "encrypted" and algorithm == "aes-256-gcm") + expected_invalid: Final = int(case in ("wrong-key", "corrupt", "invalid-shape", "invalid-scalar")) + + assert report.as_dict()["locations"]["mcp_server"] == { + "scanned": int(case not in ("empty", "null")), + "migrated": 0, + "already_v2": expected_v2, + "plaintext": 0, + "undecryptable": expected_invalid, + "legacy": expected_legacy, + } + assert report.residual_legacy == expected_legacy + assert report.total_undecryptable == expected_invalid + assert getattr(row, column) == value + client.db.litellm_mcpservertable.update.assert_not_awaited() + + @pytest.mark.asyncio async def test_check_counts_covered_table_residual(salt_key, monkeypatch): """check_encryption now scans the rotation-covered tables (model table here), diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 1225cb80224..bd59a82cbd2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -763,6 +763,7 @@ _EXPECTED_CUSTOMER = { "blocked_tools": [], "search_tools": [], "mcp_tool_search_enabled": None, + "skills": None, }, } diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 8766b1a1868..65cc23ea67f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5085,6 +5085,104 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch): assert len(deleted_keys) == 2 +class _JWTMappingRow: + def __init__(self, token, jwt_claim_name, jwt_claim_value): + self.token = token + self.jwt_claim_name = jwt_claim_name + self.jwt_claim_value = jwt_claim_value + + +class _CascadingJWTMappingTable: + """Mapping rows that LiteLLM_JWTKeyMapping_token_fkey drops when their key is deleted.""" + + def __init__(self, rows): + self.rows = rows + + async def find_many(self, where, **kwargs): + return [row for row in self.rows if row.token == where["token"]] + + def cascade(self, deleted_tokens): + self.rows = [row for row in self.rows if row.token not in deleted_tokens] + + +class _RecordingEvict: + def __init__(self): + self.cache_keys = () + + async def __call__(self, cache_keys, user_api_key_cache): + self.cache_keys = tuple(cache_keys) + + +@pytest.mark.asyncio +async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypatch): + """Deleting a key must evict its jwt_key_mapping cache entries (LIT-5380). + + The FK cascade removes the mapping rows, so a surviving cache entry would keep + resolving the deleted token hash and 401 every JWT call from that identity until + virtual_key_mapping_cache_ttl expires, instead of auto-registering again. + """ + jwt_table = _CascadingJWTMappingTable( + [_JWTMappingRow("hashed-token-1", "email", "user@example.com")] + ) + + key1 = LiteLLM_VerificationToken( + token="hashed-token-1", + user_id="user-123", + team_id=None, + key_alias="jwt-mapped-key", + spend=0.0, + max_budget=None, + models=[], + aliases={}, + config={}, + permissions={}, + metadata={}, + model_max_budget={}, + model_spend={}, + soft_budget_cooldown=False, + allowed_routes=[], + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[key1] + ) + mock_prisma_client.db.litellm_jwtkeymapping = jwt_table + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + + async def cascading_delete_data(tokens): + jwt_table.cascade(tokens) + return list(tokens) + + mock_prisma_client.delete_data = AsyncMock(side_effect=cascading_delete_data) + + recording_evict = _RecordingEvict() + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.evict_and_broadcast", + recording_evict, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", + lambda token: token, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ) + + await delete_verification_tokens( + tokens=["hashed-token-1"], + user_api_key_cache=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + ) + + assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",) + + @pytest.mark.asyncio async def test_delete_key_fn_persists_deleted_keys(monkeypatch): from litellm.proxy._types import KeyRequest @@ -17692,6 +17790,253 @@ async def test_check_project_key_limits_still_rejects_real_model_outside_project assert "Model 'gpt-5.4-mini' not in project's allowed models" in exc_info.value.detail["error"] +@pytest.mark.asyncio +async def test_update_key_soft_budget_updates_existing_budget_row(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123") + mock_db = MagicMock() + mock_db.litellm_budgettable.update = AsyncMock() + mock_db.litellm_budgettable.create = AsyncMock() + + result = await _update_key_soft_budget( + db=mock_db, + existing_key_row=existing_key, + soft_budget=25.0, + changed_by="user-1", + ) + + assert result == "budget-123" + mock_db.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "budget-123"}, + data={"soft_budget": 25.0, "updated_by": "user-1"}, + ) + + +@pytest.mark.asyncio +async def test_update_key_soft_budget_clears_existing_budget_row(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123") + mock_db = MagicMock() + mock_db.litellm_budgettable.update = AsyncMock() + mock_db.litellm_budgettable.create = AsyncMock() + + result = await _update_key_soft_budget( + db=mock_db, + existing_key_row=existing_key, + soft_budget=None, + changed_by="user-1", + ) + + assert result == "budget-123" + mock_db.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "budget-123"}, + data={"soft_budget": None, "updated_by": "user-1"}, + ) + + +@pytest.mark.asyncio +async def test_update_key_soft_budget_creates_budget_row_when_key_has_none(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + created_row = MagicMock() + created_row.budget_id = "budget-new" + mock_db = MagicMock() + mock_db.litellm_budgettable.create = AsyncMock(return_value=created_row) + mock_db.litellm_budgettable.update = AsyncMock() + + result = await _update_key_soft_budget( + db=mock_db, + existing_key_row=existing_key, + soft_budget=10.5, + changed_by="user-1", + ) + + assert result == "budget-new" + mock_db.litellm_budgettable.create.assert_awaited_once_with( + data={"soft_budget": 10.5, "created_by": "user-1", "updated_by": "user-1"} + ) + + +@pytest.mark.asyncio +async def test_update_key_soft_budget_noop_when_clearing_without_budget_row(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + mock_db = MagicMock() + mock_db.litellm_budgettable.create = AsyncMock() + mock_db.litellm_budgettable.update = AsyncMock() + + result = await _update_key_soft_budget( + db=mock_db, + existing_key_row=existing_key, + soft_budget=None, + changed_by="user-1", + ) + + assert result is None + mock_db.litellm_budgettable.create.assert_not_awaited() + mock_db.litellm_budgettable.update.assert_not_awaited() + + +def test_update_key_request_accepts_soft_budget(): + request = UpdateKeyRequest(key="sk-test", soft_budget=42.0) + assert request.soft_budget == 42.0 + assert "soft_budget" in request.model_fields_set + + +@pytest.mark.parametrize("valid_value", [None, 0.0, 25.0]) +def test_validate_soft_budget_value_accepts_valid_values(valid_value): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_soft_budget_value, + ) + + assert _validate_soft_budget_value(valid_value) is None + + +@pytest.mark.parametrize("invalid_value", [-5.0, float("nan"), float("inf")]) +def test_validate_soft_budget_value_rejects_invalid_values(invalid_value): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_soft_budget_value, + ) + + with pytest.raises(HTTPException) as exc_info: + _validate_soft_budget_value(invalid_value) + + assert exc_info.value.status_code == 400 + assert "soft_budget must be a non-negative finite number" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_apply_soft_budget_update_adds_budget_id_for_new_budget_row(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _apply_soft_budget_update, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + created_row = MagicMock() + created_row.budget_id = "budget-created-456" + mock_db = MagicMock() + mock_db.litellm_budgettable.create = AsyncMock(return_value=created_row) + mock_db.litellm_budgettable.update = AsyncMock() + + result = await _apply_soft_budget_update( + data=UpdateKeyRequest(key="sk-test", soft_budget=25.0), + non_default_values={"soft_budget": 25.0}, + db=mock_db, + existing_key_row=existing_key, + changed_by="user-1", + ) + + assert dict(result) == {"budget_id": "budget-created-456"} + + +@pytest.mark.asyncio +async def test_apply_soft_budget_update_keeps_existing_budget_id_out_of_token_update(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _apply_soft_budget_update, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123") + mock_db = MagicMock() + mock_db.litellm_budgettable.update = AsyncMock() + mock_db.litellm_budgettable.create = AsyncMock() + + result = await _apply_soft_budget_update( + data=UpdateKeyRequest(key="sk-test", soft_budget=40.0), + non_default_values={"soft_budget": 40.0, "max_budget": 100.0}, + db=mock_db, + existing_key_row=existing_key, + changed_by="user-1", + ) + + assert dict(result) == {"max_budget": 100.0} + mock_db.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "budget-123"}, + data={"soft_budget": 40.0, "updated_by": "user-1"}, + ) + + +@pytest.mark.asyncio +async def test_update_key_row_with_soft_budget_updates_budget_and_key_in_transaction(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_row_with_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + created_row = MagicMock(budget_id="budget-new") + updated_row = MagicMock() + updated_row.model_dump.return_value = {"token": "hashed", "budget_id": "budget-new"} + tx = MagicMock() + tx.litellm_budgettable.create = AsyncMock(return_value=created_row) + tx.litellm_verificationtoken.update = AsyncMock(return_value=updated_row) + tx_context = MagicMock() + tx_context.__aenter__ = AsyncMock(return_value=tx) + tx_context.__aexit__ = AsyncMock(return_value=None) + prisma_client = MagicMock() + prisma_client.tx.return_value = tx_context + prisma_client.jsonify_object = lambda data: dict(data) + + result = await _update_key_row_with_soft_budget( + prisma_client=prisma_client, + key="sk-test", + data=UpdateKeyRequest(key="sk-test", soft_budget=25.0), + non_default_values={"soft_budget": 25.0}, + existing_key_row=existing_key, + changed_by="user-1", + ) + + assert set(result) == {"token", "data"} + assert result["data"] == {"token": "hashed", "budget_id": "budget-new"} + tx.litellm_verificationtoken.update.assert_awaited_once() + update_call = tx.litellm_verificationtoken.update.await_args + assert update_call.kwargs["where"] == {"token": result["token"]} + assert update_call.kwargs["data"]["budget_id"] == "budget-new" + assert "soft_budget" not in update_call.kwargs["data"] + + +@pytest.mark.asyncio +async def test_update_key_row_with_soft_budget_propagates_transaction_error(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_row_with_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + created_row = MagicMock(budget_id="budget-new") + tx = MagicMock() + tx.litellm_budgettable.create = AsyncMock(return_value=created_row) + tx.litellm_verificationtoken.update = AsyncMock(side_effect=RuntimeError("update failed")) + tx_context = MagicMock() + tx_context.__aenter__ = AsyncMock(return_value=tx) + tx_context.__aexit__ = AsyncMock(return_value=None) + prisma_client = MagicMock() + prisma_client.tx.return_value = tx_context + prisma_client.jsonify_object = lambda data: dict(data) + + with pytest.raises(RuntimeError, match="update failed"): + await _update_key_row_with_soft_budget( + prisma_client=prisma_client, + key="sk-test", + data=UpdateKeyRequest(key="sk-test", soft_budget=25.0), + non_default_values={"soft_budget": 25.0}, + existing_key_row=existing_key, + changed_by="user-1", + ) + + tx_context.__aexit__.assert_awaited_once() + assert tx_context.__aexit__.await_args.args[0] is RuntimeError + + def test_generate_key_request_blank_team_id_is_personal(): """The UI Team-field clear submits team_id=""; it must count as no team (LIT-3925).""" from litellm.proxy._types import RegenerateKeyRequest @@ -17728,6 +18073,32 @@ def test_key_request_blank_organization_id_is_unset(): assert UpdateKeyRequest(key="sk-1", organization_id="org-1").organization_id == "org-1" +def test_update_key_request_blank_team_id_is_not_a_team_change(): + from litellm.proxy._types import UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + is_different_team, + ) + + blank = UpdateKeyRequest(key="sk-1", team_id="", key_alias="renamed") + assert blank.team_id is None + assert "team_id" not in blank.model_dump(exclude_unset=True) + assert blank.model_dump(exclude_unset=True) == {"key": "sk-1", "key_alias": "renamed"} + assert is_different_team(data=blank, existing_key_row=LiteLLM_VerificationToken(token="hashed")) is False + assert ( + is_different_team(data=blank, existing_key_row=LiteLLM_VerificationToken(token="hashed", team_id="team-1")) + is False + ) + assert "team_id" in UpdateKeyRequest(key="sk-1", team_id=None).model_dump(exclude_unset=True) + assert UpdateKeyRequest(key="sk-1", team_id="team-1").team_id == "team-1" + assert ( + is_different_team( + data=UpdateKeyRequest(key="sk-1", team_id="team-1"), + existing_key_row=LiteLLM_VerificationToken(token="hashed"), + ) + is True + ) + + def test_key_generation_check_blank_team_id_uses_personal_permissions(monkeypatch): """key_generation_check with team_id="" must take the personal-key path instead of failing the team lookup with "Unable to find team object" (LIT-3925).""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 71ff7de89b0..5e00e7d75be 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4802,6 +4802,72 @@ async def test_store_mcp_oauth_user_credential_returns_status(): assert result.expires_at == "2099-01-01T00:00:00+00:00" +@pytest.mark.asyncio +async def test_store_mcp_oauth_user_credential_blocked_when_identity_binding_enforced(): + """The direct opaque-token POST must be closed for enforce-mode identity-bound servers, + otherwise it bypasses the token-relay principal check.""" + from litellm.proxy._types import MCPOAuthUserCredentialRequest + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + store_mcp_oauth_user_credential, + ) + + server_id = "srv-binding-1" + bound_server = MCPServer( + server_id=server_id, + name=server_id, + url="https://mcp.example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", + issuer="https://idp.example.com", + audiences=["litellm-client"], + ), + ) + store_mock = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: mirrors the existing store-credential tests in this file + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( # test-quality-ok: mirrors the existing store-credential tests in this file + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + new=AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id=server_id)), + ), + patch( # test-quality-ok: mirrors the existing store-credential tests in this file + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch.object( # test-quality-ok: registry is a module-level singleton; injecting it would change the endpoint signature + manager_module.global_mcp_server_manager, + "get_mcp_server_by_id", + return_value=bound_server, + ), + patch( # test-quality-ok: asserting the DB write is never reached is the point of the test + "litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential", + new=store_mock, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await store_mcp_oauth_user_credential( + server_id=server_id, + payload=MCPOAuthUserCredentialRequest(access_token="opaque-tok", expires_in=3600), + user_api_key_dict=_make_user_auth("user-123"), + ) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_enforced" + store_mock.assert_not_called() + + @pytest.mark.asyncio async def test_delete_mcp_oauth_user_credential_only_deletes_oauth(): """delete_mcp_oauth_user_credential only deletes OAuth2 credentials, not BYOK.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 33de2a09626..5325e069813 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -30,7 +30,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( ) from litellm.proxy.utils import PrismaClient from litellm.router import Router -from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams async def _passthrough_row(update_data): @@ -392,6 +392,185 @@ class TestModelManagementAuthChecks: assert exc_info.value.code == "403" mock_update.assert_not_awaited() + def test_can_user_set_aws_session_tags_admin_success(self): + result = ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=LiteLLM_Params( + model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "genai"}] + ), + user_api_key_dict=self.admin_user, + ) + assert result is True + + def test_can_user_set_aws_session_tags_without_tags_allows_any_role(self): + result = ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=LiteLLM_Params(model="bedrock/test_model", aws_role_name="arn:aws:iam::123:role/x"), + user_api_key_dict=self.team_admin_user, + ) + assert result is True + + def test_can_user_set_aws_session_tags_team_admin_fails(self): + with pytest.raises(Exception, match="Only a proxy admin can set aws_session_tags") as exc_info: + ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=LiteLLM_Params( + model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "genai"}] + ), + user_api_key_dict=self.team_admin_user, + ) + assert exc_info.value.code == "403" + assert exc_info.value.param == "aws_session_tags" + + def test_can_user_set_aws_session_tags_unchanged_existing_allows_any_role(self): + result = ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=LiteLLM_Params( + model="bedrock/test_model", + aws_session_tags=[{"Key": "team", "Value": "genai"}, {"Key": "env", "Value": "prod"}], + ), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params( + model="bedrock/test_model", + aws_session_tags=[{"Key": "env", "Value": "prod"}, {"Key": "team", "Value": "genai"}], + ), + ) + assert result is True + + def test_can_user_set_aws_session_tags_changed_value_fails_for_team_admin(self): + with pytest.raises(Exception, match="Only a proxy admin can set aws_session_tags") as exc_info: + ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=LiteLLM_Params( + model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "platform"}] + ), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params( + model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "genai"}] + ), + ) + assert exc_info.value.code == "403" + + @pytest.mark.asyncio + async def test_add_new_model_rejects_aws_session_tags_for_non_admin(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + mock_prisma = MagicMock() + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # 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( # test-quality-ok: prior auth check needs a live DB; only the session tag check is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="tagged-bedrock", + litellm_params=LiteLLM_Params( + model="bedrock/anthropic.claude-opus-4-6-v1:0", + aws_role_name="arn:aws:iam::123456789012:role/team-role", + aws_session_tags=[{"Key": "team", "Value": "genai"}], + ), + model_info={"id": "session-tags-create-test", "team_id": "test_team"}, + ), + user_api_key_dict=self.team_admin_user, + ) + assert exc_info.value.code == "403" + assert exc_info.value.param == "aws_session_tags" + mock_prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_patch_model_rejects_aws_session_tags_for_non_admin(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + from litellm.types.router import updateLiteLLMParams + + model_id = "session-tags-patch-test" + db_model = Deployment( + model_name="tagged-bedrock", + litellm_params=LiteLLM_Params( + model="bedrock/anthropic.claude-opus-4-6-v1:0", + aws_role_name="arn:aws:iam::123456789012:role/team-role", + ), + model_info={"id": model_id, "team_id": "test_team"}, + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # 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( # test-quality-ok: stubs the DB row fetch; only the session tag check is under test + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=db_model), + ), + patch( # test-quality-ok: prior auth check needs a live DB; only the session tag check 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: asserts the DB write is never reached on rejection + "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db", + new=AsyncMock(), + ) as mock_update, + ): + with pytest.raises(ProxyException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(aws_session_tags=[{"Key": "team", "Value": "genai"}]) + ), + user_api_key_dict=self.team_admin_user, + ) + assert exc_info.value.code == "403" + assert exc_info.value.param == "aws_session_tags" + mock_update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_update_model_rejects_aws_session_tags_for_non_admin(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_model, + ) + from litellm.types.router import updateLiteLLMParams + + model_id = "session-tags-put-test" + existing = Deployment( + model_name="tagged-bedrock", + litellm_params=LiteLLM_Params( + model="bedrock/anthropic.claude-opus-4-6-v1:0", + aws_role_name="arn:aws:iam::123456789012:role/team-role", + ), + model_info={"id": model_id, "team_id": "test_team"}, + ) + existing_row = MagicMock() + existing_row.model_dump.return_value = existing.model_dump() + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # 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( # test-quality-ok: prior auth check needs a live DB; only the session tag check is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(aws_session_tags=[{"Key": "team", "Value": "genai"}]), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=self.team_admin_user, + ) + assert exc_info.value.code == "403" + assert exc_info.value.param == "aws_session_tags" + mock_prisma.db.litellm_proxymodeltable.update.assert_not_awaited() + def test_can_user_attach_credential_internal_user_fails(self): with pytest.raises(Exception, match="Only a proxy admin can attach a stored credential") as exc_info: ModelManagementAuthChecks.can_user_attach_credential( @@ -1990,7 +2169,7 @@ class TestModelInfoEndpoint: ): mock_router.get_fully_blocked_model_names.return_value = set() mock_router.get_model_list.return_value = [] - mock_router.get_configured_token_limits.return_value = (None, None) + mock_router.get_model_listing_info.return_value = None mock_router.get_deployment_by_model_group_name.return_value = Deployment( model_name="gpt-4", litellm_params=LiteLLM_Params(model="openai/gpt-4"), @@ -2067,7 +2246,7 @@ class TestModelInfoEndpoint: ): mock_router.get_fully_blocked_model_names.return_value = set() mock_router.get_model_list.return_value = [] - mock_router.get_configured_token_limits.return_value = (None, None) + mock_router.get_model_listing_info.return_value = None mock_router.get_deployment_by_model_group_name.return_value = Deployment( model_name="team-model-1", litellm_params=LiteLLM_Params(model="custom/team-model-1"), @@ -3070,6 +3249,31 @@ class TestUpdateDBModelBlocked: assert "blocked" not in result +class TestUpdateDBModelKeepsLegacyDropParams: + def test_partial_patch_keeps_encrypted_string_drop_params(self, monkeypatch): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") + legacy_row = Deployment( + model_name="gpt-5-nano", + litellm_params=LiteLLM_Params( + model="openai/gpt-5-nano", + api_key=encrypt_value_helper(value="sk-old"), + drop_params=encrypt_value_helper(value="true"), + ), + model_info=ModelInfo(id="legacy-row"), + ) + + result = update_db_model( + db_model=legacy_row, + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(api_key="sk-new")), + ) + + stored = json.loads(result["litellm_params"]) + assert decrypt_value_helper(value=stored["drop_params"], key="drop_params") == "true" + + def _build_db_model_with_pricing(): """Wildcard deployment with custom pricing in litellm_params; Deployment.__init__ mirrors SPECIAL_MODEL_INFO_PARAMS into model_info, so both blobs hold the rate.""" @@ -4440,13 +4644,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(side_effect=self._find_many), + ) + + async def _find_many(self, where: object = None) -> tuple[LiteLLM_ProxyModelTable, ...]: + json.dumps(where) # prisma serializes the filter with json.dumps and rejects a mappingproxy + return 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 +4677,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 +4838,167 @@ class TestStrategyRouterWriteValidation: assert capability is not None assert capability.sql_config_predicate.split("{config}")[-1].strip() in count_sql + _TUNED_A = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}} + _TUNED_A_EDITED = {**_TUNED_A, "dimension_weights": {"codePresence": 0.9}} + _TUNED_B = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4.1"}} + _TUNED_B_EDITED = {**_TUNED_B, "tiers": {"SIMPLE": "gpt-4o", "MEDIUM": "gpt-4.1"}} + + @staticmethod + def _db_router_row(model_id: str, config: Mapping[str, object]) -> dict[str, object]: + return { + "model_name": f"router-{model_id}", + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": dict(config)}, + "model_info": {"id": model_id, "db_model": True}, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "limit,baseline_rows,live_rows,candidate_id,candidate_config,expected", + [ + (1, ["a", "b"], {"a": "_TUNED_A", "b": "_TUNED_B"}, "a", "_TUNED_A", "allowed"), + (1, ["a", "b"], {"a": "_TUNED_A", "b": "_TUNED_B"}, "a", "_TUNED_A_EDITED", "allowed"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "a", "_TUNED_A_EDITED", "allowed"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "b", "_TUNED_B_EDITED", "refused"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "c", "_TUNED_B", "refused"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "b", "_TUNED_B", "allowed"), + (None, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "b", "_TUNED_B_EDITED", "allowed"), + (1, [], {}, "c", "_TUNED_B", "allowed"), + ], + ) + async def test_slot_enforces_baseline_relative_tuning_quota( + self, + limit: int | None, + baseline_rows: list[str], + live_rows: Mapping[str, str], + candidate_id: str, + candidate_config: str, + expected: str, + ) -> None: + """Without the license, one router may move off its recorded tuning baseline and keep being edited; + a change to a second router, or a second new tuned router, is refused. Unchanged baselines and + reverts to baseline are never counted, and a license lifts every check.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot + from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + + configs = { + "_TUNED_A": self._TUNED_A, + "_TUNED_A_EDITED": self._TUNED_A_EDITED, + "_TUNED_B": self._TUNED_B, + "_TUNED_B_EDITED": self._TUNED_B_EDITED, + } + baselines = snapshot_tuning_baselines( + [self._db_router_row(row_id, configs["_TUNED_A" if row_id == "a" else "_TUNED_B"]) for row_id in baseline_rows] + ) + effective_params = { + "model": "auto_router/complexity_router", + "complexity_router_config": configs[candidate_config], + } + # The other routers live in the DB, so the slot must read them under its own lock rather than + # trusting this pod's in-memory router: another pod's write is invisible to that list. + fake = self._FakeDb( + [], + tuning_rows=[ + { + "model_id": row_id, + "model_name": f"router-{row_id}", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": configs[name], + }, + } + for row_id, name in live_rows.items() + ], + ) + with ( + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + ): + if expected == "refused": + with pytest.raises(HTTPException) as exc_info: + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=candidate_id): + pass + assert exc_info.value.status_code == 403 + assert "changed heuristic scorer settings or tier models" in str(exc_info.value.detail) + assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) + return + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=candidate_id) as table: + assert hasattr(table, "create") + + @pytest.mark.asyncio + async def test_add_new_model_refuses_a_second_tuned_heuristic_v1_router_without_a_model_id(self) -> None: + """A create request carries no model_info at all, yet the quota still judges it: Deployment mints the + row id before the slot is entered, so a second tuned router is refused before its DB write.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + baselines = snapshot_tuning_baselines([self._db_router_row("a", self._TUNED_A)]) + fake = self._FakeDb( + [], + tuning_rows=[ + { + "model_id": "a", + "model_name": "router-a", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": self._TUNED_A_EDITED, + }, + } + ], + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the tuning quota is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="second-tuned", + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", complexity_router_config=self._TUNED_B + ), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + assert "changed heuristic scorer settings or tier models" in str(exc_info.value.message) + fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited() + fake.litellm_proxymodeltable.create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_slot_skips_tuning_quota_when_no_baseline_is_loaded(self) -> None: + """No baseline (DB-less proxy, or the startup read failed) means the gate cannot judge, so it does not.""" + from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot + + fake = self._FakeDb([]) + with ( + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", None), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + ): + async with _auto_router_capability_slot( + fake, + effective_params={"model": "auto_router/complexity_router", "complexity_router_config": self._TUNED_B}, + model_id="c", + ) as table: + assert hasattr(table, "create") + @pytest.mark.asyncio async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None: """team_model_add needs a second pool connection, so it must run only after the slot transaction diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index 4d13e054e46..7c3f4e2c6e9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -1,7 +1,8 @@ import asyncio import json from litellm._uuid import uuid -from typing import Optional, cast +from types import MappingProxyType +from typing import Final, Mapping, Optional, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1226,3 +1227,83 @@ def test_v2_update_organization_is_in_openapi_schema(): v2_path = app.openapi()["paths"]["/v2/organization/{organization_id}"] assert v2_path["patch"]["tags"] == ["organization management"] assert "OrganizationUpdateRequestV2" in json.dumps(v2_path["patch"]["requestBody"]) + + +def _organization_route_targets() -> list[tuple[str, str]]: + from fastapi.routing import APIRoute + + from litellm.proxy.management_endpoints.organization_endpoints import router + + return [ + (method, route.path.replace("{organization_id}", "org-under-test")) + for route in router.routes + if isinstance(route, APIRoute) + for method in sorted(route.methods - {"HEAD", "OPTIONS"}) + ] + + +_ORGANIZATION_ROUTE_REQUESTS: Final[Mapping[tuple[str, str], Mapping[str, object]]] = MappingProxyType( + { + ("POST", "/organization/new"): {"json": {"organization_alias": "org-under-test"}}, + ("DELETE", "/organization/delete"): {"json": {"organization_ids": ["org-under-test"]}}, + ("GET", "/organization/info"): {"params": {"organization_id": "org-under-test"}}, + ("POST", "/organization/info"): {"json": {"organizations": ["org-under-test"]}}, + ("POST", "/organization/member_add"): { + "json": {"organization_id": "org-under-test", "member": {"user_id": "user-1", "role": "internal_user"}} + }, + ("PATCH", "/organization/member_update"): {"json": {"organization_id": "org-under-test", "user_id": "user-1"}}, + ("DELETE", "/organization/member_delete"): {"json": {"organization_id": "org-under-test", "user_id": "user-1"}}, + } +) + + +def _organization_request(method: str, path: str) -> Mapping[str, object]: + return _ORGANIZATION_ROUTE_REQUESTS.get((method, path), {"json": {}}) + + +def _organization_test_client() -> TestClient: + from fastapi import FastAPI + + from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.management_endpoints.organization_endpoints import router + from litellm.proxy.proxy_server import openai_exception_handler + + app = FastAPI() + app.include_router(router) + app.add_exception_handler(ProxyException, openai_exception_handler) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="sk-test", user_role=LitellmUserRoles.PROXY_ADMIN + ) + return TestClient(app, raise_server_exceptions=False) + + +@pytest.mark.parametrize(("method", "path"), _organization_route_targets()) +def test_organization_routes_are_blocked_without_enterprise_license(monkeypatch, method, path): + """Every /organization route is enterprise-only, even for a proxy admin sending a valid request.""" + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", False, raising=False) + monkeypatch.setattr(proxy_server, "prisma_client", None, raising=False) + + response = _organization_test_client().request(method, path, **_organization_request(method, path)) + + assert response.status_code == 403 + assert "Organizations" in response.json()["detail"]["error"] + + +@pytest.mark.parametrize(("method", "path"), _organization_route_targets()) +def test_organization_routes_reach_their_handler_with_enterprise_license(monkeypatch, method, path): + """The same request a license refuses above now reaches the handler, which is the code reporting the missing database.""" + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import CommonProxyErrors + + monkeypatch.setattr(proxy_server, "premium_user", True, raising=False) + monkeypatch.setattr(proxy_server, "prisma_client", None, raising=False) + + response = _organization_test_client().request(method, path, **_organization_request(method, path)) + + assert response.status_code == 500 + assert any( + message in response.text for message in (CommonProxyErrors.db_not_connected_error.value, "No db connected") + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py index 08e931e6405..bdc12dad4bc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py @@ -19,6 +19,7 @@ from litellm.proxy._types import ( LitellmUserRoles, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.callback_config_validation import cross_entry_family_error from litellm.proxy.management_endpoints.team_callback_endpoints import ( add_team_callbacks, delete_team_callback, @@ -1443,3 +1444,118 @@ async def test_delete_team_callback_route_accepts_team_ids_containing_slashes(): assert response.json()["data"]["success_callbacks"] == ["langsmith"] written = json.loads(mock_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) assert [entry["callback_name"] for entry in written["logging"]] == ["langsmith"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_handler", + [ + lambda caller: add_team_callbacks( + data=AddTeamCallback( + callback_name="langfuse", + callback_type="success", + callback_vars={"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}, + ), + http_request=Mock(spec=Request), + team_id="team-does-not-exist", + user_api_key_dict=caller, + ), + lambda caller: get_team_callbacks( + http_request=Mock(spec=Request), + team_id="team-does-not-exist", + user_api_key_dict=caller, + ), + lambda caller: delete_team_callback( + http_request=Mock(spec=Request), + team_id="team-does-not-exist", + callback_name="langfuse", + user_api_key_dict=caller, + ), + ], + ids=["add", "get", "delete"], +) +async def test_unknown_team_is_indistinguishable_from_no_access(call_handler, unauthorized_caller): + """An unauthorized caller must not learn whether a team id exists. + + These routes are reachable by any authenticated caller so a team admin can get + as far as the access check, so a distinct "does not exist" would turn them into + a probe for valid team ids. The unknown-team response has to match the + no-access one exactly, status and body. + """ + with patch("litellm.proxy.proxy_server.prisma_client") as mock_client: # test-quality-ok: the handler imports prisma_client from proxy_server at call time, so there is no seam to inject through + mock_client.get_data = AsyncMock(return_value=None) + with pytest.raises(HTTPException) as unknown_team: + await call_handler(unauthorized_caller) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_client: # test-quality-ok: the handler imports prisma_client from proxy_server at call time, so there is no seam to inject through + mock_client.get_data = AsyncMock(return_value=_team_row()) + mock_client.db.litellm_teamtable.update = AsyncMock() + with patch( # test-quality-ok: _verify_team_access calls this module-level helper directly, so there is no seam to inject through + "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", + new_callable=AsyncMock, + return_value=False, + ): + with pytest.raises(HTTPException) as no_access: + await call_handler(unauthorized_caller) + + assert unknown_team.value.status_code == no_access.value.status_code == 403 + assert unknown_team.value.detail == no_access.value.detail + assert "does not exist" not in str(unknown_team.value.detail) + + +@pytest.mark.asyncio +async def test_proxy_admin_still_told_the_team_is_unknown(): + """The masking is only for callers who could not have managed the team; a proxy + admin keeps the diagnosable error.""" + admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin", api_key="sk-admin") + with patch("litellm.proxy.proxy_server.prisma_client") as mock_client: # test-quality-ok: the handler imports prisma_client from proxy_server at call time, so there is no seam to inject through + mock_client.get_data = AsyncMock(return_value=None) + with pytest.raises(HTTPException) as exc: + await get_team_callbacks( + http_request=Mock(spec=Request), + team_id="team-does-not-exist", + user_api_key_dict=admin, + ) + + assert exc.value.status_code == 404 + assert "does not exist" in str(exc.value.detail) + + +@pytest.mark.parametrize( + "new_vars, stored, rejected", + [ + # the redirect, in every carrier a caller could pick: an entry naming + # only a host, pairing with a key pair written on another entry + ({"langfuse_host": "http://attacker.invalid"}, [{"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], True), + # the sibling carrier -- langfuse and langfuse_otel are one account + ({"langfuse_host": "http://attacker.invalid"}, [{"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_secret_key": "sk"}], True), + # a destination variable no integration registry lists + ({"dd_agent_host": "attacker.invalid"}, [{"dd_api_key": "k", "dd_site": "us5.datadoghq.com"}], True), + # one entry owning its family end to end is the feature + ({"langfuse_host": "https://eu.cloud.langfuse.com", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}, [], False), + # a different family alongside an existing one stays fine + ({"gcs_bucket_name": "bucket"}, [{"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], False), + ({"langsmith_api_key": "k"}, [{"dd_api_key": "k"}], False), + # variables that configure no backend carry nothing to redirect + ({"turn_off_message_logging": "true"}, [{"langfuse_secret_key": "sk"}], False), + # the same integration registered for a second event: identical values + # flatten to the identical dict, so there is nothing to redirect + ({"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}, [{"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], False), + # the same credential under its other spelling is the same credential + ({"langfuse_secret": "sk"}, [{"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], False), + # a value the family already holds cannot be moved into another of its + # variables either; the exporter would address or authenticate with it + ({"langfuse_host": "pk"}, [{"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_public_key": "pk"}], True), + # the same shape with one value moved is the redirect again + ({"langfuse_host": "http://attacker.invalid", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}, [{"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], True), + ], +) +def test_one_entry_owns_a_credential_family(new_vars, stored, rejected): + """A team admin must not be able to redirect a credential they cannot read. + + The stored entries are flattened into one dict before a request reads them, + so an entry naming only a destination pairs with a key written elsewhere and + carries it to that destination. + """ + error = cross_entry_family_error(new_vars, stored) + assert (error is not None) is rejected diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 051e6bed4fd..2f6561046b1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2205,6 +2205,50 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): assert update_call_kwargs.get("include", {}).get("object_permission") is True +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint_name", ["team_model_add", "team_model_delete"]) +async def test_team_model_add_delete_keep_model_aliases_in_team_cache(endpoint_name, monkeypatch): + """LIT-5858: Prisma only returns `litellm_model_table` when the `update` asks for it, so the refreshed + cache entry lost the team's model aliases and JWT alias requests 403'd until the next DB read.""" + from litellm.proxy._types import TeamModelAddRequest, TeamModelDeleteRequest + from litellm.proxy.auth.team_grants import team_model_aliases + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import team_model_add, team_model_delete + + columns = {"team_id": "team-1234", "models": ["gpt-4o", "openai/*"]} + alias_table = {"id": 1, "model_aliases": '{"fast": "gpt-4o"}', "created_by": "admin", "updated_by": "admin"} + + async def update(where, data, include=None): + row = {**columns, "litellm_model_table": alias_table} if (include or {}).get("litellm_model_table") else columns + return SimpleNamespace(team_id="team-1234", model_dump=lambda: row) + + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=SimpleNamespace(model_dump=lambda: columns)) + prisma_client.db.litellm_teamtable.update = AsyncMock(side_effect=update) + prisma_client.db.execute_raw = AsyncMock(return_value=None) + cache = UserApiKeyCache() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None) + + admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + if endpoint_name == "team_model_add": + await team_model_add( + data=TeamModelAddRequest(team_id="team-1234", models=["team-byok-1"]), + http_request=MagicMock(), + user_api_key_dict=admin, + ) + else: + await team_model_delete( + data=TeamModelDeleteRequest(team_id="team-1234", models=["openai/*"]), + http_request=MagicMock(), + user_api_key_dict=admin, + ) + + cached_team = await cache.async_get_cache(key="team_id:team-1234", model_type=LiteLLM_TeamTableCachedObj) + assert team_model_aliases(cached_team) == {"fast": "gpt-4o"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "endpoint_name", diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 8d8bc15f9be..2050e65d2a1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2616,6 +2616,28 @@ class TestCLIKeyRegenerationFlow: ) cache.set_cache.assert_not_called() + def test_cli_sso_flow_lookup_treats_an_open_redis_breaker_as_a_miss(self): + """A Redis read refused by the open circuit breaker is a missing session, not a server error. + + The direct Redis read is what keeps the flow authoritative across workers, so the + refusal must not fall back to a possibly stale in-memory copy either. + """ + from litellm.caching.redis_cache import RedisCircuitBreakerOpenError + from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_or_raise + + redis_cache = MagicMock() + redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError("Redis circuit breaker is open") + cache = MagicMock() + cache.redis_cache = redis_cache + cache.get_cache.return_value = {"poll_secret_hash": "stale", "sso_complete": False} + + with pytest.raises(HTTPException) as exc_info: + _get_cli_sso_flow_or_raise(login_id="cli-breaker_open_1234567890", cache=cache) + + assert exc_info.value.status_code == 400 + assert "not found or expired" in exc_info.value.detail + cache.get_cache.assert_not_called() + def test_cli_sso_flow_with_enum_survives_redis_round_trip(self): """ RedisCache stores values via str(value) and reads them back through diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index d7ebb1f60bf..078315c2bf8 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -122,6 +122,29 @@ async def test_set_object_permission_persists_mcp_tool_search_enabled(): assert created_data["mcp_tool_search_enabled"] is True +@pytest.mark.asyncio +async def test_set_object_permission_persists_skills(): + mock_prisma_client = MagicMock() + mock_created_permission = MagicMock() + mock_created_permission.object_permission_id = "perm_id" + mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( + return_value=mock_created_permission + ) + + data_json = { + "object_permission": LiteLLM_ObjectPermissionBase(skills=["private-skill"]).model_dump(), + } + + await _set_object_permission(data_json=data_json, prisma_client=mock_prisma_client) + + created_data = ( + mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs[ + "data" + ] + ) + assert created_data["skills"] == ["private-skill"] + + # ---- Tests for _extract_requested_mcp_server_ids ---- diff --git a/tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py b/tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py new file mode 100644 index 00000000000..61eea00bf7d --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py @@ -0,0 +1,302 @@ +"""Tests for PerRequestRootPathMiddleware (``SERVER_ROOT_PATHS``). + +One deployment fronting several client-visible URL path prefixes: the matched +prefix becomes that request's ``root_path``, so Starlette route matching and +``request.base_url`` — and therefore every URL the proxy emits, the MCP OAuth +discovery documents among them — resolve under the prefix the client called. +""" + +import pytest +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient + +from litellm.proxy.middleware.per_request_root_path_middleware import ( + PerRequestRootPathMiddleware, + get_request_root_path, + get_server_root_paths, + normalize_root_paths, +) + + +class TestNormalizeRootPaths: + def test_strips_whitespace_and_trailing_slash(self): + assert normalize_root_paths([" /tenant-a/ ", "/tenant-b"]) == ( + "/tenant-a", + "/tenant-b", + ) + + def test_drops_empty_entries(self): + assert normalize_root_paths(["", " ", "/tenant-a"]) == ("/tenant-a",) + + def test_drops_entries_without_leading_slash(self): + # A typo'd entry must not silently match nothing at request time. + assert normalize_root_paths(["tenant-a", "/tenant-b"]) == ("/tenant-b",) + + def test_drops_bare_root(self): + # "/" would turn every request into a root_path rewrite; a + # root-mounted deployment needs no entry at all. + assert normalize_root_paths(["/", "/tenant-a"]) == ("/tenant-a",) + + def test_dedupes(self): + assert normalize_root_paths(["/t", "/t/", " /t "]) == ("/t",) + + def test_longest_first_for_nested_prefixes(self): + # Longest-first ordering is what makes the most-specific nested + # prefix win at match time. + assert normalize_root_paths(["/t", "/t/deep"]) == ("/t/deep", "/t") + + +class TestGetServerRootPaths: + def test_unset_env_is_empty(self, monkeypatch): + monkeypatch.delenv("SERVER_ROOT_PATHS", raising=False) + assert get_server_root_paths() == () + + def test_empty_env_is_empty(self, monkeypatch): + monkeypatch.setenv("SERVER_ROOT_PATHS", "") + assert get_server_root_paths() == () + + def test_comma_separated_entries(self, monkeypatch): + monkeypatch.setenv("SERVER_ROOT_PATHS", "/tenant-a, /tenant-b/") + assert get_server_root_paths() == ("/tenant-a", "/tenant-b") + + +def _capture_scope_middleware(root_paths): + """Middleware wired to a downstream that records the scope it received.""" + captured = {} + + async def downstream(scope, receive, send): + captured.update(scope) + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + return PerRequestRootPathMiddleware(downstream, root_paths=root_paths), captured + + +async def _run(mw, scope): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message): + pass + + await mw(scope, receive, send) + + +class TestPerRequestRootPathMiddleware: + @pytest.mark.asyncio + async def test_matched_prefix_becomes_root_path_path_untouched(self): + # Starlette's router strips root_path from the (unmodified) path at + # match time, so the middleware must NOT rewrite scope["path"]. + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "http", "path": "/tenant-a/mcp/x", "method": "GET", "headers": []}) + assert captured["root_path"] == "/tenant-a" + assert captured["path"] == "/tenant-a/mcp/x" + + @pytest.mark.asyncio + async def test_exact_prefix_matches(self): + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "http", "path": "/tenant-a", "method": "GET", "headers": []}) + assert captured["root_path"] == "/tenant-a" + + @pytest.mark.asyncio + async def test_segment_boundary_prevents_sibling_match(self): + # /tenant-ab must not match the /tenant-a prefix. + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "http", "path": "/tenant-ab/mcp", "method": "GET", "headers": []}) + assert "root_path" not in captured + + @pytest.mark.asyncio + async def test_unmatched_path_untouched(self): + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "http", "path": "/chat/completions", "method": "GET", "headers": []}) + assert "root_path" not in captured + assert captured["path"] == "/chat/completions" + + @pytest.mark.asyncio + async def test_longest_nested_prefix_wins(self): + mw, captured = _capture_scope_middleware(["/t", "/t/deep"]) + await _run(mw, {"type": "http", "path": "/t/deep/mcp", "method": "GET", "headers": []}) + assert captured["root_path"] == "/t/deep" + + @pytest.mark.asyncio + async def test_matched_prefix_overrides_scalar_root_path(self): + # FastAPI(root_path=SERVER_ROOT_PATH) stamps the scalar before the + # middleware stack runs; a matched dynamic prefix wins for that + # request (combining both mechanisms is warned about at startup). + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run( + mw, + {"type": "http", "path": "/tenant-a/mcp", "root_path": "/legacy", "method": "GET", "headers": []}, + ) + assert captured["root_path"] == "/tenant-a" + + @pytest.mark.asyncio + async def test_unmatched_request_keeps_scalar_root_path(self): + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run( + mw, + {"type": "http", "path": "/legacy/mcp", "root_path": "/legacy", "method": "GET", "headers": []}, + ) + assert captured["root_path"] == "/legacy" + + @pytest.mark.asyncio + async def test_websocket_scope_matched(self): + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "websocket", "path": "/tenant-a/ws", "headers": []}) + assert captured["root_path"] == "/tenant-a" + + @pytest.mark.asyncio + async def test_lifespan_scope_passes_through(self): + called = {} + + async def downstream(scope, receive, send): + called["scope"] = scope + + mw = PerRequestRootPathMiddleware(downstream, root_paths=["/tenant-a"]) + await _run(mw, {"type": "lifespan"}) + assert called["scope"] == {"type": "lifespan"} + + +def _routed_client(prefixes): + app = FastAPI() + + @app.get("/where") + def where(request: Request): + return { + "base_url": str(request.base_url), + "root_path": request.scope.get("root_path", ""), + } + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=prefixes) + return TestClient(app) + + +class TestEndToEndRouting: + def test_two_prefixes_route_on_one_app(self): + # The property a scalar SERVER_ROOT_PATH cannot provide: two + # client-visible prefixes served by the same app, each request + # reconstructing its own base URL. + client = _routed_client(["/tenant-a", "/tenant-b"]) + + resp_a = client.get("/tenant-a/where") + resp_b = client.get("/tenant-b/where") + + assert resp_a.status_code == 200 + assert resp_a.json() == { + "base_url": "http://testserver/tenant-a/", + "root_path": "/tenant-a", + } + assert resp_b.status_code == 200 + assert resp_b.json() == { + "base_url": "http://testserver/tenant-b/", + "root_path": "/tenant-b", + } + + def test_unprefixed_route_still_served(self): + client = _routed_client(["/tenant-a"]) + resp = client.get("/where") + assert resp.status_code == 200 + assert resp.json()["base_url"] == "http://testserver/" + + def test_unlisted_prefix_404s(self): + client = _routed_client(["/tenant-a"]) + assert client.get("/tenant-c/where").status_code == 404 + + +class TestGetRequestRootPath: + """``get_request_root_path`` is the accessor that plumbs the middleware's + resolved prefix to code that doesn't have scope in hand — the 401 challenge + builders in ``mcp_server_manager`` / ``server`` and ``get_custom_url`` on the + SSO callback path. Reading the SERVER_ROOT_PATH scalar there would emit URLs + under a prefix the client didn't call, and stack a second prefix onto ones it + did (the two review points this fixture pins).""" + + def test_falls_back_to_server_root_path_env_outside_a_request(self, monkeypatch): + # Outside a request the middleware's ContextVar is unset. The scalar + # env still owns the answer, so pre-middleware call sites (module-load + # UI URL builders, background tasks) behave exactly as they did. + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + assert get_request_root_path() == "/legacy" + + def test_returns_empty_string_when_no_env_and_no_request(self, monkeypatch): + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + assert get_request_root_path() == "" + + def test_returns_matched_prefix_inside_a_request(self, monkeypatch): + # With both env vars set, a request matching a SERVER_ROOT_PATHS prefix + # must see that prefix — not the SERVER_ROOT_PATH scalar — so the URL + # it emits stays under the prefix the router will resolve it against. + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + + seen: list[str] = [] + app = FastAPI() + + @app.get("/where") + def where(): + seen.append(get_request_root_path()) + return {} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a"]) + client = TestClient(app) + + assert client.get("/tenant-a/where").status_code == 200 + assert seen == ["/tenant-a"] + + def test_unmatched_request_falls_through_to_scope_scalar(self, monkeypatch): + # A request the middleware saw but did not match keeps whatever + # scope["root_path"] the app was mounted under (the scalar). The + # ContextVar still reflects the effective per-request answer, so + # emitted URLs and the router agree even on the fallback path. + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + seen: list[str] = [] + + app = FastAPI(root_path="/legacy") + + @app.get("/where") + def where(): + seen.append(get_request_root_path()) + return {} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a"]) + client = TestClient(app) + + assert client.get("/legacy/where").status_code == 200 + assert seen == ["/legacy"] + + def test_each_request_sees_its_own_prefix(self, monkeypatch): + # Two sequential requests through the same app must each see the + # prefix they arrived under, so one tenant's client is never sent + # the URL of another tenant's origin. + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + seen: list[tuple[str, str]] = [] + app = FastAPI() + + @app.get("/where") + def where(tag: str): + seen.append((tag, get_request_root_path())) + return {} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a", "/tenant-b"]) + client = TestClient(app) + assert client.get("/tenant-a/where?tag=a").status_code == 200 + assert client.get("/tenant-b/where?tag=b").status_code == 200 + assert seen == [("a", "/tenant-a"), ("b", "/tenant-b")] + + def test_context_var_reset_after_request(self, monkeypatch): + # A ContextVar left set after the request finishes would poison the + # module-load-time callers that read it lazily (they'd think they were + # inside a request under the last-seen prefix). + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + + app = FastAPI() + + @app.get("/where") + def where(): + return {"prefix": get_request_root_path()} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a"]) + client = TestClient(app) + assert client.get("/tenant-a/where").json() == {"prefix": "/tenant-a"} + # After the request finishes, the scalar-env fallback owns the answer + # again — nothing was left stashed from the last request's scope. + assert get_request_root_path() == "/legacy" diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index bca97915347..5f1e7e1fe0c 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -2920,9 +2920,9 @@ def test_unscoped_list_files_accepts_every_documented_purpose( def test_list_files_reports_a_bad_target_model_names_as_a_400( mocker: MockerFixture, monkeypatch, llm_router: Router ): - """The exception tail reports an HTTPException with its own status and error - type rather than relabelling it, so a client that branches on either keeps - reading the same thing off a bad request.""" + """The exception tail answers with the OpenAI error object a client can branch on: + the type its 400 status stands for, and a JSON null param rather than the literal + string "None" no OpenAI SDK has a case for.""" _setup_unscoped_list_files_route(mocker, monkeypatch, llm_router, _permissive_afile_list) response = _get_list_files("/v1/files?target_model_names=gpt-3.5-turbo,gpt-4o") @@ -2931,8 +2931,8 @@ def test_list_files_reports_a_bad_target_model_names_as_a_400( assert response.json() == { "error": { "message": "target_model_names on list files must be a list of one model name. Example: ['gpt-4o']", - "type": "None", - "param": "None", + "type": "invalid_request_error", + "param": None, "code": "400", } } @@ -4666,3 +4666,156 @@ def test_create_file_path_traversal_filename_rejected_before_forwarding(monkeypa assert error["param"] == "file" assert "traversal" in error["message"].lower() assert forwarded_calls == [] + + +def _setup_managed_file_route_answering_404( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + """Wire the single-file routes to a managed file store that knows no file, the way the + managed files hook answers once a file has been deleted or was never the caller's.""" + import litellm.proxy.proxy_server as ps + from fastapi import HTTPException + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + from litellm.proxy._types import LitellmUserRoles + + async def _file_not_found(file_id: str, **kwargs: object) -> None: + raise HTTPException(status_code=404, detail=f"File not found: {file_id}") + + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router) + managed_files = mocker.MagicMock(spec=BaseFileEndpoints) + managed_files.afile_retrieve = mocker.AsyncMock(side_effect=_file_not_found) + managed_files.afile_delete = mocker.AsyncMock(side_effect=_file_not_found) + managed_files.afile_content = mocker.AsyncMock(side_effect=_file_not_found) + proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + ) + + +def _call_managed_file_route(method: str, path: str) -> httpx.Response: + try: + return client.request(method, path, headers={"Authorization": "Bearer test-key"}) + finally: + import litellm.proxy.proxy_server as ps + + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def _missing_managed_file_error(file_id: str) -> dict[str, dict[str, str | None]]: + return { + "error": { + "message": f"File not found: {file_id}", + "type": "invalid_request_error", + "param": None, + "code": "404", + } + } + + +def test_create_file_reports_a_half_specified_expires_after_as_a_400( + monkeypatch: pytest.MonkeyPatch, llm_router: Router +): + """A 400 raised inside the route answers with the type a 400 stands for and a JSON null + param, not the literal string "None" in both fields, so a client can classify it.""" + setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + + response = client.post( + "/v1/files", + files={"file": ("mydata.jsonl", VALID_BATCH_LINE, "application/jsonl")}, + data={"purpose": "batch", "target_model_names": "gpt-3.5-turbo", "expires_after[anchor]": "created_at"}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert "expires_after[seconds]" in error["message"] + assert error["type"] == "invalid_request_error" + assert error["param"] is None + assert error["code"] == "400" + + +def test_get_file_reports_a_missing_managed_file_as_a_404( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +): + _setup_managed_file_route_answering_404(mocker, monkeypatch, llm_router) + file_id = _unified_managed_file_id() + + response = _call_managed_file_route("GET", f"/v1/files/{file_id}") + + assert response.status_code == 404, response.text + assert response.json() == _missing_managed_file_error(file_id) + + +def test_delete_file_reports_a_missing_managed_file_as_a_404( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +): + _setup_managed_file_route_answering_404(mocker, monkeypatch, llm_router) + file_id = _unified_managed_file_id() + + response = _call_managed_file_route("DELETE", f"/v1/files/{file_id}") + + assert response.status_code == 404, response.text + assert response.json() == _missing_managed_file_error(file_id) + + +def test_get_file_content_reports_a_missing_managed_file_as_a_404( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +): + _setup_managed_file_route_answering_404(mocker, monkeypatch, llm_router) + file_id = _unified_managed_file_id() + + response = _call_managed_file_route("GET", f"/v1/files/{file_id}/content") + + assert response.status_code == 404, response.text + assert response.json() == _missing_managed_file_error(file_id) + + +def _setup_managed_file_stored_in_an_unknown_storage_backend( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + """Wire the content route to a managed file whose row names a storage backend the + factory does not know, which is the one in-route ProxyException on these routes.""" + from types import SimpleNamespace + + import litellm.proxy.proxy_server as ps + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + from litellm.proxy._types import LitellmUserRoles + + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router) + managed_files = mocker.MagicMock(spec=BaseFileEndpoints) + managed_files.prisma_client = mocker.MagicMock() + proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files + repository = mocker.MagicMock() + repository.table.find_first = mocker.AsyncMock( + return_value=SimpleNamespace(storage_backend="ftp", storage_url="ftp://bucket/file") + ) + monkeypatch.setattr( + "litellm.proxy.openai_files_endpoints.files_endpoints.ManagedFileRepository", lambda _prisma: repository + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + ) + + +def test_get_file_content_keeps_the_status_of_a_rejection_raised_inside_the_route( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +): + """A ProxyException raised inside the route carries its status as the string ``code``, + and the tail used to rebuild it as a 500 because it only read ``status_code``.""" + _setup_managed_file_stored_in_an_unknown_storage_backend(mocker, monkeypatch, llm_router) + + response = _call_managed_file_route("GET", f"/v1/files/{_unified_managed_file_id()}/content") + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["message"].startswith("Storage backend error") + assert (error["type"], error["param"], error["code"]) == ("invalid_request_error", "file_id", "400") diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 6f9142c85df..a3d3ae32169 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -9,8 +9,10 @@ import pytest import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, + count_relayed_prompt_tokens, ) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, @@ -2037,3 +2039,56 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: if __name__ == "__main__": pytest.main([__file__]) + + +ONE_PIXEL_PNG_DATA_URL = ( + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" +) +UNREACHABLE_IMAGE_URL = "http://127.0.0.1:9/doc.png" +TEXT_ONLY_MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] + + +def _image_messages(url: str, detail: str) -> list[dict]: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": url, "detail": detail}}, + ], + } + ] + + +def test_count_relayed_prompt_tokens_counts_a_data_url_image_exactly(): + messages = _image_messages(ONE_PIXEL_PNG_DATA_URL, "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( + model="gpt-4.1-mini", messages=messages + ) + + +def test_count_relayed_prompt_tokens_keeps_a_low_detail_remote_image_at_the_base_count(): + messages = _image_messages(UNREACHABLE_IMAGE_URL, "low") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( + model="gpt-4.1-mini", messages=messages + ) + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) < high_detail_image_token_upper_bound() + + +def test_count_relayed_prompt_tokens_charges_only_the_remote_high_detail_image_at_the_upper_bound(): + messages = _image_messages(UNREACHABLE_IMAGE_URL, "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() + ) + + +@pytest.mark.parametrize("scheme", ["HTTPS://", "Http://"]) +def test_count_relayed_prompt_tokens_charges_an_uppercase_scheme_remote_high_detail_image_at_the_upper_bound(scheme): + messages = _image_messages(scheme + UNREACHABLE_IMAGE_URL.split("://", 1)[1], "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() + ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index acb45038df0..7b285674145 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1,9 +1,10 @@ +import asyncio import base64 import contextlib import json import os import traceback -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from types import MappingProxyType, SimpleNamespace from typing import Final from unittest import mock @@ -11,7 +12,9 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest +import respx from fastapi import HTTPException, Request, Response +from fastapi.routing import APIRoute from fastapi.responses import StreamingResponse from fastapi.testclient import TestClient from starlette.datastructures import FormData @@ -19,6 +22,7 @@ from starlette.datastructures import FormData import litellm from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( BaseOpenAIPassThroughHandler, @@ -43,6 +47,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( ) from litellm.proxy._types import LitellmUserRoles, SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials @@ -1852,11 +1857,11 @@ class TestBedrockAgentRuntimePassthroughToggle: return request @contextlib.contextmanager - def _patched_dispatch(self, general_settings: Mapping[str, object]): + def _patched_dispatch(self, general_settings: Mapping[str, object], credentials: object | None = None): from botocore.credentials import Credentials bedrock_llm: Final = Mock() - bedrock_llm.get_credentials = Mock(return_value=Credentials("ak", "sk")) + bedrock_llm.get_credentials = Mock(return_value=credentials or Credentials("ak", "sk")) forwarder: Final = AsyncMock(return_value="forwarded") with ( @@ -1891,6 +1896,27 @@ class TestBedrockAgentRuntimePassthroughToggle: forwarder.assert_awaited_once() assert "bedrock-agent-runtime.us-east-1.amazonaws.com" in create_route.call_args.kwargs["target"] + @pytest.mark.asyncio + async def test_agent_runtime_dispatch_signs_off_the_event_loop(self, monkeypatch): + """Regression for issue #40165: the agent-runtime pass-through signed on the loop, so botocore's + blocking credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe: Final = EventLoopProbe() + release: Final = asyncio.create_task(probe.release_refresh_from_the_loop()) + + with self._patched_dispatch(MappingProxyType({}), credentials=probe.credentials()) as (create_route, forwarder): + result: Final = await bedrock_proxy_route( + endpoint=self.AGENT_RUNTIME_ENDPOINT, + request=self._mock_request(), + fastapi_response=Mock(), + user_api_key_dict=UserAPIKeyAuth(), + ) + await release + + assert result == "forwarded" + assert create_route.call_args.kwargs["custom_headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True + @pytest.mark.asyncio @pytest.mark.parametrize("value", (True, "true", "True")) async def test_agent_runtime_dispatch_rejected_when_disabled(self, value: bool | str): @@ -3316,8 +3342,11 @@ class TestOpenAIPassthroughRoute: def _resolve_route_name(method: str, path: str) -> str | None: from starlette.routing import Match + from litellm.proxy._lazy_features import LAZY_FEATURES, _force_load from litellm.proxy.proxy_server import app + asyncio.run(_force_load(app, next(f for f in LAZY_FEATURES if f.name == "llm_passthrough"))) + scope: Final = { "type": "http", "method": method, @@ -3327,8 +3356,8 @@ def _resolve_route_name(method: str, path: str) -> str | None: "root_path": "", } for route in app.router.routes: - if route.matches(scope)[0] == Match.FULL: - return getattr(route, "name", None) + if isinstance(route, APIRoute) and route.matches(scope)[0] == Match.FULL: + return route.name return None @@ -3353,7 +3382,7 @@ def test_openai_passthrough_prefix_wins_over_native_provider_routes(method, path /{provider}/v1/files and /{provider}/v1/batches routes must never capture it with provider="openai_passthrough" (which 500s on the LlmProviders lookup). """ - assert _resolve_route_name(method, path) == "openai_proxy_route" + assert _resolve_route_name(method, path) == "openai_passthrough_route" @pytest.mark.parametrize( @@ -3370,6 +3399,41 @@ def test_native_provider_routes_are_unchanged(method, path, expected_name): assert _resolve_route_name(method, path) == expected_name +@pytest.fixture +def openai_passthrough_client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("OPENAI_API_KEY", "sk-upstream") + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual")) + yield TestClient(app) + + +@pytest.mark.parametrize( + "method, path, body", + [ + ("POST", "/v1/responses", {"model": "gpt-5.1", "input": "hi"}), + ("GET", "/v1/files", None), + ("POST", "/v1/batches", {"input_file_id": "file-abc123", "endpoint": "/v1/responses"}), + ], +) +def test_openai_passthrough_forwards_verbatim_to_openai( + openai_passthrough_client: TestClient, method: str, path: str, body: dict[str, str] | None +) -> None: + """Every /openai_passthrough request, including the /v1/files and /v1/batches + paths that native provider routes also claim, must reach OpenAI unchanged.""" + with respx.mock(assert_all_called=True) as upstream: + route = upstream.request(method, f"https://api.openai.com{path}").mock( + return_value=httpx.Response(200, json={"id": "upstream_123"}) + ) + response = openai_passthrough_client.request(method, f"/openai_passthrough{path}", json=body) + + assert (response.status_code, response.json()) == (200, {"id": "upstream_123"}) + assert route.calls.last.request.headers["authorization"] == "Bearer sk-upstream" + + class TestCursorProxyRoute: """Tests for the Cursor Cloud Agents pass-through route.""" @@ -4999,7 +5063,11 @@ class TestPassthroughRouterModelBudgetReservation: monkeypatch.setattr(proxy_server, "llm_router", RecordingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) return captured def _assert_metadata_carries_attribution(self, captured: list[dict], user_api_key_dict: UserAPIKeyAuth) -> None: @@ -5098,7 +5166,11 @@ class TestAzureRouterModelStreamingDispatch: monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) request = MagicMock(spec=Request) request.method = "POST" @@ -5158,7 +5230,11 @@ class TestAzureRouterModelStreamingKeepalive: monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) request = MagicMock(spec=Request) request.method = "POST" @@ -5206,3 +5282,213 @@ class TestAzureRouterModelStreamingKeepalive: assert result.headers["x-upstream"] == "kept" assert chunks == [b"data: hello\n\n"] + + +class TestRouterModelRelayUpstreamContract: + def _request(self, content_type: str) -> MagicMock: + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": content_type} + request.query_params = {} + return request + + def _install_router(self, monkeypatch, router, body: dict) -> None: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + async def fake_get_request_body(_request): + return body + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) + + def _recording_router(self, captured: list[dict]): + class RecordingRouter: + async def allm_passthrough_route(self, **kwargs): + captured.append(kwargs) + return httpx.Response(200, json={"ok": True}) + + return RecordingRouter() + + @pytest.mark.asyncio + async def test_azure_relay_keeps_the_json_body_when_the_content_type_carries_a_charset(self, monkeypatch): + body = {"model": "gpt-5", "messages": [{"role": "user", "content": "hi"}]} + captured: list[dict] = [] + self._install_router(monkeypatch, self._recording_router(captured), body) + + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json; charset=utf-8"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert captured[0]["json"] == body + + @pytest.mark.asyncio + async def test_vllm_relay_keeps_the_json_body_when_the_content_type_carries_a_charset(self, monkeypatch): + body = {"model": "router-model", "messages": [{"role": "user", "content": "hi"}]} + captured: list[dict] = [] + self._install_router(monkeypatch, self._recording_router(captured), body) + + await vllm_proxy_route( + endpoint="/chat/completions", + request=self._request("application/json; charset=utf-8"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert captured[0]["json"] == body + + @pytest.mark.asyncio + async def test_azure_relay_returns_the_upstream_status_and_body_when_the_deployment_rejects_the_call( + self, monkeypatch + ): + upstream_body = {"error": {"code": "DeploymentNotFound", "message": "The API deployment does not exist."}} + + class RejectingRouter: + async def allm_passthrough_route(self, **kwargs): + upstream_request = httpx.Request( + "POST", "https://my-azure.openai.azure.com/openai/deployments/gpt-5/chat/completions" + ) + upstream = httpx.Response( + 404, json=upstream_body, headers={"x-ms-request-id": "req-1"}, request=upstream_request + ) + raise httpx.HTTPStatusError("404", request=upstream_request, response=upstream) + + self._install_router(monkeypatch, RejectingRouter(), {"model": "gpt-5", "stream": False}) + + result = await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert result.status_code == 404 + assert json.loads(result.body) == upstream_body + assert result.headers["x-ms-request-id"] == "req-1" + + +@pytest.mark.asyncio +async def test_bedrock_count_tokens_error_forwards_provider_headers(): + """The count tokens route converts BedrockError into an HTTPException, and dropping the + headers there loses x-amzn-RequestId after the handler went to the trouble of keeping it.""" + from fastapi import HTTPException + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + failure = BedrockError( + status_code=500, + message="Amazon Bedrock is unable to process your request.", + headers={"x-amzn-RequestId": "req-count-tokens-500"}, + ) + + with patch( # test-quality-ok: the route's BedrockError branch is only reachable when the handler raises + "litellm.llms.bedrock.count_tokens.handler.BedrockCountTokensHandler.handle_count_tokens_request", + new=AsyncMock(side_effect=failure), + ): + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint="v1/messages/count_tokens", + request=MagicMock(), + fastapi_response=MagicMock(), + user_api_key_dict=MagicMock(), + request_body={"model": "anthropic.claude-haiku-4-5-20251001-v1:0"}, + ) + + assert exc_info.value.status_code == 500 + assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500" + + +class _AzureGroupRouter: + def __init__(self, captured: list[dict]) -> None: + self.captured = captured + + def get_model_names(self, team_id=None): + return ["gpt", "other-group"] + + def get_model_list(self, model_name=None, team_id=None): + rows = [ + {"model_name": "gpt", "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_key": "k"}}, + {"model_name": "other-group", "litellm_params": {"model": "azure/gpt-5.4", "api_key": "k"}}, + ] + return [row for row in rows if model_name is None or row["model_name"] == model_name] + + async def allm_passthrough_route(self, **kwargs): + self.captured.append(kwargs) + return httpx.Response(200, json={"ok": True}) + + +class TestAzureRelayDeploymentSegment: + """A key allowed one model group must not reach another deployment by naming it in the + ``openai/deployments/`` segment while the group segment picks the credential.""" + + def test_models_served_by_group_resolves_each_deployment_to_its_model_name(self): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _models_served_by_group + + assert _models_served_by_group(_AzureGroupRouter([]), "gpt") == frozenset({"gpt-5.4-mini"}) + assert _models_served_by_group(_AzureGroupRouter([]), "missing-group") == frozenset() + + def _install(self, monkeypatch, body: dict) -> list[dict]: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + captured: list[dict] = [] + + async def fake_get_request_body(_request): + return body + + monkeypatch.setattr(proxy_server, "llm_router", _AzureGroupRouter(captured)) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + return captured + + def _request(self) -> Request: + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": "application/json"} + request.query_params = {} + return request + + @pytest.mark.asyncio + async def test_azure_relay_rejects_a_deployment_the_group_does_not_serve(self, monkeypatch): + from fastapi import HTTPException + + captured = self._install(monkeypatch, {"model": "gpt", "messages": []}) + + with pytest.raises(HTTPException) as exc_info: + await azure_proxy_route( + endpoint="gpt/openai/deployments/gpt-5.4/chat/completions", + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token", models=["gpt"]), + ) + + assert exc_info.value.status_code == 400 + assert "gpt-5.4" in exc_info.value.detail["error"] + assert captured == [] + + @pytest.mark.asyncio + async def test_azure_relay_dispatches_the_group_and_its_own_deployment_name(self, monkeypatch): + captured = self._install(monkeypatch, {"model": "gpt", "messages": []}) + + for endpoint in ( + "gpt/openai/deployments/gpt/chat/completions", + "gpt/openai/deployments/gpt-5.4-mini/chat/completions", + ): + await azure_proxy_route( + endpoint=endpoint, + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token", models=["gpt"]), + ) + + assert [call["model"] for call in captured] == ["gpt", "gpt"] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index fb4ed3db4d2..d57bed430c1 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -11,7 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from fastapi import Request, UploadFile +from fastapi import Request, Response, UploadFile from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile @@ -22,6 +22,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( InitPassThroughEndpointHelpers, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, _registered_pass_through_routes, + chat_completion_pass_through_endpoint, create_pass_through_route, initialize_pass_through_endpoints, pass_through_request, @@ -5837,3 +5838,38 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) == "per-call-random-trace-id" ) + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_error_for_an_unknown_model( + monkeypatch: pytest.MonkeyPatch, +): + """A bare HTTPException carries no type or param, so the tail used to ship the + literal string "None" in both fields.""" + proxy_logging = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + proxy_logging.post_call_failure_hook = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + request = MagicMock(spec=Request) + request.body = AsyncMock( + return_value=json.dumps({"model": "unknown-model", "messages": [{"role": "user", "content": "hi"}]}).encode() + ) + + with pytest.raises(ProxyException) as raised: + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "400") diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index cc231e383a3..87b1bc56659 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -139,6 +139,71 @@ class TestGetAttachedPolicies: assert "gpt4-policy" in attached assert len(attached) == 3 + def test_matches_are_ordered_from_broadest_to_narrowest_scope(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "model-policy", "models": ["gpt-4"]}, + {"policy": "team-policy", "teams": ["t1"]}, + {"policy": "global-policy", "scope": "*"}, + ] + ) + + context = PolicyMatchContext(team_alias="t1", model="gpt-4") + + assert registry.get_attached_policies(context) == [ + "global-policy", + "team-policy", + "model-policy", + ] + + def test_combined_team_and_model_attachment_uses_model_specificity(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "team-policy", "teams": ["t1"]}, + {"policy": "team-model-policy", "teams": ["t1"], "models": ["gpt-4"]}, + ] + ) + + context = PolicyMatchContext(team_alias="t1", model="gpt-4") + + assert registry.get_attached_policies(context) == [ + "team-policy", + "team-model-policy", + ] + + def test_duplicate_policy_uses_broadest_matching_attachment(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "shared-policy", "models": ["gpt-4"]}, + {"policy": "model-policy", "models": ["gpt-4"]}, + {"policy": "shared-policy", "scope": "*"}, + ] + ) + + context = PolicyMatchContext(model="gpt-4") + + assert registry.get_attached_policies(context) == [ + "shared-policy", + "model-policy", + ] + assert registry.get_attached_policies_with_reasons(context)[0]["matched_via"] == "scope:*" + + def test_duplicate_policy_prefers_single_scope_over_combined_scope(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "shared-policy", "teams": ["t1"], "models": ["gpt-4"]}, + {"policy": "shared-policy", "models": ["gpt-4"]}, + ] + ) + + context = PolicyMatchContext(team_alias="t1", model="gpt-4") + + assert registry.get_attached_policies_with_reasons(context)[0]["matched_via"] == "model:gpt-4" + def test_same_policy_multiple_attachments_no_duplicates(self): """Test same policy attached multiple ways doesn't duplicate.""" registry = AttachmentRegistry() diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 4fcb7d22588..0a2641082dc 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -4,20 +4,27 @@ Tests for the pipeline executor. Uses mock guardrails to validate pipeline execution without external services. """ +import copy +import logging +from typing import Literal from unittest.mock import MagicMock import pytest import litellm +from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( CustomCodeGuardrail, ) -from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) +from litellm.types.utils import CallTypesLiteral try: from fastapi.exceptions import HTTPException @@ -158,11 +165,146 @@ class ContentCheckGuardrail(CustomGuardrail): return None +class RecordingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str, scan_raw_request: bool = False, block: bool = True): + super().__init__( + guardrail_name=guardrail_name, + event_hook="pre_call", + default_on=True, + scan_raw_request=scan_raw_request, + ) + self.block = block + + def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool: + return True + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"detected": ["aws_access_key"]}, + request_data=data, + guardrail_status="guardrail_intervened" if self.block else "success", + ) + if self.block: + raise HTTPException(status_code=400, detail="Content policy violation") + return copy.deepcopy(data) + + # ───────────────────────────────────────────────────────────────────────────── # Tests # ───────────────────────────────────────────────────────────────────────────── +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +@pytest.mark.parametrize("scan_raw_request", [False, True]) +@pytest.mark.parametrize("on_fail", ["block", "modify_response"]) +async def test_terminal_block_carries_guardrail_information_to_request( + monkeypatch: pytest.MonkeyPatch, scan_raw_request: bool, on_fail: Literal["block", "modify_response"] +): + """ + Spend logging and the Guardrails Monitor read standard_logging_guardrail_information + off the caller's request dict. A blocking step records it on the executor's + working copy (or the raw-request snapshot), so the terminal result must carry it + back onto the request or the block is never counted. + """ + guard = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=scan_raw_request) + monkeypatch.setattr(litellm, "callbacks", [guard]) + data = { + "messages": [{"role": "user", "content": "key AKIAIOSFODNN7EXAMPLE"}], + "metadata": {"user_api_key_hash": "abc"}, + } + + result = await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="credentials-api-keys", on_fail=on_fail, on_pass="next")], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="baseline-pii-protection", + raw_request_snapshot={"messages": data["messages"], "metadata": {"user_api_key_hash": "abc"}}, + ) + + assert result.terminal_action == on_fail + recorded = data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["credentials-api-keys"] + assert recorded[0]["guardrail_status"] == "guardrail_intervened" + assert data["metadata"]["user_api_key_hash"] == "abc" + assert "guardrails" not in data["metadata"] + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_terminal_block_merges_guardrail_information_without_duplicates(monkeypatch: pytest.MonkeyPatch): + """A pass_data step that returns a rewritten copy of the request, and a scan_raw_request step + that evaluates a deep copy taken before the pipeline ran, both leave earlier entries in two + dicts at once. Those must be carried back once while every step's own entry is kept.""" + first = RecordingGuardrail(guardrail_name="pii-scan", block=False) + second = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=True) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"} + data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + data["metadata"]["standard_logging_guardrail_information"] = [earlier] + + result = await PipelineExecutor.execute_steps( + steps=[ + PipelineStep(guardrail="pii-scan", on_fail="block", on_pass="next", pass_data=True), + PipelineStep(guardrail="credentials-api-keys", on_fail="block", on_pass="next"), + ], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="baseline-pii-protection", + raw_request_snapshot={ + "messages": data["messages"], + "metadata": {"standard_logging_guardrail_information": [dict(earlier)]}, + }, + ) + + assert result.terminal_action == "block" + recorded = data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "pii-scan", "credentials-api-keys"] + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_repeated_scan_raw_request_step_is_counted_once_per_evaluation(monkeypatch: pytest.MonkeyPatch): + """Running the same raw-scan guardrail twice yields two identical entries; both must reach the caller, + while the entries the raw snapshot already held before the pipeline ran are not copied again.""" + guard = RecordingGuardrail(guardrail_name="credentials-raw", scan_raw_request=True, block=False) + monkeypatch.setattr(litellm, "callbacks", [guard]) + earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"} + data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + data["metadata"]["standard_logging_guardrail_information"] = [earlier] + + result = await PipelineExecutor.execute_steps( + steps=[ + PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"), + PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"), + ], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="raw-scan-policy", + raw_request_snapshot={ + "messages": data["messages"], + "metadata": {"standard_logging_guardrail_information": [dict(earlier)]}, + }, + ) + + assert result.terminal_action == "allow" + assert result.modified_data is not None + recorded = result.modified_data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "credentials-raw", "credentials-raw"] + + @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio async def test_escalation_step1_fails_step2_blocks(monkeypatch): @@ -533,8 +675,6 @@ async def test_guardrail_not_found_uses_on_fail(monkeypatch): ], ) - monkeypatch.setattr(litellm, "callbacks", []) - result = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode=pipeline.mode, @@ -911,3 +1051,627 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch): assert outcome == "pass" assert guardrail.native_pre_call_ran is True assert "guardrail_to_apply" not in data + + +class _TextReturningGuardrail(CustomGuardrail): + def __init__(self, returned_texts): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + self.returned_texts = returned_texts + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": self.returned_texts} + + +class _TextTranslation: + delivers_ended_stream_rewrites = False + + def __init__(self): + self.seen_guardrail_names = [] + + async def process_output_streaming_response( + self, responses_so_far, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + self.seen_guardrail_names.append(guardrail_to_apply.guardrail_name) + await guardrail_to_apply.apply_guardrail( + inputs={"texts": ["hello world"]}, + request_data=request_data or {}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + +class _WritingTranslation: + """Writes the guardrail's text (and tool-call) outputs back into the buffered chunks the way the + chat/Responses/Messages handlers do on an ended stream.""" + + delivers_ended_stream_rewrites = True + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + assert deliver_ended_stream_rewrites is True + outputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}, + request_data=request_data or {}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + responses_so_far[0]["text"] = outputs["texts"][0] + if len(outputs["tool_calls"]) == 1: + responses_so_far[0]["tool_call"] = outputs["tool_calls"][0] + return responses_so_far + + +class _RefusingTranslation: + delivers_ended_stream_rewrites = True + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + responses_so_far[0]["text"] = "half-written" + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name) + + +def _chunk(): + return {"text": "hello world", "tool_call": {"function": {"name": "lookup", "arguments": '{"ssn": "123"}'}}} + + +async def _run_streaming_step(translation, streaming_chunks=None): + chunks = [object()] if streaming_chunks is None else streaming_chunks + return await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail="next", on_error="next")], + mode="post_call", + data={"model": "m"}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="p", + streaming_chunks=chunks, + endpoint_translation=translation, + ) + + +def _assert_passed_with_discard_warning(result, caplog): + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records) + assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_text_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) + translation = _TextTranslation() + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(translation, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + assert translation.seen_guardrail_names == ["masker"] + + +@pytest.mark.asyncio +async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(("hello world",))]) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_TextTranslation()) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +class _InPlaceMutatingGuardrail(CustomGuardrail): + """Rewrites like bedrock/presidio do: rebinds inputs["texts"] on the dict it was handed + and returns that same dict, so a post-call comparison against inputs sees no change.""" + + def __init__(self): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + inputs["texts"] = ["hello [MASKED]"] + return inputs + + +@pytest.mark.asyncio +async def test_streaming_step_in_place_rewrite_is_discarded_without_write_back(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_InPlaceMutatingGuardrail()]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_TextTranslation(), chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +class _TextAndToolCallRewritingGuardrail(CustomGuardrail): + def __init__(self, rewrite_tool_call): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + self.rewrite_tool_call = rewrite_tool_call + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + tool_calls = ( + [{"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}] + if self.rewrite_tool_call + else inputs["tool_calls"] + ) + return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": tool_calls} + + +@pytest.mark.asyncio +async def test_streaming_step_delivers_text_rewrite_through_writing_translation(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=False)]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_WritingTranslation(), chunks) + + assert result.terminal_action == "allow" + assert chunks[0]["text"] == "hello [MASKED]" + assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "123"}' + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_delivers_tool_call_rewrite_through_writing_translation(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_WritingTranslation(), chunks) + + assert result.terminal_action == "allow" + assert chunks[0]["text"] == "hello [MASKED]" + assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "[MASKED]"}' + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +class _ToolCallDroppingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": []} + + +@pytest.mark.asyncio +async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool_call(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_ToolCallDroppingGuardrail()]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_WritingTranslation(), chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_TextTranslation(), chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +class _BlockingStreamGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + +def _recorded_guardrail_statuses(result): + return [ + entry["guardrail_status"] + for entry in result.modified_data["metadata"]["standard_logging_guardrail_information"] + ] + + +@pytest.mark.asyncio +async def test_streaming_step_records_guardrail_information_once_on_mask(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert result.terminal_action == "allow" + assert _recorded_guardrail_statuses(result) == ["success"] + + +@pytest.mark.asyncio +async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_header(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + + +@pytest.mark.asyncio +async def test_streaming_step_records_guardrail_information_once_on_block(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert [step.outcome for step in result.step_results] == ["fail"] + assert _recorded_guardrail_statuses(result) == ["guardrail_intervened"] + + +@pytest.mark.asyncio +async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewrite(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_RefusingTranslation(), chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +class _LegacyHookGuardrail(CustomGuardrail): + """A guardrail with only the legacy post-call hook: it never defines apply_guardrail.""" + + def __init__(self, replacement=None, raises=None, guardrail_name="masker", rewrite_in_place=None): + super().__init__(guardrail_name=guardrail_name, event_hook="post_call", default_on=True) + self.replacement = replacement + self.raises = raises + self.rewrite_in_place = rewrite_in_place + self.calls = [] + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.calls.append({"data": data, "user_api_key_dict": user_api_key_dict, "response": response}) + if self.raises is not None: + raise self.raises + if self.rewrite_in_place is not None: + response["text"] = self.rewrite_in_place + return self.replacement + + +class _NativeHooksGuardrail(_LegacyHookGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail") + + +class _LegacyScanningTranslation: + """Stores the assembled response under request_data["response"] before scanning, like the + chat, Responses, and Messages handlers, hands hooks a route-native shape, and re-extracts one + text per entry of a replacement's "texts".""" + + delivers_ended_stream_rewrites = True + + def post_call_hook_response(self, response): + return {"native": True, "text": response["text"], "tool_calls": response["tool_calls"]} + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + request_data.setdefault( + "response", {"text": responses_so_far[0]["text"], "tool_calls": [dict(responses_so_far[0]["tool_call"])]} + ) + outputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + responses_so_far[0]["text"] = outputs["texts"][0] + return responses_so_far + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + inputs = {"texts": [response["text"]] if "text" in response else list(response["texts"])} + if response.get("tool_calls"): + inputs["tool_calls"] = list(response["tool_calls"]) + await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data={"response": response}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return response + + +class _ToolOnlyLegacyScanningTranslation(_LegacyScanningTranslation): + """Like the Messages handler on a tool-only message: the ended-stream scan omits "texts" from + the inputs, while the non-streaming scan of the same response sends an empty list.""" + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + request_data.setdefault("response", {"text": "", "tool_calls": [dict(responses_so_far[0]["tool_call"])]}) + await guardrail_to_apply.apply_guardrail( + inputs={"tool_calls": [dict(responses_so_far[0]["tool_call"])]}, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + await guardrail_to_apply.apply_guardrail( + inputs={"texts": [], "tool_calls": list(response.get("tool_calls") or [])}, + request_data={"response": response}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return response + + +def _tool_only_chunk(): + return {"text": "", "tool_call": _chunk()["tool_call"]} + + +def _native(text): + return {"native": True, "text": text, "tool_calls": [_chunk()["tool_call"]]} + + +def _legacy_replacement(*texts, tool_calls=None): + return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls} + + +async def _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, on_fail="block", on_error="next", translation=None +): + return await _run_legacy_streaming_steps( + monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error, translation=translation + ) + + +async def _run_legacy_streaming_steps( + monkeypatch, guardrails, chunks, on_fail="block", on_error="next", translation=None +): + monkeypatch.setattr(litellm, "callbacks", list(guardrails)) + return await PipelineExecutor.execute_steps( + steps=[ + PipelineStep( + guardrail=guardrail.guardrail_name, + on_pass="next" if position + 1 < len(guardrails) else "allow", + on_fail=on_fail, + on_error=on_error, + ) + for position, guardrail in enumerate(guardrails) + ], + mode="post_call", + data={"model": "m"}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="p", + streaming_chunks=chunks, + endpoint_translation=_LegacyScanningTranslation() if translation is None else translation, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("guardrail_class", [_LegacyHookGuardrail, _NativeHooksGuardrail]) +async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypatch, caplog, guardrail_class): + guardrail = guardrail_class(replacement=_legacy_replacement("[REWRITTEN] hello world")) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert [call["response"] for call in guardrail.calls] == [_native("hello world")] + assert guardrail.calls[0]["data"]["model"] == "m" + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_delivers_a_legacy_rewrite_made_in_place(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(rewrite_in_place="[REWRITTEN] hello world") + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_passes_untouched_when_legacy_hook_returns_none(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=None) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert len(guardrail.calls) == 1 + assert chunks == [_chunk()] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_streaming_step_blocks_with_the_legacy_hook_exception(monkeypatch): + exc = HTTPException(status_code=400, detail={"error": "output blocked"}) + chunks = [_chunk()] + + result = await _run_legacy_streaming_step(monkeypatch, _LegacyHookGuardrail(raises=exc), chunks) + + assert result.terminal_action == "block" + assert [step.outcome for step in result.step_results] == ["fail"] + assert result.original_exception is exc + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_takes_on_error_when_legacy_hook_crashes(monkeypatch): + chunks = [_chunk()] + + result = await _run_legacy_streaming_step( + monkeypatch, _LegacyHookGuardrail(raises=ValueError("boom")), chunks, on_error="block" + ) + + assert result.terminal_action == "block" + assert [step.outcome for step in result.step_results] == ["error"] + assert result.step_results[0].error_detail == "boom" + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("split", "in two")) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(monkeypatch, caplog): + masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} + guardrail = _LegacyHookGuardrail( + replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[masked_tool_call]) + ) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[])) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alone(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=None) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert chunks == [_tool_only_chunk()] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog): + masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call])) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_tool_only_chunk()] + + +class _NoHooksGuardrail(CustomGuardrail): + pass + + +class _IteratorAndLegacyHookGuardrail(_LegacyHookGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for item in response: + yield item + + +class _UnscannableRewriteTranslation(_LegacyScanningTranslation): + """Like the chat handler on a response whose choices are plain dicts: the non-streaming scan + never hands anything to the guardrail.""" + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + return response + + +def test_streaming_execution_runs_legacy_hooks_only_when_that_hook_is_their_only_streaming_path(): + assert PipelineExecutor.supports_streaming_execution(_LegacyHookGuardrail()) is True + assert PipelineExecutor.supports_streaming_execution(_NativeHooksGuardrail()) is True + assert PipelineExecutor.supports_streaming_execution(_IteratorAndLegacyHookGuardrail()) is False + assert PipelineExecutor.supports_streaming_execution(_NoHooksGuardrail(guardrail_name="neither")) is False + + +@pytest.mark.asyncio +async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_rescan(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("hello [MASKED]")) + chunks = [_chunk()] + + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation()) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): + masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world")) + auditor = _LegacyHookGuardrail(replacement=None, guardrail_name="auditor") + chunks = [_chunk()] + + result = await _run_legacy_streaming_steps(monkeypatch, [masker, auditor], chunks, on_fail="next") + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass", "pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert [call["response"] for call in masker.calls] == [_native("hello world")] + assert [call["response"] for call in auditor.calls] == [_native("[REWRITTEN] hello world")] diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py new file mode 100644 index 00000000000..37849101b3a --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -0,0 +1,262 @@ +import logging +from collections.abc import Iterator, Mapping + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.router import Deployment, LiteLLM_Params + +GOVERNED_MODEL_GROUP = "gpt-5.4-mini" +GOVERNED_MODEL_ID = "deployment-governed" +UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini" +UNGOVERNED_MODEL_ID = "deployment-ungoverned" +WILDCARD_MODEL_GROUP = "openai/*" +WILDCARD_MODEL_ID = "deployment-wildcard" + + +class FakeRouter: + def __init__(self, deployments: dict[str, Deployment], model_group_alias: dict[str, object] | None = None): + self._deployments = deployments + self.model_group_alias = model_group_alias or {} + + def get_deployment(self, model_id: str) -> Deployment | None: + return self._deployments.get(model_id) + + +def _deployment(model_group: str, model_id: str) -> Deployment: + return Deployment( + model_name=model_group, + litellm_params=LiteLLM_Params(model=f"openai/{model_group}"), + model_info={"id": model_id}, + ) + + +def _router(model_group_alias: dict[str, object] | None = None) -> FakeRouter: + return FakeRouter( + { + GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID), + UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID), + WILDCARD_MODEL_ID: _deployment(WILDCARD_MODEL_GROUP, WILDCARD_MODEL_ID), + }, + model_group_alias, + ) + + +def _encoded_response_id(model_id: str) -> str: + return ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id=model_id, response_id="resp_upstream" + ) + + +def _pipeline_policy(guardrail: str, mode: str = "post_call") -> dict[str, object]: + return { + "guardrails": {"add": [guardrail]}, + "pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": "block"}]}, + } + + +@pytest.fixture +def policy_engine() -> Iterator[None]: + policy_registry = get_policy_registry() + attachment_registry = get_attachment_registry() + policy_registry.load_policies( + { + "response-governance": _pipeline_policy("output-word-filter"), + "input-governance": _pipeline_policy("input-word-filter", mode="pre_call"), + "team-governance": _pipeline_policy("team-word-filter"), + "tag-governance": _pipeline_policy("tag-word-filter"), + } + ) + attachment_registry.load_attachments( + [ + {"policy": "response-governance", "models": [GOVERNED_MODEL_GROUP]}, + {"policy": "input-governance", "models": [GOVERNED_MODEL_GROUP]}, + {"policy": "team-governance", "teams": ["governed-team"]}, + {"policy": "tag-governance", "tags": ["governed"]}, + ] + ) + yield + policy_registry.clear() + attachment_registry.clear() + + +def _retrieval_data(model_id: str) -> dict[str, object]: + return {"response_id": _encoded_response_id(model_id), "litellm_metadata": {}} + + +def _attached_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, str], ...]: + bucket = data["litellm_metadata"] + assert isinstance(bucket, dict) + return tuple( + (policy_name, ",".join(step.guardrail for step in pipeline.steps)) + for policy_name, pipeline in bucket["_guardrail_pipelines"] + ) + + +def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine: None) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert data["litellm_metadata"]["_pipeline_managed_guardrails"] == frozenset({"output-word-filter"}) + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + assert data["litellm_metadata"]["applied_guardrails"] == ["output-word-filter"] + assert data["litellm_metadata"]["policy_sources"] == {"response-governance": "model:gpt-5.4-mini"} + assert "model" not in data + assert "guardrails" not in data["litellm_metadata"] + + +def test_key_and_team_context_also_governs_retrieval(policy_engine: None) -> None: + data = _retrieval_data(UNGOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval( + data=data, user_api_key_dict=UserAPIKeyAuth(team_alias="governed-team"), llm_router=_router() + ) + + assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),) + + +def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine: None) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + + +def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine: None) -> None: + data: dict[str, object] = { + "response_id": _encoded_response_id(UNGOVERNED_MODEL_ID), + "litellm_metadata": {"tags": ["governed"]}, + } + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("tag-governance", "tag-word-filter"),) + assert data["litellm_metadata"]["policy_sources"] == {"tag-governance": "tag:governed"} + + +def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine: None) -> None: + data = _retrieval_data(UNGOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(UNGOVERNED_MODEL_ID) + + +def test_already_attached_policy_is_not_attached_twice(policy_engine: None) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + router = _router() + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + + +def _hidden_submit_model_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "the model name it was submitted as" in record.getMessage() + ] + + +def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: + data = _retrieval_data(WILDCARD_MODEL_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(WILDCARD_MODEL_ID) + assert [ + "as model group openai/* (a wildcard deployment)" in message + for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_aliased_model_group_still_attaches_its_own_policies_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + router = _router({"gpt-mini": GOVERNED_MODEL_GROUP, "gpt-hidden": {"model": GOVERNED_MODEL_GROUP, "hidden": True}}) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert [ + "(the target of model_group_alias gpt-mini, gpt-hidden)" in message + for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval( + data=_retrieval_data(GOVERNED_MODEL_ID), + user_api_key_dict=UserAPIKeyAuth(), + llm_router=_router({"other-alias": UNGOVERNED_MODEL_GROUP}), + ) + + assert _hidden_submit_model_warnings(caplog) == [] + + +def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING + and "retrieved without its post_call policy pipelines" in record.getMessage() + ] + + +@pytest.mark.parametrize( + ("response_id", "reason"), + [ + ("resp_plain_upstream_id", "response id names no deployment"), + (_encoded_response_id("deployment-missing-from-router"), "deployment no longer in the router"), + (None, "response id names no deployment"), + ], +) +def test_unresolvable_response_id_attaches_nothing_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture, response_id: str, reason: str +) -> None: + data = {"response_id": response_id, "litellm_metadata": {}} + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == {"response_id": response_id, "litellm_metadata": {}} + assert [message.endswith(f"({reason})") for message in _ungoverned_retrieval_warnings(caplog)] == [True] + + +def test_without_a_router_attaches_nothing_and_warns(policy_engine: None, caplog: pytest.LogCaptureFixture) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) + + assert data == _retrieval_data(GOVERNED_MODEL_ID) + assert [message.endswith("(no router)") for message in _ungoverned_retrieval_warnings(caplog)] == [True] + + +def test_without_policy_engine_attaches_nothing_quietly(caplog: pytest.LogCaptureFixture) -> None: + get_policy_registry().clear() + data = _retrieval_data(GOVERNED_MODEL_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) + + assert data == _retrieval_data(GOVERNED_MODEL_ID) + assert _ungoverned_retrieval_warnings(caplog) == [] diff --git a/tests/test_litellm/proxy/proxy_server/test_background_health.py b/tests/test_litellm/proxy/proxy_server/test_background_health.py index 990844369f7..d99aa637955 100644 --- a/tests/test_litellm/proxy/proxy_server/test_background_health.py +++ b/tests/test_litellm/proxy/proxy_server/test_background_health.py @@ -20,6 +20,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest import litellm.proxy.proxy_server as proxy_server +from litellm.constants import BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME from litellm.proxy.proxy_server import ( _adaptive_router_flusher_loop, _get_endpoint_exception_status, @@ -111,13 +112,11 @@ async def test_run_direct_health_check_with_instrumentation_returns_results( lambda _gs: {}, ) - healthy, unhealthy, exceptions = ( - await _run_direct_health_check_with_instrumentation( - model_list=[{"model_name": "gpt-4"}], - details=False, - max_concurrency=1, - instrumentation_context={"source": "test"}, - ) + healthy, unhealthy, exceptions = await _run_direct_health_check_with_instrumentation( + model_list=[{"model_name": "gpt-4"}], + details=False, + max_concurrency=1, + instrumentation_context={"source": "test"}, ) assert normalize( @@ -245,6 +244,137 @@ async def test_schedule_background_health_check_db_save_invalid_no_event_loop_ra ) +def _lock_manager(redis_cache, acquired): + manager = MagicMock() + manager.redis_cache = redis_cache + manager.acquire_lock = AsyncMock(return_value=acquired) + manager.release_lock = AsyncMock() + return manager + + +def _capture_saves(monkeypatch, persisted=True): + saves = [] + + async def _fake_save(*_args, **kwargs): + saves.append(kwargs) + return persisted + + import litellm.proxy.health_endpoints._health_endpoints as he + + monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save) + return saves + + +def _cancel_during_save(monkeypatch): + async def _fake_save(*_args, **_kwargs): + raise asyncio.CancelledError() + + import litellm.proxy.health_endpoints._health_endpoints as he + + monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save) + + +def _schedule_with(lock_manager): + _schedule_background_health_check_db_save( + prisma_client=MagicMock(), + shared_health_manager=None, + model_list=[], + healthy_endpoints=[], + unhealthy_endpoints=[], + pod_lock_manager=lock_manager, + lock_ttl=300, + ) + + +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_skips_a_window_another_pod_persisted(monkeypatch): + saves = _capture_saves(monkeypatch) + lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=False) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert saves == [] + + +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_holds_the_window_lock_for_the_whole_interval(monkeypatch): + """The lock is the "saved this window" marker: never reentrant, TTL = interval, and never released.""" + saves = _capture_saves(monkeypatch) + lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert normalize( + { + "saves": len(saves), + "lock_request": lock_manager.acquire_lock.await_args.kwargs, + "released": lock_manager.release_lock.await_count, + } + ) == { + "saves": 1, + "lock_request": { + "cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, + "ttl": 300, + "allow_reentrant": False, + }, + "released": 0, + } + + +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_reports_failure( + monkeypatch, +): + """A failed save must not burn the window: release the lock so another pod's cycle can retry.""" + saves = _capture_saves(monkeypatch, persisted=False) + lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert normalize( + { + "saves": len(saves), + "release_request": lock_manager.release_lock.await_args.kwargs, + "release_count": lock_manager.release_lock.await_count, + } + ) == { + "saves": 1, + "release_request": {"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME}, + "release_count": 1, + } + + +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_is_cancelled( + monkeypatch, +): + """A pod shutting down mid-save releases the lock instead of holding it until the TTL.""" + _cancel_during_save(monkeypatch) + lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert ( + lock_manager.release_lock.await_args.kwargs, + lock_manager.release_lock.await_count, + ) == ({"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME}, 1) + + +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_runs_ungated_without_redis(monkeypatch): + saves = _capture_saves(monkeypatch, persisted=False) + lock_manager = _lock_manager(redis_cache=None, acquired=True) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert (len(saves), lock_manager.acquire_lock.await_count, lock_manager.release_lock.await_count) == (1, 0, 0) + + # --------------------------------------------------------------------------- # _get_endpoint_exception_status # --------------------------------------------------------------------------- @@ -319,13 +449,9 @@ def test_write_health_state_to_router_cache_sets_states(monkeypatch): _write_health_state_to_router_cache(healthy, unhealthy, exceptions) - fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with( - fake_states - ) + fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(fake_states) - call_args = fake_router.health_state_cache.set_deployment_health_states.call_args[ - 0 - ][0] + call_args = fake_router.health_state_cache.set_deployment_health_states.call_args[0][0] assert normalize( { "states_keys": sorted(call_args.keys()), @@ -367,9 +493,7 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp fake_router.cooldown_time = 30 monkeypatch.setattr(proxy_server, "llm_router", fake_router) - monkeypatch.setattr( - proxy_server, "general_settings", {"model_list_healthy_only": True} - ) + monkeypatch.setattr(proxy_server, "general_settings", {"model_list_healthy_only": True}) fake_states = {"m1": {"is_healthy": True}, "m2": {"is_healthy": False}} @@ -403,9 +527,7 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp {"m2": SimpleNamespace(status_code=500)}, ) - fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with( - fake_states - ) + fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(fake_states) assert cooldowns == [] assert failures == [] @@ -415,9 +537,7 @@ def test_write_health_state_to_router_cache_swallows_internal_failures(monkeypat fake_router = MagicMock() fake_router.enable_health_check_routing = True fake_router.health_check_ignore_transient_errors = False - fake_router.health_state_cache.set_deployment_health_states.side_effect = ( - RuntimeError("cache exploded") - ) + fake_router.health_state_cache.set_deployment_health_states.side_effect = RuntimeError("cache exploded") monkeypatch.setattr(proxy_server, "llm_router", fake_router) @@ -447,9 +567,7 @@ async def test_adaptive_router_flusher_loop_flushes_each_router(monkeypatch): from litellm.types.router import TaggedPreRoutingStrategy fake_router = MagicMock() - fake_router.adaptive_routers = { - "alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)] - } + fake_router.adaptive_routers = {"alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)]} monkeypatch.setattr(proxy_server, "llm_router", fake_router) monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) @@ -547,12 +665,8 @@ async def test_run_background_health_check_runs_one_cycle_then_cancels(monkeypat "_run_direct_health_check_with_instrumentation", _fake_direct, ) - monkeypatch.setattr( - proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None - ) - monkeypatch.setattr( - proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None - ) + monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None) + monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None) monkeypatch.setattr( proxy_server, "health_check_filter_kwargs_from_general_settings", @@ -630,12 +744,8 @@ async def test_run_background_health_check_probes_only_listed_model_groups(monke "_run_direct_health_check_with_instrumentation", _fake_direct, ) - monkeypatch.setattr( - proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None - ) - monkeypatch.setattr( - proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None - ) + monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None) + monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None) monkeypatch.setattr( proxy_server, "health_check_filter_kwargs_from_general_settings", diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a06e7142122..deb7289d2d1 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -22,6 +22,7 @@ import inspect import json import logging import os +import subprocess from collections.abc import Awaitable, Callable from typing import List, Optional, Union from unittest.mock import AsyncMock, MagicMock, patch @@ -61,6 +62,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 +72,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 +80,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, } @@ -249,6 +253,38 @@ async def test_flush_spend_logs_queue_on_shutdown_swallows_drain_errors(monkeypa await ps._flush_spend_logs_queue_on_shutdown() +@pytest.mark.asyncio +async def test_flush_spend_counters_on_shutdown_commits_buffered_spend(monkeypatch): + fake_prisma = MagicMock() + monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) + commit = AsyncMock() + monkeypatch.setattr(ps.proxy_logging_obj.db_spend_update_writer, "db_update_spend_transaction_handler", commit) + + await ps.flush_spend_counters_on_shutdown() + + observed = { + "commit_calls": commit.await_count, + "commit_prisma": commit.await_args.kwargs["prisma_client"] is fake_prisma, + "commit_proxy_logging": commit.await_args.kwargs["proxy_logging_obj"] is ps.proxy_logging_obj, + } + assert observed == {"commit_calls": 1, "commit_prisma": True, "commit_proxy_logging": True} + + +@pytest.mark.asyncio +async def test_flush_spend_counters_on_shutdown_logs_and_swallows_commit_errors(monkeypatch, caplog): + monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False) + monkeypatch.setattr( + ps.proxy_logging_obj.db_spend_update_writer, + "db_update_spend_transaction_handler", + AsyncMock(side_effect=RuntimeError("db gone")), + ) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + await ps.flush_spend_counters_on_shutdown() + + assert "Error flushing spend counters on shutdown: db gone" in caplog.text + + # --------------------------------------------------------------------------- # _initialize_shared_aiohttp_session # --------------------------------------------------------------------------- @@ -746,6 +782,30 @@ async def test_proxy_startup_event_invalid_missing_app_arg_raises(): pass +@pytest.mark.asyncio +async def test_proxy_startup_event_prunes_dead_workers_live_gauges(tmp_path): + """With PROMETHEUS_MULTIPROC_DIR set, a booting worker drops the live-gauge files of pids that no longer + exist, so a crashed worker's in-flight samples leave the aggregate as soon as its replacement starts.""" + exited = subprocess.Popen(["true"]) + assert exited.wait(timeout=30) == 0 + stale = tmp_path / f"gauge_livesum_{exited.pid}.db" + stale.touch() + counter = tmp_path / f"counter_{exited.pid}.db" + counter.touch() + + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} + clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path) + with patch.dict(os.environ, clean_env, clear=True): + try: + async with proxy_startup_event(app=None): + pass + except Exception: + pass + + assert not stale.exists() + assert counter.exists() + + def test_otel_global_provider_published_after_callback_init(): """The OTel V2 global-provider publish must run after callback initialization in ``proxy_startup_event``. @@ -818,6 +878,42 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): ) +@pytest.mark.asyncio +async def test_tuning_baseline_v2_is_created_alongside_the_legacy_row(): + from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT + + prisma_client = MagicMock() + prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_config.create = AsyncMock() + deployment = { + "model_name": "a", + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": {}}, + } + + result = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, [deployment]) + + assert result == {'yaml:["a",[]]': DEFAULT_TUNING_FINGERPRINT} + assert prisma_client.db.litellm_config.create.await_args.kwargs["data"] == { + "param_name": "auto_router_tuning_baseline_v2", + "param_value": json.dumps(dict(result)), + } + + +@pytest.mark.asyncio +async def test_tuning_baseline_waits_for_a_complete_db_model_census(monkeypatch): + prisma_client = MagicMock() + monkeypatch.setattr(ps.proxy_config, "_get_models_from_db", AsyncMock(return_value=None)) + + result = await ProxyStartupEvent.enforce_heuristic_v1_tuning_baseline( + prisma_client=prisma_client, + llm_router=None, + limit=1, + ) + + assert result is None + prisma_client.db.litellm_config.find_unique.assert_not_called() + + # --------------------------------------------------------------------------- # _initialize_slack_alerting_jobs — spend-report pod locking (issue #14809) # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 2babfe432f3..e79448d0620 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -9,6 +9,7 @@ Pins covered: from __future__ import annotations import json +import logging import os import re from types import SimpleNamespace @@ -19,6 +20,7 @@ import pytest import litellm from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.proxy_server import ( ProxyConfig, _is_remote_module_url, @@ -1633,6 +1635,32 @@ async def test_ProxyConfig_load_config_minimal_yaml(tmp_path, monkeypatch): } +@pytest.mark.asyncio +@pytest.mark.parametrize("setting", ["true", "false", "null", "'true'", None]) +async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monkeypatch, caplog, setting): + config_file = tmp_path / "budget.yaml" + flag = f" disable_budget_reservation: {setting}\n" if setting is not None else "" + config_file.write_text( + "model_list: []\nlitellm_settings: {}\ngeneral_settings:\n" + " master_key: null\n" + flag + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr("litellm.constants.budget_reservation_disabled_info_emitted", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + config = ProxyConfig() + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + for _ in range(3): + await config.load_config(router=None, config_file_path=str(config_file)) + + records = [ + record for record in caplog.records + if "disable_budget_reservation is enabled" in record.message + ] + assert [record.levelno for record in records] == ([logging.INFO] if setting == "true" else []) + + @pytest.mark.asyncio async def test_ProxyConfig_load_config_resolves_router_settings_plugins(tmp_path, monkeypatch): """Regression: router_settings.plugins dotted-path strings must be resolved to @@ -2401,6 +2429,111 @@ def test_ProxyConfig__add_deployment_resolves_env_refs_on_arbitrary_field(monkey assert deployment.litellm_params.some_future_field == "resolved-custom-value" +@pytest.mark.parametrize( + "stored_drop_params", + ["true", "os.environ/DROP_PARAMS_FLAG"], +) +def test_ProxyConfig__add_deployment_turns_stored_drop_params_string_into_bool(monkeypatch, stored_drop_params): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") + monkeypatch.setenv("DROP_PARAMS_FLAG", "true") + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + db_model = SimpleNamespace( + model_id="model-1", + model_name="gpt-5-nano", + model_info={"id": "model-1"}, + litellm_params={ + "model": encrypt_value_helper(value="openai/gpt-5-nano"), + "drop_params": encrypt_value_helper(value=stored_drop_params), + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model]) + deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"] + + assert added == 1 + assert deployment.litellm_params.drop_params is True + + +def test_ProxyConfig__add_deployment_keeps_loading_rows_after_a_non_flag_drop_params(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + + def db_model(model_id, drop_params): + return SimpleNamespace( + model_id=model_id, + model_name="gpt-5-nano", + model_info={"id": model_id}, + litellm_params={ + "model": encrypt_value_helper(value="openai/gpt-5-nano"), + "drop_params": encrypt_value_helper(value=drop_params), + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model("bad-row", 2), db_model("good-after", "true")]) + deployments = [call.kwargs["deployment"] for call in fake_router.upsert_deployment.call_args_list] + + assert added == 2 + assert [d.litellm_params.drop_params for d in deployments] == [None, True] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("configured, expected", [("true", True), ("false", False)]) +async def test_ProxyConfig_load_config_turns_litellm_settings_drop_params_string_into_bool( + tmp_path, monkeypatch, configured, expected +): + f = tmp_path / "c.yaml" + f.write_text(f'model_list: []\nlitellm_settings:\n drop_params: "{configured}"\n') + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "drop_params", not expected) + + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert litellm.drop_params is expected + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_resolves_a_litellm_settings_drop_params_env_ref(tmp_path, monkeypatch): + f = tmp_path / "c.yaml" + f.write_text("model_list: []\nlitellm_settings:\n drop_params: os.environ/DROP_PARAMS_FROM_ENV\n") + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setenv("DROP_PARAMS_FROM_ENV", "true") + monkeypatch.setattr(litellm, "drop_params", False) + + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert litellm.drop_params is True + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_warns_and_turns_off_a_non_flag_litellm_settings_drop_params( + tmp_path, monkeypatch, caplog +): + f = tmp_path / "c.yaml" + f.write_text("model_list: []\nlitellm_settings:\n drop_params: ture\n") + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "drop_params", True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert litellm.drop_params is False + assert "litellm_settings.drop_params='ture' is not a flag value, treating it as off" in caplog.text + + # --------------------------------------------------------------------------- # ProxyConfig.decrypt_model_list_from_db # --------------------------------------------------------------------------- @@ -2912,6 +3045,129 @@ async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch fake_router.update_settings.assert_called_once_with(routing_strategy="latency-based-routing") +def _stub_add_deployment_collaborators( + monkeypatch: pytest.MonkeyPatch, pc: ProxyConfig, fake_prisma: MagicMock +) -> None: + from litellm.proxy import proxy_server + + fake_router = MagicMock() + fake_router.get_model_list = MagicMock(return_value=[]) + + async def fake_get_config(*args: object, **kwargs: object) -> dict[str, object]: + return {} + + monkeypatch.setattr(litellm, "credential_list", []) + monkeypatch.setattr(pc, "get_config", fake_get_config) + monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", AsyncMock()) + monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock()) + monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + monkeypatch.setattr(proxy_server, "master_key", "sk-master") + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "proxy_config", pc) + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + + +def _encrypted_credential_row(credential_name: str, api_key: str) -> dict[str, object]: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + return { + "credential_name": credential_name, + "credential_values": {"api_key": encrypt_value_helper(api_key, new_encryption_key="sk-master")}, + "credential_info": {"custom_llm_provider": "openai"}, + } + + +def _fake_prisma_with_encrypted_credential(credential_name: str, api_key: str) -> MagicMock: + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.find_many = AsyncMock( + return_value=[_encrypted_credential_row(credential_name, api_key)] + ) + return fake_prisma + + +@pytest.mark.asyncio +async def test_ProxyConfig_add_deployment_loads_db_credentials_before_reconciling_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy import proxy_server + from litellm.utils import load_credentials_from_list + + pc = ProxyConfig() + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + monkeypatch.setattr(proxy_server, "general_settings", {}) + installed = MagicMock() + + async def read_models_while_a_credential_lands(prisma_client: object) -> list[MagicMock]: + fake_prisma.db.litellm_credentialstable.find_many.return_value = [ + _encrypted_credential_row("openai-cred", "sk-from-db") + ] + return [MagicMock()] + + async def install_models(new_models: object, proxy_logging_obj: object) -> None: + installed(credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(pc, "_get_models_from_db", read_models_while_a_credential_lands) + monkeypatch.setattr(pc, "_update_llm_router", install_models) + + await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) + + installed.assert_called_once_with(credential={"api_key": "sk-from-db"}) + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"} + request_kwargs = {"litellm_credential_name": "openai-cred"} + load_credentials_from_list(request_kwargs) + assert request_kwargs == {"litellm_credential_name": "openai-cred", "api_key": "sk-from-db"} + + +@pytest.mark.asyncio +async def test_ProxyConfig_add_deployment_loads_db_credentials_even_when_models_are_not_db_objects( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy import proxy_server + + pc = ProxyConfig() + fake_prisma = _fake_prisma_with_encrypted_credential("openai-cred", "sk-from-db") + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["mcp"]}) + models_fetch = AsyncMock(return_value=[]) + monkeypatch.setattr(pc, "_get_models_from_db", models_fetch) + + await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) + + models_fetch.assert_not_awaited() + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"} + + +@pytest.mark.asyncio +async def test_ProxyConfig_get_credentials_reads_from_writer_not_replica(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + pc = ProxyConfig() + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + writer_inner.litellm_credentialstable.find_many = AsyncMock( + return_value=[_encrypted_credential_row("openai-cred", "sk-from-writer")] + ) + reader_inner.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + fake_prisma = MagicMock() + fake_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + + await pc.get_credentials(prisma_client=fake_prisma) + + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-writer"} + reader_inner.litellm_credentialstable.find_many.assert_not_awaited() + + # --------------------------------------------------------------------------- # ProxyConfig._add_general_settings_from_db_config # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index dcb63b8ca82..2d9c1bd8b46 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -13,9 +13,12 @@ Routes covered: from __future__ import annotations +import asyncio import json from unittest.mock import AsyncMock, MagicMock +import pytest + from .conftest import VOLATILE_KEYS, normalize @@ -229,10 +232,7 @@ def test_config_update_no_db_error(client, auth_as, monkeypatch): json={"general_settings": {"alerting": ["slack"]}}, ) assert response.status_code != 200 - assert ( - "db" in str(response.json()).lower() - or "connect" in str(response.json()).lower() - ) + assert "db" in str(response.json()).lower() or "connect" in str(response.json()).lower() # --------------------------------------------------------------------------- @@ -273,9 +273,7 @@ def test_config_field_update_happy_admin(client, auth_as, mock_prisma, monkeypat } -def test_config_field_update_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_update_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin cannot update config fields — returns 400 with not-allowed detail (handler uses 400 for the auth gate, not 403).""" from litellm.proxy import proxy_server as ps @@ -335,9 +333,7 @@ def test_config_field_info_happy_admin(client, auth_as, mock_prisma, monkeypatch monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert response.status_code == 200 assert normalize(response.json()) == { "field_name": "max_parallel_requests", @@ -345,9 +341,7 @@ def test_config_field_info_happy_admin(client, auth_as, mock_prisma, monkeypatch } -def test_config_field_info_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin (INTERNAL_USER) is denied — admin-view gate fires.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -356,9 +350,7 @@ def test_config_field_info_non_admin_rejected( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.INTERNAL_USER): - response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert response.status_code == 400 assert "error" in response.json().get("detail", {}) @@ -375,16 +367,12 @@ def test_config_field_info_field_not_in_db(client, auth_as, mock_prisma, monkeyp monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert response.status_code == 400 assert "not in DB" in response.json().get("detail", {}).get("error", "") -def test_config_field_info_redacts_nested_secret_for_view_only_admin( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_redacts_nested_secret_for_view_only_admin(client, auth_as, mock_prisma, monkeypatch): """A view-only admin reading a structured field must not receive nested credentials. database_args carries aws_web_identity_token (a DynamoDB role-assumption credential); it must come back redacted while non-secret @@ -405,9 +393,7 @@ def test_config_field_info_redacts_nested_secret_for_view_only_admin( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): - response = client.get( - "/config/field/info", params={"field_name": "database_args"} - ) + response = client.get("/config/field/info", params={"field_name": "database_args"}) assert response.status_code == 200 value = response.json()["field_value"] assert value["aws_web_identity_token"] == "REDACTED" @@ -415,9 +401,7 @@ def test_config_field_info_redacts_nested_secret_for_view_only_admin( assert value["user_table_name"] == "LiteLLM_UserTable" -def test_config_field_info_full_admin_sees_nested_secret( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_full_admin_sees_nested_secret(client, auth_as, mock_prisma, monkeypatch): """The redaction must not over-redact for a full PROXY_ADMIN, who needs the real nested value to populate the edit form.""" from litellm.proxy import proxy_server as ps @@ -435,18 +419,14 @@ def test_config_field_info_full_admin_sees_nested_secret( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/field/info", params={"field_name": "database_args"} - ) + response = client.get("/config/field/info", params={"field_name": "database_args"}) assert response.status_code == 200 value = response.json()["field_value"] assert value["aws_web_identity_token"] == "sk-super-secret-token" assert value["region_name"] == "us-east-1" -def test_config_field_info_redacts_top_level_scalar_for_view_only( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_redacts_top_level_scalar_for_view_only(client, auth_as, mock_prisma, monkeypatch): """The top-level scalar branch must also redact for a view-only admin. database_url carries DB credentials and is not caught by the name masker, so it is in the explicit secret set.""" @@ -460,9 +440,7 @@ def test_config_field_info_redacts_top_level_scalar_for_view_only( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): - response = client.get( - "/config/field/info", params={"field_name": "database_url"} - ) + response = client.get("/config/field/info", params={"field_name": "database_url"}) assert response.status_code == 200 assert response.json()["field_value"] == "REDACTED" @@ -476,17 +454,12 @@ def test_redact_general_setting_value_recurses_list_of_dicts(): {"path": "/foo", "headers": {"Authorization": "Bearer sk-x"}}, {"path": "/bar", "client_secret": "sk-y"}, ] - redacted = ps._redact_general_setting_value( - "some_list_field", value, is_full_admin=False - ) + redacted = ps._redact_general_setting_value("some_list_field", value, is_full_admin=False) assert redacted[0]["headers"]["Authorization"] == "REDACTED" assert redacted[0]["path"] == "/foo" assert redacted[1]["client_secret"] == "REDACTED" assert redacted[1]["path"] == "/bar" - assert ( - ps._redact_general_setting_value("some_list_field", value, is_full_admin=True) - == value - ) + assert ps._redact_general_setting_value("some_list_field", value, is_full_admin=True) == value def test_redact_secret_values_in_obj_fails_closed_at_max_depth(): @@ -504,22 +477,16 @@ def test_redact_secret_values_in_obj_fails_closed_at_max_depth(): for _ in range(ps._REDACT_SECRET_MAX_DEPTH + 2): nested = {"wrap": nested} - out = ps._redact_general_setting_value( - "some_struct_field", nested, is_full_admin=False - ) + out = ps._redact_general_setting_value("some_struct_field", nested, is_full_admin=False) # the secret must not survive anywhere in the returned tree assert "sk-leak-bottom" not in repr(out) # full admin is unaffected by the cap — the value comes back untouched - admin_out = ps._redact_general_setting_value( - "some_struct_field", nested, is_full_admin=True - ) + admin_out = ps._redact_general_setting_value("some_struct_field", nested, is_full_admin=True) assert admin_out is nested -def test_config_list_redacts_pass_through_secret_for_view_only( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_list_redacts_pass_through_secret_for_view_only(client, auth_as, mock_prisma, monkeypatch): """/config/list must not leak pass_through_endpoints upstream credentials to a view-only admin. pass_through_endpoints is a known secret-bearing field, so a non-admin gets it redacted; a full admin still sees it.""" @@ -546,24 +513,16 @@ def test_config_list_redacts_pass_through_secret_for_view_only( ) def _pass_through_value(body): - return next( - entry["field_value"] - for entry in body - if entry["field_name"] == "pass_through_endpoints" - ) + return next(entry["field_value"] for entry in body if entry["field_name"] == "pass_through_endpoints") with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): - view_resp = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + view_resp = client.get("/config/list", params={"config_type": "general_settings"}) assert view_resp.status_code == 200 assert "sk-UPSTREAM-SECRET" not in view_resp.text assert _pass_through_value(view_resp.json()) == "REDACTED" with auth_as(LitellmUserRoles.PROXY_ADMIN): - admin_resp = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + admin_resp = client.get("/config/list", params={"config_type": "general_settings"}) assert admin_resp.status_code == 200 admin_value = _pass_through_value(admin_resp.json()) assert admin_value[0]["headers"]["Authorization"] == "Bearer sk-UPSTREAM-SECRET" @@ -587,9 +546,7 @@ def test_config_list_happy_admin(client, auth_as, mock_prisma, monkeypatch): monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + response = client.get("/config/list", params={"config_type": "general_settings"}) assert response.status_code == 200 body = response.json() assert isinstance(body, list) @@ -695,9 +652,7 @@ def test_config_list_non_admin_rejected(client, auth_as, mock_prisma, monkeypatc monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.INTERNAL_USER): - response = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + response = client.get("/config/list", params={"config_type": "general_settings"}) assert response.status_code == 400 assert "role" in response.json().get("detail", {}).get("error", "").lower() @@ -710,9 +665,7 @@ def test_config_list_no_db_error(client, auth_as, monkeypatch): monkeypatch.setattr(ps, "prisma_client", None) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + response = client.get("/config/list", params={"config_type": "general_settings"}) assert response.status_code == 400 assert "error" in response.json().get("detail", {}) @@ -756,9 +709,7 @@ def test_config_field_delete_happy_admin(client, auth_as, mock_prisma, monkeypat } -def test_config_field_delete_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_delete_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin caller hits the 400 not-allowed branch with role in detail.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -778,9 +729,7 @@ def test_config_field_delete_non_admin_rejected( assert "role" in response.json().get("detail", {}).get("error", "").lower() -def test_config_field_delete_field_not_in_config( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_delete_field_not_in_config(client, auth_as, mock_prisma, monkeypatch): """If there is no general_settings row at all, returns 400 'not in config'.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -825,9 +774,7 @@ def test_config_callback_delete_happy_admin(client, auth_as, mock_prisma, monkey monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.post( - "/config/callback/delete", json={"callback_name": "langfuse"} - ) + response = client.post("/config/callback/delete", json={"callback_name": "langfuse"}) assert response.status_code == 200 # `deleted_at` is an ISO timestamp generated at request time — extend # the volatile set just for this assertion so dict-equality still works. @@ -840,9 +787,7 @@ def test_config_callback_delete_happy_admin(client, auth_as, mock_prisma, monkey } -def test_config_callback_delete_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_callback_delete_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin caller is rejected with 400 not-allowed.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -852,9 +797,7 @@ def test_config_callback_delete_non_admin_rejected( monkeypatch.setattr(ps, "store_model_in_db", True) with auth_as(LitellmUserRoles.INTERNAL_USER): - response = client.post( - "/config/callback/delete", json={"callback_name": "langfuse"} - ) + response = client.post("/config/callback/delete", json={"callback_name": "langfuse"}) assert response.status_code == 400 assert "role" in response.json().get("detail", {}).get("error", "").lower() @@ -869,22 +812,15 @@ def test_config_callback_delete_not_found(client, auth_as, mock_prisma, monkeypa monkeypatch.setattr(ps, "store_model_in_db", True) fake_proxy_config = MagicMock() - fake_proxy_config.get_config = AsyncMock( - return_value={"litellm_settings": {"success_callback": ["slack"]}} - ) + fake_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {"success_callback": ["slack"]}}) monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.post( - "/config/callback/delete", json={"callback_name": "langfuse"} - ) + response = client.post("/config/callback/delete", json={"callback_name": "langfuse"}) # The handler re-raises HTTPException(404) verbatim (only generic # `Exception` becomes a 500 ProxyException), so pin 404 strictly. assert response.status_code == 404 - assert ( - "langfuse" in str(response.json()).lower() - or "not found" in str(response.json()).lower() - ) + assert "langfuse" in str(response.json()).lower() or "not found" in str(response.json()).lower() # --------------------------------------------------------------------------- @@ -948,10 +884,7 @@ def test_get_config_callbacks_internal_error(client, auth_as, mock_prisma, monke with auth_as(LitellmUserRoles.PROXY_ADMIN): response = client.get("/get/config/callbacks") assert response.status_code >= 400 - assert ( - "boom" in str(response.json()).lower() - or "error" in str(response.json()).lower() - ) + assert "boom" in str(response.json()).lower() or "error" in str(response.json()).lower() _CALLBACK_ENV_FIXTURE = { @@ -985,14 +918,10 @@ def _install_callbacks_config(monkeypatch, mock_prisma): def _callback_variables(body: dict, name: str) -> dict: - return next( - cb["variables"] for cb in body["callbacks"] if cb["name"] == name - ) + return next(cb["variables"] for cb in body["callbacks"] if cb["name"] == name) -def test_get_config_callbacks_redacts_secret_env_vars_for_view_only_admin( - client, auth_as, mock_prisma, monkeypatch -): +def test_get_config_callbacks_redacts_secret_env_vars_for_view_only_admin(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy._types import LitellmUserRoles _install_callbacks_config(monkeypatch, mock_prisma) @@ -1024,9 +953,7 @@ def test_get_config_callbacks_redacts_secret_env_vars_for_view_only_admin( assert otel_vars["OTEL_ENDPOINT"] == _CALLBACK_ENV_FIXTURE["OTEL_ENDPOINT"] -def test_get_config_callbacks_full_admin_still_sees_secret_env_vars( - client, auth_as, mock_prisma, monkeypatch -): +def test_get_config_callbacks_full_admin_still_sees_secret_env_vars(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy._types import LitellmUserRoles _install_callbacks_config(monkeypatch, mock_prisma) @@ -1047,9 +974,7 @@ def test_get_config_callbacks_full_admin_still_sees_secret_env_vars( assert otel_vars["OTEL_HEADERS"] == _CALLBACK_ENV_FIXTURE["OTEL_HEADERS"] -def test_get_config_callbacks_redacts_slack_webhook_urls_for_view_only_admin( - client, auth_as, mock_prisma, monkeypatch -): +def test_get_config_callbacks_redacts_slack_webhook_urls_for_view_only_admin(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -1170,6 +1095,413 @@ def test_get_config_callbacks_redacts_email_alerting_vars_for_view_only_admin( assert admin_email["SMTP_HOST"] == "smtp.resend.com" +def test_get_config_callbacks_appends_runtime_only_callbacks(client, auth_as, mock_prisma, monkeypatch): + """LIT-5281: a YAML callback that the DB callback list replaced in the merged config still runs, so it must + show up as a read_only row next to the editable DB-configured one.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": ["langfuse"]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.langsmith import LangsmithLogger + from litellm.integrations.opentelemetry import OpenTelemetry + + monkeypatch.setattr(litellm, "success_callback", ["langfuse", LangsmithLogger()]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", [OpenTelemetry()]) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse", "success", False), + ("langsmith", "success", True), + ("otel", "success_and_failure", True), + ] + + +def test_get_config_callbacks_accepts_scalar_and_null_yaml_callbacks(client, auth_as, mock_prisma, monkeypatch): + """`success_callback: langfuse` (a YAML scalar) is one configured callback, not eight single-letter ones, and a + `callbacks: null` key contributes nothing.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": "langfuse", "failure_callback": None, "callbacks": None}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.langsmith import LangsmithLogger + + monkeypatch.setattr(litellm, "success_callback", ["langfuse", LangsmithLogger()]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse", "success", False), + ("langsmith", "success", True), + ] + + +def test_get_config_callbacks_deduplicates_configured_and_runtime(client, auth_as, mock_prisma, monkeypatch): + """A configured callback shows once as editable, whether the runtime holds its string or an initialized instance + (arize initializes an ArizeLogger, logfire a bare OpenTelemetry that only its class identifies).""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": ["langfuse", "arize", "logfire"]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.arize.arize import ArizeLogger + from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig + + arize_logger = ArizeLogger(config=OpenTelemetryConfig(exporter="console"), callback_name="arize") + monkeypatch.setattr(litellm, "success_callback", ["langfuse", arize_logger, OpenTelemetry()]) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse", "success", False), + ("arize", "success", False), + ("logfire", "success", False), + ] + + +def test_get_config_callbacks_keeps_yaml_otel_family_callbacks_next_to_configured_one( + client, auth_as, mock_prisma, monkeypatch +): + """LIT-5281: arize, weave_otel and langfuse_otel all initialize OpenTelemetry subclasses. Saving one of them + from the dashboard replaces the YAML `callbacks` list, so the YAML siblings keep running and must stay listed + under their own names instead of being hidden as duplicates of the configured OTel callback.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"callbacks": ["langfuse_otel"]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.arize.arize import ArizeLogger + from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger + from litellm.integrations.langsmith import LangsmithLogger + from litellm.integrations.opentelemetry import OpenTelemetryConfig + from litellm.integrations.weave.weave_otel import WeaveOtelLogger + + console_config = OpenTelemetryConfig(exporter="console") + monkeypatch.setattr(litellm, "success_callback", [LangsmithLogger()]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr( + litellm, + "callbacks", + [ + ArizeLogger(config=console_config, callback_name="arize"), + WeaveOtelLogger(config=console_config), + LangfuseOtelLogger(config=console_config, callback_name="langfuse_otel"), + ], + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse_otel", "success_and_failure", False), + ("arize", "success_and_failure", True), + ("langsmith", "success", True), + ("weave_otel", "success_and_failure", True), + ] + + +def _dotted_path_test_function(*args, **kwargs): + pass + + +@pytest.mark.parametrize("handler_kind", ["instance", "function"]) +@pytest.mark.parametrize( + "config_key,expected_type", + [ + ("success_callback", "success"), + ("failure_callback", "failure"), + ("callbacks", "success_and_failure"), + ], +) +def test_get_config_callbacks_deduplicates_dotted_path_callback( + client, auth_as, mock_prisma, monkeypatch, config_key, expected_type, handler_kind +): + """A dotted-path callback stays a single editable row instead of duplicating under its class or function name.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + import litellm + from litellm.integrations.custom_logger import CustomLogger + + class _DottedPathTestHandler(CustomLogger): + pass + + dotted_handler = _DottedPathTestHandler() if handler_kind == "instance" else _dotted_path_test_function + dotted_path = f"{__name__}.dotted_handler" + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {config_key: [dotted_path]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + monkeypatch.setattr(litellm, "callbacks", [dotted_handler]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + + assert response.status_code == 200 + callbacks = response.json()["callbacks"] + assert [(callback["name"], callback["type"], callback.get("read_only", False)) for callback in callbacks] == [ + (dotted_path, expected_type, False) + ] + + +def test_get_config_callbacks_lists_dict_shaped_config_callbacks(client, auth_as, mock_prisma, monkeypatch): + """Dict-shaped success_callback config values list their keys as editable rows.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": {"langsmith": {"batch_size": 1}}}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_callbacks_by_type", + MagicMock(return_value={"success": ["langsmith"], "failure": [], "success_and_failure": []}), + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + + assert response.status_code == 200 + callbacks = response.json()["callbacks"] + assert [(callback["name"], callback.get("read_only", False)) for callback in callbacks] == [("langsmith", False)] + + +def test_get_config_callbacks_excludes_internal_runtime_callbacks(client, auth_as, mock_prisma, monkeypatch): + """Proxy infrastructure callbacks are excluded from callback inventory.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": []}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + import litellm + from litellm._service_logger import ServiceLogging + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_logger import CustomLogger + from litellm.integrations.langsmith import LangsmithLogger + from litellm.integrations.s3_v2 import S3Logger + from litellm.integrations.sqs import SQSLogger + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import VectorStorePreCallHook + from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter + from litellm.router import Router + + class _InventoryTestGuardrail(CustomGuardrail): + pass + + class _UserCodeLogger(CustomLogger): + pass + + def user_code_function(*args, **kwargs): + pass + + async def build_aws_loggers() -> tuple[S3Logger, SQSLogger]: + return S3Logger(s3_bucket_name="inventory-bucket"), SQSLogger(sqs_queue_url="https://sqs.example/inventory") + + s3_logger, sqs_logger = asyncio.run(build_aws_loggers()) + router = Router(model_list=[]) + monkeypatch.setattr(litellm, "input_callback", []) + monkeypatch.setattr( + litellm, "success_callback", [LangsmithLogger(), s3_logger, router.sync_deployment_callback_on_success] + ) + monkeypatch.setattr(litellm, "_async_success_callback", [sqs_logger, router.deployment_callback_on_success]) + monkeypatch.setattr(litellm, "failure_callback", [user_code_function]) + monkeypatch.setattr(litellm, "_async_failure_callback", [router.async_deployment_callback_on_failure]) + monkeypatch.setattr( + litellm, + "callbacks", + [ + _PROXY_MaxBudgetLimiter(), + _PROXY_LiteLLMManagedFiles(internal_usage_cache=MagicMock(), prisma_client=MagicMock()), + ServiceLogging(), + VectorStorePreCallHook(), + _InventoryTestGuardrail(guardrail_name="inventory-test-guardrail"), + _UserCodeLogger(), + ], + ) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + assert "cache" in litellm.success_callback and "cache" in litellm._async_success_callback + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + + assert response.status_code == 200 + assert [ + (callback["name"], callback["type"], callback["read_only"]) for callback in response.json()["callbacks"] + ] == [ + ("_UserCodeLogger", "success_and_failure", True), + ("langsmith", "success", True), + ("s3", "success", True), + ("sqs", "success", True), + ("user_code_function", "failure", True), + ] + + +def test_get_config_callbacks_redacts_runtime_only_row_secrets_for_view_only_admin( + client, auth_as, mock_prisma, monkeypatch +): + """Runtime-only callback rows are subject to the same redaction gate as configured.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": []}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", ["otel"]) + + with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + body = response.json() + + callbacks = body["callbacks"] + otel_cb = next((cb for cb in callbacks if cb["name"] == "otel"), None) + assert otel_cb is not None + assert otel_cb["type"] == "success_and_failure" + assert otel_cb["read_only"] is True + assert otel_cb["variables"]["OTEL_HEADERS"] == "REDACTED" + assert otel_cb["variables"]["OTEL_ENDPOINT"] == _CALLBACK_ENV_FIXTURE["OTEL_ENDPOINT"] + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + admin_response = client.get("/get/config/callbacks") + assert admin_response.status_code == 200 + admin_body = admin_response.json() + admin_otel = next((cb for cb in admin_body["callbacks"] if cb["name"] == "otel"), None) + assert admin_otel is not None + assert admin_otel["variables"]["OTEL_HEADERS"] == _CALLBACK_ENV_FIXTURE["OTEL_HEADERS"] + + # --------------------------------------------------------------------------- # GET /config/yaml # --------------------------------------------------------------------------- @@ -1183,9 +1515,7 @@ def test_config_yaml_returns_demo_payload(client, auth_as): response = client.request("GET", "/config/yaml", json={}) shape = { "status": response.status_code, - "media_type_yaml": response.headers.get("content-type", "").startswith( - "application/json" - ), + "media_type_yaml": response.headers.get("content-type", "").startswith("application/json"), "has_body": len(response.content) > 0, } assert shape == { diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py index b75ee1caccf..40bd66ea91c 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py @@ -11,15 +11,21 @@ Routes covered: from __future__ import annotations import json +from pathlib import Path from unittest.mock import AsyncMock, MagicMock +from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map_provenance + from .conftest import VOLATILE_KEYS, normalize # Some response bodies include a "timestamp" — extend the volatile set so # dict-equality assertions remain stable. _VOLATILE = VOLATILE_KEYS | frozenset({"timestamp"}) +_SERVED_ETAG = 'W/"cost-map-etag"' +_ROOT_COST_MAP = Path(__file__).resolve().parents[4] / "model_prices_and_context_window.json" + # --------------------------------------------------------------------------- # Helpers @@ -83,6 +89,7 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma): "status": "success", "models_count": 2, "timestamp": "", + **get_model_cost_map_provenance(), } assert table.upsert.await_count == 1 update_payload = table.upsert.await_args.kwargs["data"]["update"] @@ -90,6 +97,54 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma): assert update_payload["reload_revision"] == {"increment": 1} +def test_reload_model_cost_map_surfaces_the_blob_id_of_the_bytes_served_on_every_status_surface( + client, auth_as, monkeypatch, mock_prisma +): + import httpx + + import litellm + from litellm.litellm_core_utils.get_model_cost_map import git_blob_id + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _attach_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + body = _ROOT_COST_MAP.read_bytes() + expected = {"source_revision": git_blob_id(body), "etag": _SERVED_ETAG} + served = httpx.Response(200, headers={"ETag": _SERVED_ETAG}, content=body) + monkeypatch.setattr( + "litellm.litellm_core_utils.get_model_cost_map._default_reload_client", + lambda: httpx.AsyncClient(transport=httpx.MockTransport(lambda request: served)), + ) + monkeypatch.setattr("litellm.add_known_models", lambda model_cost_map=None: None) + monkeypatch.setattr("litellm.model_cost", {}, raising=False) + + async def _fake_invalidate(name): + return None + + monkeypatch.setattr(ps, "invalidate_config_param", _fake_invalidate) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + reload_response = client.post("/reload/model_cost_map") + source_response = client.get("/model/cost_map/source") + status_response = client.get("/schedule/model_cost_map_reload/status") + public_response = client.get("/public/litellm_model_cost_map") + + assert reload_response.status_code == 200 + reload_body = reload_response.json() + assert {key: reload_body[key] for key in expected} == expected + assert source_response.status_code == 200 + source_body = source_response.json() + assert {key: source_body[key] for key in expected} == expected + assert source_body["source"] == "remote" + assert status_response.status_code == 200 + assert {key: status_response.json()[key] for key in expected} == expected + assert public_response.status_code == 200 + assert "gpt-4o" in public_response.json() + assert reload_body["models_count"] == len(litellm.model_cost) + + def test_reload_model_cost_map_fetch_failure_502_keeps_map( client, auth_as, monkeypatch, mock_prisma ): @@ -270,7 +325,7 @@ def test_cancel_model_cost_map_reload_no_db_500(client, auth_as, monkeypatch): def test_get_model_cost_map_reload_status_no_db_not_scheduled( client, auth_as, monkeypatch ): - """No prisma client → returns the not-scheduled shape (4 keys, all-null).""" + """No prisma client → returns the not-scheduled shape (all-null) plus the cost map provenance.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -283,6 +338,7 @@ def test_get_model_cost_map_reload_status_no_db_not_scheduled( "interval_hours": None, "last_run": None, "next_run": None, + **get_model_cost_map_provenance(), } @@ -309,6 +365,7 @@ def test_get_model_cost_map_reload_status_scheduled( "interval_hours": 12, "last_run": None, "next_run": None, + **get_model_cost_map_provenance(), } @@ -337,6 +394,7 @@ def test_get_model_cost_map_reload_status_reports_persisted_last_run( "interval_hours": 6, "last_run": "2024-01-01T06:00:00+00:00", "next_run": "2024-01-01T12:00:00+00:00", + **get_model_cost_map_provenance(), } @@ -365,6 +423,7 @@ def test_get_model_cost_map_reload_status_no_config_not_scheduled( "interval_hours": None, "last_run": None, "next_run": None, + **get_model_cost_map_provenance(), } @@ -391,6 +450,8 @@ def test_get_model_cost_map_source_happy(client, auth_as, monkeypatch): "url": "https://example.invalid/cost_map.json", "is_env_forced": False, "fallback_reason": None, + "loaded_at": "2026-09-07T01:02:03+00:00", + **get_model_cost_map_provenance(), } monkeypatch.setattr( "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map_source_info", @@ -406,6 +467,8 @@ def test_get_model_cost_map_source_happy(client, auth_as, monkeypatch): "url": "https://example.invalid/cost_map.json", "is_env_forced": False, "fallback_reason": None, + "loaded_at": "2026-09-07T01:02:03+00:00", + **get_model_cost_map_provenance(), "model_count": 3, } diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py index bc6106a06f8..d31d952a03e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_models.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py @@ -17,6 +17,7 @@ import litellm from litellm.proxy import proxy_server from litellm.proxy import utils as proxy_utils from litellm.proxy.utils import create_model_info_response +from litellm.types.router import DeploymentModelListingInfo from .conftest import normalize # type: ignore[import-not-found] @@ -166,13 +167,16 @@ def test_anthropic_format_carries_router_configured_token_limits(client, auth_as built from another entry's lookup shows up as the wrong numbers.""" def _configured(model_name): - return (300000, 32000) if model_name == "gpt-4" else (500000, 4096) + max_input, max_output = (300000, 32000) if model_name == "gpt-4" else (500000, 4096) + return DeploymentModelListingInfo( + cost_map_keys=(model_name,), max_input_tokens=max_input, max_output_tokens=max_output + ) def _cost_map_lookup(model_id): max_input, max_output = (200000, 64000) if model_id == "gpt-4" else (100000, 8000) return {"max_input_tokens": max_input, "max_output_tokens": max_output, "mode": "chat"} - patched_models.get_configured_token_limits = MagicMock(side_effect=_configured) + patched_models.get_model_listing_info = MagicMock(side_effect=_configured) def _resolved(**kwargs): return create_model_info_response(**kwargs, get_model_info=_cost_map_lookup) @@ -343,3 +347,53 @@ def test_anthropic_format_returns_public_team_model_name( assert response.status_code == 200 assert [m["id"] for m in response.json()["data"]] == ["gpt-4-team"] assert internal_name not in response.text + + +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +@pytest.mark.parametrize( + "caller_headers", + [ + {"anthropic-version": "2023-06-01", "user-agent": "claude-code/2.1.267"}, + {"anthropic-version": "2023-06-01", "user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, + {"anthropic-version": "2023-06-01", "x-gateway-client": "claude-code"}, + ], +) +def test_anthropic_format_lists_claude_code_view_ids_for_claude_code( + client, auth_as, patched_models, monkeypatch, path, caller_headers +): + """Claude Code drops every id without claude/anthropic in it and reads [1m] as its 1M marker, so for Claude + Code (its discovery fetch's own user agent, its SDK's, or the gateway-client header a launcher sends) every + group is listed under a Claude-shaped id with the marker where the window reaches 1M; the display name stays + the served name.""" + + def _create_model_info_response(model_id, provider="openai", **kwargs): + if model_id != "claude-sonnet": + return _stub_model_info_response(model_id=model_id, provider=provider) + return {**_stub_model_info_response(model_id=model_id, provider=provider), "max_input_tokens": 1000000} + + patched_models.model_group_alias = {} + patched_models.has_model_id.return_value = False + patched_models.get_candidate_model_ids_for_route.side_effect = lambda name, team_id=None: frozenset({name}) if name in ("gpt-4", "claude-sonnet") else frozenset() + monkeypatch.setattr(proxy_utils, "create_model_info_response", _create_model_info_response) + + with auth_as(): + response = client.get(path, headers=caller_headers) + + assert response.status_code == 200 + body = response.json() + assert [(m["id"], m["display_name"]) for m in body["data"]] == [ + ("claude-router-6770742d34", "gpt-4"), + ("claude-sonnet[1m]", "claude-sonnet"), + ] + assert (body["first_id"], body["last_id"]) == ("claude-router-6770742d34", "claude-sonnet[1m]") + assert [row["source_model"] for row in body["data"]] == ["gpt-4", "claude-sonnet"] + + +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +def test_anthropic_format_keeps_served_ids_for_other_anthropic_clients(client, auth_as, patched_models, path): + """An Anthropic SDK asking for the vendor shape gets the served ids: the view is Claude Code's alone.""" + with auth_as(): + response = client.get(path, headers={"anthropic-version": "2023-06-01", "user-agent": "anthropic-sdk-python/0.40"}) + + assert response.status_code == 200 + assert [m["id"] for m in response.json()["data"]] == ["gpt-4", "claude-sonnet"] diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 86e97a334df..2f47736a398 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -4,11 +4,12 @@ Pins covered: - ``get_current_spend`` - ``increment_spend_counters`` - ``_reconcile_budget_reservation_for_counter_update`` -- ``_increment_end_user_and_tag_spend_counters`` -- ``_increment_org_spend_counter`` -- ``_init_and_increment_unreserved_spend_counter`` -- ``_init_and_increment_spend_counter`` -- ``_init_and_increment_window_spend_counter`` +- ``_prepare_end_user_and_tag_spend_increments`` +- ``_prepare_org_spend_increment`` +- ``_prepare_unreserved_spend_counter_increment`` +- ``_prepare_spend_counter_increment`` +- ``_prepare_window_spend_counter_increment`` +- ``_apply_spend_counter_increments`` - ``_ensure_spend_counter_initialized`` - ``_get_source_cache_base_spend`` - ``_ensure_window_spend_counter_initialized`` @@ -48,9 +49,7 @@ def _make_spend_counter_cache( cache.in_memory_cache.delete_cache = MagicMock() if with_redis: cache.redis_cache = MagicMock() - cache.redis_cache.async_get_cache = AsyncMock( - return_value=redis_get_value, side_effect=redis_get_side_effect - ) + cache.redis_cache.async_get_cache = AsyncMock(return_value=redis_get_value, side_effect=redis_get_side_effect) cache.redis_cache.async_increment = AsyncMock( return_value=redis_increment_value, side_effect=redis_increment_side_effect, @@ -58,6 +57,8 @@ def _make_spend_counter_cache( cache.redis_cache.async_delete_cache = AsyncMock() cache.redis_cache.async_set_cache = AsyncMock() cache.redis_cache.async_set_max = AsyncMock() + cache.redis_cache.async_increment_pipeline = AsyncMock(return_value=None) + cache.redis_cache.get_ttl = MagicMock(return_value=None) else: cache.redis_cache = None cache.async_increment_cache = AsyncMock(return_value=redis_increment_value) @@ -70,9 +71,7 @@ def _make_spend_counter_cache( def _make_user_api_key_cache(get_value=None, get_side_effect=None): cache = MagicMock() - cache.async_get_cache = AsyncMock( - return_value=get_value, side_effect=get_side_effect - ) + cache.async_get_cache = AsyncMock(return_value=get_value, side_effect=get_side_effect) cache.async_set_cache_pipeline = AsyncMock() return cache @@ -109,9 +108,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(monkeypatch ) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=99.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=99.0) assert result == 17.0 @@ -136,9 +133,7 @@ async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch # the stale counter is repaired up to the authoritative DB value via a # monotonic set-max so other workers read the corrected total, and a # concurrent increment cannot be clobbered - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key="spend:key:abc", value=12.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:key:abc", value=12.0) @pytest.mark.asyncio @@ -169,9 +164,7 @@ async def test_get_current_spend_no_floor_without_max_budget(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0) assert result == 2.0 assert from_db.await_count == 0 @@ -210,12 +203,8 @@ async def test_get_current_spend_floor_caches_db_read(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - first = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) - second = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) + first = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) + second = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) assert first == 12.0 assert second == 12.0 @@ -336,9 +325,7 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): assert result == 15.0 assert wfsl.await_count == 1 - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) def _make_window_spend_prisma(row=None, spend_logs_total=0.0): @@ -379,9 +366,7 @@ async def test_get_current_spend_floors_window_against_maintained_row(monkeypatc assert result == 15.0 fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) @pytest.mark.asyncio @@ -393,9 +378,7 @@ async def test_get_current_spend_floors_window_against_logs_when_row_stale(monke window_start = datetime(2026, 1, 8, tzinfo=timezone.utc) fake_prisma = _make_window_spend_prisma( - row=SimpleNamespace( - window_start=window_start - timedelta(days=7), spend=999.0 - ), + row=SimpleNamespace(window_start=window_start - timedelta(days=7), spend=999.0), spend_logs_total=15.0, ) fake_cache = _make_spend_counter_cache(redis_get_value=2.0) @@ -423,21 +406,13 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat rather than admitted on an unverifiable budget.""" from fastapi import HTTPException - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) with pytest.raises(HTTPException) as exc: - await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert exc.value.status_code == 503 @@ -445,18 +420,12 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat async def test_get_current_spend_fail_closed_off_admits_when_unverifiable(monkeypatch): """Default (flag off): an unverifiable read keeps the existing behavior and admits using the cached fallback — no new rejection.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "general_settings", {}) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -466,13 +435,9 @@ async def test_get_current_spend_fail_closed_admits_when_redis_verified(monkeypa authoritative, so an under-budget request is admitted normally.""" fake_cache = _make_spend_counter_cache(redis_get_value=1.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -481,16 +446,10 @@ async def test_get_current_spend_fail_closed_allows_authoritative_fallback(monke """End-user/tag callers pass fallback_authoritative=True (their spend is loaded fresh from the DB in auth), so fail-closed does not reject them even when the counter path is unreadable.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) result = await ps.get_current_spend( counter_key="spend:end_user:e1", @@ -508,9 +467,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa re-checks the authoritative DB and enforces against it.""" fake_cache = _make_spend_counter_cache(redis_get_value=0.00001) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) from_db = AsyncMock(return_value=0.5) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -532,9 +489,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa @pytest.mark.asyncio async def test_increment_spend_counters_increments_all_buckets(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -543,9 +498,7 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): async def _fake_coalesced(**kwargs): return None - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced)) await ps.increment_spend_counters( token="hashed-tok", @@ -554,25 +507,36 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): response_cost=5.0, ) + pipeline = fake_cache.redis_cache.async_increment_pipeline + pipeline.assert_awaited_once() + increment_list = pipeline.await_args.kwargs["increment_list"] + assert {op["key"] for op in increment_list} == { + "spend:key:hashed-tok", + "spend:team:t1", + "spend:team_member:u1:t1", + "spend:user:u1", + } + assert all(op["increment_value"] == 5.0 for op in increment_list) observed = { "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, + "pipeline_calls": pipeline.await_count, "user_cache_used": fake_user_cache.async_get_cache.called, } assert normalize(observed) == { - "redis_increment_called": True, - "increment_calls": 4, + "redis_increment_called": False, + "pipeline_calls": 1, "user_cache_used": True, } class _ConcurrencyProbe: - """Stand-in for redis_cache.async_increment that pins concurrency. + """Stand-in for redis_cache.async_get_cache that pins concurrency. - Each call registers itself as in-flight and blocks on ``release`` until the - test lets it proceed. ``all_arrived`` fires once ``expected`` distinct scope - increments are simultaneously suspended here, which can only happen if the - per-scope increments are gathered rather than awaited one after another. + Each warm-check read registers itself as in-flight and blocks on ``release`` + until the test lets it proceed. ``all_arrived`` fires once ``expected`` + distinct scope warm-checks are simultaneously suspended here, which can only + happen if the per-scope prepares are gathered rather than awaited one after + another. """ def __init__(self, expected_concurrency: int): @@ -581,36 +545,45 @@ class _ConcurrencyProbe: self.max_in_flight = 0 self.all_arrived = asyncio.Event() self.release = asyncio.Event() - self.values: dict[str, float] = {} + self.keys: list[str] = [] - async def async_increment(self, *, key, value, refresh_ttl=True): + async def async_get_cache(self, *, key, **kwargs): self.in_flight += 1 self.max_in_flight = max(self.max_in_flight, self.in_flight) + self.keys.append(key) if self.in_flight >= self.expected: self.all_arrived.set() if not self.release.is_set(): await self.release.wait() self.in_flight -= 1 - self.values[key] = self.values.get(key, 0.0) + value - return self.values[key] + return 1.0 @pytest.mark.asyncio async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): """The six independent scopes (key, team, team_member, user, end_user+tags, - org) must be incremented concurrently. The probe only fires once all six are - suspended in async_increment at the same time, which is impossible if the + org) must prepare their increments concurrently. The probe only fires once + all eight warm-check reads (one per counter: 6 scopes + 2 tags) are + suspended in async_get_cache at the same time, which is impossible if the awaits are chained sequentially.""" - probe = _ConcurrencyProbe(expected_concurrency=6) - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = probe.async_increment + probe = _ConcurrencyProbe(expected_concurrency=8) + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_get_cache = probe.async_get_cache + recorded: dict[str, float] = {} + + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] + results.append(recorded[op["key"]]) + return results + + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) task = asyncio.create_task( ps.increment_spend_counters( @@ -630,16 +603,16 @@ async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): probe.release.set() await task pytest.fail( - "scope increments did not run concurrently; sequential awaits " - f"detected (peak in-flight was {probe.max_in_flight}, expected 6)" + "scope prepares did not run concurrently; sequential awaits " + f"detected (peak in-flight was {probe.max_in_flight}, expected 8)" ) - assert probe.in_flight == 6 - assert probe.max_in_flight == 6 + assert probe.in_flight == 8 + assert probe.max_in_flight == 8 probe.release.set() await task - assert probe.values == { + assert recorded == { "spend:key:hashed-tok": 5.0, "spend:team:t1": 5.0, "spend:team_member:u1:t1": 5.0, @@ -659,26 +632,25 @@ async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch) import litellm.proxy.spend_tracking.budget_reservation as br reserved = {"spend:key:hashed-tok", "spend:org:org1"} - monkeypatch.setattr( - br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved)) - ) + monkeypatch.setattr(br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved))) monkeypatch.setattr(br, "reconcile_budget_reservation", AsyncMock()) recorded: dict[str, float] = {} - async def _record_increment(*, key, value, refresh_ttl=True): - recorded[key] = recorded.get(key, 0.0) + value - return recorded[key] + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] + results.append(recorded[op["key"]]) + return results fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = _record_increment + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) reservation = {"finalized": False} await ps.increment_spend_counters( @@ -708,27 +680,46 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ ): """A failure in one scope must propagate to the caller (so it can invalidate reserved counters) while every other scope still settles rather than being - left as an orphaned background task, and the reservation is not finalized.""" - recorded: dict[str, float] = {} + left as an orphaned background task, and the reservation is not finalized. + The surviving scopes' increments are still applied in the single pipeline: + dropping them would under-count spend, the unsafe direction for budget + enforcement.""" + warmed_keys: list[str] = [] - async def _increment(*, key, value, refresh_ttl=True): + async def _warm_check(*, key, **kwargs): + warmed_keys.append(key) if key == "spend:team:t1": - raise RuntimeError("redis increment failed") - recorded[key] = recorded.get(key, 0.0) + value - return recorded[key] + raise RuntimeError("redis get failed") + return 1.0 - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = _increment + async def _reseed_fails(*, counter_key, **kwargs): + if counter_key == "spend:team:t1": + raise RuntimeError("reseed failed") + + applied: dict[str, float] = {} + + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + applied[op["key"]] = op["increment_value"] + results.append(op["increment_value"]) + return results + + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_get_cache = AsyncMock(side_effect=_warm_check) + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) + ps.SpendCounterReseed, + "coalesced", + AsyncMock(side_effect=_reseed_fails), ) reservation = {"finalized": False} - with pytest.raises(RuntimeError, match="redis increment failed"): + with pytest.raises(RuntimeError, match="reseed failed"): await ps.increment_spend_counters( token="hashed-tok", team_id="t1", @@ -741,7 +732,19 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ ) assert reservation["finalized"] is False - assert recorded == { + # every sibling scope settled (its warm-check ran) before the error propagated + assert set(warmed_keys) == { + "spend:key:hashed-tok", + "spend:team:t1", + "spend:team_member:u1:t1", + "spend:user:u1", + "spend:end_user:eu1", + "spend:tag:a", + "spend:org:org1", + } + # the surviving scopes' increments were still applied, in one pipeline call + fake_cache.redis_cache.async_increment_pipeline.assert_awaited_once() + assert applied == { "spend:key:hashed-tok": 5.0, "spend:team_member:u1:t1": 5.0, "spend:user:u1": 5.0, @@ -749,6 +752,7 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ "spend:tag:a": 5.0, "spend:org:org1": 5.0, } + fake_cache.redis_cache.async_increment.assert_not_awaited() @pytest.mark.asyncio @@ -772,6 +776,108 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( assert reservation == {"finalized": True} assert fake_cache.redis_cache.async_increment.called is False + fake_cache.redis_cache.async_increment_pipeline.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( + monkeypatch, +): + """Every scope's increment must go out in a single async_increment_pipeline + call, not one INCRBYFLOAT round-trip per scope.""" + counter_cache = ps.DualCache() + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm + + async def _pipeline(increment_list, **_): + return [1.5] * len(increment_list) + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=_pipeline) + fake_redis.async_increment = AsyncMock() + fake_redis.get_ttl = MagicMock(return_value=None) + counter_cache.redis_cache = fake_redis + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", None) + + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + fake_redis.async_increment_pipeline.assert_awaited_once() + assert fake_redis.async_increment.await_count == 0 + increment_list = fake_redis.async_increment_pipeline.await_args.kwargs["increment_list"] + expected_keys = { + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", + } + assert {op["key"] for op in increment_list} == expected_keys + assert all(op["increment_value"] == 0.5 for op in increment_list) + for key in expected_keys: + assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters( + monkeypatch, +): + """A failing pipeline must invalidate every pending counter so the next + request reseeds from the DB (which already holds this request's cost) + instead of trusting a value the write may have partially applied.""" + from redis.exceptions import MaxConnectionsError + + counter_cache = ps.DualCache() + pending_keys = ( + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", + ) + for key in pending_keys: + counter_cache.in_memory_cache.set_cache(key=key, value=1.0) + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm + fake_redis.async_increment_pipeline = AsyncMock(side_effect=MaxConnectionsError()) + fake_redis.async_increment = AsyncMock() + fake_redis.async_delete_cache = AsyncMock() + fake_redis.get_ttl = MagicMock(return_value=None) + counter_cache.redis_cache = fake_redis + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", None) + + with pytest.raises(MaxConnectionsError): + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + assert fake_redis.async_increment.await_count == 0 + deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list} + assert deleted_keys == set(pending_keys) + for key in pending_keys: + assert counter_cache.in_memory_cache.get_cache(key=key) is None # --------------------------------------------------------------------------- @@ -781,9 +887,7 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( @pytest.mark.asyncio async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set_when_none(): - result = await ps._reconcile_budget_reservation_for_counter_update( - budget_reservation=None, response_cost=1.0 - ) + result = await ps._reconcile_budget_reservation_for_counter_update(budget_reservation=None, response_cost=1.0) assert result == set() @@ -818,179 +922,151 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat # --------------------------------------------------------------------------- -# _increment_end_user_and_tag_spend_counters +# _prepare_end_user_and_tag_spend_increments # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_increment_end_user_and_tag_spend_counters_increments_each_unique_tag( +async def test_prepare_end_user_and_tag_spend_increments_returns_each_unique_tag( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=3.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=1.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - await ps._increment_end_user_and_tag_spend_counters( + pending = await ps._prepare_end_user_and_tag_spend_increments( end_user_id="eu1", tags=["a", "b", "a", "", None], response_cost=3.0, reserved_counter_keys=set(), ) - observed = { - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count, - "called": fake_cache.redis_cache.async_increment.called, - } - assert normalize(observed) == { - "increment_calls": 3, - "in_memory_set_calls": 3, - "called": True, + assert {item.counter_key for item in pending} == { + "spend:end_user:eu1", + "spend:tag:a", + "spend:tag:b", } + assert all(item.increment == 3.0 for item in pending) @pytest.mark.asyncio -async def test_increment_end_user_and_tag_spend_counters_no_end_user_no_tags_invalid_input_noop( +async def test_prepare_end_user_and_tag_spend_increments_no_end_user_no_tags_invalid_input_noop( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._increment_end_user_and_tag_spend_counters( + pending = await ps._prepare_end_user_and_tag_spend_increments( end_user_id=None, tags=None, response_cost=1.0, reserved_counter_keys=set(), ) + assert pending == () assert fake_cache.redis_cache.async_increment.called is False # --------------------------------------------------------------------------- -# _increment_org_spend_counter +# _prepare_org_spend_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_increment_org_spend_counter_increments_when_org_present(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=10.0 - ) +async def test_prepare_org_spend_increment_returns_pending_when_org_present(monkeypatch): + fake_cache = _make_spend_counter_cache(redis_get_value=1.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - await ps._increment_org_spend_counter( + pending = await ps._prepare_org_spend_increment( org_id="org-1", response_cost=10.0, reserved_counter_keys=set(), ) - observed = { - "increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "counter_key_arg": fake_cache.redis_cache.async_increment.call_args.kwargs[ - "key" - ], - } - assert normalize(observed) == { - "increment_called": True, - "increment_calls": 1, - "counter_key_arg": "spend:org:org-1", - } + assert len(pending) == 1 + assert pending[0].counter_key == "spend:org:org-1" + assert pending[0].increment == 10.0 @pytest.mark.asyncio -async def test_increment_org_spend_counter_no_org_is_noop_invalid_id(monkeypatch): +async def test_prepare_org_spend_increment_no_org_is_noop_invalid_id(monkeypatch): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._increment_org_spend_counter( + pending = await ps._prepare_org_spend_increment( org_id=None, response_cost=1.0, reserved_counter_keys=set(), ) + assert pending == () assert fake_cache.redis_cache.async_increment.called is False # --------------------------------------------------------------------------- -# _init_and_increment_unreserved_spend_counter +# _prepare_unreserved_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_unreserved_spend_counter_skips_reserved_keys( +async def test_prepare_unreserved_spend_counter_increment_skips_reserved_keys( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._init_and_increment_unreserved_spend_counter( + pending = await ps._prepare_unreserved_spend_counter_increment( counter_key="spend:tag:x", source_cache_key="tag:x", increment=1.0, reserved_counter_keys={"spend:tag:x"}, ) + assert pending is None assert fake_cache.redis_cache.async_increment.called is False @pytest.mark.asyncio -async def test_init_and_increment_unreserved_spend_counter_proceeds_when_not_reserved( +async def test_prepare_unreserved_spend_counter_increment_proceeds_when_not_reserved( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=2.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None) fake_user_cache = _make_user_api_key_cache() + reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) - await ps._init_and_increment_unreserved_spend_counter( + pending = await ps._prepare_unreserved_spend_counter_increment( counter_key="spend:tag:y", source_cache_key="tag:y", increment=2.0, reserved_counter_keys=set(), ) - observed = { - "increment_called": fake_cache.redis_cache.async_increment.called, - "redis_get_called": fake_cache.redis_cache.async_get_cache.called, - "reseed_consulted": True, - } - assert observed == { - "increment_called": True, - "redis_get_called": True, - "reseed_consulted": True, - } + assert pending is not None + assert pending.counter_key == "spend:tag:y" + assert pending.increment == 2.0 + assert fake_cache.redis_cache.async_get_cache.called is True + assert reseed.called is True # --------------------------------------------------------------------------- -# _init_and_increment_spend_counter +# _prepare_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=11.0, redis_increment_value=14.0 - ) +async def test_prepare_spend_counter_increment_warm_cache_skips_reseed(monkeypatch): + fake_cache = _make_spend_counter_cache(redis_get_value=11.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -998,12 +1074,14 @@ async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypa reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) - await ps._init_and_increment_spend_counter( + pending = await ps._prepare_spend_counter_increment( counter_key="spend:key:k", source_cache_key="k", increment=3.0, ) + assert pending.counter_key == "spend:key:k" + assert pending.increment == 3.0 observed = { "reseed_called": reseed.called, "increment_called": fake_cache.redis_cache.async_increment.called, @@ -1011,23 +1089,21 @@ async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypa } assert normalize(observed) == { "reseed_called": False, - "increment_called": True, + "increment_called": False, "in_memory_seeded_from_redis": True, } # --------------------------------------------------------------------------- -# _init_and_increment_window_spend_counter +# _prepare_window_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_window_spend_counter_increments_when_initialized( +async def test_prepare_window_spend_counter_increment_returns_pending_when_initialized( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=0.0, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=0.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( @@ -1036,7 +1112,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ AsyncMock(return_value=0.0), ) - await ps._init_and_increment_window_spend_counter( + pending = await ps._prepare_window_spend_counter_increment( counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", @@ -1045,26 +1121,19 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ increment=5.0, ) - observed = { - "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count, - } - assert normalize(observed) == { - "redis_increment_called": True, - "increment_calls": 1, - "in_memory_set_calls": 2, - } + assert pending is not None + assert pending.counter_key == "spend:key:k:window:1d" + assert pending.increment == 5.0 @pytest.mark.asyncio -async def test_init_and_increment_window_spend_counter_missing_window_start_invalid_skips( +async def test_prepare_window_spend_counter_increment_missing_window_start_invalid_skips( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._init_and_increment_window_spend_counter( + pending = await ps._prepare_window_spend_counter_increment( counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", @@ -1073,9 +1142,58 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva increment=5.0, ) + assert pending is None assert fake_cache.redis_cache.async_increment.called is False +# --------------------------------------------------------------------------- +# _apply_spend_counter_increments +# --------------------------------------------------------------------------- + + +def _two_pending_increments() -> tuple[ps._PendingSpendIncrement, ...]: + return ( + ps._PendingSpendIncrement(counter_key="spend:key:k", increment=1.5), + ps._PendingSpendIncrement(counter_key="spend:team:t", increment=1.5), + ) + + +@pytest.mark.asyncio +async def test_apply_spend_counter_increments_open_breaker_invalidates_and_returns(monkeypatch): + """An open Redis circuit breaker is a known, already-logged state, not a per-request tracking failure. + + Re-raising the refusal sent every request through the cost callback's error path, which + logged an ERROR and fired the failed-tracking alert once per request for the whole outage. + """ + from litellm.caching.redis_cache import RedisCircuitBreakerOpenError + + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_increment_pipeline = AsyncMock( + side_effect=RedisCircuitBreakerOpenError("Redis circuit breaker is open") + ) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + + await ps._apply_spend_counter_increments(_two_pending_increments()) + + deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) + assert deleted_keys == ["spend:key:k", "spend:team:t"] + fake_cache.in_memory_cache.set_cache.assert_not_called() + + +@pytest.mark.asyncio +async def test_apply_spend_counter_increments_other_redis_error_invalidates_and_raises(monkeypatch): + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("redis down")) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + + with pytest.raises(ConnectionError, match="redis down"): + await ps._apply_spend_counter_increments(_two_pending_increments()) + + deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) + assert deleted_keys == ["spend:key:k", "spend:team:t"] + fake_cache.in_memory_cache.set_cache.assert_not_called() + + # --------------------------------------------------------------------------- # _ensure_spend_counter_initialized # --------------------------------------------------------------------------- @@ -1114,16 +1232,12 @@ async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source( async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=7.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=7.0) fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0}) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) await ps._ensure_spend_counter_initialized( counter_key="spend:user:u", @@ -1163,9 +1277,7 @@ async def test_get_source_cache_base_spend_reads_first_hit_from_list(monkeypatch fake_user_cache.async_get_cache = AsyncMock(side_effect=_get) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) - result = await ps._get_source_cache_base_spend( - source_cache_key=["miss", "hit-obj", "miss2"] - ) + result = await ps._get_source_cache_base_spend(source_cache_key=["miss", "hit-obj", "miss2"]) observed = { "result": result, @@ -1294,9 +1406,7 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey fake_cache = _make_spend_counter_cache(redis_increment_value=44.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=4.0 - ) + result = await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=4.0) observed = { "result": result, @@ -1314,15 +1424,11 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_increment_side_effect=RuntimeError("incr fail") - ) + fake_cache = _make_spend_counter_cache(redis_increment_side_effect=RuntimeError("incr fail")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) with pytest.raises(RuntimeError): - await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=1.0 - ) + await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=1.0) assert fake_cache.in_memory_cache.delete_cache.called is True assert fake_cache.redis_cache.async_delete_cache.called is True @@ -1343,9 +1449,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) observed = { "in_memory_delete_called": fake_cache.in_memory_cache.delete_cache.called, "redis_delete_called": fake_cache.redis_cache.async_delete_cache.called, - "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs[ - "key" - ], + "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs["key"], } assert normalize(observed) == { "in_memory_delete_called": True, @@ -1357,9 +1461,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) @pytest.mark.asyncio async def test_invalidate_spend_counter_swallows_redis_failure_no_raise(monkeypatch): fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_delete_cache = AsyncMock( - side_effect=RuntimeError("redis down") - ) + fake_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) await ps._invalidate_spend_counter(counter_key="spend:key:k") diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 0fb9b1a6d88..bc346874e0d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -27,6 +27,7 @@ from litellm.proxy.proxy_server import ( _get_proxy_model_info, _translate_model_name_for_response, ) +from litellm.types.router import DeploymentModelListingInfo def _team_row() -> dict: @@ -880,7 +881,7 @@ async def test_v1_models_translates_team_model_for_access_group_key(monkeypatch) router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] @@ -922,7 +923,7 @@ async def test_v1_models_keeps_internal_names_when_public_name_flag_disabled( router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] @@ -957,7 +958,7 @@ async def test_v1_models_translates_team_model_with_metadata(monkeypatch): router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] router.get_model_group_info.return_value = None @@ -1003,7 +1004,7 @@ async def test_v1_models_metadata_fallbacks_use_internal_routing_key(monkeypatch router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] # Fallbacks are keyed on the internal routing name, as the router stores them. @@ -1060,7 +1061,7 @@ async def test_v1_models_metadata_does_not_leak_other_team_fallbacks(monkeypatch router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_x, team_y] router.get_model_list.return_value = [team_x, team_y] router.fallbacks = [ @@ -1089,6 +1090,101 @@ async def test_v1_models_metadata_does_not_leak_other_team_fallbacks(monkeypatch ] +@pytest.mark.asyncio +async def test_v1_models_team_alias_inherits_token_limits_and_chat_mode(monkeypatch): + team_dep = { + "model_name": "model_name_teamX_terra_uuid", + "litellm_params": {"model": "azure/gpt-4.1"}, + "model_info": { + "id": "id-terra", + "team_id": "teamX", + "team_public_model_name": "GPT Terra", + "access_groups": ["grp-a"], + "mode": "chat", + "max_input_tokens": 876000, + "max_output_tokens": 128000, + }, + } + router = MagicMock() + router.get_model_names.return_value = ["model_name_teamX_terra_uuid"] + router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_terra_uuid"]} + router.get_fully_blocked_model_names.return_value = set() + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_keys=("azure/gpt-4.1",), + max_input_tokens=876000, + max_output_tokens=128000, + ) + router.get_configured_mode.return_value = "chat" + router.model_list = [team_dep] + router.get_model_list.return_value = [team_dep] + router.get_model_group_info.return_value = None + + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "general_settings", {"use_team_public_model_name": True}) + + key = UserAPIKeyAuth(user_id="user", api_key="***", models=["grp-a"], team_models=[]) + response = await ps.model_list(user_api_key_dict=key, include_metadata=True) + + assert response["data"] == [ + { + "id": "GPT Terra", + "object": "model", + "created": 1677610602, + "owned_by": "openai", + "mode": "chat", + "max_input_tokens": 876000, + "max_output_tokens": 128000, + "metadata": {"fallbacks": []}, + } + ] + + +@pytest.mark.asyncio +async def test_v1_models_team_image_alias_inherits_image_generation_mode(monkeypatch): + team_dep = { + "model_name": "model_name_teamX_image_uuid", + "litellm_params": {"model": "openai/gpt-image-1"}, + "model_info": { + "id": "id-image", + "team_id": "teamX", + "team_public_model_name": "image", + "access_groups": ["grp-a"], + "mode": "image_generation", + }, + } + router = MagicMock() + router.get_model_names.return_value = ["model_name_teamX_image_uuid"] + router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_image_uuid"]} + router.get_fully_blocked_model_names.return_value = set() + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_keys=("openai/gpt-image-1",), + max_input_tokens=None, + max_output_tokens=None, + ) + router.get_configured_mode.return_value = "image_generation" + router.model_list = [team_dep] + router.get_model_list.return_value = [team_dep] + router.get_model_group_info.return_value = None + + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "general_settings", {"use_team_public_model_name": True}) + + key = UserAPIKeyAuth(user_id="user", api_key="***", models=["grp-a"], team_models=[]) + response = await ps.model_list(user_api_key_dict=key) + + assert response["data"] == [ + { + "id": "image", + "object": "model", + "created": 1677610602, + "owned_by": "openai", + "mode": "image_generation", + } + ] + + def test_translate_team_model_names_for_listing_swaps_and_dedupes(): """Internal team routing keys -> public name; sibling deployments sharing a public name collapse to one entry (order preserved); globals untouched.""" @@ -1315,7 +1411,7 @@ def test_translate_team_model_names_for_listing_respects_legacy_flag(): def _public_named_router(*team_rows: dict) -> MagicMock: router = MagicMock() router.get_model_list.return_value = list(team_rows) - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None return router @@ -1427,8 +1523,13 @@ async def test_retrieve_model_by_public_name_returns_200(monkeypatch): team_row = _team_row() router = _public_named_router(team_row) deployment = MagicMock() - deployment.litellm_params.model = "azure/gpt-5.2-low-rpm-testing" + deployment.litellm_params.model = "azure/gpt-4.1" router.get_deployment_by_model_group_name.return_value = deployment + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_keys=("azure/gpt-4.1",), + max_input_tokens=16384, + max_output_tokens=4096, + ) monkeypatch.setattr(ps, "llm_router", router) monkeypatch.setattr(ps, "general_settings", {}) @@ -1445,6 +1546,9 @@ async def test_retrieve_model_by_public_name_returns_200(monkeypatch): resp = await ps.model_info(model_id="team-claude-sonnet", user_api_key_dict=key) assert resp["id"] == "team-claude-sonnet" + assert resp.get("mode") == "chat" + assert resp.get("max_input_tokens") == 16384 + assert resp.get("max_output_tokens") == 4096 # lookup happened by the internal routing key, not the public name router.get_deployment_by_model_group_name.assert_called_once_with( "model_name_team-abc-123_4a6b8" diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 4a19ad3541c..0d82ed778f5 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1,5 +1,6 @@ import re from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -327,6 +328,85 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False +def test_chatgpt_provider_fields(): + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/providers/fields") + assert response.status_code == 200 + providers = response.json() + + chatgpt = next((p for p in providers if p["provider"] == "CHATGPT"), None) + assert chatgpt is not None, "ChatGPT provider entry not found" + + assert chatgpt["provider_display_name"] == "ChatGPT Subscription" + assert chatgpt["litellm_provider"] == LlmProviders.CHATGPT.value + assert chatgpt["default_model_placeholder"].startswith("chatgpt/") + assert chatgpt["credential_fields"] == [] + + +ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( + { + "a2a", + "a2a_agent", + "amazon_nova", + "apertis", + "aws_polly", + "black_forest_labs", + "charity_engine", + "chutes", + "darkbloom", + "gdc", + "helicone", + "inception", + "langflow", + "langgraph", + "libertai", + "litellm_agent", + "manus", + "meta", + "modelscope", + "mongodb", + "nano-gpt", + "neosantara", + "parasail", + "pinstripes", + "poe", + "publicai", + "ragflow", + "reducto", + "s3_vectors", + "sagemaker_nova", + "scaleway", + "stability", + "synthetic", + "tencent", + "tensormesh", + "text-completion-inception", + "valkey", + "xiaomi_mimo", + "zai", + } +) + + +def test_every_backend_provider_is_listed_in_add_model_or_frozen_as_unlisted(): + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/providers/fields") + assert response.status_code == 200 + listed = {p["litellm_provider"] for p in response.json()} + + unlisted = {provider.value for provider in LlmProviders} - listed + assert unlisted == ADD_MODEL_UNLISTED_PROVIDERS, ( + "Add Model dropdown drift: give the new provider an entry in provider_create_fields.json " + "rather than adding it to ADD_MODEL_UNLISTED_PROVIDERS" + ) + + def test_google_ai_studio_provider_fields_expose_api_base(): """The Google AI Studio (gemini) credential form must let admins set a custom api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted diff --git a/tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py b/tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py index 66eeb3cef34..82f2ef097aa 100644 --- a/tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py +++ b/tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py @@ -6,10 +6,13 @@ Tests for LiteLLM proxy realtime WebRTC HTTP endpoints: import json import time +from collections.abc import Awaitable +from typing import Protocol from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient @@ -159,17 +162,27 @@ def mock_route_request_realtime_calls(): return _mock_route +class AddLitellmDataToRequest(Protocol): + def __call__(self, data: dict[str, object], **kwargs: object) -> Awaitable[dict[str, object]]: ... + + +class PreCallHook(Protocol): + def __call__( + self, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str + ) -> Awaitable[dict[str, object]]: ... + + @pytest.fixture -def mock_add_litellm_data(): - async def _mock(data, **kwargs): +def mock_add_litellm_data() -> AddLitellmDataToRequest: + async def _mock(data: dict[str, object], **kwargs: object) -> dict[str, object]: return data return _mock @pytest.fixture -def mock_pre_call_hook(): - async def _mock(user_api_key_dict, data, call_type): +def mock_pre_call_hook() -> PreCallHook: + async def _mock(user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> dict[str, object]: return data return _mock @@ -1199,3 +1212,78 @@ async def test_transcription_sessions_wraps_route_exception( assert "Model not allowed" in response.text finally: proxy_app.dependency_overrides.pop(user_api_key_auth, None) + + +def test_realtime_calls_upstream_rejection_answers_an_openai_typed_error( + proxy_app: FastAPI, + mock_add_litellm_data: AddLitellmDataToRequest, + mock_pre_call_hook: PreCallHook, + monkeypatch: pytest.MonkeyPatch, +): + """A bare HTTPException carries no type or param, so the tail used to ship the + literal string "None" in both fields of the error the browser client reads.""" + token_payload = _encode_realtime_token_payload( + ephemeral_key="fake_upstream_epk", + model_id="gpt-4o-realtime-preview", + user_id=None, + team_id=None, + expires_at=int(time.time()) + 3600, + ) + encrypted_token = encrypt_value_helper(token_payload) + + async def failing_route_request(*args: object, **kwargs: object) -> None: + raise HTTPException( + status_code=404, + detail={"error": "realtime: Invalid model name passed in model=gpt-4o-realtime-preview"}, + ) + + proxy_logging = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + proxy_logging.post_call_failure_hook = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.route_request", failing_route_request) + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", mock_add_litellm_data) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + + response = TestClient(proxy_app).post( + "/v1/realtime/calls", + headers={"Authorization": f"Bearer {encrypted_token}"}, + content=b"v=0\r\no=- 0 0 IN IP4 0.0.0.0\r\ns=-\r\n", + ) + + assert response.status_code == 404 + assert (response.json()["error"]["type"], response.json()["error"]["param"]) == ("invalid_request_error", None) + + +def test_transcription_sessions_rejection_answers_an_openai_typed_error( + proxy_app: FastAPI, + mock_add_litellm_data: AddLitellmDataToRequest, + mock_pre_call_hook: PreCallHook, + monkeypatch: pytest.MonkeyPatch, +): + """A model the router cannot serve surfaces as a bare HTTPException, which this tail + used to relabel with the literal string "None" for both type and param.""" + + async def failing_route_request(*args: object, **kwargs: object) -> None: + raise HTTPException( + status_code=400, + detail={"error": "realtime: Invalid model name passed in model=no-such-transcribe"}, + ) + + proxy_logging = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + proxy_logging.post_call_failure_hook = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.route_request", failing_route_request) + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", mock_add_litellm_data) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + proxy_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="test-user") + try: + response = TestClient(proxy_app, raise_server_exceptions=False).post( + "/v1/realtime/transcription_sessions", + headers={"Authorization": "Bearer sk-test-master-key"}, + json={"input_audio_transcription": {"model": "no-such-transcribe"}}, + ) + finally: + proxy_app.dependency_overrides.pop(user_api_key_auth, None) + + assert response.status_code == 400 + assert (response.json()["error"]["type"], response.json()["error"]["param"]) == ("invalid_request_error", None) diff --git a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py index 9f11ff6f20d..ea858e04e0f 100644 --- a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py @@ -6,11 +6,11 @@ import json from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import Request, Response +from fastapi import HTTPException, Request, Response import litellm.proxy.common_request_processing as common_request_processing_mod import litellm.proxy.proxy_server as proxy_server_mod -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.rerank_endpoints.endpoints import rerank from litellm.types.utils import RerankResponse @@ -118,3 +118,55 @@ async def test_rerank_omits_detailed_timing_headers_when_disabled(): fastapi_response = await _call_rerank() assert "x-litellm-timing-llm-api-ms" not in fastapi_response.headers + + +async def _rerank_failure( + failure: Exception, *, raised_before_routing: bool, monkeypatch: pytest.MonkeyPatch +) -> ProxyException: + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock( + side_effect=failure if raised_before_routing else lambda **kwargs: kwargs["data"] + ) + proxy_logging_obj.post_call_failure_hook = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + async def failing_route_request(**kwargs: object) -> None: + raise failure + + monkeypatch.setattr(proxy_server_mod, "add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr(proxy_server_mod, "route_request", failing_route_request) + monkeypatch.setattr(proxy_server_mod, "proxy_logging_obj", proxy_logging_obj) + monkeypatch.setattr(proxy_server_mod, "llm_router", MagicMock()) + monkeypatch.setattr(proxy_server_mod, "version", "1.2.3") + + with pytest.raises(ProxyException) as raised: + await rerank( + request=_build_request(), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + return raised.value + + +@pytest.mark.asyncio +async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(monkeypatch: pytest.MonkeyPatch): + """A bare HTTPException carries no type or param, so the tail used to ship the + literal string "None" in both fields.""" + failure = HTTPException(status_code=404, detail={"error": "rerank: Invalid model name passed in model=rerank-model"}) + + error = await _rerank_failure(failure, raised_before_routing=False, monkeypatch=monkeypatch) + + assert (error.type, error.param, error.code) == ("invalid_request_error", None, "404") + + +@pytest.mark.asyncio +async def test_a_rejection_raised_before_routing_keeps_its_own_status(monkeypatch: pytest.MonkeyPatch): + """A ProxyException stores its status as the string ``code``, which the tail used to + miss and rewrap as a 500 while keeping the 4xx type and param.""" + rejection = ProxyException(message="session_id is required", type="bad_request_error", param="session_id", code=400) + + error = await _rerank_failure(rejection, raised_before_routing=True, monkeypatch=monkeypatch) + + assert (error.type, error.param, error.code) == ("bad_request_error", "session_id", "400") diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index f65f68812a2..ee062b14b96 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -1,18 +1,37 @@ +import json +import math from typing import Final import pytest +import litellm from litellm.caching import DualCache +from litellm.proxy import proxy_server from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request +from litellm.proxy.spend_tracking.budget_reservation import ( + count_request_input_tokens, + estimate_request_max_cost, + reserve_budget_for_request, +) from litellm.proxy.utils import ProxyLogging +from litellm.router import Router +from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge import token_counter as rust_token_counter +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo TOKEN_COUNTING_ROUTES: Final = ( "/responses/input_tokens", "/v1/responses/input_tokens", "/openai/v1/responses/input_tokens", "/utils/token_counter", + "/v1/messages/count_tokens", + "/v1beta/models/gemini-3.8-flash:countTokens", + "/models/gemini-3.8-flash:countTokens", + "/bedrock/v1/messages/count-tokens", + "/bedrock/model/us.anthropic.claude-sonnet-4-6/count-tokens", + "/vertex_ai/v1/projects/p/locations/us-east5/publishers/anthropic/models/count-tokens:rawPredict", + "/vertex-ai/v1/projects/p/locations/us-east5/publishers/anthropic/models/count-tokens:rawPredict", ) @@ -46,3 +65,267 @@ async def test_non_exempt_llm_route_still_reserves_budget(): assert reservation is not None assert reservation["reserved_cost"] > 0 + + +@pytest.mark.asyncio +async def test_reservation_carries_the_admission_input_token_count(): + reservation: Final = await _reserve("/v1/responses") + expected: Final = litellm.token_counter(model="gpt-4o", text="hello") + + assert reservation is not None + assert expected > 0 + assert reservation["input_tokens"] == expected + + +ANTHROPIC_MESSAGES: Final = [{"role": "user", "content": "hello!!!"}] +COUNT_TOKENS_REQUESTS: Final[tuple[tuple[str, dict[str, object]], ...]] = ( + ("/v1/messages/count_tokens", {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}), + ("/v1beta/models/gemini-3.8-flash:countTokens", {"contents": [{"role": "user", "parts": [{"text": "hello!!!"}]}]}), + ( + "/vertex_ai/v1/projects/p/locations/us-east5/publishers/anthropic/models/count-tokens:rawPredict", + {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}, + ), + ("/bedrock/v1/messages/count-tokens", {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}), +) +TINY_BUDGET_KEY_TOKEN: Final = "hashed-count-tokens-key" + + +@pytest.fixture +def spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> DualCache: + cache: Final = DualCache() + monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", None) + return cache + + +async def _reserve_for_tiny_budget_key(route: str, request_body: dict[str, object]) -> dict[str, object] | None: + return await reserve_budget_for_request( + request_body=request_body, + route=route, + llm_router=None, + valid_token=UserAPIKeyAuth(token=TINY_BUDGET_KEY_TOKEN, max_budget=0.01, spend=0.0), + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("route", "request_body"), COUNT_TOKENS_REQUESTS) +async def test_repeated_token_counting_never_touches_a_tiny_budget( + spend_counter_cache: DualCache, route: str, request_body: dict[str, object] +): + counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" + + assert await _reserve_for_tiny_budget_key(route, request_body) is None + assert await _reserve_for_tiny_budget_key(route, request_body) is None + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is None + + completion: Final = await _reserve_for_tiny_budget_key( + "/v1/messages", {"model": "claude-sonnet-5", "max_tokens": 16, "messages": ANTHROPIC_MESSAGES} + ) + assert completion is not None + reserved_cost: Final = completion["reserved_cost"] + assert isinstance(reserved_cost, float) + assert reserved_cost > 0 + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reserved_cost) + + +BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6" +CONVERSE_BODY: Final = { + "messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}], + "inferenceConfig": {"maxTokens": 5}, +} +INVOKE_BODY: Final = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 5, + "messages": [{"role": "user", "content": "Reply with one word: pong"}], +} + + +def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window(): + converse_cost: Final = estimate_request_max_cost( + request_body=CONVERSE_BODY, + route=f"/bedrock/model/{BEDROCK_SONNET}/converse", + llm_router=None, + input_token_counts={}, + ) + invoke_cost: Final = estimate_request_max_cost( + request_body=INVOKE_BODY, + route=f"/bedrock/model/{BEDROCK_SONNET}/invoke", + llm_router=None, + input_token_counts={}, + ) + assert converse_cost is not None and invoke_cost is not None + assert invoke_cost < converse_cost < 2 * invoke_cost + + +def _tiered_deployment(input_cost_per_token: float) -> Deployment: + return Deployment( + model_name="tiered-group", + litellm_params=LiteLLM_Params(model="dashscope/qwen3-max", api_key="sk-fake"), + model_info=ModelInfo( + id="tiered-deployment", + max_output_tokens=1000, + tiered_pricing=[ + { + "input_cost_per_token": input_cost_per_token, + "output_cost_per_token": input_cost_per_token, + "range": [0, 128000], + } + ], + ), + ) + + +TIERED_BODY: Final = {"model": "tiered-group", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10} + + +def test_repeated_estimates_reuse_cached_model_cost_info() -> None: + router: Final = Router(model_list=[_tiered_deployment(1e-06).model_dump()]) + first: Final = estimate_request_max_cost(request_body=TIERED_BODY, route="/chat/completions", llm_router=router) + hits_before: Final = router.cached_deployment_model_info.cache_info().hits + + second: Final = estimate_request_max_cost(request_body=TIERED_BODY, route="/chat/completions", llm_router=router) + + assert second == first + assert router.cached_deployment_model_info.cache_info().hits == hits_before + 1 + + +def test_deployment_pricing_update_invalidates_cached_estimate() -> None: + router: Final = Router(model_list=[_tiered_deployment(1e-06).model_dump()]) + before: Final = estimate_request_max_cost(request_body=TIERED_BODY, route="/chat/completions", llm_router=router) + assert before is not None + + router.upsert_deployment(_tiered_deployment(1e-03)) + + after: Final = estimate_request_max_cost(request_body=TIERED_BODY, route="/chat/completions", llm_router=router) + assert after is not None + assert math.isclose(after, before * 1000) + + +ANTHROPIC_TOKENIZER_MODEL: Final = "claude-sonnet-4-5-20250929" +RUST_COUNTED_BODY: Final = {"model": ANTHROPIC_TOKENIZER_MODEL, "max_tokens": 16, "messages": ANTHROPIC_MESSAGES} +RUST_INPUT_TOKENS: Final = 4_321 + + +class _FakeDeclined(Exception): + pass + + +class _FakeUpstream(Exception): + pass + + +class _FakeNative: + RustBridgeDeclined = _FakeDeclined + RustUpstreamError = _FakeUpstream + + +class _RecordingCounter: + bodies: Final[list[bytes]] = [] + + def __init__(self, tokenizer_json: str) -> None: + pass + + async def acount_request(self, body: bytes) -> object: + self.bodies.append(body) + return {"model": ANTHROPIC_TOKENIZER_MODEL, "input_tokens": RUST_INPUT_TOKENS} + + +class _DecliningCounter: + def __init__(self, tokenizer_json: str) -> None: + pass + + async def acount_request(self, body: bytes) -> object: + raise _FakeDeclined("unsupported content block") + + +@pytest.fixture +def rust_counter(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) + rust_token_counter._anthropic_counter.cache_clear() + configuration.reset_rust_configuration() + _RecordingCounter.bodies.clear() + yield + rust_token_counter.TOKEN_COUNTER.reset() + rust_token_counter._anthropic_counter.cache_clear() + configuration.reset_rust_configuration() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "request_body"), + ( + ("/v1/messages", RUST_COUNTED_BODY), + ("/v1/chat/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "messages": ANTHROPIC_MESSAGES}), + ("/v1/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "prompt": "hi"}), + ("/v1/responses", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": "hi"}), + ("/v1/embeddings", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": ["hi"]}), + ("/v1/rerank", {"model": ANTHROPIC_TOKENIZER_MODEL, "query": "hi", "documents": ["a"]}), + ), +) +async def test_rust_count_replaces_python_tokenizing_on_every_llm_route( + rust_counter: None, route: str, request_body: dict +) -> None: + litellm.rust(True) + rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter) + raw_body: Final = json.dumps(request_body).encode() + + counts: Final = await count_request_input_tokens( + request_body=request_body, route=route, llm_router=None, raw_body=raw_body + ) + + assert dict(counts) == {ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS} + assert _RecordingCounter.bodies == [raw_body] + + +@pytest.mark.asyncio +async def test_rust_decline_falls_back_to_python_count(rust_counter: None) -> None: + litellm.rust(True) + rust_token_counter.TOKEN_COUNTER.override(_DecliningCounter) + python_counts: Final = await count_request_input_tokens( + request_body=RUST_COUNTED_BODY, route="/v1/messages", llm_router=None + ) + + counts: Final = await count_request_input_tokens( + request_body=RUST_COUNTED_BODY, + route="/v1/messages", + llm_router=None, + raw_body=json.dumps(RUST_COUNTED_BODY).encode(), + ) + + assert dict(counts) == dict(python_counts) + assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS + + +@pytest.mark.asyncio +async def test_disabled_rust_never_sees_the_raw_body(rust_counter: None) -> None: + litellm.rust(False) + rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter) + + counts: Final = await count_request_input_tokens( + request_body=RUST_COUNTED_BODY, + route="/v1/messages", + llm_router=None, + raw_body=json.dumps(RUST_COUNTED_BODY).encode(), + ) + + assert _RecordingCounter.bodies == [] + assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS + + +@pytest.mark.asyncio +async def test_non_anthropic_tokenizer_models_stay_in_python(rust_counter: None) -> None: + litellm.rust(True) + rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter) + body: Final = {"model": "gpt-4o", "messages": ANTHROPIC_MESSAGES} + + counts: Final = await count_request_input_tokens( + request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=json.dumps(body).encode() + ) + + assert _RecordingCounter.bodies == [] + assert counts["gpt-4o"] != RUST_INPUT_TOKENS diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py index c123eeeed36..6165af4920d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py @@ -41,6 +41,15 @@ class _FlakyRedisCache: self._store[key] = float(value) return True + async def async_increment_pipeline(self, increment_list, **kwargs): + results = [] + for op in increment_list: + results.append(await self.async_increment(op["key"], op["increment_value"])) + return results + + def get_ttl(self, **kwargs): + return None + @pytest.mark.asyncio async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failure( diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 7a80319239d..89be341c87b 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -1,21 +1,60 @@ +import asyncio +import time from collections.abc import Sequence +from datetime import datetime, timedelta from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest from prisma.errors import PrismaError +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import ( + SPEND_LOG_KEY_METADATA_CACHE_TTL, + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, + SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, +) from litellm.proxy.spend_tracking.key_metadata_recovery import ( fill_missing_api_key_aliases, recover_double_hashed_key_metadata, + recover_key_metadata_from_spend_logs, ) from litellm.proxy.utils import hash_token -def _digest_row(digest: str, key_alias: str, team_id: str | None, user_id: str | None) -> dict[str, str | None]: +def _digest_row(digest: str, key_alias: str | None, team_id: str | None, user_id: str | None) -> dict[str, str | None]: return {"digest": digest, "key_alias": key_alias, "team_id": team_id, "user_id": user_id} +def _query_raw_spend_logs(rows: Sequence[dict[str, str | None]]) -> AsyncMock: + async def query_raw(sql: str, *params: object) -> list[dict[str, str | None]]: + if '"LiteLLM_SpendLogs"' in sql: + return list(rows) + raise AssertionError(f"unexpected query: {sql}") + + return AsyncMock(side_effect=query_raw) + + +def _spend_log_row(digest: str, key_alias: str | None, team_id: str | None, user_id: str | None) -> dict[str, str | None]: + return { + "digest": digest, + "first_alias": key_alias, + "last_alias": key_alias, + "first_team": team_id, + "last_team": team_id, + "first_owner": user_id, + "last_owner": user_id, + } + + +def _spend_log_transaction(mock_prisma: MagicMock, query_raw: AsyncMock) -> AsyncMock: + transaction = MagicMock() + transaction.execute_raw = AsyncMock(return_value=0) + transaction.query_raw = query_raw + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + return query_raw + + def _query_raw_by_table( active_rows: Sequence[dict[str, str | None]], deleted_rows: Sequence[dict[str, str | None]], @@ -218,3 +257,334 @@ async def test_fill_missing_api_key_aliases_skips_named_keys_that_have_no_email( assert filled == rows mock_prisma.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_resolves_session_token_from_metadata(): + session_digest = hash_token("cli-session-repro-user-6852") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [_spend_log_row(session_digest, "cli-session-repro-user-6852", None, "repro-user-6852")] + ), + ) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {session_digest}, window, cache=InMemoryCache()) + + assert result[session_digest]["key_alias"] == "cli-session-repro-user-6852" + assert result[session_digest]["user_id"] == "repro-user-6852" + ((_, digests, start, end),) = [call.args for call in query_raw.call_args_list] + assert digests == [session_digest] + assert (start, end) == window + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_skips_query_when_no_missing_keys(): + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(return_value=[])) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, set(), window, cache=InMemoryCache()) + + assert result == {} + query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_returns_empty_on_prisma_error(): + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(side_effect=PrismaError("db down"))) + + result = await recover_key_metadata_from_spend_logs( + mock_prisma, {hash_token("cli-session-x")}, window, cache=InMemoryCache() + ) + + assert result == {} + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_ignores_foreign_and_all_null_rows(): + wanted = hash_token("cli-session-wanted") + all_null = hash_token("cli-session-null") + foreign = hash_token("cli-session-foreign") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [ + _spend_log_row(wanted, "kept-alias", None, "owner-1"), + _spend_log_row(all_null, None, None, None), + _spend_log_row(foreign, "foreign-alias", None, "owner-2"), + ] + ), + ) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {wanted, all_null}, window, cache=InMemoryCache()) + + assert set(result) == {wanted} + assert result[wanted]["key_alias"] == "kept-alias" + assert result[wanted]["user_id"] == "owner-1" + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_skips_non_sha256_keys(): + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(return_value=[])) + + result = await recover_key_metadata_from_spend_logs( + mock_prisma, {"cli-session-raw-1798", "key-hash-short"}, window, cache=InMemoryCache() + ) + + assert result == {} + query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_accepts_hashed_jwt_digests(): + jwt_digest = f"hashed-jwt-{hash_token('jwt-subject-1')}" + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(jwt_digest, None, "team-jwt", "jwt-user")])) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {jwt_digest}, window, cache=InMemoryCache()) + + assert result[jwt_digest]["team_id"] == "team-jwt" + assert result[jwt_digest]["user_id"] == "jwt-user" + ((_, digests, _, _),) = [call.args for call in query_raw.call_args_list] + assert digests == [jwt_digest] + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_serves_repeat_lookups_from_the_cache(): + found = hash_token("cli-session-found") + unknown = hash_token("cli-session-unknown") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache() + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(found, "found-alias", None, "owner-1")])) + + first = await recover_key_metadata_from_spend_logs(mock_prisma, {found, unknown}, window, cache=cache) + second = await recover_key_metadata_from_spend_logs(mock_prisma, {found, unknown}, window, cache=cache) + + assert first == second + assert set(first) == {found} + assert first[found]["key_alias"] == "found-alias" + assert query_raw.await_count == 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_only_queries_digests_the_cache_has_not_seen(): + cached_digest = hash_token("cli-session-cached") + new_digest = hash_token("cli-session-new") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache() + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(cached_digest, "cached-alias", None, None)])) + await recover_key_metadata_from_spend_logs(mock_prisma, {cached_digest}, window, cache=cache) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(new_digest, "new-alias", None, None)])) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {cached_digest, new_digest}, window, cache=cache) + + assert result[cached_digest]["key_alias"] == "cached-alias" + assert result[new_digest]["key_alias"] == "new-alias" + ((_, digests, _, _),) = [call.args for call in query_raw.call_args_list] + assert digests == [new_digest] + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_rescans_when_the_window_changes(): + digest = hash_token("cli-session-windowed") + cache = InMemoryCache() + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([])) + await recover_key_metadata_from_spend_logs( + mock_prisma, {digest}, (datetime(2026, 9, 1), datetime(2026, 9, 4)), cache=cache + ) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(digest, "later-alias", None, None)])) + + result = await recover_key_metadata_from_spend_logs( + mock_prisma, {digest}, (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=cache + ) + + assert result[digest]["key_alias"] == "later-alias" + assert query_raw.await_count == 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_retries_a_failed_query_only_after_the_miss_ttl(): + digest = hash_token("cli-session-retry") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(side_effect=PrismaError("statement timeout"))) + started = time.time() + assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(digest, "back-online", None, None)])) + + assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} + query_raw.assert_not_awaited() + miss_key = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before")) + assert cache.ttl_dict[miss_key] - started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1 + cache.ttl_dict[miss_key] = time.time() - 1 + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) + + assert result[digest]["key_alias"] == "back-online" + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_drops_the_owner_of_a_digest_shared_by_several_users(): + shared_ui_digest = hash_token("ui-token") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [{**_spend_log_row(shared_ui_digest, "ui-token", "litellm-dashboard", None), "first_owner": "alice", "last_owner": "bob"}] + ), + ) + + result = await recover_key_metadata_from_spend_logs( + mock_prisma, {shared_ui_digest}, window, cache=InMemoryCache() + ) + + assert result[shared_ui_digest] == {"key_alias": "ui-token", "team_id": "litellm-dashboard", "user_id": None} + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_keeps_the_owner_when_every_named_row_agrees(): + digest = hash_token("cli-session-one-owner") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [_spend_log_row(digest, None, None, "carol")] + ), + ) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) + + assert result[digest]["user_id"] == "carol" + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_forgets_a_miss_long_before_a_hit(): + found = hash_token("cli-session-found") + unknown = hash_token("cli-session-unknown") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(found, "found-alias", None, None)])) + started = time.time() + + await recover_key_metadata_from_spend_logs(mock_prisma, {found, unknown}, window, cache=cache) + + hit_expires = next(deadline for key, deadline in cache.ttl_dict.items() if found in key) + miss_expires = next(deadline for key, deadline in cache.ttl_dict.items() if unknown in key) + assert miss_expires - started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1 + assert hit_expires - started >= SPEND_LOG_KEY_METADATA_CACHE_TTL - 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_runs_one_query_for_concurrent_lookups(): + digest = hash_token("cli-session-shared") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache() + lock = asyncio.Lock() + mock_prisma = MagicMock() + + async def slow_query_raw(sql: str, *params: object) -> list[dict[str, str | None]]: + await asyncio.sleep(0.01) + return [_spend_log_row(digest, "shared-alias", None, None)] + + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(side_effect=slow_query_raw)) + + results = await asyncio.gather( + *( + recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache, lock=lock) + for _ in range(9) + ) + ) + + assert all(result[digest]["key_alias"] == "shared-alias" for result in results) + assert query_raw.await_count == 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_keeps_a_repeated_miss_as_long_as_a_hit(): + unknown = hash_token("cli-session-never-named") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) + mock_prisma = MagicMock() + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([])) + await recover_key_metadata_from_spend_logs(mock_prisma, {unknown}, window, cache=cache) + first_miss_key = next(key for key in cache.ttl_dict if unknown in key and not key.endswith(":missed-before")) + cache.ttl_dict[first_miss_key] = time.time() - 1 + started = time.time() + + await recover_key_metadata_from_spend_logs(mock_prisma, {unknown}, window, cache=cache) + + assert query_raw.await_count == 2 + assert cache.ttl_dict[first_miss_key] - started >= SPEND_LOG_KEY_METADATA_CACHE_TTL - 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_keeps_the_owner_older_rows_agree_on_when_the_newest_is_nameless(): + digest = hash_token("cli-session-owner-from-older-rows") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(digest, None, "team-x", "alice")])) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) + + assert result[digest] == {"key_alias": None, "team_id": "team-x", "user_id": "alice"} + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_names_nothing_for_a_field_whose_rows_disagree(): + digest = hash_token("cli-session-disagreeing-rows") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [ + { + **_spend_log_row(digest, None, None, "carol"), + "first_alias": "old-alias", + "last_alias": "renamed-alias", + "first_team": "team-a", + "last_team": "team-b", + } + ] + ), + ) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) + + assert result[digest] == {"key_alias": None, "team_id": None, "user_id": "carol"} + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_bounds_the_scan_with_a_statement_timeout(): + digest = hash_token("cli-session-bounded-scan") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + calls: list[str] = [] + transaction = MagicMock() + transaction.execute_raw = AsyncMock(side_effect=lambda sql: calls.append(sql) or 0) + transaction.query_raw = AsyncMock(side_effect=lambda sql, *args: calls.append("scan") or []) + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + + await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) + + assert calls == [f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}", "scan"] + assert mock_prisma.db.tx.call_args.kwargs["timeout"] == timedelta( + milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS + ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index 3f775d82b7f..e466edab131 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -1,4 +1,4 @@ - +from typing import Final import pytest @@ -121,6 +121,147 @@ def _caching_usage(read: int, written: int, text: int = 10, out: int = 100) -> d } +@pytest.mark.parametrize( + "model,provider,prompt,reads,writes_5m,writes_1h,tier,region,location", + [ + pytest.param("claude-sonnet-4-5", "anthropic", 50000, 20000, 20000, 0, None, None, None, id="5m"), + pytest.param("claude-sonnet-4-5", "anthropic", 50000, 20000, 0, 20000, None, None, None, id="1h"), + pytest.param("claude-sonnet-4-5", "anthropic", 50000, 20000, 12000, 8000, None, None, None, id="mixed-ttl"), + pytest.param("claude-sonnet-4-5", "anthropic", 199999, 80000, 20000, 0, None, None, None, id="below-200k"), + pytest.param("claude-sonnet-4-5", "anthropic", 200000, 80000, 20000, 0, None, None, None, id="exactly-200k"), + pytest.param("claude-sonnet-4-5", "anthropic", 200001, 80000, 20000, 0, None, None, None, id="above-200k"), + pytest.param("claude-sonnet-4-5", "anthropic", 250000, 80000, 12000, 8000, None, None, None, id="ttl-and-200k"), + pytest.param( + "claude-sonnet-4-5", "anthropic", 250000, 80000, 0, 20000, "priority", None, None, id="absent-tier" + ), + pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "priority", None, None, id="priority"), + pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "flex", None, None, id="flex"), + pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "batch", None, None, id="batch-resolver-fallback"), + pytest.param("gpt-5.5", "openai", 300000, 80000, 0, 0, "flex", None, None, id="flex-and-272k"), + pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "priority", "eu", None, id="priority-and-eu"), + pytest.param("gemini-2.5-pro", "vertex_ai", 250000, 80000, 20000, 0, None, None, None, id="variant-only-write"), + pytest.param("gemini-3.5-flash", "vertex_ai", 100000, 80000, 0, 0, None, None, "us-east5", id="vertex-region"), + pytest.param("gpt-5.5", "openai", 300000, 0, 0, 0, "priority", "eu", None, id="no-cache"), + ], +) +def test_caching_savings_agree_with_biller_on_the_request_pricing_basis( + model: str, + provider: str, + prompt: int, + reads: int, + writes_5m: int, + writes_1h: int, + tier: str | None, + region: str | None, + location: str | None, +) -> None: + pricing: Final = litellm.get_model_info(model=model, custom_llm_provider=provider) + usage: Final = Usage( + prompt_tokens=prompt, + completion_tokens=100, + total_tokens=prompt + 100, + prompt_tokens_details={ + "cached_tokens": reads, + "cache_creation_tokens": writes_5m + writes_1h, + "text_tokens": prompt - reads - writes_5m - writes_1h, + "cache_creation_token_details": { + "ephemeral_5m_input_tokens": writes_5m, + "ephemeral_1h_input_tokens": writes_1h, + }, + }, + ) + uncached: Final = Usage( + prompt_tokens=prompt, + completion_tokens=100, + total_tokens=prompt + 100, + prompt_tokens_details={"text_tokens": prompt, "cached_tokens": 0, "cache_creation_tokens": 0}, + ) + costs: Final = tuple( + sum( + generic_cost_per_token( + model=model, + usage=arm, + custom_llm_provider=provider, + model_info=pricing, + service_tier=tier, + data_residency=region, + vertex_location=location, + ) + ) + for arm in (uncached, usage) + ) + expected: Final = costs[0] - costs[1] + for attributed in (False, True): + result: Final = compute_savings_spend( + model=model, + custom_llm_provider=provider, + compression_saved_tokens=4389, + gateway_injected_cache=attributed, + usage_object=usage.model_dump(), + cost_breakdown={"service_tier": tier, "data_residency": region, "vertex_location": location}, + billed_at="2026-09-07T12:00:00+00:00", + ) + assert result.prompt_caching == pytest.approx(expected) + assert result.gateway_injected_caching == pytest.approx(expected if attributed else 0.0) + assert result.compression == pytest.approx(4389 * (pricing["input_cost_per_token"] or 0.0)) + assert result.autorouter == 0.0 + if reads + writes_5m + writes_1h == 0: + assert expected == 0.0 + + +def test_negative_ttl_counts_do_not_become_cache_write_credits() -> None: + results: Final = tuple( + compute_savings_spend( + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + gateway_injected_cache=True, + usage_object={ + "prompt_tokens": 6000, + "completion_tokens": 100, + "prompt_tokens_details": { + "text_tokens": 1000, + "cache_creation_tokens": 5000, + "cache_creation_token_details": { + "ephemeral_5m_input_tokens": short_count, + "ephemeral_1h_input_tokens": 5000, + }, + }, + }, + ) + for short_count in (-5000, 0) + ) + assert results[0] == results[1] + assert results[0].prompt_caching < 0 + + +def test_unpublished_one_hour_price_uses_the_ordinary_write_price() -> None: + model: Final = "claude-4-opus-20250514" + pricing: Final = litellm.get_model_info(model=model, custom_llm_provider="anthropic") + assert pricing.get("cache_creation_input_token_cost_above_1hr") is None + assert pricing["cache_creation_input_token_cost"] > pricing["input_cost_per_token"] + results: Final = tuple( + compute_savings_spend( + model=model, + custom_llm_provider="anthropic", + compression_saved_tokens=0, + gateway_injected_cache=True, + usage_object={ + "prompt_tokens": 6000, + "completion_tokens": 100, + "prompt_tokens_details": { + "text_tokens": 1000, + "cache_creation_tokens": 5000, + "cache_creation_token_details": ttl, + }, + }, + ) + for ttl in (None, {"ephemeral_1h_input_tokens": 5000}) + ) + assert results[0] == results[1] + assert results[0].prompt_caching < 0 + + def test_prompt_caching_savings_nets_out_the_cache_write_premium(): """A cache-writing request is only credited the read discount minus the write premium.""" input_cost, cache_read_cost = _anthropic_costs("claude-sonnet-5") @@ -1021,6 +1162,104 @@ def test_prompt_caching_prices_at_the_deployment_rate_not_the_public_one(): assert result.prompt_caching > at_public_rates.prompt_caching +@pytest.mark.parametrize( + "baseline_id, selected_id, selected_multiplier, billed_input, classifier_cost, expected", + [ + ("baseline", "selected", 0.1, None, 0.0, 0.0135), + ("baseline", "selected", 2.0, None, 0.0, -0.015), + ("baseline", "selected", 1.0, None, 0.0, 0.0), + ("baseline", "selected", 0.1, 0.004, 0.001, 0.01), + ("baseline", "baseline", 0.1, 0.004, 0.001, -0.001), + (None, "selected", 0.1, None, 0.0, 0.0), + ("baseline", None, 0.1, None, 0.0, 0.0), + (None, None, 0.1, None, 0.0, 0.0), + ("", "selected", 0.1, None, 0.0, 0.0), + ("baseline", "", 0.1, None, 0.0, 0.0), + ], +) +def test_autorouter_savings_distinguishes_priced_deployments( + baseline_id: str | None, + selected_id: str | None, + selected_multiplier: float, + billed_input: float | None, + classifier_cost: float, + expected: float, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": "anthropic/claude-opus-5", + "api_key": "test-key", + "input_cost_per_token": 1e-5 * multiplier, + "output_cost_per_token": 5e-5 * multiplier, + }, + "model_info": {"id": name}, + } + for name, multiplier in (("baseline", 1.0), ("selected", selected_multiplier)) + ] + ) + result: Final = compute_savings_spend( + model="claude-opus-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + gateway_injected_cache=False, + model_id=selected_id, + llm_router=lambda: router, + routing_decision={ + "savings_baseline_model": "anthropic/claude-opus-5", + "savings_baseline_deployment_id": baseline_id, + "conversation_continuing": False, + "classifier_cost": classifier_cost, + }, + usage_object={"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100}, + cost_breakdown=None if billed_input is None else {"input_cost": billed_input, "output_cost": 0.0}, + ) + assert result.autorouter == pytest.approx(expected) + + +@pytest.mark.parametrize("selected_model", ["azure/contract-deployment", "contract-deployment"]) +def test_autorouter_savings_recognizes_one_deployment_under_its_base_model(selected_model: str) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "contract", + "litellm_params": { + "model": "azure/contract-deployment", + "api_key": "test-key", + "api_base": "https://example.openai.azure.com", + "input_cost_per_token": 0.0001, + "output_cost_per_token": 0.0002, + "cache_read_input_token_cost": 0.00001, + }, + "model_info": {"id": "contract", "base_model": "azure/gpt-5.5"}, + } + ] + ) + result: Final = compute_savings_spend( + model=selected_model, + custom_llm_provider="azure", + compression_saved_tokens=0, + gateway_injected_cache=False, + model_id="contract", + llm_router=lambda: router, + routing_decision={ + "savings_baseline_model": "azure/gpt-5.5", + "savings_baseline_deployment_id": "contract", + "conversation_continuing": True, + }, + usage_object={ + "prompt_tokens": 21000, + "completion_tokens": 100, + "total_tokens": 21100, + "prompt_tokens_details": {"text_tokens": 1000, "cached_tokens": 0, "cache_creation_tokens": 20000}, + }, + cost_breakdown={"input_cost": 2.1, "output_cost": 0.02}, + ) + assert result.autorouter == 0.0 + + def test_a_recorded_baseline_deployment_prices_at_its_configured_rate(): """A hardest-tier deployment with a negotiated rate is what the traffic would really have cost; pricing its model publicly misstates the saving.""" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_event.py b/tests/test_litellm/proxy/spend_tracking/test_spend_event.py new file mode 100644 index 00000000000..ff449235582 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_event.py @@ -0,0 +1,213 @@ +import json +from datetime import datetime +from typing import Final + +import pytest + +import litellm +from litellm.caching.caching import Cache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.spend_tracking.spend_event import ( + CACHE_OFF_KEY, + SpendEventBuildError, + SpendEventDecodeError, + build_spend_event, + decode_spend_event, + is_offloadable_success, + spend_event_callback_args, +) +from litellm.types.utils import LiteLLMBatch, ModelResponse, Usage + +_BIG_PROMPT: Final = "x" * 20_000 +_RESERVATION: Final = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "key:hash", "reserved_cost": 0.5}], + "finalized": False, + "input_cost": 0.1, + "input_tokens": 5000, +} + + +def _response(tool_name: str | None = None) -> ModelResponse: + tool_calls: Final = ( + [{"id": "call-1", "type": "function", "function": {"name": tool_name, "arguments": "{}"}}] + if tool_name is not None + else None + ) + return ModelResponse( + id="chatcmpl-1", + model="gpt-4o-2024-08-06", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "y" * 20_000, "tool_calls": tool_calls}, + "finish_reason": "tool_calls" if tool_name else "stop", + } + ], + usage=Usage(prompt_tokens=5000, completion_tokens=4000, total_tokens=9000), + ) + + +def _success_kwargs(preset_cache_key: str | None = "preset-key") -> dict: + return { + "litellm_call_id": "call-1", + "call_type": "acompletion", + "model": "gpt-4o", + "custom_llm_provider": "openai", + "stream": False, + "cache_hit": None, + "response_cost": 0.0125, + "completion_start_time": datetime(2026, 1, 1, 0, 0, 1), + "messages": [{"role": "user", "content": _BIG_PROMPT}], + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + "litellm_params": { + "api_base": "https://api.openai.com", + "preset_cache_key": preset_cache_key, + "proxy_server_request": {"body": {"messages": [{"role": "user", "content": _BIG_PROMPT}]}}, + "metadata": { + "user_api_key": "hash-1", + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-1", + "user_api_key_org_id": "org-1", + "user_api_key_end_user_id": "end-user-1", + "user_api_key_auth": UserAPIKeyAuth(api_key="hash-1", budget_reservation=dict(_RESERVATION)), + "model_group": "gpt-4o", + "model_info": {"id": "deployment-1"}, + "tags": ["tag-a"], + "litellm_parent_otel_span": object(), + }, + }, + "standard_logging_object": { + "response_cost": 0.0125, + "model": "gpt-4o-2024-08-06", + "model_id": "deployment-1", + "request_tags": ["tag-a"], + "request_model_access_groups": ["premium"], + "messages": [{"role": "user", "content": _BIG_PROMPT}], + "response": {"choices": [{"message": {"content": "y" * 20_000}}]}, + "model_parameters": {"temperature": 0.1}, + "metadata": {"user_api_key_hash": "hash-1", "usage_object": {"prompt_tokens": 5000}}, + "hidden_params": {"litellm_overhead_time_ms": 3}, + "model_map_information": {}, + }, + } + + +def _build(kwargs: dict, response: object, store_bodies: bool = False) -> bytes: + line: Final = build_spend_event( + kwargs, response, datetime(2026, 1, 1), datetime(2026, 1, 1, 0, 0, 2), store_bodies=store_bodies + ) + assert isinstance(line, bytes) + return line + + +def test_event_is_compact_and_omits_bodies_by_default(): + line: Final = _build(_success_kwargs(), _response(tool_name="get_weather")) + assert line.endswith(b"\n") + assert len(line) < 4_000 + assert _BIG_PROMPT.encode() not in line + assert b"yyyy" not in line + decoded: Final = json.loads(line) + assert "messages" not in decoded["standard_logging_object"] + assert "response" not in decoded["standard_logging_object"] + assert decoded["litellm_params"]["proxy_server_request"] is None + + +def test_event_carries_bodies_when_spend_logs_store_them(): + line: Final = _build(_success_kwargs(), _response(), store_bodies=True) + decoded: Final = json.loads(line) + assert decoded["standard_logging_object"]["messages"][0]["content"] == _BIG_PROMPT + assert decoded["standard_logging_object"]["response"]["choices"][0]["message"]["content"] == "y" * 20_000 + assert decoded["litellm_params"]["proxy_server_request"]["body"]["messages"][0]["content"] == _BIG_PROMPT + + +def test_round_trip_preserves_identity_usage_reservation_and_tools(): + line: Final = _build(_success_kwargs(), _response(tool_name="get_weather")) + event: Final = decode_spend_event(line) + assert not isinstance(event, SpendEventDecodeError) + args: Final = spend_event_callback_args(event) + + metadata: Final = args.kwargs["litellm_params"]["metadata"] + assert metadata is not None + assert (metadata["user_api_key"], metadata["user_api_key_team_id"], metadata["user_api_key_org_id"]) == ( + "hash-1", + "team-1", + "org-1", + ) + assert metadata["user_api_key_budget_reservation"] == _RESERVATION + assert "user_api_key_auth" not in metadata + assert "litellm_parent_otel_span" not in metadata + assert args.kwargs["standard_logging_object"]["request_model_access_groups"] == ["premium"] + assert args.kwargs["standard_logging_object"]["response_cost"] == 0.0125 + assert args.kwargs["tools"] == ({"type": "function", "function": {"name": "get_weather"}},) + assert args.kwargs["completion_start_time"] == datetime(2026, 1, 1, 0, 0, 1) + assert (args.start_time, args.end_time) == (datetime(2026, 1, 1), datetime(2026, 1, 1, 0, 0, 2)) + assert args.response_obj is not None + assert args.response_obj["id"] == "chatcmpl-1" + assert args.response_obj["usage"]["prompt_tokens"] == 5000 + assert args.response_obj["usage"]["completion_tokens"] == 4000 + tool_calls: Final = args.response_obj["choices"][0]["message"]["tool_calls"] + assert [call["function"]["name"] for call in tool_calls] == ["get_weather"] + assert "complete_streaming_response" not in args.kwargs + + +def test_streaming_event_reconstructs_complete_streaming_response(): + kwargs: Final = {**_success_kwargs(), "stream": True, "complete_streaming_response": _response()} + event: Final = decode_spend_event(_build(kwargs, _response())) + assert not isinstance(event, SpendEventDecodeError) + args: Final = spend_event_callback_args(event) + assert args.kwargs["stream"] is True + assert args.kwargs["complete_streaming_response"] == args.response_obj + + +class _HashingCache(Cache): + def __init__(self) -> None: + pass + + def get_cache_key(self, **kwargs) -> str: + raise AssertionError("the fast path must not hash the request body") + + +@pytest.mark.parametrize( + ("cache", "preset", "expected"), + [ + (None, "preset-key", CACHE_OFF_KEY), + (_HashingCache(), "preset-key", "preset-key"), + (_HashingCache(), None, None), + ], +) +def test_event_reuses_preset_cache_key_and_never_hashes(monkeypatch, cache, preset, expected): + monkeypatch.setattr(litellm, "cache", cache) + decoded: Final = json.loads(_build(_success_kwargs(preset_cache_key=preset), _response())) + assert decoded["litellm_params"]["preset_cache_key"] == expected + + +def test_unbuildable_kwargs_fall_back_to_in_process_tracking(): + kwargs: Final = {**_success_kwargs(), "response_cost": "not-a-number"} + assert isinstance( + build_spend_event(kwargs, _response(), datetime.now(), datetime.now(), False), SpendEventBuildError + ) + + +def test_undecodable_line_is_an_error_value(): + assert isinstance(decode_spend_event(b'{"version": 2}\n'), SpendEventDecodeError) + assert isinstance(decode_spend_event(b"not json\n"), SpendEventDecodeError) + + +def test_batch_retrieves_stay_in_process(): + assert is_offloadable_success(_response()) is True + assert is_offloadable_success(None) is True + assert ( + is_offloadable_success( + LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + ) + ) + is False + ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_event_producer.py b/tests/test_litellm/proxy/spend_tracking/test_spend_event_producer.py new file mode 100644 index 00000000000..adfb1251d1c --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_event_producer.py @@ -0,0 +1,359 @@ +import asyncio +from collections.abc import Awaitable, Callable +from pathlib import Path +from typing import Final + +import pytest +import uvloop + +from litellm.proxy.spend_tracking.spend_event_producer import ( + AddressError, + CollectorAddress, + CollectorSettings, + SpendEventProducer, + TcpAddress, + UnixAddress, + build_spend_event_producer, + open_collector_connection, + parse_collector_address, +) + + +class _Sidecar: + """A unix-socket server that records every line it receives, standing in for the collector.""" + + def __init__(self, path: Path, reads: bool = True, limit: int = 2**16) -> None: + self.path = path + self.reads = reads + self.limit = limit + self.lines: list[bytes] = [] # mutable-ok: test double records what the producer sent + self._server: asyncio.Server | None = None + self._stopped = asyncio.Event() + self._connections: list[asyncio.StreamWriter] = [] # mutable-ok: test double tracks peers to hang up on + + async def __aenter__(self) -> "_Sidecar": + self._server = await asyncio.start_unix_server(self._on_connection, path=str(self.path), limit=self.limit) + return self + + async def __aexit__(self, *exc: object) -> None: + self._stopped.set() + await self.hang_up() + + async def hang_up(self) -> None: + """Exit the way a stopped sidecar does: stop listening and close every producer connection.""" + assert self._server is not None + self._server.close() + for connection in self._connections: + connection.close() + await connection.wait_closed() + await self._server.wait_closed() + + async def _on_connection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + self._connections.append(writer) + if not self.reads: + await self._stopped.wait() + return + while line := await reader.readline(): + self.lines.append(line) + writer.close() + + +class _CrashingSidecar(_Sidecar): + """Bills a few lines, then dies mid-stream with the producer's backlog still queued behind them.""" + + def __init__(self, path: Path, lines_before_crash: int) -> None: + super().__init__(path, limit=2**20) + self._lines_before_crash = lines_before_crash + + async def _on_connection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + self._connections.append(writer) + for _ in range(self._lines_before_crash): + self.lines.append(await reader.readline()) + writer.transport.abort() + + +class _Fallback: + def __init__(self) -> None: + self.lines: list[bytes] = [] # mutable-ok: test double records what fell back to in-process + + async def __call__(self, line: bytes) -> None: + self.lines.append(line) + + +class _GatedFallback(_Fallback): + """A fallback that blocks, like a slow database write, until the test releases it.""" + + def __init__(self) -> None: + super().__init__() + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def __call__(self, line: bytes) -> None: + self.started.set() + await self.release.wait() + await super().__call__(line) + + +class _StalledDrainWriter(asyncio.StreamWriter): + """Hands bytes to the real transport but never wakes ``drain()``: the loop iteration between a flush + completing and the writer task resuming, frozen in place.""" + + def __init__(self, real: asyncio.StreamWriter, reader: asyncio.StreamReader) -> None: + super().__init__(real.transport, real.transport.get_protocol(), reader, asyncio.get_running_loop()) + self._real_writer_whose_finalizer_would_close_the_transport = real + + async def drain(self) -> None: + await asyncio.Event().wait() + + +async def _open_with_stalled_drain( + address: CollectorAddress, timeout: float +) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + reader, writer = await open_collector_connection(address, timeout) + return reader, _StalledDrainWriter(writer, reader) + + +def _producer( + path: Path, + fallback: _Fallback, + on_unavailable="fallback", + buffer_size: int = 100, + open_connection: Callable[ + [CollectorAddress, float], Awaitable[tuple[asyncio.StreamReader, asyncio.StreamWriter]] + ] = open_collector_connection, +) -> SpendEventProducer: + return SpendEventProducer( + address=UnixAddress(path=str(path)), + on_unavailable=on_unavailable, + buffer_size=buffer_size, + connect_timeout=1.0, + fallback=fallback, + open_connection=open_connection, + ) + + +def test_parse_collector_address(): + assert parse_collector_address("unix:///var/run/litellm/collector.sock") == UnixAddress( + path="/var/run/litellm/collector.sock" + ) + assert parse_collector_address("tcp://127.0.0.1:4100") == TcpAddress(host="127.0.0.1", port=4100) + assert parse_collector_address("tcp://localhost:4100") == TcpAddress(host="localhost", port=4100) + assert parse_collector_address("tcp://[::1]:4100") == TcpAddress(host="::1", port=4100) + assert isinstance(parse_collector_address("redis://localhost:6379"), AddressError) + assert isinstance(parse_collector_address("tcp://127.0.0.1"), AddressError) + + +@pytest.mark.parametrize("address", ["tcp://0.0.0.0:4100", "tcp://10.0.0.5:4100", "tcp://collector.svc:4100"]) +def test_tcp_address_outside_loopback_is_refused(address: str): + """The socket has no authentication, so anything reachable from outside the pod would accept forged spend.""" + error: Final = parse_collector_address(address) + assert isinstance(error, AddressError) + assert "loopback" in error.reason + assert build_spend_event_producer(CollectorSettings(enabled=True, address=address), _Fallback()) is None + + +def test_gateway_produces_only_when_enabled_and_not_the_sidecar_itself(): + fallback: Final = _Fallback() + assert build_spend_event_producer(CollectorSettings(enabled=False), fallback) is None + assert build_spend_event_producer(CollectorSettings(enabled=True, job_role="collector"), fallback) is None + assert build_spend_event_producer(CollectorSettings(enabled=True, address="redis://x"), fallback) is None + assert isinstance(build_spend_event_producer(CollectorSettings(enabled=True), fallback), SpendEventProducer) + + +def test_settings_read_the_documented_env(monkeypatch): + monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "true") + monkeypatch.setenv("LITELLM_COLLECTOR_ADDRESS", "tcp://127.0.0.1:4100") + monkeypatch.setenv("LITELLM_COLLECTOR_BUFFER_SIZE", "50") + monkeypatch.setenv("LITELLM_COLLECTOR_ON_UNAVAILABLE", "drop") + monkeypatch.setenv("LITELLM_JOB_ROLE", "collector") + settings: Final = CollectorSettings() + assert (settings.enabled, settings.address, settings.buffer_size, settings.on_unavailable) == ( + True, + "tcp://127.0.0.1:4100", + 50, + "drop", + ) + assert settings.produces is False + + +@pytest.mark.asyncio +async def test_events_reach_the_sidecar_once_and_in_order(tmp_path: Path): + fallback: Final = _Fallback() + async with _Sidecar(tmp_path / "spend.sock") as sidecar: + producer: Final = _producer(sidecar.path, fallback) + outcomes: Final = [await producer.publish(f"event-{i}\n".encode()) for i in range(20)] + await producer.close(drain_timeout=5.0) + await asyncio.sleep(0.05) + + assert outcomes == ["queued"] * 20 + assert sidecar.lines == [f"event-{i}\n".encode() for i in range(20)] + assert fallback.lines == [] + stats: Final = producer.stats() + assert (stats.queued, stats.sent, stats.fallback, stats.dropped) == (20, 20, 0, 0) + + +@pytest.mark.asyncio +async def test_unreachable_sidecar_falls_back_in_process_and_backs_off(tmp_path: Path): + fallback: Final = _Fallback() + producer: Final = _producer(tmp_path / "missing.sock", fallback) + first: Final = await producer.publish(b"event-1\n") + await asyncio.sleep(0.05) + second: Final = await producer.publish(b"event-2\n") + await producer.close(drain_timeout=5.0) + + assert first == "queued" + assert second == "fallback" + assert fallback.lines == [b"event-1\n", b"event-2\n"] + stats: Final = producer.stats() + assert (stats.sent, stats.fallback, stats.dropped, stats.connected) == (0, 2, 0, False) + + +@pytest.mark.parametrize("loop_factory", [asyncio.new_event_loop, uvloop.new_event_loop], ids=["asyncio", "uvloop"]) +def test_sidecar_hang_up_falls_back_instead_of_losing_events( + tmp_path: Path, loop_factory: Callable[[], asyncio.AbstractEventLoop] +): + async def scenario() -> tuple[list[bytes], list[bytes], tuple[int, int, int]]: + fallback: Final = _Fallback() + sidecar: Final = _Sidecar(tmp_path / "spend.sock") + async with sidecar: + producer: Final = _producer(sidecar.path, fallback) + await producer.publish(b"event-1\n") + await asyncio.sleep(0.05) + await sidecar.hang_up() + await asyncio.sleep(0.05) + await producer.publish(b"event-2\n") + await producer.close(drain_timeout=5.0) + stats: Final = producer.stats() + return sidecar.lines, fallback.lines, (stats.sent, stats.fallback, stats.dropped) + + with asyncio.Runner(loop_factory=loop_factory) as runner: + sidecar_lines, fallback_lines, counts = runner.run(scenario()) + + assert sidecar_lines == [b"event-1\n"] + assert fallback_lines == [b"event-2\n"] + assert counts == (1, 1, 0) + + +@pytest.mark.parametrize("loop_factory", [asyncio.new_event_loop, uvloop.new_event_loop], ids=["asyncio", "uvloop"]) +def test_mid_stream_crash_never_bills_an_event_on_both_sides( + tmp_path: Path, loop_factory: Callable[[], asyncio.AbstractEventLoop] +): + """Events large enough to straddle the kernel buffer, a sidecar that reads some and then drops the socket: a + failed write may only fall back when the sidecar cannot have read the whole line.""" + events: Final = tuple(f"event-{i:03d}-".encode() + b"x" * 65536 + b"\n" for i in range(64)) + + async def scenario() -> tuple[list[bytes], list[bytes], tuple[int, int, int]]: + fallback: Final = _Fallback() + sidecar: Final = _CrashingSidecar(tmp_path / "spend.sock", lines_before_crash=3) + async with sidecar: + producer: Final = _producer(sidecar.path, fallback) + for event in events: + assert await producer.publish(event) == "queued" + await asyncio.sleep(0.2) + await producer.close(drain_timeout=5.0) + stats: Final = producer.stats() + return sidecar.lines, fallback.lines, (stats.sent, stats.fallback, stats.dropped) + + with asyncio.Runner(loop_factory=loop_factory) as runner: + sidecar_lines, fallback_lines, counts = runner.run(scenario()) + + assert sidecar_lines == list(events[:3]) + assert set(sidecar_lines).isdisjoint(fallback_lines) + assert len(fallback_lines) == len(set(fallback_lines)) + assert fallback_lines[-1] == events[-1] + assert counts[0] + counts[1] == len(events) and counts[2] == 0 + assert counts[0] >= len(sidecar_lines) + + +@pytest.mark.asyncio +async def test_drain_timeout_hands_the_in_flight_event_to_fallback(tmp_path: Path): + """A sidecar that stops reading leaves one event half-written; cancelling the writer must not lose it.""" + fallback: Final = _Fallback() + stuck: Final = b"x" * (4 * 1024 * 1024) + b"\n" + async with _Sidecar(tmp_path / "spend.sock", reads=False) as sidecar: + producer: Final = _producer(sidecar.path, fallback) + assert await producer.publish(stuck) == "queued" + await asyncio.sleep(0.1) + await producer.close(drain_timeout=0.2) + + assert fallback.lines == [stuck] + stats: Final = producer.stats() + assert (stats.sent, stats.fallback, stats.connected) == (0, 1, False) + + +@pytest.mark.asyncio +async def test_shutdown_lets_the_writer_finish_a_fallback_already_in_progress(tmp_path: Path): + """Cancelling the writer while it runs the pipeline in-process must neither lose nor repeat that event.""" + fallback: Final = _GatedFallback() + producer: Final = _producer(tmp_path / "missing.sock", fallback) + assert await producer.publish(b"event-1\n") == "queued" + await asyncio.wait_for(fallback.started.wait(), 5.0) + closing: Final = asyncio.ensure_future(producer.close(drain_timeout=0.05)) + await asyncio.sleep(0.2) + assert fallback.lines == [] + fallback.release.set() + await asyncio.wait_for(closing, 5.0) + + assert fallback.lines == [b"event-1\n"] + assert producer.stats().fallback == 1 + + +@pytest.mark.asyncio +async def test_shutdown_does_not_replay_an_event_the_kernel_already_took(tmp_path: Path): + """Cancelling a drain whose bytes already left the process must not run the event a second time in-process.""" + fallback: Final = _Fallback() + async with _Sidecar(tmp_path / "spend.sock") as sidecar: + producer: Final = _producer(sidecar.path, fallback, open_connection=_open_with_stalled_drain) + assert await producer.publish(b"event-1\n") == "queued" + await asyncio.sleep(0.1) + await producer.close(drain_timeout=0.2) + await asyncio.sleep(0.05) + + assert sidecar.lines == [b"event-1\n"] + assert fallback.lines == [] + stats: Final = producer.stats() + assert (stats.fallback, stats.dropped, stats.connected) == (0, 0, False) + + +@pytest.mark.asyncio +async def test_drop_policy_counts_instead_of_running_in_process(tmp_path: Path): + fallback: Final = _Fallback() + producer: Final = _producer(tmp_path / "missing.sock", fallback, on_unavailable="drop") + await producer.publish(b"event-1\n") + await producer.close(drain_timeout=5.0) + assert await producer.publish(b"event-2\n") == "dropped" + + assert fallback.lines == [] + assert producer.stats().dropped == 2 + + +@pytest.mark.asyncio +async def test_full_buffer_applies_the_unavailable_policy_immediately(tmp_path: Path): + fallback: Final = _Fallback() + async with _Sidecar(tmp_path / "spend.sock") as sidecar: + producer: Final = _producer(sidecar.path, fallback, buffer_size=2) + outcomes: Final = [await producer.publish(f"event-{i}\n".encode()) for i in range(3)] + await producer.close(drain_timeout=5.0) + await asyncio.sleep(0.05) + + assert outcomes == ["queued", "queued", "fallback"] + assert fallback.lines == [b"event-2\n"] + assert sidecar.lines == [b"event-0\n", b"event-1\n"] + + +@pytest.mark.asyncio +async def test_close_flushes_buffered_events_then_refuses_new_ones(tmp_path: Path): + fallback: Final = _Fallback() + async with _Sidecar(tmp_path / "spend.sock") as sidecar: + producer: Final = _producer(sidecar.path, fallback) + for i in range(50): + await producer.publish(f"event-{i}\n".encode()) + assert sidecar.lines == [] + await producer.close(drain_timeout=5.0) + await asyncio.sleep(0.05) + after_close: Final = await producer.publish(b"late\n") + + assert len(sidecar.lines) == 50 + assert after_close == "fallback" + assert fallback.lines == [b"late\n"] + assert producer.stats().connected is False diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 329a33eb440..6e43ac4a12b 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -455,6 +455,34 @@ async def test_assert_user_can_view_request_id_rejects_both_users_none(): assert exc_info.value.status_code == 403 +@pytest.mark.asyncio +async def test_assert_user_can_view_request_id_rejects_missing_row(): + """ + A request_id with no spend-log row (e.g. pruned by retention) must not + authorize reading the payload from cold storage; a missing row is not + the same as an owned row. + """ + + class MockSpendLogs: + async def find_unique(self, where, include=None): + return None + + class MockDB: + def __init__(self): + self.litellm_spendlogs = MockSpendLogs() + + class MockPrisma: + def __init__(self): + self.db = MockDB() + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1") + with pytest.raises(HTTPException) as exc_info: + await spend_management_endpoints._assert_user_can_view_request_id( + MockPrisma(), auth, "req-missing-row" + ) + assert exc_info.value.status_code == 403 + + def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypatch): """ Without prisma, non-admins cannot be authorized to read request/response @@ -5703,6 +5731,29 @@ def _cold_storage_handler(payload): return ColdStorageHandler(cold_storage_logger=logger), logger +@pytest.mark.asyncio +@pytest.mark.parametrize("cold_has_audit", [False, True]) +async def test_resolve_payload_recovers_truncated_classifier_audit_without_losing_existing_fields(cold_has_audit): + full_audit = {"classifier_input": {"system": "full rubric"}, "originating_request_masked": {"input": "source"}} + truncated_request = {"model": "classifier", "classifier_input": {"system": "litellm_truncated"}} + handler, logger = _cold_storage_handler({ + "proxy_server_request": {"body": {}}, **(full_audit if cold_has_audit else {}), + }) + row = { + "messages": '[{"role":"user","content":"ask"}]', "response": '{"tier":"SIMPLE"}', + "proxy_server_request": json.dumps(truncated_request), "metadata": {"cold_storage_object_key": "k/audit.json"}, + } + resolved = await spend_management_endpoints._resolve_request_response_payload(row, cold_storage_handler=handler) + assert logger.requested_object_keys == ["k/audit.json"] + assert resolved.messages == row["messages"] + assert resolved.response == row["response"] + if cold_has_audit: + assert resolved.proxy_server_request["classifier_input"] == full_audit["classifier_input"] + assert resolved.proxy_server_request["originating_request_masked"] == full_audit["originating_request_masked"] + else: + assert resolved.proxy_server_request == row["proxy_server_request"] + + @pytest.mark.parametrize( "value, expected", [ @@ -6578,6 +6629,30 @@ def _session_page_row(session_key, last_activity): return {"session_key": session_key, "api_key": "hashed-key", "last_activity": last_activity} +def _session_grouped_paginating_prisma(sessions, counted_total=None): + """Mock prisma serving the grouped page query out of ``sessions``, honoring the LIMIT and OFFSET it asks for.""" + + async def mock_query_raw(sql_query, *params): + if "COUNT(*) AS total_count" in sql_query: + return [{"total_count": min(len(sessions) if counted_total is None else counted_total, params[-1])}] + if "DISTINCT ON" in sql_query: + return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]] + if "COALESCE(SUM(spend)" in sql_query: + return [] + bounds = re.search(r"LIMIT \$(\d+)(?: OFFSET \$(\d+))?", sql_query) + limit = params[int(bounds.group(1)) - 1] + offset = params[int(bounds.group(2)) - 1] if bounds.group(2) else 0 + return [ + _session_page_row(session_key, last_activity) + for session_key, last_activity in sessions[offset : offset + limit] + ] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + return mock_prisma + + @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatch): """One row per (session, api_key), session-count total, and a keyset cursor for the next page.""" @@ -6691,6 +6766,203 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(client, monkeypatch): + """page > 1 with no session_cursor (the UI's last-page jump) serves the sessions that page starts at.""" + sessions = tuple((f"sess-{index:02d}", f"2026-08-29 10:{59 - index:02d}:00") for index in range(60)) + mock_prisma = _session_grouped_paginating_prisma(sessions) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page": 3, + "page_size": 25, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["page"] == 3 + assert data["total"] == 60 + assert data["has_more"] is False + assert data["next_session_cursor"] is None + assert [row["request_id"] for row in data["data"]] == [f"req-sess-{index:02d}" for index in range(50, 60)] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_short_page_totals_itself(client, monkeypatch): + """A page that runs out of sessions is the end of the list, so the total comes from it and nothing is counted.""" + sessions = tuple((f"sess-{index:02d}", f"2026-08-29 10:{59 - index:02d}:00") for index in range(10)) + mock_prisma = _session_grouped_paginating_prisma(sessions, counted_total=999) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page": 1, + "page_size": 25, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["total"] == 10, "the count query's 999 would have won if it had been asked" + assert data["total_is_capped"] is False + assert data["total_pages"] == 1 + assert len(data["data"]) == 10 + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_page_past_the_end_keeps_the_real_total(client, monkeypatch): + """An empty page past the last one says nothing about the total, so it is counted rather than inferred.""" + sessions = tuple((f"sess-{index:02d}", f"2026-08-29 10:{59 - index:02d}:00") for index in range(100)) + mock_prisma = _session_grouped_paginating_prisma(sessions) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page": 4, + "page_size": 50, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["data"] == [] + assert data["total"] == 100, "the empty page's offset is not a total" + assert data["total_pages"] == 2 + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_page_past_count_cap_is_empty(client, monkeypatch): + """The last page inside the capped total still lists sessions; the page after it is empty and costs no query.""" + cap = spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + sessions = tuple((f"sess-{index:06d}", "2026-08-29 10:00:00") for index in range(cap + 50)) + mock_prisma = _session_grouped_paginating_prisma(sessions) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + params = { + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page_size": 25, + } + last_page = client.get( + "/spend/logs/ui", + params={**params, "page": cap // 25}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert last_page.status_code == 200, last_page.text + last_page_data = last_page.json() + assert last_page_data["total"] == cap + assert last_page_data["total_is_capped"] is True + assert last_page_data["data"][0]["request_id"] == f"req-sess-{cap - 25:06d}" + assert len(last_page_data["data"]) == 25 + + mock_prisma.db.query_raw.reset_mock() + past_cap = client.get( + "/spend/logs/ui", + params={**params, "page": cap // 25 + 1}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert past_cap.status_code == 200, past_cap.text + past_cap_data = past_cap.json() + assert past_cap_data["data"] == [] + assert past_cap_data["has_more"] is False + assert past_cap_data["total"] == cap + assert mock_prisma.db.query_raw.await_count == 1, "only the bounded count query runs past the capped window" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_last_page_stops_at_the_capped_total(client, monkeypatch): + """A page size that does not divide the cap still ends the last page at the capped total it reports.""" + cap = spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP + sessions = tuple((f"sess-{index:06d}", "2026-08-29 10:00:00") for index in range(cap + 50)) + mock_prisma = _session_grouped_paginating_prisma(sessions) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page": cap // 7 + 1, + "page_size": 7, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["total"] == cap + assert [row["request_id"] for row in data["data"]] == [ + f"req-sess-{index:06d}" for index in range(cap - cap % 7, cap) + ] + assert data["has_more"] is True + assert data["next_session_cursor"] is not None + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort( client, monkeypatch @@ -6777,3 +7049,161 @@ async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_sess assert "next_session_cursor" not in data finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def _fake_prisma_with_owned_spend_log(owner_user_id, messages_json, response_json): + class _Row: + user = owner_user_id + team_id = None + + class _SpendLogs: + async def find_unique(self, where, include=None): + return _Row() + + class _DB: + def __init__(self): + self.litellm_spendlogs = _SpendLogs() + + async def query_raw(self, _sql, *_args): + return [ + { + "messages": messages_json, + "response": response_json, + "proxy_server_request": "{}", + "metadata": "{}", + } + ] + + class _Prisma: + def __init__(self): + self.db = _DB() + + return _Prisma() + + +def test_ui_view_request_response_internal_user_owner_gets_payload(client, monkeypatch): + """ + An internal_user who owns the spend-log row can fetch the Logs drawer + detail payload for their own request (regression for #34099, where the + route was blocked for INTERNAL_USER before reaching this ownership check). + """ + messages_json = json.dumps([{"role": "user", "content": "hi"}]) + response_json = json.dumps({"choices": [{"message": {"content": "hello"}}]}) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + _fake_prisma_with_owned_spend_log("user_a", messages_json, response_json), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_a" + ) + try: + response = client.get( + "/spend/logs/ui/req-owned-by-user-a", + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + body = response.json() + assert json.loads(body["messages"]) == [{"role": "user", "content": "hi"}] + assert json.loads(body["response"]) == { + "choices": [{"message": {"content": "hello"}}] + } + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +class _RecordingAdditionalLoggingUtils: + """Injectable custom logger that records every request_id it's asked for.""" + + def __init__(self, payload): + self._payload = payload + self.requested_ids = [] + + async def get_request_response_payload(self, request_id, start_time_utc, end_time_utc): + self.requested_ids.append(request_id) + return self._payload + + +def test_ui_view_request_response_internal_user_non_owner_forbidden(client, monkeypatch): + """ + A different internal_user requesting someone else's row is forbidden; + guards against _assert_user_can_view_request_id being skipped in the + detail-drawer handler. Also proves the handler stops before it ever asks + a custom logger or the DB for the payload. + """ + messages_json = json.dumps([{"role": "user", "content": "hi"}]) + response_json = json.dumps({"choices": [{"message": {"content": "hello"}}]}) + fake_prisma = _fake_prisma_with_owned_spend_log("user_a", messages_json, response_json) + original_query_raw = fake_prisma.db.query_raw + query_raw_calls = [] + + async def _spy_query_raw(*args, **kwargs): + query_raw_calls.append((args, kwargs)) + return await original_query_raw(*args, **kwargs) + + fake_prisma.db.query_raw = _spy_query_raw + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma) + + custom_logger = _RecordingAdditionalLoggingUtils({"messages": "should-not-be-returned"}) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_additional_logging_utils_from_custom_logger", + lambda: [custom_logger], + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_b" + ) + try: + response = client.get( + "/spend/logs/ui/req-owned-by-user-a", + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + assert custom_logger.requested_ids == [] + assert query_raw_calls == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def test_ui_view_request_response_internal_user_missing_row_forbidden(client, monkeypatch): + """ + Regression for the fail-open in _assert_user_can_view_request_id: a + request_id with no spend-log row (e.g. pruned by retention) must be + denied before the handler ever consults a custom logger, otherwise a + non-admin who guesses/obtains a request_id could read another tenant's + payload out of cold storage. Fails if `if row is None: return` is + reintroduced. + """ + + class _SpendLogs: + async def find_unique(self, where, include=None): + return None + + class _DB: + def __init__(self): + self.litellm_spendlogs = _SpendLogs() + + from types import SimpleNamespace + + fake_prisma = SimpleNamespace(db=_DB()) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma) + + custom_logger = _RecordingAdditionalLoggingUtils({"messages": "should-not-be-returned"}) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_additional_logging_utils_from_custom_logger", + lambda: [custom_logger], + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_a" + ) + try: + response = client.get( + "/spend/logs/ui/req-pruned", + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + assert custom_logger.requested_ids == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index a7de3f1d8d6..54e5a6d5385 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -528,8 +528,8 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key" session_rows = [ - {"session_key": "req-1", "api_key": "k", "last_activity": "2026-02-16 10:00:00"}, - {"session_key": "req-2", "api_key": "k", "last_activity": "2026-02-16 09:00:00"}, + {"session_key": f"req-{index}", "api_key": "k", "last_activity": f"2026-02-16 10:{59 - index:02d}:00"} + for index in range(51) ] representative_rows = [ {"request_id": "req-1", "api_key": "k", "metadata": "{}", "session_id": None}, @@ -538,7 +538,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): async def mock_query_raw(sql_query, *params): if "COUNT(*) AS total_count" in sql_query: - return [{"total_count": 12}] + return [{"total_count": 60}] if "DISTINCT ON" in sql_query: return representative_rows return session_rows @@ -590,11 +590,11 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): assert "COUNT(*) OVER ()" not in rep_sql assert [row["request_id"] for row in response["data"]] == ["req-1", "req-2"] - assert response["total"] == 12 + assert response["total"] == 60 assert response["total_is_capped"] is False - assert response["total_pages"] == 1 - assert response["has_more"] is False - assert response["next_session_cursor"] is None + assert response["total_pages"] == 2 + assert response["has_more"] is True + assert response["next_session_cursor"] == "2026-02-16 10:10:00|k|req-49" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 95ddc4477e1..dc8a87a97ad 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1,8 +1,8 @@ import asyncio import datetime import json -from datetime import timezone from collections.abc import Mapping +from datetime import timezone from typing import Any, Final, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -15,12 +15,15 @@ from litellm.constants import ( LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, + MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, + UNKNOWN_MODEL_SPEND_LOG_MODEL, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_messages_for_spend_logs_payload, _get_proxy_server_request_for_spend_logs_payload, @@ -35,11 +38,10 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, _sanitize_guardrail_information_for_spend_logs, _sanitize_request_body_for_spend_logs_payload, - _should_store_prompts_and_responses_in_spend_logs, get_logging_payload, get_spend_logs_id, + should_store_prompts_and_responses_in_spend_logs, ) -from litellm.proxy._types import SpendLogsPayload from litellm.proxy.utils import hash_token from litellm.types.utils import ( StandardLoggingHiddenParams, @@ -67,6 +69,30 @@ def _get_additional_usage_values_for_usage(usage: litellm.Usage) -> dict: return metadata["additional_usage_values"] +@pytest.mark.parametrize("store_prompts,redact", [(True, False), (False, False), (True, True)]) +def test_classifier_audit_spend_storage_obeys_privacy_and_truncation(monkeypatch, store_prompts, redact): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": store_prompts}) + audit: Final = { + "classifier_input": {"system": "rubric" * 1000, "messages": [{"role": "user", "content": "ask"}]}, + "originating_request_masked": {"input": "source-only", "api_key": "REDACTED"}, + } + stored: Final = json.loads(_get_proxy_server_request_for_spend_logs_payload( + metadata={}, litellm_params={"proxy_server_request": {"body": {"model": "classifier"}}}, + kwargs={"standard_logging_object": audit, "standard_callback_dynamic_params": {"turn_off_message_logging": redact}}, + )) + if not store_prompts or redact: + assert "classifier_input" not in stored + assert "originating_request_masked" not in stored + else: + assert stored["classifier_input"]["messages"] == audit["classifier_input"]["messages"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in json.dumps(stored["classifier_input"]) + assert stored["originating_request_masked"]["input"] == "source-only" + assert stored["model"] == "classifier" + assert audit["classifier_input"]["system"] == "rubric" * 1000 + + def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_tokens(): additional_usage_values = _get_additional_usage_values_for_usage( litellm.Usage( @@ -81,6 +107,48 @@ def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_token assert additional_usage_values["prompt_tokens_details"]["cached_tokens"] == 123 +class _HashingCache(litellm.Cache): + def __init__(self) -> None: + pass + + def get_cache_key(self, **kwargs) -> str: + raise AssertionError("a preset cache key must be reused instead of hashing the request") + + +def _cache_key_in_spend_log(monkeypatch: pytest.MonkeyPatch, cache: litellm.Cache | None, preset: str | None) -> str: + monkeypatch.setattr(litellm, "cache", cache) + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "x" * 10_000}], + "litellm_params": {"metadata": {"user_api_key": "test-key"}, "preset_cache_key": preset}, + }, + response_obj=litellm.ModelResponse(id="chatcmpl-test", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + return payload["cache_key"] + + +def test_get_logging_payload_reuses_the_preset_cache_key_instead_of_hashing_the_body(monkeypatch): + assert _cache_key_in_spend_log(monkeypatch, _HashingCache(), "preset-key") == "preset-key" + + +def test_get_logging_payload_records_cache_off_without_hashing(monkeypatch): + assert _cache_key_in_spend_log(monkeypatch, None, None) == "Cache OFF" + + +def test_get_logging_payload_still_hashes_when_caching_is_on_and_no_preset_key_exists(monkeypatch): + class _RecordingCache(litellm.Cache): + def __init__(self) -> None: + pass + + def get_cache_key(self, **kwargs) -> str: + return "hashed-from-" + kwargs["model"] + + assert _cache_key_in_spend_log(monkeypatch, _RecordingCache(), None) == "hashed-from-gpt-4o-mini" + + _TRACE_ONLY_STANDARD_LOGGING: Final = cast( StandardLoggingPayload, { @@ -132,6 +200,68 @@ def test_legacy_policy_keeps_trace_id_fallback(): assert len(str(generated)) == 36 +def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id(): + """The create call's request id IS the batch id and the poller's cost row appends + _batch_cost to it, so deriving the session from the request id lands both rows in one + trace on the logs UI even though the poller builds a fresh logging context per cycle.""" + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id + + create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1") + cost_session: Final = _get_batch_trace_session_id( + call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost" + ) + assert create_session == cost_session == "batch-uid-1" + + +def test_non_batch_call_types_derive_no_batch_session(): + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id + + assert _get_batch_trace_session_id(call_type="acompletion", request_id="chatcmpl-1") is None + + +def test_batch_session_outranks_the_per_request_trace_id(): + """Each batch lifecycle call carries its own auto-generated trace id, so letting the + trace id win would scatter the rows across sessions again.""" + session_id: Final = _get_session_id_for_spend_log( + kwargs={"litellm_trace_id": "trace-abc"}, + metadata={"trace_id": "trace-abc"}, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=False, + batch_trace_session_id="batch-uid-1", + ) + assert session_id == "batch-uid-1" + + +def test_omit_policy_still_suppresses_batch_sessions(): + session_id: Final = _get_session_id_for_spend_log( + kwargs={}, + metadata=None, + standard_logging_payload=None, + omit_when_missing=True, + batch_trace_session_id="batch-uid-1", + ) + assert session_id is None + + +def test_get_logging_payload_groups_batch_create_and_cost_rows_in_one_session(): + def _payload(call_type: str) -> SpendLogsPayload: + return get_logging_payload( + kwargs={ + "call_type": call_type, + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=litellm.ModelResponse(id="batch-uid-1", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + create_payload: Final = _payload("acreate_batch") + cost_payload: Final = _payload("aretrieve_batch") + assert cost_payload["request_id"] == "batch-uid-1_batch_cost" + assert create_payload["session_id"] == cost_payload["session_id"] == "batch-uid-1" + + @pytest.mark.parametrize( ("request_metadata", "expected"), [ @@ -525,11 +655,11 @@ def test_sanitize_request_body_for_spend_logs_payload_circular_reference(): assert sanitized == {"b": {"a": {}}} # Should return empty dict for circular reference -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_store_prompts_true( mock_should_store, ): - # When _should_store_prompts_and_responses_in_spend_logs returns True + # When should_store_prompts_and_responses_in_spend_logs returns True mock_should_store.return_value = True # Sample vector store request metadata @@ -543,11 +673,11 @@ def test_get_vector_store_request_for_spend_logs_payload_store_prompts_true( assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == "sensitive information" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_store_prompts_false( mock_should_store, ): - # When _should_store_prompts_and_responses_in_spend_logs returns False + # When should_store_prompts_and_responses_in_spend_logs returns False mock_should_store.return_value = False # Sample vector store request metadata @@ -563,7 +693,7 @@ def test_get_vector_store_request_for_spend_logs_payload_store_prompts_false( assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["type"] == "text" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_null_input(mock_should_store): # When input is None mock_should_store.return_value = False @@ -571,7 +701,7 @@ def test_get_vector_store_request_for_spend_logs_payload_null_input(mock_should_ assert result is None -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_realtime_returns_messages(mock_should_store): """ Test that _get_messages_for_spend_logs_payload returns messages @@ -598,7 +728,7 @@ def test_get_messages_for_spend_logs_realtime_returns_messages(mock_should_store assert parsed[1]["content"] == "What is the weather today?" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from messages.""" mock_should_store.return_value = True @@ -615,7 +745,7 @@ def test_get_messages_for_spend_logs_strips_null_bytes(mock_should_store): assert parsed[0]["content"] == "helloworld" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_realtime_empty_when_disabled(mock_should_store): """ Test that _get_messages_for_spend_logs_payload returns '{}' for realtime calls @@ -633,7 +763,7 @@ def test_get_messages_for_spend_logs_realtime_empty_when_disabled(mock_should_st assert result == "{}" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_non_realtime_returns_empty(mock_should_store): """ Test that _get_messages_for_spend_logs_payload returns '{}' for non-realtime @@ -651,7 +781,7 @@ def test_get_messages_for_spend_logs_non_realtime_returns_empty(mock_should_stor assert result == "{}" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_store): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB @@ -679,7 +809,7 @@ def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_ assert parsed["data"][0]["other_field"] == "value" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from response.""" mock_should_store.return_value = True @@ -692,7 +822,7 @@ def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store assert json.loads(response_json)["content"] == "answerhere" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_truncates_large_embedding( mock_should_store, ): @@ -747,7 +877,7 @@ def test_truncation_includes_db_safeguard_note(): ) -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_response_truncation_logs_info_message(mock_should_store): """ Test that when response is truncated before DB storage, an info log is emitted @@ -769,7 +899,7 @@ def test_response_truncation_logs_info_message(mock_should_store): assert "response was truncated" in log_msg -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_request_body_truncation_logs_info_message(mock_should_store): """ Test that when request body is truncated before DB storage, an info log is emitted. @@ -860,6 +990,93 @@ def test_safe_dumps_complex_metadata_like_object(): assert parsed["model"] == "gpt-4" +_RAW_MODEL_WITH_PROMPT: Final = "opus-4.6 Please summarize my medical records\nPatient has diabetes" + + +_BEDROCK_INFERENCE_PROFILE_ARN: Final = ( + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/claude-sonnet-4-5" +) +_OVERLONG_MODEL: Final = "m" * (MAX_SPEND_LOG_MODEL_NAME_LENGTH + 1) + + +@pytest.mark.parametrize( + ("requested_model", "failure", "expected_model"), + [ + ( + _RAW_MODEL_WITH_PROMPT, + ProxyModelNotFoundError(route="acompletion", model_name=_RAW_MODEL_WITH_PROMPT), + UNKNOWN_MODEL_SPEND_LOG_MODEL, + ), + ( + _RAW_MODEL_WITH_PROMPT, + ValueError("Upstream passthrough request failed with status 404"), + UNKNOWN_MODEL_SPEND_LOG_MODEL, + ), + (_OVERLONG_MODEL, ValueError("provider timed out"), UNKNOWN_MODEL_SPEND_LOG_MODEL), + ( + "gpt-5.2", + ProxyModelNotFoundError(route="acompletion", model_name="gpt-5.2"), + UNKNOWN_MODEL_SPEND_LOG_MODEL, + ), + ("gpt-5.2", ValueError("provider timed out"), "gpt-5.2"), + (_BEDROCK_INFERENCE_PROFILE_ARN, ValueError("provider timed out"), _BEDROCK_INFERENCE_PROFILE_ARN), + ( + "MCP: deepwiki-ask_question", + ValueError("Content blocked: keyword 'confidential' detected"), + "MCP: deepwiki-ask_question", + ), + ], +) +def test_get_logging_payload_replaces_rejected_or_prompt_shaped_models_with_the_placeholder( + requested_model: str, failure: Exception, expected_model: str +): + kwargs: Final = { + "model": requested_model, + "messages": [{"role": "user", "content": "hi"}], + "call_type": "acompletion", + "litellm_params": {"metadata": {"user_api_key": "sk-test", "status": "failure"}}, + } + + payload: Final = get_logging_payload( + kwargs=kwargs, + response_obj=failure, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["model"] == expected_model + + +@pytest.mark.parametrize( + ("metadata", "response_obj"), + [ + ({"user_api_key": "sk-test"}, litellm.ModelResponse(id="chatcmpl-test", choices=[])), + ( + {"user_api_key": "sk-test", "model_group": "team alias", "status": "failure"}, + ValueError("provider timed out"), + ), + ], +) +def test_get_logging_payload_keeps_a_whitespace_model_name_on_success_or_a_routed_failure( + metadata: dict[str, str], response_obj: litellm.ModelResponse | Exception +): + kwargs: Final = { + "model": _RAW_MODEL_WITH_PROMPT, + "messages": [{"role": "user", "content": "hi"}], + "call_type": "acompletion", + "litellm_params": {"metadata": metadata}, + } + + payload: Final = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["model"] == _RAW_MODEL_WITH_PROMPT + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_api_key_preserved_when_standard_logging_payload_is_none(): @@ -1350,7 +1567,7 @@ def test_get_logging_payload_handles_missing_overhead_gracefully(): ) -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_spend_logs_redacts_request_and_response_when_turn_off_message_logging_enabled( mock_should_store, ): @@ -1414,7 +1631,7 @@ def test_should_store_prompts_and_responses_in_spend_logs_case_insensitive_strin mock_get_secret_bool, ): """ - Test that _should_store_prompts_and_responses_in_spend_logs handles + Test that should_store_prompts_and_responses_in_spend_logs handles case-insensitive string values for store_prompts_in_spend_logs in general_settings. """ # Test case-insensitive string "true" variations @@ -1424,7 +1641,7 @@ def test_should_store_prompts_and_responses_in_spend_logs_case_insensitive_strin {"store_prompts_in_spend_logs": true_value}, ): mock_get_secret_bool.return_value = False # Ensure env var is False - result = _should_store_prompts_and_responses_in_spend_logs() + result = should_store_prompts_and_responses_in_spend_logs() assert result is True, f"Expected True for '{true_value}', got {result}" # Test boolean True @@ -1433,7 +1650,7 @@ def test_should_store_prompts_and_responses_in_spend_logs_case_insensitive_strin {"store_prompts_in_spend_logs": True}, ): mock_get_secret_bool.return_value = False - result = _should_store_prompts_and_responses_in_spend_logs() + result = should_store_prompts_and_responses_in_spend_logs() assert result is True, f"Expected True for boolean True, got {result}" # Test that non-true values fall back to environment variable @@ -1444,22 +1661,22 @@ def test_should_store_prompts_and_responses_in_spend_logs_case_insensitive_strin ): # When env var is True, should return True mock_get_secret_bool.return_value = True - result = _should_store_prompts_and_responses_in_spend_logs() + result = should_store_prompts_and_responses_in_spend_logs() assert result is True, f"Expected True (from env var) for '{false_value}', got {result}" # When env var is False, should return False mock_get_secret_bool.return_value = False - result = _should_store_prompts_and_responses_in_spend_logs() + result = should_store_prompts_and_responses_in_spend_logs() assert result is False, f"Expected False (from env var) for '{false_value}', got {result}" # Test when general_settings doesn't have the key at all with patch("litellm.proxy.proxy_server.general_settings", {}): mock_get_secret_bool.return_value = True - result = _should_store_prompts_and_responses_in_spend_logs() + result = should_store_prompts_and_responses_in_spend_logs() assert result is True, "Expected True (from env var) when key missing, got False" mock_get_secret_bool.return_value = False - result = _should_store_prompts_and_responses_in_spend_logs() + result = should_store_prompts_and_responses_in_spend_logs() assert result is False, "Expected False (from env var) when key missing, got True" @@ -1492,7 +1709,7 @@ def test_get_spend_logs_metadata_guardrail_info_fallback_from_metadata(): assert result["guardrail_information"] is None -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_redacts_all_prompt_carrying_fields_when_flag_false( mock_should_store, ): @@ -1528,7 +1745,7 @@ def test_sanitize_guardrail_information_redacts_all_prompt_carrying_fields_when_ assert entry["guardrail_action"] == "NONE" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_redacts_prompt_fields_when_flag_false( mock_should_store, ): @@ -1592,7 +1809,7 @@ def test_sanitize_guardrail_information_redacts_prompt_fields_when_flag_false( } -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_preserves_guardrail_usage_when_flag_false( mock_should_store, ): @@ -1624,7 +1841,7 @@ def test_sanitize_guardrail_information_preserves_guardrail_usage_when_flag_fals assert entry["guardrail_usage"] == {"topicPolicyUnits": 1, "contentPolicyUnits": 1, "wordPolicyUnits": 0} -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_passthrough_when_flag_true( mock_should_store, ): @@ -1647,13 +1864,13 @@ def test_sanitize_guardrail_information_passthrough_when_flag_true( assert result == guardrail_info -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_none_passthrough(mock_should_store): mock_should_store.return_value = False assert _sanitize_guardrail_information_for_spend_logs(None) is None -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_normalizes_bare_dict_input(mock_should_store): """ Regression: xecguard (xecguard.py:246) assigns a bare dict to @@ -1685,7 +1902,7 @@ def test_sanitize_guardrail_information_normalizes_bare_dict_input(mock_should_s assert entry["start_time"] == 1.0 -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_drops_non_dict_items_in_list(mock_should_store): """ A stray non-dict item in the list (e.g. from a buggy caller that @@ -1704,7 +1921,7 @@ def test_sanitize_guardrail_information_drops_non_dict_items_in_list(mock_should assert result == [{"guardrail_name": "x", "guardrail_response": REDACTED_BY_LITELM_STRING}] -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_preserves_absent_prompt_fields(mock_should_store): """ Entries that never carried guardrail_request or guardrail_response must @@ -2120,7 +2337,7 @@ def test_sanitize_request_body_strips_secret_fields(): assert sanitized["messages"] == [{"role": "user", "content": "hi"}] -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_proxy_server_request_payload_excludes_secret_fields(mock_should_store): """ End-to-end test: when the proxy_server_request body contains @@ -2200,7 +2417,7 @@ def test_redact_prompt_leaks_empty_string(): assert _redact_prompt_leaks_in_error_string("") == "" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_redacts_when_not_storing_prompts( mock_should_store, ): @@ -2228,7 +2445,7 @@ def test_sanitize_error_information_redacts_when_not_storing_prompts( assert sanitized["llm_provider"] == "openai" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_skips_redaction_when_storing_prompts( mock_should_store, ): @@ -2250,7 +2467,7 @@ def test_sanitize_error_information_skips_redaction_when_storing_prompts( assert REDACTED_BY_LITELM_STRING not in sanitized["error_message"] -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_caps_size_regardless_of_prompt_flag( mock_should_store, ): @@ -2281,7 +2498,7 @@ def test_sanitize_error_information_none_passthrough(): assert _sanitize_error_information_for_spend_logs(None) is None -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_reproduces_lit_2992(mock_should_store): # Mirrors the reproduced row body from LIT-2992 — a RateLimitError whose # message embeds 178 pydantic validation errors, each carrying a full @@ -2366,7 +2583,7 @@ def test_redact_prompt_leaks_handles_unterminated_value(): assert REDACTED_BY_LITELM_STRING in redacted -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_redacts_traceback_when_not_storing_prompts( mock_should_store, ): @@ -2398,7 +2615,7 @@ def test_sanitize_error_information_redacts_traceback_when_not_storing_prompts( assert "ValueError: invalid request" in sanitized["traceback"] -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_skips_traceback_redaction_when_storing_prompts( mock_should_store, ): @@ -2506,7 +2723,7 @@ def test_redact_prompt_leaks_combined_quoted_key_and_pydantic_assignment(): assert redacted.count(REDACTED_BY_LITELM_STRING) >= 2 -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_redacts_pydantic_assignment_form( mock_should_store, ): @@ -3102,7 +3319,7 @@ def test_get_logging_payload_hashes_bearer_prefixed_api_key(): ) -@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_preserves_headroom_compression_token_stats( mock_should_store, ): diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index b8fb6170d34..40ebc03781c 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2198,6 +2198,83 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): assert reservation["finalized"] is True +class _ExpiringRedisCache: + def __init__(self) -> None: + self.store: dict[str, float] = {} + + async def async_get_cache(self, key: str, *args: object, **kwargs: object) -> float | None: + return self.store.get(key) + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self.store[key] = self.store.get(key, 0.0) + float(value) + return self.store[key] + + async def async_set_max(self, key: str, value: float, **kwargs: object) -> float: + self.store[key] = max(self.store.get(key, float("-inf")), float(value)) + return self.store[key] + + async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool: + self.store[key] = float(value) + return True + + async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: + self.store.pop(key, None) + + async def async_increment_pipeline(self, increment_list, **kwargs): + results = [] + for op in increment_list: + results.append(await self.async_increment(op["key"], op["increment_value"])) + return results + + def get_ttl(self, **kwargs) -> None: + return None + + +@pytest.mark.asyncio +async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( + spend_counter_state, +): + """Redis key expired mid-stream while the pod's in-memory copy still holds the + reserved value: reconcile must reseed from the DB floor plus the settled cost + instead of applying ``actual - reserved`` to the empty key.""" + import litellm.proxy.proxy_server as ps + + counter_cache, _ = spend_counter_state + counter_key = "spend:team_member:user-expiry:team-expiry" + redis_cache = _ExpiringRedisCache() + counter_cache.redis_cache = redis_cache + counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.6) + + reservation = { + "reserved_cost": 0.6, + "entries": [ + { + "counter_key": counter_key, + "entity_type": "TeamMember", + "entity_id": "user-expiry:team-expiry", + "reserved_cost": 0.6, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + with patch.object( # test-quality-ok: the reseed reads the DB floor through a Prisma client the test has no seam for + ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.3) + ): + await ps.increment_spend_counters( + token="key-expiry", + team_id="team-expiry", + user_id="user-expiry", + response_cost=0.05, + budget_reservation=reservation, + ) + + assert redis_cache.store[counter_key] == pytest.approx(0.35) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.35) + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_should_invalidate_reserved_counters_after_persisted_spend_failure( spend_counter_state, diff --git a/tests/test_litellm/proxy/test_collector.py b/tests/test_litellm/proxy/test_collector.py new file mode 100644 index 00000000000..ac3e1f5e566 --- /dev/null +++ b/tests/test_litellm/proxy/test_collector.py @@ -0,0 +1,229 @@ +import asyncio +import logging +from collections.abc import Callable, Iterator +from pathlib import Path +from typing import Final + +import pytest +import uvloop + +from litellm._logging import verbose_logger, verbose_proxy_logger, verbose_router_logger +from litellm.proxy.collector import ( + SpendEventConsumer, + address_argument, + apply_log_level, + pod_pgbouncer_database_url, +) +from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings +from litellm.proxy.spend_tracking.spend_event_producer import ( + AddressError, + SpendEventProducer, + TcpAddress, + UnixAddress, + open_collector_connection, +) + + +class _Handler: + def __init__(self, fail_on: bytes | None = None) -> None: + self.lines: list[bytes] = [] # mutable-ok: test double records the events the consumer handed over + self._fail_on = fail_on + + async def __call__(self, line: bytes) -> None: + if line == self._fail_on: + raise RuntimeError("pipeline failed") + self.lines.append(line) + + +async def _no_fallback(line: bytes) -> None: + raise AssertionError(f"unexpected fallback for {line!r}") + + +class _Fallback: + def __init__(self) -> None: + self.lines: list[bytes] = [] # mutable-ok: test double records the events run in-process + + async def __call__(self, line: bytes) -> None: + self.lines.append(line) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["unix", "tcp"]) +async def test_consumer_handles_each_producer_line_once_in_order(tmp_path: Path, transport: str): + handler: Final = _Handler(fail_on=b"event-3\n") + consumer: Final = SpendEventConsumer(handler) + server: Final = await consumer.serve( + UnixAddress(path=str(tmp_path / "spend.sock")) if transport == "unix" else TcpAddress("127.0.0.1", 0) + ) + address: Final = ( + UnixAddress(path=str(tmp_path / "spend.sock")) + if transport == "unix" + else TcpAddress("127.0.0.1", server.sockets[0].getsockname()[1]) + ) + producer: Final = SpendEventProducer( + address=address, on_unavailable="fallback", buffer_size=100, connect_timeout=1.0, fallback=_no_fallback + ) + for i in range(6): + await producer.publish(f"event-{i}\n".encode()) + await producer.close(drain_timeout=5.0) + + server.close() + assert await consumer.drain(timeout=5.0) == 0 + assert handler.lines == [f"event-{i}\n".encode() for i in range(6) if i != 3] + assert (consumer.received, consumer.handled, consumer.failed) == (6, 5, 1) + + +@pytest.mark.asyncio +async def test_consumer_discards_a_truncated_trailing_event(tmp_path: Path): + handler: Final = _Handler() + consumer: Final = SpendEventConsumer(handler) + address: Final = UnixAddress(path=str(tmp_path / "spend.sock")) + server: Final = await consumer.serve(address) + _, writer = await open_collector_connection(address, timeout=1.0) + writer.write(b"whole\npartial-without-newline") + await writer.drain() + writer.close() + await writer.wait_closed() + await asyncio.sleep(0.05) + + server.close() + assert await consumer.drain(timeout=5.0) == 0 + assert handler.lines == [b"whole\n"] + assert consumer.received == 1 + + +@pytest.mark.asyncio +async def test_drain_reports_producers_still_connected_after_the_timeout(tmp_path: Path): + consumer: Final = SpendEventConsumer(_Handler()) + address: Final = UnixAddress(path=str(tmp_path / "spend.sock")) + server: Final = await consumer.serve(address) + _, writer = await open_collector_connection(address, timeout=1.0) + await asyncio.sleep(0.05) + + server.close() + assert await consumer.drain(timeout=0.1) == 1 + writer.close() + await writer.wait_closed() + assert await consumer.drain(timeout=5.0) == 0 + + +@pytest.mark.asyncio +async def test_graceful_stop_hands_the_producer_over_to_its_fallback_without_losing_events(tmp_path: Path): + handler: Final = _Handler() + fallback: Final = _Fallback() + consumer: Final = SpendEventConsumer(handler) + address: Final = UnixAddress(path=str(tmp_path / "spend.sock")) + server: Final = await consumer.serve(address) + producer: Final = SpendEventProducer( + address=address, on_unavailable="fallback", buffer_size=100, connect_timeout=1.0, fallback=fallback + ) + await producer.publish(b"event-1\n") + await asyncio.sleep(0.05) + + server.close() + draining: Final = asyncio.ensure_future(consumer.drain(timeout=5.0)) + await asyncio.sleep(0.05) + await producer.publish(b"event-2\n") + await producer.close(drain_timeout=5.0) + + assert await draining == 0 + assert handler.lines == [b"event-1\n"] + assert fallback.lines == [b"event-2\n"] + assert (producer.stats().sent, producer.stats().fallback) == (1, 1) + + +@pytest.mark.parametrize("loop_factory", [asyncio.new_event_loop, uvloop.new_event_loop], ids=["asyncio", "uvloop"]) +def test_drain_still_hands_over_live_producers_when_another_connection_already_died( + tmp_path: Path, loop_factory: Callable[[], asyncio.AbstractEventLoop] +): + """A transport the loop force-closed under a busy handler must not abort the half-close of the others.""" + + async def scenario() -> tuple[int, list[bytes]]: + release: Final = asyncio.Event() + + async def slow_handler(line: bytes) -> None: + await release.wait() + + consumer: Final = SpendEventConsumer(slow_handler) + address: Final = UnixAddress(path=str(tmp_path / "spend.sock")) + server: Final = await consumer.serve(address) + _, dead = await open_collector_connection(address, timeout=1.0) + dead.write(b"stuck\n") + await dead.drain() + await asyncio.sleep(0.05) + for connection in consumer._open_connections: # pyright: ignore[reportPrivateUsage] # force-close like uvloop does on a socket error + connection.transport.close() + dead.close() + fallback: Final = _Fallback() + producer: Final = SpendEventProducer( + address=address, on_unavailable="fallback", buffer_size=100, connect_timeout=1.0, fallback=fallback + ) + await producer.publish(b"event-1\n") + await asyncio.sleep(0.05) + + server.close() + draining: Final = asyncio.ensure_future(consumer.drain(timeout=0.5)) + await asyncio.sleep(0.05) + await producer.publish(b"event-2\n") + await producer.close(drain_timeout=5.0) + still_open: Final = await draining + release.set() + await asyncio.sleep(0.05) + return still_open, fallback.lines + + with asyncio.Runner(loop_factory=loop_factory) as runner: + still_open, fallback_lines = runner.run(scenario()) + + assert still_open == 2 + assert fallback_lines == [b"event-2\n"] + + +def test_address_argument(): + assert address_argument((), default="unix:///tmp/x.sock") == "unix:///tmp/x.sock" + assert address_argument(("--address", "tcp://127.0.0.1:4100"), default="unix:///tmp/x.sock") == ( + "tcp://127.0.0.1:4100" + ) + assert isinstance(address_argument(("--listen", "x"), default="unix:///tmp/x.sock"), AddressError) + + +def test_pod_pgbouncer_database_url_points_at_the_proxy_containers_pooler(): + """With pgbouncer on, the sidecar must not open its own upstream connections but share the pod's pooler.""" + upstream: Final = "postgresql://u:p@db.internal:5432/litellm?schema=public" + environ: Final = {"DATABASE_URL": upstream} + assert pod_pgbouncer_database_url(PgBouncerSettings(enabled=False), environ, token_auth=False) is None + assert ( + pod_pgbouncer_database_url(PgBouncerSettings(enabled=True, port=6543), environ, token_auth=False) + == "postgresql://u:p@127.0.0.1:6543/litellm?schema=public&pgbouncer=true" + ) + assert isinstance(pod_pgbouncer_database_url(PgBouncerSettings(enabled=True), {}, token_auth=False), PgBouncerError) + + +def test_pod_pgbouncer_database_url_goes_direct_under_token_auth(): + """The proxy's pgbouncer only knows the token that container minted, so the sidecar must mint its own upstream.""" + iam_upstream: Final = "postgresql://u@db.internal:5432/litellm?schema=public" + assert ( + pod_pgbouncer_database_url(PgBouncerSettings(enabled=True), {"DATABASE_URL": iam_upstream}, token_auth=True) + is None + ) + assert pod_pgbouncer_database_url(PgBouncerSettings(enabled=True), {}, token_auth=True) is None + + +@pytest.fixture +def restore_log_levels() -> Iterator[None]: + loggers: Final = (verbose_logger, verbose_router_logger, verbose_proxy_logger) + levels: Final = tuple(logger.level for logger in loggers) + yield + for logger, level in zip(loggers, levels, strict=True): + logger.setLevel(level) + + +@pytest.mark.usefixtures("restore_log_levels") +@pytest.mark.parametrize( + ("litellm_log", "expected"), + [("DEBUG", logging.DEBUG), ("info", logging.INFO), (None, logging.WARNING), ("loud", logging.WARNING)], +) +def test_apply_log_level_mirrors_the_proxy_env_contract(litellm_log: str | None, expected: int): + verbose_proxy_logger.setLevel(logging.WARNING) + apply_log_level(litellm_log) + assert verbose_proxy_logger.isEnabledFor(expected) + assert not verbose_proxy_logger.isEnabledFor(expected - 10) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 96d9b0c5a26..50b26577e5c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2,8 +2,8 @@ import asyncio import copy import datetime import json -from types import SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Optional +from types import MappingProxyType, SimpleNamespace +from typing import AsyncGenerator, Callable, Final, Iterator, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -417,7 +417,7 @@ class TestProxyBaseLLMRequestProcessing: ) fake_llm_router = MagicMock() - fake_llm_router.get_model_list.return_value = [ + fake_llm_router.deployments_for_request.return_value = [ { "model_name": "smart-router", "litellm_params": { @@ -1809,6 +1809,146 @@ class TestCommonRequestProcessingHelpers: response = await create_response(mock_generator(), "text/event-stream", custom_headers) assert response.headers["x-custom-header"] == "TestValue" + async def test_create_streaming_response_refresh_headers_after_first_chunk(self): + """LIT-6767: headers a caller can only resolve once the first chunk exists. + + A pre-first-chunk fallback replaces the deployment while the response + headers are still uncommitted, so ``refresh_headers`` is consulted after + the first chunk is buffered and its result wins. + """ + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + refresh_headers: Final = AsyncMock( + return_value={"x-litellm-model-id": "fallback-deployment", "llm_provider-x-request-id": "req-FALLBACK"} + ) + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment", "llm_provider-x-request-id": "req-FAILED"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, StreamingResponse) + assert refresh_headers.await_count == 1 + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + assert response.headers["llm_provider-x-request-id"] == "req-FALLBACK" + # the buffering headers are still applied on top of the refreshed set + assert response.headers["x-accel-buffering"] == "no" + assert response.headers["cache-control"] == "no-cache" + + async def test_create_streaming_response_refreshes_only_after_the_first_chunk(self): + """LIT-6767: the refresh has to be consulted after the generator produced a chunk. + + A pre-first-chunk fallback only repoints the response while that first chunk is + being produced, so a refresh consulted any earlier still describes the attempt + that failed and the headers go out wrong. + """ + first_chunk_produced: Final = asyncio.Event() + + async def mock_generator(): + first_chunk_produced.set() + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + served = "fallback-deployment" if first_chunk_produced.is_set() else "failed-deployment" + return {"x-litellm-model-id": served} + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + + async def test_create_streaming_response_empty_stream_uses_refreshed_headers(self): + """LIT-6767: a fallback that served nothing still gets to name itself. + + The empty-generator branch returns its own StreamingResponse, so it needs the + refreshed headers too or the client is told the failed deployment answered. + """ + + async def mock_generator(): + return + yield # make it an async generator + + async def refresh_headers(): + return {"x-litellm-model-id": "fallback-deployment"} + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, StreamingResponse) + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + assert response.headers["x-accel-buffering"] == "no" + + async def test_create_streaming_response_without_refresh_headers_is_unchanged(self): + """LIT-6767: the default keeps the caller-supplied headers verbatim.""" + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + ) + assert response.headers["x-litellm-model-id"] == "failed-deployment" + + async def test_create_streaming_response_refresh_headers_failure_keeps_stream(self): + """LIT-6767: the first chunk is already paid for, so a failing refresh + falls back to the caller's headers instead of erroring the stream.""" + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + raise RuntimeError("boom") + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, StreamingResponse) + assert response.status_code == status.HTTP_200_OK + assert response.headers["x-litellm-model-id"] == "failed-deployment" + assert await self.consume_stream(response) == [ + 'data: {"content": "data"}\n\n', + "data: [DONE]\n\n", + ] + + async def test_create_response_first_chunk_error_uses_refreshed_headers(self): + """LIT-6767: the JSON error response built from a bad first chunk carries + the refreshed headers too, so it cannot describe a deployment that no + longer served the request.""" + + async def mock_generator(): + yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + return {"x-litellm-model-id": "fallback-deployment"} + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, JSONResponse) + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + async def test_create_streaming_response_disables_proxy_buffering(self): """Regression for #28384: every StreamingResponse create_response returns must carry the headers that stop nginx/ingress/Envoy from buffering the @@ -2062,6 +2202,19 @@ class TestGuardrailBlockErrorPayloadNeverStringifiesNone: assert frame["error"]["param"] is None assert frame["error"]["code"] == "400" + def test_a_streaming_frame_keeps_the_status_a_proxy_exception_was_raised_with(self): + """ProxyException stores its status as the string ``code``, so a 429 raised before the + first chunk used to reach the SSE frame as a 500.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.common_request_processing import sse_error_payload + + error_status, error_obj = sse_error_payload( + ProxyException(message="Rate limit reached", type="rate_limit_error", param=None, code=429) + ) + + assert error_status == 429 + assert (error_obj["type"], error_obj["code"]) == ("rate_limit_error", "429") + @pytest.mark.parametrize( "status_code, expected_type", [ @@ -8014,3 +8167,256 @@ class TestDetachedStreamFailureHook: await logging_obj._on_detached_stream_failure(failure) assert [call["original_exception"] for call in recorder.calls] == [failure] + + +class TestStreamingResponseHeadersFollowFallback: + """LIT-6767: the streaming branch has to publish the deployment that served the stream.""" + + @staticmethod + def _fallback_adopting_stream(): + class _Stream: + def __init__(self) -> None: + self._hidden_params = { + "model_id": "failed-deployment", + "api_base": "http://127.0.0.1:20769/v1", + "additional_headers": {"llm_provider-stale-marker": "failed-deployment"}, + } + self.fallback_headers_adopted = False + + def adopt(self) -> None: + self._hidden_params = { + "model_id": "served-deployment", + "api_base": "https://api.openai.com", + "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, + } + self.fallback_headers_adopted = True + + return _Stream() + + @pytest.mark.asyncio + async def test_streaming_headers_name_the_deployment_that_served(self, monkeypatch): + """A pre-first-chunk fallback repoints the stream while the headers are still + uncommitted, so the published headers must describe the fallback, not the attempt + the Router picked first.""" + stream = self._fallback_adopting_stream() + + def select_data_generator(**kwargs): + async def generator(): + stream.adopt() + yield 'data: {"choices": [{"delta": {"content": "OK"}}]}\n\n' + yield "data: [DONE]\n\n" + + return generator() + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "lit-6767-call" + logging_obj._defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + + processor = ProxyBaseLLMRequestProcessing( + data={"model": "oa-midfail", "stream": True, "litellm_logging_obj": logging_obj} + ) + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock( + side_effect=lambda data, user_api_key_dict, response: response + ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-header": "kept"} + ) + + async def fake_route_request(**kwargs): + async def call(): + return stream + + return call() + + monkeypatch.setattr( + litellm.proxy.common_request_processing, "route_request", fake_route_request + ) + + result = await processor.base_process_llm_request( + request=Request(scope={"type": "http", "headers": []}), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + is_streaming_request=True, + skip_pre_call_logic=True, + ) + + assert isinstance(result, StreamingResponse) + assert result.headers["x-litellm-model-id"] == "served-deployment" + assert result.headers["x-litellm-model-api-base"] == "https://api.openai.com" + assert result.headers["llm_provider-x-request-id"] == "req-SERVED" + assert "llm_provider-stale-marker" not in result.headers + assert result.headers["x-callback-header"] == "kept" + + +class TestPassthroughHeadersAcceptImmutableMappings: + """LIT-6767: the streaming branch now hands the passthrough helpers an immutable mapping.""" + + def test_merge_passthrough_streaming_headers_accepts_a_read_only_mapping(self): + merged = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers( + response_headers=httpx.Headers({"content-type": "text/event-stream", "transfer-encoding": "chunked"}), + custom_headers=MappingProxyType({"x-litellm-model-id": "served-deployment"}), + ) + + assert merged["x-litellm-model-id"] == "served-deployment" + assert merged["content-type"] == "text/event-stream" + # the excluded hop-by-hop header is still dropped + assert "transfer-encoding" not in merged + + +@pytest.mark.asyncio +async def test_handle_llm_api_exception_forwards_provider_headers_on_http_status_error(): + """The httpx.HTTPStatusError branch dropped the headers its sibling branches forward. + + A Bedrock passthrough failure reaches this branch, so the request id was gone + before the client saw the response. + """ + import httpx + + from litellm.proxy._types import UserAPIKeyAuth + + request = httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse") + response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-passthrough-500"}, + content=b'{"message": "Amazon Bedrock is unable to process your request."}', + request=request, + ) + + processor = ProxyBaseLLMRequestProcessing(data={}) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + with pytest.raises(HTTPException) as exc_info: + await processor._handle_llm_api_exception( + e=httpx.HTTPStatusError("boom", request=request, response=response), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + assert exc_info.value.headers is not None + assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-passthrough-500" + + +class TestBackgroundResponseRetrievalGovernance: + """LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines.""" + + GOVERNED_MODEL_GROUP = "gpt-5.4-mini" + GOVERNED_MODEL_ID = "deployment-governed" + + @pytest.fixture + def policy_engine(self) -> Iterator[None]: + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + get_policy_registry().load_policies( + { + "response-governance": { + "guardrails": {"add": ["output-word-filter"]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": "output-word-filter", "on_pass": "allow", "on_fail": "block"}], + }, + } + } + ) + get_attachment_registry().load_attachments( + [{"policy": "response-governance", "models": [self.GOVERNED_MODEL_GROUP]}] + ) + yield + get_policy_registry().clear() + get_attachment_registry().clear() + + def _router(self) -> MagicMock: + from litellm.types.router import Deployment, LiteLLM_Params + + router = MagicMock() + router.get_deployment.side_effect = lambda model_id: ( + Deployment( + model_name=self.GOVERNED_MODEL_GROUP, + litellm_params=LiteLLM_Params(model=f"openai/{self.GOVERNED_MODEL_GROUP}"), + model_info={"id": model_id}, + ) + if model_id == self.GOVERNED_MODEL_ID + else None + ) + return router + + async def _pre_call(self, route_type: str, monkeypatch: pytest.MonkeyPatch) -> dict[str, object]: + from litellm.responses.utils import ResponsesAPIRequestUtils + + client_facing_response_id = "resp_opaque-client-facing-id" + encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id=self.GOVERNED_MODEL_ID, response_id="resp_upstream" + ) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"response_id": client_facing_response_id, "litellm_metadata": {}} + ) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def passthrough_add_litellm_data_to_request( + data: dict[str, object], **kwargs: object + ) -> dict[str, object]: + return data + + async def decrypting_pre_call_hook( + user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str + ) -> dict[str, object]: + if data.get("response_id") == client_facing_response_id: + data["response_id"] = encoded_response_id + return data + + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + passthrough_add_litellm_data_to_request, + ) + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=decrypting_pre_call_hook) + proxy_config = MagicMock(spec=ProxyConfig) + proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(), + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + route_type=route_type, + llm_router=self._router(), + ) + return returned_data + + @pytest.mark.asyncio + async def test_retrieving_a_background_response_attaches_its_model_post_call_pipeline( + self, policy_engine: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + data = await self._pre_call("aget_responses", monkeypatch) + + assert data["response_id"].startswith("resp_bGl0ZWxsbTpjdXN0b21f") + pipelines = data["litellm_metadata"]["_guardrail_pipelines"] + assert [(policy_name, [step.guardrail for step in pipeline.steps]) for policy_name, pipeline in pipelines] == [ + ("response-governance", ["output-word-filter"]) + ] + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + assert data["model"] is None + + @pytest.mark.asyncio + async def test_submitting_a_response_does_not_attach_pipelines_from_its_response_id( + self, policy_engine: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + data = await self._pre_call("aresponses", monkeypatch) + + assert "_guardrail_pipelines" not in data["litellm_metadata"] + assert "applied_policies" not in data["litellm_metadata"] diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index fdae11d517a..5aa80213134 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -6,6 +6,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.health_check_latest import LatestHealthCheckRow from litellm.proxy.health_endpoints._health_endpoints import ( _aggregate_health_check_results, _build_model_param_to_info_mapping, @@ -13,6 +15,7 @@ from litellm.proxy.health_endpoints._health_endpoints import ( _save_background_health_checks_to_db, _save_health_check_results_if_changed, _save_health_check_to_db, + latest_health_checks_endpoint, ) from litellm.proxy.utils import PrismaClient @@ -21,12 +24,8 @@ from litellm.proxy.utils import PrismaClient def mock_prisma(): """Simplified mock PrismaClient with bound methods""" client = MagicMock() - client.db.litellm_healthchecktable.create = AsyncMock( - return_value={"id": "test-id"} - ) - client.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[{"id": "1", "model_name": "test"}] - ) + client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "test-id"}) + client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[{"id": "1", "model_name": "test"}]) # Bind actual methods import types @@ -52,14 +51,10 @@ def mock_prisma(): ("healthy", 1, 0, False), # Database error case ], ) -async def test_save_health_check_result( - mock_prisma, status, healthy, unhealthy, should_succeed -): +async def test_save_health_check_result(mock_prisma, status, healthy, unhealthy, should_succeed): """Test health check result saving with various scenarios""" if not should_succeed: - mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception( - "DB Error" - ) + mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception("DB Error") result = await mock_prisma.save_health_check_result( model_name="test-model", @@ -187,9 +182,7 @@ def test_aggregate_health_check_results(): {"model": "gpt-4", "error": "Rate limit exceeded"}, ] - result = _aggregate_health_check_results( - model_param_to_info, healthy_endpoints, unhealthy_endpoints - ) + result = _aggregate_health_check_results(model_param_to_info, healthy_endpoints, unhealthy_endpoints) # Check gpt-3.5-turbo is healthy gpt35_key = ("model-123", "gpt-3.5-turbo") @@ -220,9 +213,7 @@ def test_aggregate_health_check_results_multiple_endpoints(): ] unhealthy_endpoints = [] - result = _aggregate_health_check_results( - model_param_to_info, healthy_endpoints, unhealthy_endpoints - ) + result = _aggregate_health_check_results(model_param_to_info, healthy_endpoints, unhealthy_endpoints) key = ("model-123", "gpt-3.5-turbo") assert result[key]["healthy_count"] == 2 @@ -398,7 +389,7 @@ async def test_save_background_health_checks_to_db(): start_time = 1234567890.0 - await _save_background_health_checks_to_db( + persisted = await _save_background_health_checks_to_db( mock_prisma, model_list, healthy_endpoints, @@ -407,7 +398,8 @@ async def test_save_background_health_checks_to_db(): "background_health_check", ) - # Should call get_all_latest_health_checks and save_health_check_result + # Should call get_all_latest_health_checks and save_health_check_result, and report completion + assert persisted is True mock_prisma.get_all_latest_health_checks.assert_called_once() mock_prisma.save_health_check_result.assert_called_once() @@ -418,22 +410,112 @@ async def test_save_background_health_checks_to_db(): assert call_kwargs["checked_by"] == "background_health_check" +def _two_model_results(): + return { + ("model-1", "gpt-4"): { + "model_name": "gpt-4", + "model_id": "model-1", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + }, + ("model-2", "gpt-4o"): { + "model_name": "gpt-4o", + "model_id": "model-2", + "healthy_count": 0, + "unhealthy_count": 1, + "error_message": "boom", + }, + } + + +@pytest.mark.asyncio +async def test_save_health_check_results_if_changed_awaits_every_write_and_reports_success(): + """Writes are awaited, not detached, so the caller can tell the cycle's persistence completed.""" + mock_prisma = MagicMock() + mock_prisma.save_health_check_result = AsyncMock(return_value={"id": "row"}) + + persisted = await _save_health_check_results_if_changed( + mock_prisma, _two_model_results(), {}, 1234567890.0, "background_health_check" + ) + + assert (persisted, mock_prisma.save_health_check_result.await_count) == (True, 2) + + +@pytest.mark.asyncio +async def test_save_health_check_results_if_changed_reports_failure_when_a_write_returns_none(): + """save_health_check_result swallows DB errors and returns None; that must surface as False.""" + mock_prisma = MagicMock() + mock_prisma.save_health_check_result = AsyncMock(side_effect=[{"id": "row"}, None]) + + persisted = await _save_health_check_results_if_changed( + mock_prisma, _two_model_results(), {}, 1234567890.0, "background_health_check" + ) + + assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 2) + + +@pytest.mark.asyncio +async def test_save_health_check_results_if_changed_reports_success_when_nothing_needed_writing(): + mock_prisma = MagicMock() + mock_prisma.save_health_check_result = AsyncMock() + model_results = { + ("model-1", "gpt-4"): { + "model_name": "gpt-4", + "model_id": "model-1", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + }, + } + latest_checks_map = { + "model-1": MagicMock(status="healthy", checked_at=datetime.now(timezone.utc) - timedelta(minutes=5)), + } + + persisted = await _save_health_check_results_if_changed( + mock_prisma, model_results, latest_checks_map, 1234567890.0, "background_health_check" + ) + + assert (persisted, mock_prisma.save_health_check_result.await_count) == (True, 0) + + +def _one_model_setup(): + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "model_info": {"id": "model-123"}, + "litellm_params": {"model": "gpt-3.5-turbo"}, + }, + ] + return model_list, [{"model": "gpt-3.5-turbo"}], [] + + +@pytest.mark.asyncio +async def test_save_background_health_checks_to_db_returns_false_when_a_write_fails(): + mock_prisma = MagicMock() + mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[]) + mock_prisma.save_health_check_result = AsyncMock(return_value=None) + model_list, healthy_endpoints, unhealthy_endpoints = _one_model_setup() + + persisted = await _save_background_health_checks_to_db( + mock_prisma, model_list, healthy_endpoints, unhealthy_endpoints, 1234567890.0, "background_health_check" + ) + + assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 1) + + @pytest.mark.asyncio async def test_save_background_health_checks_to_db_no_prisma(): """Test graceful handling when no prisma client""" - result = await _save_background_health_checks_to_db( - None, [], [], [], 0.0, "background_health_check" - ) - assert result is None + result = await _save_background_health_checks_to_db(None, [], [], [], 0.0, "background_health_check") + assert result is False @pytest.mark.asyncio async def test_save_background_health_checks_to_db_exception_handling(): """Test exception handling in background health check save""" mock_prisma = MagicMock() - mock_prisma.get_all_latest_health_checks = AsyncMock( - side_effect=Exception("DB Error") - ) + mock_prisma.get_all_latest_health_checks = AsyncMock(side_effect=Exception("DB Error")) model_list = [ { @@ -443,104 +525,134 @@ async def test_save_background_health_checks_to_db_exception_handling(): }, ] - # Should not raise exception, should handle gracefully - await _save_background_health_checks_to_db( + # Must not raise (the health check loop has to survive a DB outage) but must report + # the failure, so the window lock can be released for another pod to retry + persisted = await _save_background_health_checks_to_db( mock_prisma, model_list, [], [], 0.0, "background_health_check" ) - # Function should complete without raising + assert persisted is False + + +def _raw_latest_row(model_name: str, model_id, checked_at: datetime) -> dict: + return { + "health_check_id": f"hc-{model_id or 'no-id'}-{model_name}", + "model_name": model_name, + "model_id": model_id, + "status": "healthy", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + "response_time_ms": 10.0, + "details": None, + "checked_by": "pod-1", + "checked_at": checked_at.isoformat(), + "created_at": checked_at.isoformat(), + "updated_at": checked_at.isoformat(), + } @pytest.mark.asyncio -async def test_get_all_latest_health_checks_with_model_id(mock_prisma): - """Test get_all_latest_health_checks properly groups by model_id""" - mock_check2 = MagicMock() - mock_check2.model_id = "model-456" - mock_check2.model_name = "gpt-3.5-turbo" - mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=5) - - mock_check3 = MagicMock() - mock_check3.model_id = "model-123" - mock_check3.model_name = "gpt-3.5-turbo" - mock_check3.checked_at = datetime.now(timezone.utc) - timedelta( - minutes=1 - ) # Latest for model-123 - - # Order by checked_at desc - mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[mock_check3, mock_check2] - ) - - result = await mock_prisma.get_all_latest_health_checks() - - # Should return 2 unique models (by model_id) - assert len(result) == 2 - - # Should have latest check for each model_id - model_ids = {check.model_id for check in result} - assert "model-123" in model_ids - assert "model-456" in model_ids - - # model-123 should have the latest check (1 minute ago) - model123_check = next(c for c in result if c.model_id == "model-123") - assert model123_check.checked_at == mock_check3.checked_at - - -@pytest.mark.asyncio -async def test_get_all_latest_health_checks_without_model_id(mock_prisma): - """Test get_all_latest_health_checks groups by model_name when model_id is None""" - mock_check2 = MagicMock() - mock_check2.model_id = None - mock_check2.model_name = "gpt-3.5-turbo" - mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=1) # Latest - - mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[mock_check2] - ) - - result = await mock_prisma.get_all_latest_health_checks() - - # Should return 1 unique model (by model_name) - assert len(result) == 1 - assert result[0].model_name == "gpt-3.5-turbo" - assert result[0].checked_at == mock_check2.checked_at # Latest - - -@pytest.mark.asyncio -async def test_get_all_latest_health_checks_same_name_with_and_without_model_id( - mock_prisma, -): +async def test_get_all_latest_health_checks_keeps_every_distinct_group_with_its_own_checked_at(mock_prisma): """ - Same model_name can appear twice after DISTINCT ON: once keyed by (model_id, name) - and once by (NULL, name) — different Postgres groups than a single row with id. + Postgres owns the dedup. (id, name), (other id, name) and (NULL, name) are distinct groups and each row + must arrive typed, with its own checked_at, for the 1h re-save compare and the id-or-name lookup key. """ now = datetime.now(timezone.utc) - with_id = MagicMock() - with_id.model_id = "deployment-abc" - with_id.model_name = "gpt-4" - with_id.checked_at = now - timedelta(minutes=2) - - without_id = MagicMock() - without_id.model_id = None - without_id.model_name = "gpt-4" - without_id.checked_at = now - timedelta(minutes=1) - - mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[without_id, with_id] + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + _raw_latest_row("gpt-3.5-turbo", "model-123", now - timedelta(minutes=1)), + _raw_latest_row("gpt-3.5-turbo", "model-456", now - timedelta(minutes=5)), + _raw_latest_row("gpt-4", "deployment-abc", now - timedelta(minutes=2)), + _raw_latest_row("gpt-4", None, now - timedelta(minutes=3)), + ] ) result = await mock_prisma.get_all_latest_health_checks() - assert len(result) == 2 - names = {r.model_name for r in result} - assert names == {"gpt-4"} - ids = {r.model_id for r in result} - assert "deployment-abc" in ids - assert None in ids + assert {(check.model_id, check.model_name): check.checked_at for check in result} == { + ("model-123", "gpt-3.5-turbo"): now - timedelta(minutes=1), + ("model-456", "gpt-3.5-turbo"): now - timedelta(minutes=5), + ("deployment-abc", "gpt-4"): now - timedelta(minutes=2), + (None, "gpt-4"): now - timedelta(minutes=3), + } - by_key = {(r.model_id, r.model_name): r for r in result} - assert by_key[("deployment-abc", "gpt-4")].checked_at == with_id.checked_at - assert by_key[(None, "gpt-4")].checked_at == without_id.checked_at + +@pytest.mark.asyncio +async def test_save_background_health_checks_compares_raw_checked_at_against_utc_now(mock_prisma): + """ + Raw rows carry ISO strings and the engine may omit the offset. A naive checked_at would TypeError + inside the 1h compare, be swallowed, and silently stop every save; a stale row must still re-save. + """ + stale = (datetime.now(timezone.utc) - timedelta(hours=2)).replace(tzinfo=None) + fresh = datetime.now(timezone.utc) - timedelta(minutes=5) + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + _raw_latest_row("stale-model", "stale-id", stale), + _raw_latest_row("fresh-model", "fresh-id", fresh), + ] + ) + mock_prisma.save_health_check_result = AsyncMock() + model_list = [ + {"model_name": "stale-model", "model_info": {"id": "stale-id"}, "litellm_params": {"model": "openai/stale"}}, + {"model_name": "fresh-model", "model_info": {"id": "fresh-id"}, "litellm_params": {"model": "openai/fresh"}}, + ] + + await _save_background_health_checks_to_db( + mock_prisma, + model_list, + [{"model": "openai/stale"}, {"model": "openai/fresh"}], + [], + time.time(), + "pod-1", + ) + await asyncio.sleep(0) + + assert [call.kwargs["model_id"] for call in mock_prisma.save_health_check_result.await_args_list] == ["stale-id"] + + +@pytest.mark.asyncio +async def test_latest_health_checks_endpoint_serialises_raw_rows(monkeypatch): + row = LatestHealthCheckRow( + health_check_id="hc-1", + model_name="gpt-4", + model_id="deployment-abc", + status="healthy", + healthy_count=1, + unhealthy_count=0, + error_message=None, + response_time_ms=12.5, + details='{"region": "eu"}', + checked_by="pod-1", + checked_at=datetime(2026, 8, 25), + created_at=datetime(2026, 8, 25, tzinfo=timezone.utc), + updated_at=datetime(2026, 8, 25, tzinfo=timezone.utc), + ) + prisma = MagicMock() + prisma.get_all_latest_health_checks = AsyncMock(return_value=(row,)) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + + response = await latest_health_checks_endpoint(user_api_key_dict=UserAPIKeyAuth()) + + assert response == { + "latest_health_checks": { + "deployment-abc": { + "health_check_id": "hc-1", + "model_name": "gpt-4", + "model_id": "deployment-abc", + "status": "healthy", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + "response_time_ms": 12.5, + "details": {"region": "eu"}, + "checked_by": "pod-1", + "checked_at": "2026-08-25T00:00:00+00:00", + "created_at": "2026-08-25T00:00:00+00:00", + } + }, + "total_models": 1, + } @pytest.mark.asyncio @@ -628,12 +740,7 @@ def test_parse_background_health_check_model_groups_unset_returns_none(): assert parse_background_health_check_model_groups(None) is None assert parse_background_health_check_model_groups({}) is None - assert ( - parse_background_health_check_model_groups( - {"background_health_check_model_groups": None} - ) - is None - ) + assert parse_background_health_check_model_groups({"background_health_check_model_groups": None}) is None def test_parse_background_health_check_model_groups_list_returns_frozenset(): @@ -650,9 +757,7 @@ def test_parse_background_health_check_model_groups_malformed_raises(bad_value): from litellm.proxy.health_check import parse_background_health_check_model_groups with pytest.raises(ValueError, match="must be a list of model group names"): - parse_background_health_check_model_groups( - {"background_health_check_model_groups": bad_value} - ) + parse_background_health_check_model_groups({"background_health_check_model_groups": bad_value}) def test_filter_deployments_to_model_groups(): @@ -665,9 +770,7 @@ def test_filter_deployments_to_model_groups(): ] assert filter_deployments_to_model_groups(model_list, None) == tuple(model_list) - assert filter_deployments_to_model_groups( - model_list, frozenset({"prod-openai"}) - ) == (model_list[0], model_list[2]) + assert filter_deployments_to_model_groups(model_list, frozenset({"prod-openai"})) == (model_list[0], model_list[2]) assert filter_deployments_to_model_groups(model_list, frozenset()) == () diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 7070617ce3e..37b983d709a 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -23,6 +23,7 @@ from litellm.proxy.litellm_pre_call_utils import ( _get_dynamic_logging_metadata, _get_enforced_params, _get_metadata_variable_name, + _match_and_track_policies, _promoted_trace_control_fields, _resolve_credential_from_model_config, _resolve_provider_from_deployment, @@ -35,6 +36,7 @@ from litellm.proxy.litellm_pre_call_utils import ( ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY +from litellm.litellm_core_utils.redact_messages import _get_turn_off_message_logging_from_dynamic_params from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) @@ -767,6 +769,7 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], + "api_key": "request-key", } user_api_key_dict = UserAPIKeyAuth( @@ -794,6 +797,8 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r assert "proxy_server_request" not in snapshot_body, ( "proxy_server_request must be excluded from its own body snapshot to prevent the body from self-referencing" ) + assert "api_key" not in snapshot_body + assert updated["proxy_server_request"]["credential_fields"] == ("api_key",) def test_refresh_proxy_server_request_body_snapshot_picks_up_guardrail_masking(): @@ -4148,6 +4153,72 @@ async def test_add_guardrails_from_policy_engine(): attachment_registry._initialized = False +def test_match_and_track_policies_preserves_attachment_and_request_body_order(): + from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry + from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext + + attachment_policy_names = [f"attachment-policy-{index}" for index in range(8)] + request_body_policy_names = ["body-policy-1", "body-policy-2"] + policy_names = [*attachment_policy_names, *request_body_policy_names] + policies = {policy_name: Policy() for policy_name in policy_names} + attachment_registry = AttachmentRegistry() + attachment_registry.load_attachments( + [{"policy": policy_name, "scope": "*"} for policy_name in attachment_policy_names] + ) + + applied_policy_names, _ = _match_and_track_policies( + data={"metadata": {}}, + context=PolicyMatchContext(model="gpt-4"), + request_body_policies=request_body_policy_names, + policies_override=policies, + attachment_registry_override=attachment_registry, + ) + + assert applied_policy_names == policy_names + + +@pytest.mark.asyncio +async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_its_pipeline_also_steps(): + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import ( + GuardrailPipeline, + PipelineStep, + Policy, + PolicyAttachment, + PolicyGuardrails, + ) + + data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], "metadata": {}} + policy_registry = get_policy_registry() + policy_registry._policies = { + "response-governance": Policy( + guardrails=PolicyGuardrails(add=["pii_blocker"]), + pipeline=GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="pii_blocker")]), + ), + } + policy_registry._initialized = True + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [PolicyAttachment(policy="response-governance", scope="*")] + attachment_registry._initialized = True + + try: + await add_guardrails_from_policy_engine( + data=data, + metadata_variable_name="metadata", + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + assert data["metadata"]["guardrails"] == ["pii_blocker"] + assert data["metadata"]["_pipeline_managed_guardrails"] == {"pii_blocker"} + assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_from_data(): """ @@ -7272,7 +7343,12 @@ def _reserved_stamp_key(key_metadata: dict | None = None) -> UserAPIKeyAuth: ) -_PLANTED_STAMPS = {"attempted_fallbacks": 99, "original_model_group": "spoofed-group", "client_key": "client_value"} +_PLANTED_STAMPS = { + "attempted_fallbacks": 99, + "original_model_group": "spoofed-group", + "_client_output_ceiling": {"api_base": "https://attacker.example"}, + "client_key": "client_value", +} @pytest.mark.asyncio @@ -7301,6 +7377,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "litellm_metadata" not in updated assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] + assert "_client_output_ceiling" not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" @@ -7630,6 +7707,107 @@ async def test_missing_session_id_omit_keeps_client_supplied_session_id(): assert _spend_log_session_id(updated) == "client-session-1" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_body", + [ + {"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1"}, + {"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1", "metadata": {"trace_id": "trace-1"}}, + ], +) +async def test_missing_session_id_omit_keeps_body_litellm_session_id( + monkeypatch: pytest.MonkeyPatch, client_body: dict[str, object] +): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + updated = await add_litellm_data_to_request( + data=client_body, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + callback_session_id = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=SimpleNamespace(litellm_session_id=""), + litellm_params=get_litellm_params(litellm_session_id="cust-sess-1", metadata=updated["metadata"]), + ) + assert callback_session_id == "cust-sess-1" + assert updated["metadata"]["session_id"] == "cust-sess-1" + assert _spend_log_session_id(updated) == "cust-sess-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_body_litellm_session_id_does_not_override_metadata_session_id(): + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-4o", + "messages": [], + "litellm_session_id": "cust-sess-1", + "metadata": {"session_id": "meta-sess-1"}, + }, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["metadata"]["session_id"] == "meta-sess-1" + assert _spend_log_session_id(updated) == "meta-sess-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +async def test_missing_session_id_omit_keeps_metadata_session_id_on_litellm_metadata_routes(path: str): + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-4o", + "input": "hi", + "litellm_session_id": "cust-sess-1", + "metadata": {"session_id": "meta-sess-1"}, + }, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "meta-sess-1" + assert _spend_log_session_id(updated, "litellm_metadata") == "meta-sess-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +async def test_missing_session_id_omit_keeps_body_litellm_session_id_on_litellm_metadata_routes(path: str): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "input": "hi", "litellm_session_id": "cust-sess-1"}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "cust-sess-1" + assert _spend_log_session_id(updated, "litellm_metadata") == "cust-sess-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_ignores_empty_body_litellm_session_id(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], "litellm_session_id": ""}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert "session_id" not in updated["metadata"] + assert _spend_log_session_id(updated) is None + + @pytest.mark.asyncio async def test_missing_session_id_generate_reuses_traceparent_trace_id(): """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" @@ -7760,3 +7938,36 @@ async def test_client_supplied_omit_marker_never_reaches_the_spend_log( if general_settings.get("missing_session_id") == "generate" else "per-call-random-trace-id" ) + + +def test_default_team_settings_bool_turn_off_message_logging_redacts(): + from litellm.proxy.proxy_server import ProxyConfig + + pc = ProxyConfig() + pc.config = { + "litellm_settings": { + "default_team_settings": [ + { + "team_id": "team-redact", + "success_callback": ["gcs_bucket"], + "failure_callback": ["gcs_bucket"], + "turn_off_message_logging": True, + } + ] + } + } + + callback_metadata = LiteLLMProxyRequestSetup.add_team_based_callbacks_from_config( + team_id="team-redact", + proxy_config=pc, + ) + + assert callback_metadata is not None + assert callback_metadata.success_callback == ["gcs_bucket"] + assert callback_metadata.callback_vars == {"turn_off_message_logging": "True"} + assert ( + _get_turn_off_message_logging_from_dynamic_params( + {"standard_callback_dynamic_params": dict(callback_metadata.callback_vars)} + ) + is True + ) diff --git a/tests/test_litellm/proxy/test_pointfive_dashboard_config.py b/tests/test_litellm/proxy/test_pointfive_dashboard_config.py new file mode 100644 index 00000000000..9e46e7f0b3e --- /dev/null +++ b/tests/test_litellm/proxy/test_pointfive_dashboard_config.py @@ -0,0 +1,47 @@ +import json +from pathlib import Path + +import litellm +from litellm.integrations.custom_logger import CustomLogger + + +def _dashboard_configs() -> tuple[dict, ...]: + path = Path(litellm.__file__).parent / "integrations" / "callback_configs.json" + return tuple(json.loads(path.read_text())) + + +def _pointfive_config() -> dict: + return next(config for config in _dashboard_configs() if config["id"] == "pointfive") + + +def test_pointfive_appears_in_the_dashboard_callback_dropdown(): + """The dropdown is served from callback_configs.json, so an entry only in the dashboard source is invisible.""" + entry = _pointfive_config() + + assert entry["displayName"] == "PointFive" + assert entry["supports_key_team_logging"] is False + assert entry["dynamic_params"]["POINTFIVE_API_KEY"]["required"] is True + assert entry["dynamic_params"]["POINTFIVE_API_KEY"]["type"] == "password" + assert entry["dynamic_params"]["POINTFIVE_API_URL"]["required"] is False + + +def test_the_dropdown_logo_asset_exists(): + """A logo the dashboard cannot resolve degrades silently to a letter tile.""" + logo = _pointfive_config()["logo"] + repo_root = Path(litellm.__file__).parent.parent + asset = repo_root / "ui" / "litellm-dashboard" / "public" / "assets" / "logos" / logo + + assert logo == "pointfive.png" + assert asset.is_file() + + +def test_the_dropdown_fields_are_the_env_vars_the_logger_reads(): + """ + The field names are the environment variables verbatim. + + The proxy would uppercase them either way, but naming them as stored means the edit + form finds the saved values and prefills them instead of showing blanks. + """ + fields = tuple(_pointfive_config()["dynamic_params"]) + + assert fields == tuple(CustomLogger.get_callback_env_vars("pointfive")) diff --git a/tests/test_litellm/proxy/test_pointfive_ui_callback.py b/tests/test_litellm/proxy/test_pointfive_ui_callback.py new file mode 100644 index 00000000000..13591d92005 --- /dev/null +++ b/tests/test_litellm/proxy/test_pointfive_ui_callback.py @@ -0,0 +1,15 @@ +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import AllCallbacks + + +def test_pointfive_is_offered_in_the_ui_callback_registry(): + """The proxy ui builds its form from this registry, so an absent entry is an absent form.""" + entry = AllCallbacks().pointfive + + assert entry.litellm_callback_name == "pointfive" + assert entry.ui_callback_name == "PointFive" + + +def test_the_ui_offers_the_two_settings_the_plugin_reads(): + """get_callback_env_vars is what the ui renders; it must match what the logger looks up.""" + assert tuple(CustomLogger.get_callback_env_vars("pointfive")) == ("POINTFIVE_API_KEY", "POINTFIVE_API_URL") diff --git a/tests/test_litellm/proxy/test_prisma_engine_watchdog.py b/tests/test_litellm/proxy/test_prisma_engine_watchdog.py index d73f74c5cd2..ae750b0770a 100644 --- a/tests/test_litellm/proxy/test_prisma_engine_watchdog.py +++ b/tests/test_litellm/proxy/test_prisma_engine_watchdog.py @@ -18,12 +18,15 @@ import asyncio import os import threading import time +from typing import Final from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest from litellm.proxy.utils import PrismaClient, ProxyLogging +WRITER_PROBE_SQL: Final = "SELECT current_setting('transaction_read_only') AS transaction_read_only" + @pytest.fixture(autouse=True) def mock_prisma_binary(): @@ -260,7 +263,7 @@ async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive( """Direct reconnect (engine alive) probes the writer first and skips the recreate when the probe is healthy. - The engine-alive path now runs a SELECT 1 probe before recreating. A + The engine-alive path now runs a writability probe before recreating. A healthy probe means the connection is fine — e.g. an IAM token refresh already replaced the engine (issue #29176) — so recreating would kill a working engine. Recreate happens only when the probe fails (covered in @@ -278,7 +281,7 @@ async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive( await engine_client._run_reconnect_cycle(timeout_seconds=5.0) engine_client.db.recreate_prisma_client.assert_not_awaited() - engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.query_raw.assert_awaited_once_with(WRITER_PROBE_SQL) engine_client.db.disconnect.assert_not_awaited() engine_client._start_engine_watcher.assert_awaited_once() @@ -296,7 +299,7 @@ async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown( await engine_client._run_reconnect_cycle(timeout_seconds=5.0) engine_client.db.recreate_prisma_client.assert_not_awaited() - engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.query_raw.assert_awaited_once_with(WRITER_PROBE_SQL) engine_client.db.disconnect.assert_not_awaited() engine_client._start_engine_watcher.assert_awaited_once() diff --git a/tests/test_litellm/proxy/test_prometheus_cleanup.py b/tests/test_litellm/proxy/test_prometheus_cleanup.py index 93b9b694c2c..6a1b95c51ff 100644 --- a/tests/test_litellm/proxy/test_prometheus_cleanup.py +++ b/tests/test_litellm/proxy/test_prometheus_cleanup.py @@ -6,13 +6,77 @@ ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir. from __future__ import annotations import os +import subprocess +import sys +from pathlib import Path +from typing import Final from unittest.mock import patch import pytest +from prometheus_client import CollectorRegistry, multiprocess -from litellm.proxy.prometheus_cleanup import mark_worker_exit, wipe_directory +from litellm.proxy.prometheus_cleanup import mark_dead_workers, mark_worker_exit, wipe_directory from litellm.proxy.proxy_cli import ProxyInitializationHelpers +_WORKER: Final = """ +import sys, time +from prometheus_client import Gauge +Gauge("litellm_in_flight", "", multiprocess_mode="livesum").set(float(sys.argv[1])) +print("ready", flush=True) +if sys.argv[2] == "stay": + time.sleep(120) +""" + + +def _spawn_worker(directory: Path, in_flight: str, lifetime: str) -> subprocess.Popen[str]: + env = {**os.environ, "PROMETHEUS_MULTIPROC_DIR": str(directory)} + worker = subprocess.Popen( + [sys.executable, "-c", _WORKER, in_flight, lifetime], env=env, stdout=subprocess.PIPE, text=True + ) + assert worker.stdout is not None and worker.stdout.readline() == "ready\n" + return worker + + +def _livesum(directory: Path) -> float: + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry, path=str(directory)) + value = registry.get_sample_value("litellm_in_flight") + return 0.0 if value is None else value + + +class TestMarkDeadWorkers: + def test_drops_live_gauges_of_exited_workers_and_keeps_running_ones(self, tmp_path: Path) -> None: + """A worker that died mid-request leaves its livesum file behind; the replacement worker's startup prune + must remove exactly that file so the aggregate stops counting requests nobody is serving.""" + dead = _spawn_worker(tmp_path, "3", "exit") + assert dead.wait(timeout=30) == 0 + alive = _spawn_worker(tmp_path, "2", "stay") + try: + assert (tmp_path / f"gauge_livesum_{dead.pid}.db").exists() + assert _livesum(tmp_path) == 5.0 + + assert mark_dead_workers(str(tmp_path)) == (dead.pid,) + + assert not (tmp_path / f"gauge_livesum_{dead.pid}.db").exists() + assert (tmp_path / f"gauge_livesum_{alive.pid}.db").exists() + assert _livesum(tmp_path) == 2.0 + assert mark_dead_workers(str(tmp_path)) == () + finally: + alive.kill() + alive.wait(timeout=30) + + def test_leaves_counters_of_exited_workers_alone(self, tmp_path: Path) -> None: + (tmp_path / "counter_424242.db").touch() + (tmp_path / "histogram_424242.db").touch() + assert mark_dead_workers(str(tmp_path)) == () + assert sorted(p.name for p in tmp_path.glob("*.db")) == ["counter_424242.db", "histogram_424242.db"] + + def test_keeps_live_gauges_of_workers_it_may_not_signal(self, tmp_path: Path) -> None: + """Signal 0 to pid 1 raises PermissionError for an unprivileged proxy; that pid is alive, not dead.""" + (tmp_path / "gauge_livesum_1.db").touch() + assert mark_dead_workers(str(tmp_path)) == () + assert (tmp_path / "gauge_livesum_1.db").exists() + class TestWipeDirectory: def test_deletes_all_db_files(self, tmp_path): diff --git a/tests/test_litellm/proxy/test_prometheus_metrics_server.py b/tests/test_litellm/proxy/test_prometheus_metrics_server.py index fc1fa381fa4..e2f461e61a6 100644 --- a/tests/test_litellm/proxy/test_prometheus_metrics_server.py +++ b/tests/test_litellm/proxy/test_prometheus_metrics_server.py @@ -1,5 +1,5 @@ -"""The separate metrics server must aggregate PROMETHEUS_MULTIPROC_DIR, expose only /metrics, and follow its -parent's lifetime. +"""The separate metrics server must aggregate PROMETHEUS_MULTIPROC_DIR, expose /metrics plus a probe-friendly +/health, and follow its parent's lifetime. Everything here runs on loopback against a child of this test process; no LLM keys or external network. """ @@ -91,7 +91,10 @@ def test_metrics_app_aggregates_multiproc_dir_and_reports_pid(tmp_path: Path, mo assert metrics.headers[PID_HEADER] == str(os.getpid()) assert 'litellm_requests_metric_total{model="gpt-5"} 5.0' in metrics.text - assert client.get("/health").status_code == 404 + health: Final = client.get("/health") + assert health.status_code == 200 + assert health.json() == {"status": "healthy", "multiproc_dir": str(tmp_path)} + assert client.get("/docs").status_code == 404 empty: Final = TestClient(build_metrics_app(str(other_dir))).get("/metrics") assert empty.status_code == 200 diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 0c20d5e0ff0..c76ff189a8a 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1525,7 +1525,12 @@ class TestProxyInitializationHelpers: def capture_run(self): captured["options"] = dict(self.options) - with patch("gunicorn.app.base.BaseApplication.run", capture_run): + with ( + patch("gunicorn.app.base.BaseApplication.run", capture_run), + patch( # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership + "litellm.proxy.proxy_cli.start_query_engine_reaper" + ), + ): ProxyInitializationHelpers._run_gunicorn_server( host="127.0.0.1", port=4010, @@ -1553,6 +1558,9 @@ class TestProxyInitializationHelpers: with ( patch("gunicorn.app.base.BaseApplication.run", capture_run), patch("builtins.print") as mock_print, + patch( # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership + "litellm.proxy.proxy_cli.start_query_engine_reaper" + ), ): ProxyInitializationHelpers._run_gunicorn_server( host="127.0.0.1", diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 9f1321aec2c..28ff4571b44 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -637,14 +637,14 @@ def test_callback_capabilities_excludes_opted_out_guardrail_from_iterator_overri assert [cb for cb, _ in caps.iterator_overrides if cb is opted_out] == [] -def test_deployment_pre_call_target_stays_native_when_opted_out(): +def test_deployment_hook_target_stays_native_when_opted_out(): """Model-level guardrails resolve their target here rather than through ProxyLogging.""" - assert _KeepsNativeHooks()._deployment_pre_call_target() is not None + assert _KeepsNativeHooks()._deployment_hook_target() is not None opted_out = _KeepsNativeHooks() - assert opted_out._deployment_pre_call_target() is opted_out - assert _AppliesGuardrail()._deployment_pre_call_target() is not None + assert opted_out._deployment_hook_target() is opted_out + assert _AppliesGuardrail()._deployment_hook_target() is not None routed = _AppliesGuardrail() - assert routed._deployment_pre_call_target() is not routed + assert routed._deployment_hook_target() is not routed @pytest.mark.asyncio @@ -671,6 +671,95 @@ async def test_deferred_stream_guardrails_run_native_hook_when_opted_out(monkeyp assert routed.native_hooks_ran == [] +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monkeypatch): + """A post_call pipeline step already ran the opted-out guardrail's own hook against + the buffered stream, so the deferred audit must not run it a second time.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed]) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")]) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == [] + + +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_whose_pipeline_could_not_stream(monkeypatch): + """A pipeline step with neither streaming interface keeps the whole pipeline off the + stream, so the deferred audit is the only place the opted-out guardrail's own hook + still runs, the way it did before pipelines ran on streams.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + class NeitherHookGuardrail(CustomGuardrail): + pass + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + neither = NeitherHookGuardrail(guardrail_name="gr-neither", event_hook=GuardrailEventHooks.post_call) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed, neither]) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="keeps_native", on_fail="next"), + PipelineStep(guardrail="gr-neither", on_fail="block"), + ], + ) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_on_route_without_translation(monkeypatch): + """A route with no endpoint guardrail translation cannot gate the stream through its + pipelines, so the deferred audit still owes the opted-out guardrail its own hook.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed]) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")]) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/custom/stream"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + @pytest.mark.asyncio async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch): """The realtime path calls apply_guardrail directly, so the opt-out has to be diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 9393ec0f8e6..3b3031647dc 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6,6 +6,7 @@ import os import re import socket import subprocess +import time import types from datetime import datetime, timedelta, timezone from pathlib import Path @@ -28,7 +29,7 @@ from litellm.caching.caching import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app, initialize from litellm.utils import _invalidate_model_cost_lowercase_map @@ -822,16 +823,16 @@ def _mock_scheduled_proxy_config() -> MagicMock: @pytest.mark.asyncio -async def test_initialize_scheduled_jobs_credentials(monkeypatch): - """ - Test that get_credentials is only called when store_model_in_db is True - """ +async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_deployment( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from apscheduler.schedulers.asyncio import AsyncIOScheduler + from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging - # Mock dependencies mock_prisma_client = MagicMock() mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() @@ -841,25 +842,6 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch("litellm.proxy.proxy_server.store_model_in_db", False), - ): # set store_model_in_db to False - # Test when store_model_in_db is False - await ProxyStartupEvent.initialize_scheduled_background_jobs( - general_settings={}, - prisma_client=mock_prisma_client, - proxy_budget_rescheduler_min_time=1, - proxy_budget_rescheduler_max_time=2, - proxy_batch_write_at=5, - proxy_logging_obj=mock_proxy_logging, - ) - - # Verify get_credentials was not called - mock_proxy_config.get_credentials.assert_not_called() - - # Now test with store_model_in_db = True - with ( - patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), - patch("litellm.proxy.proxy_server.store_model_in_db", True), - patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True), ): await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings={}, @@ -870,12 +852,31 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): proxy_logging_obj=mock_proxy_logging, ) - # Verify get_credentials was called both directly and scheduled - assert mock_proxy_config.get_credentials.call_count == 1 # Direct call + mock_proxy_config.get_credentials.assert_not_called() + mock_proxy_config.add_deployment.assert_not_called() - # Verify a scheduled job was added for get_credentials - mock_scheduler_calls = [call[0] for call in mock_proxy_config.get_credentials.mock_calls] - assert len(mock_scheduler_calls) > 0 + scheduler = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + assert scheduler.get_job("get_credentials_job") is None + assert scheduler.get_job("add_deployment_job") is not None + mock_proxy_config.get_credentials.assert_not_called() + assert mock_proxy_config.add_deployment.call_count == 1 + finally: + scheduler.shutdown(wait=False) @pytest.mark.asyncio @@ -924,7 +925,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat @pytest.mark.asyncio async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch): """ - The DB config-reload jobs (add_deployment, get_credentials) that keep multi-pod + The DB config-reload job (add_deployment) that keeps multi-pod deployments in sync must be scheduled at the configured proxy_config_reload_interval_seconds, not a hardcoded value. """ @@ -967,7 +968,7 @@ async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval( if "id" in job_call.kwargs } assert scheduled_seconds["add_deployment_job"] == configured_interval - assert scheduled_seconds["get_credentials_job"] == configured_interval + assert "get_credentials_job" not in scheduled_seconds @pytest.mark.asyncio @@ -1011,7 +1012,7 @@ async def test_initialize_scheduled_jobs_rejects_non_positive_config_reload_inte if "id" in job_call.kwargs } assert scheduled_seconds["add_deployment_job"] == 30 - assert scheduled_seconds["get_credentials_job"] == 30 + assert "get_credentials_job" not in scheduled_seconds @pytest.mark.asyncio @@ -3166,6 +3167,47 @@ async def test_custom_ui_sso_sign_in_handler_config_loading(): os.unlink(config_file_path) +@pytest.mark.asyncio +async def test_startup_initializes_string_callbacks_after_all_litellm_settings_load(tmp_path, monkeypatch): + from litellm.integrations.s3_v2 import S3Logger + from litellm.litellm_core_utils import litellm_logging + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import ProxyLogging + + config_file = tmp_path / "config.yaml" + config_file.write_text( + "model_list: []\n" + "litellm_settings:\n" + " success_callback:\n" + " - s3_v2\n" + " failure_callback:\n" + " - s3_v2\n" + " s3_callback_params:\n" + " s3_bucket_name: ordering-regression-bucket\n" + " s3_region_name: us-west-2\n" + ) + + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "s3_callback_params", None) + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", []) + + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + ProxyLogging(user_api_key_cache=MagicMock())._init_litellm_callbacks(llm_router=None) + + success_loggers = [cb for cb in litellm._async_success_callback if isinstance(cb, S3Logger)] + failure_loggers = [cb for cb in litellm._async_failure_callback if isinstance(cb, S3Logger)] + assert len(success_loggers) == 1 + assert len(failure_loggers) == 1 + assert success_loggers[0].s3_bucket_name == "ordering-regression-bucket" + assert success_loggers[0].s3_region_name == "us-west-2" + assert "s3_v2" not in litellm.success_callback + assert "s3_v2" not in litellm.failure_callback + + @pytest.mark.asyncio async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeypatch): """ @@ -7271,6 +7313,88 @@ async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins( assert ps.general_settings["apply_user_budget_to_team_keys"] is True +def _fill_user_api_key_cache(cache: DualCache, count: int) -> None: + for index in range(count): + cache.set_cache(key=f"key-{index}", value={"token": f"key-{index}"}, local_only=True) + + +@pytest.mark.asyncio +async def test_update_general_settings_user_api_key_cache_max_size_resizes_the_running_cache(monkeypatch): + """The Admin UI writes the capacity to the DB config, so the running cache has + to pick it up on reload; otherwise the knob only works after a restart.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + cache = UserApiKeyCache() + monkeypatch.setattr(proxy_server_module, "general_settings", {}) + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await ProxyConfig()._update_general_settings(db_general_settings={"user_api_key_cache_max_size": 300}) + + assert proxy_server_module.general_settings["user_api_key_cache_max_size"] == 300 + + _fill_user_api_key_cache(cache, 250) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + +@pytest.mark.asyncio +async def test_update_general_settings_clearing_user_api_key_cache_max_size_restores_the_default(monkeypatch): + """Blanking the field in the dashboard deletes the key, so the cache must fall + back to the default capacity rather than keep the last configured size.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + cache = UserApiKeyCache() + cache.update_in_memory_max_size(5000) + monkeypatch.setattr(proxy_server_module, "general_settings", {"user_api_key_cache_max_size": 5000}) + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await ProxyConfig()._update_general_settings(db_general_settings={"store_model_in_db": True}) + + assert "user_api_key_cache_max_size" not in proxy_server_module.general_settings + + _fill_user_api_key_cache(cache, 201) + assert cache.get_cache(key="key-0", local_only=True) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("db_value", [0, -5, "lots"]) +async def test_update_general_settings_ignores_an_invalid_user_api_key_cache_max_size(db_value, monkeypatch): + """A non-positive capacity would make the eviction loop pop an empty heap on the + next write, so a bad DB value must leave the running cache untouched.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + cache = UserApiKeyCache() + cache.update_in_memory_max_size(300) + monkeypatch.setattr(proxy_server_module, "general_settings", {}) + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await ProxyConfig()._update_general_settings(db_general_settings={"user_api_key_cache_max_size": db_value}) + + assert "user_api_key_cache_max_size" not in proxy_server_module.general_settings + + _fill_user_api_key_cache(cache, 250) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + +@pytest.mark.asyncio +async def test_update_general_settings_user_api_key_cache_max_size_yaml_wins(monkeypatch): + """A DB value must not silently override an explicit YAML capacity on reload.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_general_settings_keys = {"user_api_key_cache_max_size"} + cache = UserApiKeyCache() + cache.update_in_memory_max_size(300) + monkeypatch.setattr(proxy_server_module, "general_settings", {"user_api_key_cache_max_size": 300}) + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await proxy_config._update_general_settings(db_general_settings={"user_api_key_cache_max_size": 10}) + + assert proxy_server_module.general_settings["user_api_key_cache_max_size"] == 300 + + _fill_user_api_key_cache(cache, 250) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "db_value,expected", @@ -7446,10 +7570,8 @@ async def test_store_model_in_db_db_override_when_config_false(): # store_model_in_db should now be True (overridden by DB) assert ps.store_model_in_db is True - # add_deployment and get_credentials should have been called - # since store_model_in_db is now True assert mock_proxy_config.add_deployment.call_count == 1 - assert mock_proxy_config.get_credentials.call_count == 1 + mock_proxy_config.get_credentials.assert_not_called() @pytest.mark.asyncio @@ -7710,7 +7832,7 @@ async def test_increment_spend_counters_team_and_member(): @pytest.mark.asyncio -async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss(): +async def test_prepare_spend_counter_increment_reseeds_from_db_on_counter_miss(): """When the Redis counter is missing, the reseed path reads the authoritative spend from the DB (not a stale cache), so the next increment continues from the correct base value.""" @@ -7723,8 +7845,17 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( recorded_increments.append({"key": key, "value": value, "ttl": ttl}) return value + async def record_pipeline(increment_list, **kwargs): + results = [] + for op in increment_list: + await record_increment(key=op["key"], value=op["increment_value"], ttl=op["ttl"]) + results.append(op["increment_value"]) + return results + fake_redis = AsyncMock() fake_redis.async_increment = AsyncMock(side_effect=record_increment) + fake_redis.async_increment_pipeline = AsyncMock(side_effect=record_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) fake_redis.async_get_cache = AsyncMock(return_value=None) # counter missing fake_redis.async_set_cache = AsyncMock(return_value=True) # SET NX wins counter_cache.redis_cache = fake_redis @@ -7743,7 +7874,10 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( stale_cache.in_memory_cache.set_cache(key="team_id:team-9", value=stale_team) import litellm.proxy.proxy_server as ps - from litellm.proxy.proxy_server import _init_and_increment_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_spend_counter_increment, + ) orig_user, orig_counter, orig_prisma = ( ps.user_api_key_cache, @@ -7754,11 +7888,12 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_spend_counter( + pending = await _prepare_spend_counter_increment( counter_key="spend:team:team-9", source_cache_key="team_id:team-9", increment=1.5, ) + await _apply_spend_counter_increments(pending=(pending,)) fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-9"}) # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. @@ -7937,7 +8072,10 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): @pytest.mark.asyncio async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() window_start = datetime.now(timezone.utc) - timedelta(hours=1) @@ -7953,7 +8091,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key="spend:key:key-window:window:1h", entity_type="Key", entity_id="key-window", @@ -7961,6 +8099,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( by=["api_key"], @@ -7976,7 +8115,10 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): @pytest.mark.asyncio async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:team:team-stale-local" @@ -7998,6 +8140,15 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis db_row = MagicMock() @@ -8016,11 +8167,12 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): ps.prisma_client = fake_prisma ps.user_api_key_cache = DualCache() try: - await _init_and_increment_spend_counter( + pending = await _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key="team_id:team-stale-local", increment=1.5, ) + await _apply_spend_counter_increments(pending=(pending,)) fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-stale-local"}) # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. @@ -8035,7 +8187,10 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): @pytest.mark.asyncio async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:key:key-window-stale-local:window:1h" @@ -8058,6 +8213,15 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -8072,7 +8236,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key=counter_key, entity_type="Key", entity_id="key-window-stale-local", @@ -8080,6 +8244,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( by=["api_key"], @@ -8099,7 +8264,10 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): @pytest.mark.asyncio async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:key:key-window-concurrent-seed:window:1h" @@ -8122,6 +8290,15 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) fake_redis.async_set_cache = AsyncMock(return_value=False) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -8136,7 +8313,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key=counter_key, entity_type="Key", entity_id="key-window-concurrent-seed", @@ -8144,6 +8321,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_redis.async_set_cache.assert_awaited_once_with( key=counter_key, @@ -8160,7 +8338,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() @pytest.mark.asyncio async def test_window_spend_counter_skips_invalid_window_start(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import _prepare_window_spend_counter_increment counter_cache = DualCache() @@ -8169,7 +8347,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): orig_counter = ps.spend_counter_cache ps.spend_counter_cache = counter_cache try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key="spend:key:key-invalid-window:window:not-a-duration", entity_type="Key", entity_id="key-invalid-window", @@ -8177,6 +8355,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): window_start=None, increment=0.5, ) + assert pending is None assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-invalid-window:window:not-a-duration") is None finally: @@ -8240,6 +8419,9 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): async def assert_reservation_not_finalized_yet(**kwargs): assert budget_reservation["finalized"] is False incremented_counters.append(kwargs["counter_key"]) + return ps._PendingSpendIncrement( + counter_key=kwargs["counter_key"], increment=kwargs["increment"] + ) import litellm.proxy.proxy_server as ps @@ -8248,7 +8430,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): ps.user_api_key_cache = DualCache() try: with patch( - "litellm.proxy.proxy_server._init_and_increment_spend_counter", + "litellm.proxy.proxy_server._prepare_spend_counter_increment", new=AsyncMock(side_effect=assert_reservation_not_finalized_yet), ): await increment_spend_counters( @@ -8317,8 +8499,8 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter( """When the reservation reconcile finds the counter in an inconsistent state (here: missing), it must NOT delete the counter and fail open (the old behavior, which left the counter unenforced after a Redis reload). It reseeds - from the authoritative DB so the counter reflects the recorded total and - budget gating continues.""" + from the authoritative DB and adds this request's settled cost, which the + async spend flush has not written yet, so budget gating continues.""" from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import increment_spend_counters from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed @@ -8355,9 +8537,7 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter( ) assert budget_reservation["finalized"] is True - # counter reseeded to the authoritative DB value, not deleted/left None - # and not double-counted via a direct increment - assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-bad-reserved-counter") == pytest.approx(0.6) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-bad-reserved-counter") == pytest.approx(0.85) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -8583,7 +8763,7 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): async def test_concurrent_read_and_write_paths_share_one_db_query(): """ The read path (`get_current_spend`) and the write path - (`_init_and_increment_spend_counter`) both reseed cold counters from + (`_prepare_spend_counter_increment`) both reseed cold counters from the DB. They must share the per-counter lock so a concurrent pre-call enforcement read and post-call increment for the same counter collapse to one DB query, not two. @@ -8592,7 +8772,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import ( - _init_and_increment_spend_counter, + _prepare_spend_counter_increment, get_current_spend, ) @@ -8646,7 +8826,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): try: results = await _asyncio.gather( get_current_spend(counter_key=counter_key, fallback_spend=0.0), - _init_and_increment_spend_counter( + _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key="ignored", increment=1.5, @@ -9085,6 +9265,126 @@ class TestLazyFeaturesNotImportedAtStartup: class TestLazyFeatureMiddleware: """Behavior of the middleware itself, exercised in isolation.""" + @pytest.mark.asyncio + async def test_llm_passthrough_loads_on_first_provider_request(self, monkeypatch): + """An app that never registered the provider passthrough routes 404s a + provider request; behind the middleware the same request registers the + routes and is forwarded to the provider with the configured key.""" + import respx + from fastapi import FastAPI + + from litellm.proxy._lazy_features import LAZY_FEATURES, LazyFeatureMiddleware + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-upstream") + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + feat = next(f for f in LAZY_FEATURES if f.name == "llm_passthrough") + target_app = FastAPI() + target_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(api_key="sk-virtual") + mw = LazyFeatureMiddleware(target_app, fastapi_app=target_app, features=(feat,)) + + with respx.mock() as upstream: + route = upstream.get("https://api.mistral.ai/v1/models").mock( + return_value=httpx.Response(200, json={"object": "list", "data": []}) + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=target_app), base_url="http://t") as bare: + assert (await bare.get("/mistral/v1/models")).status_code == 404 + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=mw), base_url="http://t") as lazy: + response = await lazy.get("/mistral/v1/models") + + assert (response.status_code, response.json()) == (200, {"object": "list", "data": []}) + assert route.calls.last.request.headers["authorization"] == "Bearer sk-upstream" + + def test_llm_passthrough_prefixes_cover_every_route_the_module_registers(self): + """A route the module registers under a prefix the feature does not claim + would 404 until an unrelated provider request happens to load the module.""" + from litellm.proxy._lazy_features import LAZY_FEATURES + + feat = next(f for f in LAZY_FEATURES if f.name == "llm_passthrough") + paths = [r.path for r in importlib.import_module(feat.module_path).router.routes] + + assert {"/mistral/{endpoint:path}", "/openai/{endpoint:path}"} <= set(paths) + unreachable = [p for p in paths if not feat.matches(p.replace("{endpoint:path}", "x"))] + assert unreachable == [], f"routes the middleware would never load: {unreachable}" + + @pytest.mark.asyncio + @pytest.mark.parametrize("first_hit", ["/v1/realtime/calls", "/openai/v1/models"]) + async def test_lazy_routes_land_in_registry_order_not_first_hit_order(self, first_hit): + """Two lazy features with overlapping paths must answer a request with the + same handler no matter which one a deployment happens to hit first.""" + from fastapi import APIRouter, FastAPI + + from litellm.proxy._lazy_features import LazyFeature, LazyFeatureMiddleware + + def make_register(path, handler): + def register(app, module): + router = APIRouter() + router.add_api_route(path, lambda: {"handler": handler}, methods=["POST"]) + app.include_router(router) + + return register + + catch_all = LazyFeature( + name="catch_all", + module_path="json", + path_prefixes=("/openai/",), + register_fn=make_register("/openai/{endpoint:path}", "catch_all"), + ) + specific = LazyFeature( + name="specific", + module_path="base64", + path_prefixes=("/openai/v1/realtime", "/v1/realtime"), + register_fn=make_register("/openai/v1/realtime/calls", "specific"), + ) + + target_app = FastAPI() + mw = LazyFeatureMiddleware(target_app, fastapi_app=target_app, features=(catch_all, specific)) + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=mw), base_url="http://t") as client: + await client.post(first_hit) + await client.post("/openai/v1/models") + response = await client.post("/openai/v1/realtime/calls") + + assert response.json() == {"handler": "catch_all"} + + @pytest.mark.asyncio + @pytest.mark.parametrize("root_path", ["", "/api"]) + async def test_reserved_slot_keeps_lazy_catch_all_ahead_of_later_eager_routes(self, root_path): + """/{mcp_server_name}/mcp is registered after the provider passthrough router + at startup, so /mistral/mcp must keep reaching the provider catch-all once + that router loads lazily instead of being swallowed by the MCP route. The + native /mistral/v1/files route sits ahead of it, so that path neither loads + the feature nor changes owner, with or without a SERVER_ROOT_PATH prefix.""" + from fastapi import APIRouter, FastAPI + + from litellm.proxy._lazy_features import LazyFeature, LazyFeatureMiddleware, reserve_lazy_slot + + def register(app, module): + router = APIRouter() + router.add_api_route("/mistral/{endpoint:path}", lambda: {"handler": "passthrough"}, methods=["POST"]) + app.include_router(router) + + passthrough = LazyFeature( + name="llm_passthrough", module_path="json", path_prefixes=("/mistral/",), register_fn=register + ) + target_app = FastAPI(root_path=root_path) + target_app.add_api_route("/mistral/v1/files", lambda: {"handler": "files"}, methods=["POST"]) + reserve_lazy_slot(target_app, "llm_passthrough", features=(passthrough,)) + target_app.add_api_route("/{mcp_server_name}/mcp", lambda: {"handler": "mcp"}, methods=["POST"]) + target_app.add_middleware(LazyFeatureMiddleware, fastapi_app=target_app, features=(passthrough,)) + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=target_app), base_url="http://t") as client: + files_first = (await client.post(f"{root_path}/mistral/v1/files")).json()["handler"] + loaded_after_files = frozenset(target_app.state.lazy_loaded) + handlers = [ + (await client.post(f"{root_path}{path}")).json()["handler"] + for path in ("/mistral/mcp", "/mistral/v1/files") + ] + + assert (files_first, loaded_after_files) == ("files", frozenset()) + assert handlers == ["passthrough", "files"] + @pytest.mark.asyncio async def test_first_request_triggers_load_subsequent_does_not(self): from fastapi import FastAPI @@ -9224,6 +9524,82 @@ class TestLazyFeatureMiddleware: else: assert loads == [], f"{case}: feature must not load" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "env_root_path,scope_root_path,request_path,should_load,case", + [ + # Per-request root_path (PerRequestRootPathMiddleware under + # SERVER_ROOT_PATHS) with no scalar env: strip and match. + ("", "/tenant-a", "/tenant-a/dummy/x", True, "per-request root_path strip"), + # scope root_path is authoritative over the cached env scalar. + ("/api/v1", "/tenant-a", "/tenant-a/dummy/x", True, "scope wins over env scalar"), + # Boundary check still applies to the per-request value. + ("", "/tenant-a", "/tenant-ab/dummy/x", False, "boundary check on scope root_path"), + # Empty scope root_path falls back to the env scalar. + ("/api/v1", "", "/api/v1/dummy/x", True, "empty scope falls back to env"), + ], + ) + async def test_per_request_root_path_handling( + self, monkeypatch, env_root_path, scope_root_path, request_path, should_load, case + ): + """ + ``scope["root_path"]`` must be stripped before prefix matching when + set — the scalar SERVER_ROOT_PATH lands there via + ``FastAPI(root_path=...)``, and PerRequestRootPathMiddleware + (SERVER_ROOT_PATHS) resolves a per-request prefix there. Otherwise + lazily-registered features — the MCP OAuth discovery router among + them — stay unloaded under a client-visible prefix and 404. + """ + from fastapi import FastAPI + + from litellm.proxy._lazy_features import ( + LazyFeature, + LazyFeatureMiddleware, + ) + + monkeypatch.setenv("SERVER_ROOT_PATH", env_root_path) + + loads = [] + + def fake_register(app, module): + loads.append(getattr(module, "__name__", "?")) + + feat = LazyFeature( + name=f"dummy_prr_{case}", + module_path="json", + path_prefixes=("/dummy",), + register_fn=fake_register, + ) + + async def downstream(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + target_app = FastAPI() + mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message): + pass + + await mw( + { + "type": "http", + "path": request_path, + "root_path": scope_root_path, + "method": "GET", + "headers": [], + }, + receive, + send, + ) + if should_load: + assert loads == ["json"], f"{case}: expected feature to load" + else: + assert loads == [], f"{case}: feature must not load" + @pytest.mark.asyncio async def test_concurrent_first_requests_only_register_once(self): """ @@ -10050,6 +10426,27 @@ def test_get_config_list_includes_apply_user_budget_to_team_keys(monkeypatch): app.dependency_overrides.clear() +def test_get_config_list_includes_user_api_key_cache_max_size(monkeypatch): + """The Admin UI General Settings table renders whatever /config/list returns, + so the cache capacity has to be exposed there as an Integer to be editable.""" + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(proxy_server_module, "prisma_client", mock_prisma) + app.dependency_overrides[proxy_server_module.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + assert resp.status_code == 200, resp.text + fields = {item["field_name"]: item for item in resp.json()} + assert fields["user_api_key_cache_max_size"]["field_type"] == "Integer" + finally: + app.dependency_overrides.clear() + + def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch): """The throttle fraction is a litellm_settings scalar surfaced on the General Settings table as a Float field so it sits with the other global limits; it @@ -12652,6 +13049,48 @@ async def test_load_config_router_authorizes_fallback_targets_against_the_callin assert router.fallback_access_check is router_fallback_access_check +@pytest.mark.asyncio +async def test_load_config_user_api_key_cache_max_size_keeps_more_than_200_entries(tmp_path, monkeypatch): + """The auth cache used to be pinned at InMemoryCache's 200 entry default, so a + deployment with more keys than that evicted constantly and every request + fell through to the DB. The YAML knob has to raise the cap on the live cache.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump({"general_settings": {"user_api_key_cache_max_size": "1000"}})) + + cache = UserApiKeyCache() + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + _fill_user_api_key_cache(cache, 999) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_value", [0, -1, "unbounded"]) +async def test_load_config_rejects_a_non_positive_user_api_key_cache_max_size(tmp_path, bad_value, monkeypatch): + """InMemoryCache treats 0 as 'cache nothing' and a negative cap makes eviction + pop an empty heap, so the proxy must refuse to boot with such a value instead + of silently disabling auth caching.""" + from pydantic import ValidationError + + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump({"general_settings": {"user_api_key_cache_max_size": bad_value}})) + + cache = UserApiKeyCache() + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + with pytest.raises(ValidationError): + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + _fill_user_api_key_cache(cache, 150) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + def test_docs_redoc_openapi_are_reachable_by_default(): """ LIT-6745: the interactive/machine-readable docs surfaces are on by @@ -12781,3 +13220,59 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() import litellm.proxy.proxy_server as ps assert ps.general_settings["enable_openai_websocket_passthrough"] is False + + +async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + warm_tokenizer("claude-fable-5") + + response, took, lags = await timed_with_loop_lags( + lambda: proxy_server_module.token_counter(TokenCountRequest(model="claude-fable-5", prompt=text * 100)) + ) + + assert response.total_tokens > 0 + assert_loop_stayed_free(took, lags) + + +async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch): + from tokenizers import Tokenizer + + from litellm import Router + from tests.test_litellm.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags + + claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] + + class SlowHubTokenizer: + @staticmethod + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: + time.sleep(0.3) + return claude_tokenizer + + monkeypatch.setattr(litellm.utils, "Tokenizer", SlowHubTokenizer) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + Router( + model_list=[ + { + "model_name": "self-hosted", + "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"}, + "model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}}, + } + ] + ), + ) + + response, took, lags = await timed_with_loop_lags( + lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop")) + ) + + assert response.tokenizer_type == "huggingface_tokenizer" + assert response.total_tokens > 0 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index dcf18e773e8..9def21c0573 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1,5 +1,6 @@ import datetime as real_datetime import smtplib +from typing import Final import pytest from fastapi import HTTPException @@ -8,7 +9,7 @@ from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.types.guardrails import GuardrailEventHooks @@ -580,6 +581,75 @@ class TestPostCallFailureHookLiftsStandardLoggingObject: assert "standard_logging_object" not in request_data +class TestPostCallFailureHookLiftsCallTypeAndStartTime: + """A guardrail-blocked MCP tool call fails before any LLM call. The failure + spend row is built from request_data after ``litellm_logging_obj`` is popped, + so ``call_type`` and the request ``start_time`` must be lifted off the logging + object first, or the Logs page shows the row as an LLM call with a blank call + type and a 0s duration (LIT-7453). + """ + + @pytest.mark.asyncio + async def test_failed_mcp_tool_call_spend_row_keeps_call_type_model_and_duration(self): + import traceback + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + request_start = real_datetime.datetime.now() - real_datetime.timedelta(seconds=2) + logging_obj = Logging( + model="MCP: deepwiki-ask_question", + messages=[], + stream=False, + call_type="call_mcp_tool", + start_time=request_start, + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + model="MCP: deepwiki-ask_question", + user="", + optional_params={}, + litellm_params={"metadata": {"user_api_key_hash": "hashed"}}, + ) + blocked = Exception("Content blocked: keyword 'confidential' detected") + logging_obj.failure_handler(blocked, traceback.format_exc(), request_start, real_datetime.datetime.now()) + request_data = { + "name": "deepwiki-ask_question", + "arguments": {"question": "confidential"}, + "litellm_logging_obj": logging_obj, + } + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + spend_writer = SimpleNamespace(update_database=AsyncMock()) + original_callbacks = list(litellm.callbacks) + litellm.callbacks = [_ProxyDBLogger(spend_writer=lambda: spend_writer)] + try: + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=blocked, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + finally: + litellm.callbacks = original_callbacks + ProxyLogging._callback_capabilities_cache.clear() + + db_call = spend_writer.update_database.call_args.kwargs + payload = get_logging_payload( + kwargs=db_call["kwargs"], + response_obj=db_call["completion_response"], + start_time=db_call["start_time"], + end_time=db_call["end_time"], + ) + assert payload["call_type"] == "call_mcp_tool" + assert payload["model"] == "MCP: deepwiki-ask_question" + assert payload["endTime"] - payload["startTime"] >= real_datetime.timedelta(seconds=2) + + class TestPostCallFailureHookEstimatesDispatchedInputTokens: """A non-stream request that failed after dispatch (timeout, provider error) consumed provider-billed input tokens but recovered no usage. @@ -887,6 +957,7 @@ from typing import cast import litellm from litellm.proxy.utils import create_model_info_response +from litellm.types.router import DeploymentModelListingInfo from litellm.types.utils import ModelInfo @@ -914,7 +985,7 @@ def test_create_model_info_response_includes_max_tokens_from_lookup(): def test_create_model_info_response_does_not_call_router_group_info(): router = MagicMock() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None response = create_model_info_response( model_id="some-model", @@ -929,7 +1000,9 @@ def test_create_model_info_response_does_not_call_router_group_info(): def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map(): router = MagicMock() - router.get_configured_token_limits.return_value = (32000, 8000) + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_keys=("my-custom-deployment",), max_input_tokens=32000, max_output_tokens=8000 + ) response = create_model_info_response( model_id="my-custom-deployment", @@ -986,7 +1059,9 @@ def test_create_model_info_response_uses_deployment_mode_for_auto_router(): def test_create_model_info_response_deployment_limits_override_cost_map(): router = MagicMock() - router.get_configured_token_limits.return_value = (200000, None) + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_keys=("gpt-4o",), max_input_tokens=200000, max_output_tokens=None + ) response = create_model_info_response( model_id="gpt-4o", @@ -999,6 +1074,54 @@ def test_create_model_info_response_deployment_limits_override_cost_map(): assert response["max_output_tokens"] == 16384 +def test_create_model_info_response_reports_widest_window_in_a_mixed_group(): + """A group mixing models advertises the widest window, not whichever is listed first.""" + limits = { + "small-model": _fake_model_info(max_input_tokens=200000, max_output_tokens=4096, mode="chat"), + "large-model": _fake_model_info(max_input_tokens=1000000, max_output_tokens=128000, mode="chat"), + } + + for keys in (("small-model", "large-model"), ("large-model", "small-model")): + router = MagicMock() + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_keys=keys, max_input_tokens=None, max_output_tokens=None + ) + + response = create_model_info_response( + model_id="house-claude", + provider="openai", + llm_router=router, + get_model_info=lambda model: limits[model], + ) + + assert response["max_input_tokens"] == 1000000, keys + assert response["max_output_tokens"] == 128000, keys + + +def test_create_model_info_response_resolves_alias_once_per_listing(): + """The alias is the same for every deployment in the group, so it is looked up once.""" + seen: list[str] = [] + + def _tracking_get_model_info(model: str) -> ModelInfo: + seen.append(model) + return _fake_model_info(max_input_tokens=128000) + + router = MagicMock() + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_keys=("model-a", "model-b"), max_input_tokens=None, max_output_tokens=None + ) + + create_model_info_response( + model_id="house-model", + provider="openai", + llm_router=router, + get_model_info=_tracking_get_model_info, + ) + + assert seen.count("house-model") == 1 + assert sorted(seen) == ["house-model", "model-a", "model-b"] + + def test_create_model_info_response_survives_malformed_configured_limits(): from litellm import Router @@ -1794,13 +1917,14 @@ def test_a_failure_with_no_logging_object_lifts_nothing(): assert dict(_failure_fields_to_lift({"litellm_logging_obj": _LoggingObj({})})) == {} -def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs(): +def test_a_dispatched_failure_lifts_the_fields_the_spend_log_needs(): from litellm.proxy.utils import _failure_fields_to_lift lifted = _failure_fields_to_lift( { "litellm_logging_obj": _LoggingObj( { + "start_time": 1699999999.0, "first_api_call_start_time": 1700000000.0, "call_type": "acompletion", "model": FAILURE_USAGE_MODEL, @@ -1812,17 +1936,59 @@ def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs(): ) assert set(lifted) == { + "start_time", "first_api_call_start_time", + "call_type", "combined_usage_object", "response_cost", "standard_logging_object", } + assert lifted["start_time"] == 1699999999.0 assert lifted["first_api_call_start_time"] == 1700000000.0 + assert lifted["call_type"] == "acompletion" assert lifted["response_cost"] == 0.0 assert lifted["combined_usage_object"].prompt_tokens > 0 assert lifted["standard_logging_object"] == {"id": "log-1"} +@pytest.mark.asyncio +async def test_a_dispatched_failure_is_counted_off_the_event_loop(): + from unittest.mock import AsyncMock, patch + + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("claude-fable-5") + request_data = { + "litellm_logging_obj": _LoggingObj( + { + "first_api_call_start_time": 1700000000.0, + "call_type": "acompletion", + "model": "claude-fable-5", + "messages": [{"role": "user", "content": text * 100}], + } + ), + "metadata": {}, + } + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + _, took, lags = await timed_with_loop_lags( + lambda: proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=Exception("boom"), + user_api_key_dict=UserAPIKeyAuth(), + ) + ) + + assert request_data["combined_usage_object"].prompt_tokens > 0 + assert_loop_stayed_free(took, lags) + + @pytest.mark.asyncio async def test_proxy_only_error_expected_4xx_skips_traceback_for_both_handlers(monkeypatch): """Regression for LIT-6043: an expected 4xx must not format a traceback for @@ -1922,6 +2088,150 @@ async def test_proxy_only_error_5xx_keeps_traceback_and_runs_sync_callbacks(monk assert "test_proxy_utils" in captured["async_traceback"] +def test_create_model_info_response_resolves_alias_to_deployment_model(): + """A public model name that is not itself a cost-map key must not be resolved through + the fallback-generalization rules: `bedrock-claude-opus-5` matches the generic + claude-family baseline (200k/64k) by substring, while the deployment it fronts really + accepts 1M/128k. Regression for the /v1/models alias resolution introduced in v1.94.0.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "bedrock-claude-opus-5", + "litellm_params": { + "custom_llm_provider": "bedrock", + "model": "bedrock/eu.anthropic.claude-opus-5", + }, + "model_info": {"base_model": "eu.anthropic.claude-opus-5"}, + } + ] + ) + + response = create_model_info_response( + model_id="bedrock-claude-opus-5", provider="openai", llm_router=router + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert response["max_input_tokens"] == 1000000 + assert response["max_output_tokens"] == 128000 + + +def test_create_model_info_response_keeps_exact_alias_over_generalized_deployment_model(): + """Mirror of the alias bug: when the deployment points at a custom backend name that + only matches a generalization rule, the listed name's exact cost-map entry is the + better answer and must win.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "claude-opus-5", + "litellm_params": { + "custom_llm_provider": "bedrock", + "model": "bedrock/my-claude-opus-5-provisioned", + }, + } + ] + ) + + response = create_model_info_response( + model_id="claude-opus-5", provider="openai", llm_router=router + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert response["max_input_tokens"] == 1000000 + + +def test_create_model_info_response_falls_back_to_alias_for_opaque_deployment_name(): + """An Azure deployment named after the resource rather than the model has no cost-map + entry; the listed name still does, and must keep answering.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "azure/my-gpt4o-deployment"}, + } + ] + ) + + response = create_model_info_response( + model_id="gpt-4o", provider="openai", llm_router=router + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert response["max_input_tokens"] == 128000 + assert response["max_output_tokens"] == 16384 + + +def test_create_model_info_response_resolves_mode_through_deployment_model(): + """`mode` is derived from the same lookup, so an aliased embedding deployment + currently reports no mode at all; it must report `embedding`.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "my-embeddings", + "litellm_params": {"model": "openai/text-embedding-3-small"}, + } + ] + ) + + response = create_model_info_response( + model_id="my-embeddings", provider="openai", llm_router=router + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert response["mode"] == "embedding" + + +@pytest.mark.parametrize( + "key_metadata, team_metadata, expected_to_run", + [ + ({"guardrails": ["key-scoped-guardrail"]}, None, True), + ({}, {"guardrails": ["key-scoped-guardrail"]}, True), + ({"guardrails": ["some-other-guardrail"]}, None, False), + ({}, None, False), + ], +) +def test_convert_mcp_to_llm_format_carries_key_and_team_guardrails(key_metadata, team_metadata, expected_to_run): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + guardrail = CustomGuardrail(guardrail_name="key-scoped-guardrail", event_hook="pre_mcp_call", default_on=False) + kwargs = { + "name": "ask_question", + "arguments": {"question": "hello"}, + "server_name": "deepwiki", + "user_api_key_auth": UserAPIKeyAuth(metadata=key_metadata, team_metadata=team_metadata), + } + request_obj = proxy_logging._create_mcp_request_object_from_kwargs(kwargs) + + with patch( # test-quality-ok: the key-guardrail premium gate reads this proxy_server module global and has no injection seam + "litellm.proxy.proxy_server.premium_user", True + ): + synthetic = proxy_logging._convert_mcp_to_llm_format(request_obj, kwargs) + + assert guardrail.should_run_guardrail(synthetic, GuardrailEventHooks.pre_mcp_call) is expected_to_run + + class _TracebackRecordingLogger(CustomLogger): def __init__(self) -> None: super().__init__() @@ -1972,3 +2282,49 @@ async def test_post_call_failure_hook_redacts_traceback_before_callbacks(monkeyp assert recorder.received_traceback is not None assert provider_key not in recorder.received_traceback assert "REDACTED" in recorder.received_traceback + + +class TestPrismaClientTokenAuthBehindThePool: + """Behind the in-container pool the supervisor renews the writer's database + token and hands the workers a loopback URL with a static password, so the + writer wrapper must not run its own refresh loop. The reader is not pooled + and keeps refreshing its own token.""" + + UPSTREAM: Final = "postgresql://litellm:TOKEN@db.internal:5432/litellm" + READER: Final = "postgresql://litellm:TOKEN@reader.internal:5432/litellm" + + def _client(self, monkeypatch: pytest.MonkeyPatch, pooled: bool) -> PrismaClient: + from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR + + monkeypatch.delenv("AZURE_POSTGRESQL_AUTH", raising=False) + monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + monkeypatch.setenv("DATABASE_URL", self.UPSTREAM) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", self.READER) + if pooled: + monkeypatch.setenv(PGBOUNCER_POOLED_ENV_VAR, "true") + else: + monkeypatch.delenv(PGBOUNCER_POOLED_ENV_VAR, raising=False) + rds: Final = MagicMock() + rds.generate_db_auth_token.return_value = "TOKEN" + with patch("boto3.client", return_value=rds): + return PrismaClient(database_url=self.UPSTREAM, proxy_logging_obj=MagicMock(spec=ProxyLogging)) + + def test_a_pooled_writer_leaves_token_refresh_to_the_pooler_while_the_reader_keeps_its_own( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + client = self._client(monkeypatch, pooled=True) + assert isinstance(client.db, RoutingPrismaWrapper) + assert client.db.writer.iam_token_db_auth is False + assert client.db.reader.iam_token_db_auth is True + assert client.token_auth is not None + + def test_an_unpooled_writer_still_refreshes_its_own_token(self, monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + client = self._client(monkeypatch, pooled=False) + assert isinstance(client.db, RoutingPrismaWrapper) + assert client.db.writer.iam_token_db_auth is True + assert client.db.reader.iam_token_db_auth is True diff --git a/tests/test_litellm/proxy/test_route_priority.py b/tests/test_litellm/proxy/test_route_priority.py new file mode 100644 index 00000000000..dfdc816f4b4 --- /dev/null +++ b/tests/test_litellm/proxy/test_route_priority.py @@ -0,0 +1,171 @@ +import sys +from types import ModuleType + +import httpx +import pytest +from fastapi import APIRouter, FastAPI +from fastapi.testclient import TestClient +from starlette.routing import Match + +from litellm.proxy.route_priority import HOT_ROUTE_PATHS, hot_routes_first + +FILLER_COUNT = 300 + + +def _routes_scanned_before_dispatch(app: FastAPI, method: str, path: str) -> int: + """Number of route.matches() calls Starlette's Router.app makes before it finds a full match.""" + scope = {"type": "http", "method": method, "path": path, "root_path": "", "headers": [], "query_string": b""} + for i, route in enumerate(app.router.routes): + match, _ = route.matches(dict(scope)) + if match == Match.FULL: + return i + 1 + raise AssertionError(f"{method} {path} has no route") + + +def _hot_router() -> APIRouter: + router = APIRouter() + + @router.get("/health/liveliness") + @router.get("/health/liveness") + async def liveliness(): + return "I'm alive!" + + @router.post("/v1/chat/completions") + @router.post("/chat/completions") + async def chat(): + return {"object": "chat.completion"} + + return router + + +def _app_with_filler_then_hot_routes() -> FastAPI: + app = FastAPI() + for i in range(FILLER_COUNT): + + @app.get(f"/filler/{i}") + async def filler(i: int = i): + return {"filler": i} + + app.include_router(_hot_router()) + return app + + +def test_hot_routes_first_puts_hot_routes_ahead_of_everything_else(): + app = _app_with_filler_then_hot_routes() + assert _routes_scanned_before_dispatch(app, "GET", "/health/liveliness") > FILLER_COUNT + + app.router.routes = hot_routes_first(app.router.routes) + + hot_count = sum(1 for r in app.router.routes if getattr(r, "path", None) in HOT_ROUTE_PATHS) + assert _routes_scanned_before_dispatch(app, "GET", "/health/liveliness") <= hot_count + assert _routes_scanned_before_dispatch(app, "GET", "/health/liveness") <= hot_count + assert _routes_scanned_before_dispatch(app, "POST", "/v1/chat/completions") <= hot_count + assert _routes_scanned_before_dispatch(app, "POST", "/chat/completions") <= hot_count + + +def test_hot_routes_first_keeps_the_other_routes_in_order_and_dispatching(): + app = _app_with_filler_then_hot_routes() + before = [r.path for r in app.router.routes if getattr(r, "path", "").startswith("/filler/")] + + app.router.routes = hot_routes_first(app.router.routes) + + after = [r.path for r in app.router.routes if getattr(r, "path", "").startswith("/filler/")] + assert after == before + client = TestClient(app) + assert client.get("/health/liveliness").json() == "I'm alive!" + assert client.get("/filler/7").json() == {"filler": 7} + assert client.post("/v1/chat/completions").json() == {"object": "chat.completion"} + assert client.get("/v1/chat/completions").status_code == 405 + assert client.get("/does/not/exist").status_code == 404 + + +def test_hot_routes_first_is_idempotent(): + app = _app_with_filler_then_hot_routes() + once = hot_routes_first(app.router.routes) + assert hot_routes_first(once) == once + + +@pytest.mark.asyncio +async def test_lazy_loaded_hot_route_moves_to_the_front(monkeypatch): + from litellm.proxy._lazy_features import LazyFeature, LazyFeatureMiddleware + + messages_router = APIRouter() + + @messages_router.post("/v1/messages") + async def messages(): + return {"type": "message"} + + fake_module = ModuleType("fake_anthropic_endpoints") + fake_module.router = messages_router + monkeypatch.setitem(sys.modules, fake_module.__name__, fake_module) + + target_app = _app_with_filler_then_hot_routes() + target_app.router.routes = hot_routes_first(target_app.router.routes) + + async def downstream(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + feat = LazyFeature(name="anthropic", module_path=fake_module.__name__, path_prefixes=("/v1/messages",)) + mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message): + pass + + await mw({"type": "http", "path": "/v1/messages", "method": "POST", "headers": []}, receive, send) + + hot_count = sum(1 for r in target_app.router.routes if getattr(r, "path", None) in HOT_ROUTE_PATHS) + assert _routes_scanned_before_dispatch(target_app, "POST", "/v1/messages") <= hot_count + assert TestClient(target_app).post("/v1/messages").json() == {"type": "message"} + + +@pytest.mark.asyncio +async def test_hot_routes_first_keeps_reserved_lazy_slot_ahead_of_later_eager_routes(): + """Liveness is registered after the provider passthrough slot, so pulling it to the + front must not shift where the lazily loaded catch-all is spliced back in.""" + from litellm.proxy._lazy_features import LazyFeature, LazyFeatureMiddleware, reserve_lazy_slot + + def register(app, module): + router = APIRouter() + router.add_api_route("/mistral/{endpoint:path}", lambda: {"handler": "passthrough"}, methods=["POST"]) + app.include_router(router) + + passthrough = LazyFeature( + name="llm_passthrough", module_path="json", path_prefixes=("/mistral/",), register_fn=register + ) + target_app = FastAPI() + target_app.add_api_route("/mistral/v1/files", lambda: {"handler": "files"}, methods=["POST"]) + target_app.add_api_route("/mistral/v1/batches", lambda: {"handler": "batches"}, methods=["POST"]) + reserve_lazy_slot(target_app, "llm_passthrough", features=(passthrough,)) + target_app.include_router(_hot_router()) + target_app.add_api_route("/{mcp_server_name}/mcp", lambda: {"handler": "mcp"}, methods=["POST"]) + target_app.router.routes = hot_routes_first(target_app.router.routes) + target_app.add_middleware(LazyFeatureMiddleware, fastapi_app=target_app, features=(passthrough,)) + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=target_app), base_url="http://t") as client: + batches_first = (await client.post("/mistral/v1/batches")).json()["handler"] + loaded_after_batches = frozenset(target_app.state.lazy_loaded) + handlers = [ + (await client.post(path)).json()["handler"] + for path in ("/mistral/mcp", "/mistral/v1/files", "/mistral/v1/batches") + ] + + assert (batches_first, loaded_after_batches) == ("batches", frozenset()) + assert handlers == ["passthrough", "files", "batches"] + hot_count = sum(1 for r in target_app.router.routes if getattr(r, "path", None) in HOT_ROUTE_PATHS) + assert _routes_scanned_before_dispatch(target_app, "GET", "/health/liveliness") <= hot_count + + +def test_proxy_app_dispatches_liveness_and_chat_completions_before_the_rest(): + from litellm.proxy.proxy_server import app + + hot_count = sum(1 for r in app.router.routes if getattr(r, "path", None) in HOT_ROUTE_PATHS) + assert hot_count >= 4 + assert len(app.router.routes) > 100 + assert _routes_scanned_before_dispatch(app, "GET", "/health/liveliness") <= hot_count + assert _routes_scanned_before_dispatch(app, "GET", "/health/liveness") <= hot_count + assert _routes_scanned_before_dispatch(app, "POST", "/v1/chat/completions") <= hot_count + assert _routes_scanned_before_dispatch(app, "POST", "/chat/completions") <= hot_count diff --git a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py index e73e3c151e0..117c5aa3081 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py @@ -1,7 +1,10 @@ +import asyncio import json +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException +from prisma.errors import DataError from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.utils import get_error_message_str, handle_exception_on_proxy @@ -135,7 +138,7 @@ def test_handle_exception_on_proxy_happy_path_generic_exception_defaults_to_500( "message": "kaboom", "type": ProxyErrorTypes.internal_server_error.value, "code": "500", - "param": "None", + "param": None, } @@ -171,3 +174,45 @@ def test_handle_exception_on_proxy_error_path_none_input_wraps_as_500(): "code": "500", "type": ProxyErrorTypes.internal_server_error.value, } + + +@pytest.mark.asyncio +async def test_handle_exception_on_proxy_read_only_transaction_forces_writer_recreate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + prisma_client = MagicMock() + prisma_client.recreate_read_only_writer = AsyncMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + exc = DataError( + data={ + "user_facing_error": { + "message": 'PostgresError { code: "25006", message: "cannot execute UPDATE in a read-only transaction" }' + } + } + ) + + result = handle_exception_on_proxy(exc) + await asyncio.sleep(0) + + snapshot = { + "code": result.code, + "recreate_kwargs": prisma_client.recreate_read_only_writer.await_args.kwargs, + } + assert snapshot == { + "code": "500", + "recreate_kwargs": {"reason": "postgres_read_only_transaction"}, + } + + +@pytest.mark.asyncio +async def test_handle_exception_on_proxy_leaves_writer_alone_for_other_db_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + prisma_client = MagicMock() + prisma_client.recreate_read_only_writer = AsyncMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + + handle_exception_on_proxy(DataError(data={"user_facing_error": {"message": "deadlock detected"}})) + await asyncio.sleep(0) + + assert prisma_client.recreate_read_only_writer.await_count == 0 diff --git a/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py index 31ea1bdce74..5f23c5fc20e 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py @@ -314,3 +314,76 @@ def test_normalize_route_for_root_path_error_path_when_route_not_under_root( _clear_url_env(monkeypatch) monkeypatch.setenv("SERVER_ROOT_PATH", "/proxy") assert normalize_route_for_root_path("/other/v1/chat") is None + + +# --------------------------------------------------------------------------- +# get_custom_url under a per-request prefix (SERVER_ROOT_PATHS): +# when a request lives under a dynamic prefix, request.base_url already +# contains it. Appending the SERVER_ROOT_PATH scalar on top produced +# double-prefixed SSO callback / login URLs (a path that doesn't exist on +# the deployment). These pin that the emitted URL now stays under one +# prefix — the one the request actually arrived on. +# --------------------------------------------------------------------------- + + +def test_get_custom_url_uses_per_request_prefix_when_middleware_ran(monkeypatch): + """The middleware stashes the effective per-request prefix in a ContextVar. + ``get_custom_url`` reads that in preference to the SERVER_ROOT_PATH scalar, + and ``join_paths``'s tail-dedup collapses the append so a request whose + ``base_url`` already ends in ``/tenant-a`` does not become + ``/tenant-a/legacy/route``.""" + _clear_url_env(monkeypatch) + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + + from litellm.proxy.middleware.per_request_root_path_middleware import ( + _request_root_path_var, + ) + + token = _request_root_path_var.set("/tenant-a") + try: + # request.base_url already carries the tenant prefix; the scalar + # SERVER_ROOT_PATH must not be re-appended on top. + result = get_custom_url( + request_base_url="https://request.example.com/tenant-a/", + route="/v1/chat", + ) + finally: + _request_root_path_var.reset(token) + + assert result == "https://request.example.com/tenant-a/v1/chat" + + +def test_get_custom_url_no_double_prefix_when_both_env_vars_configured(monkeypatch): + """Regression guard for the review point: with SERVER_ROOT_PATH also set + (scalar-legacy) and the request matched by SERVER_ROOT_PATHS (per-request), + the emitted URL is under one prefix — the per-request one — never both.""" + _clear_url_env(monkeypatch) + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + monkeypatch.setenv("SERVER_ROOT_PATHS", "/tenant-a") + + from litellm.proxy.middleware.per_request_root_path_middleware import ( + _request_root_path_var, + ) + + token = _request_root_path_var.set("/tenant-a") + try: + # No "/legacy" ever appears — the fix pins the reviewer's expected + # behavior: only one prefix should apply per request. + result = get_custom_url("https://api.example.com/tenant-a", "/sso/callback") + finally: + _request_root_path_var.reset(token) + + assert result == "https://api.example.com/tenant-a/sso/callback" + assert "/legacy" not in result + + +def test_get_custom_url_scalar_only_still_stamps_root_path(monkeypatch): + """Pre-middleware deployments have not opted into SERVER_ROOT_PATHS at all; + the ContextVar stays unset and the SERVER_ROOT_PATH scalar owns the answer + — the behavior every existing scalar-only deployment relies on.""" + _clear_url_env(monkeypatch) + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + + result = get_custom_url("https://request.example.com", "/v1/chat") + + assert result == "https://request.example.com/legacy/v1/chat" diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py index 9f48ba68b4f..fbe9934f06c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py @@ -22,6 +22,10 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from litellm.proxy.db.health_check_latest import ( + LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, + LATEST_HEALTH_CHECKS_SQL, +) from litellm.proxy.utils import PrismaClient @@ -261,36 +265,49 @@ async def test_get_health_check_history_db_error_returns_empty_list( assert await prisma_client.get_health_check_history() == [] +def _raw_health_check_row(model_name: str = "gpt-4", model_id: str | None = "deployment-abc") -> dict[str, Any]: + return { + "health_check_id": f"hc-{model_name}", + "model_name": model_name, + "model_id": model_id, + "status": "healthy", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + "response_time_ms": 12.5, + "details": None, + "checked_by": "pod-1", + "checked_at": "2026-08-25T00:00:00+00:00", + "created_at": "2026-08-25T00:00:00+00:00", + "updated_at": "2026-08-25T00:00:00+00:00", + } + + @pytest.mark.asyncio -async def test_get_all_latest_health_checks_uses_distinct( +async def test_get_all_latest_health_checks_dedups_in_postgres( prisma_client: PrismaClient, ) -> None: - rows = [MagicMock(name=f"row-{i}") for i in range(3)] - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=rows) + """A revert to prisma find_many(distinct=...) streams the whole history table; the SQL must own the DISTINCT.""" + prisma_client.db.query_raw = AsyncMock(return_value=[_raw_health_check_row()]) result = await prisma_client.get_all_latest_health_checks() - kwargs = prisma_client.db.litellm_healthchecktable.find_many.await_args.kwargs actual = { - "len": len(result), - "distinct": kwargs["distinct"], - "order_len": len(kwargs["order"]), - "first_order": kwargs["order"][0], + "query": prisma_client.db.query_raw.await_args.args, + "rows": [(row.model_id, row.model_name, row.status) for row in result], + "row_type": type(result[0]).__name__, } assert actual == { - "len": 3, - "distinct": ["model_id", "model_name"], - "order_len": 3, - "first_order": {"model_id": "asc"}, + "query": (LATEST_HEALTH_CHECKS_SQL,), + "rows": [("deployment-abc", "gpt-4", "healthy")], + "row_type": "LatestHealthCheckRow", } @pytest.mark.asyncio -async def test_get_all_latest_health_checks_db_error_returns_empty_list( +async def test_get_all_latest_health_checks_db_error_returns_no_rows( prisma_client: PrismaClient, ) -> None: - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock( - side_effect=RuntimeError("oops") - ) - assert await prisma_client.get_all_latest_health_checks() == [] + prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("oops")) + assert await prisma_client.get_all_latest_health_checks() == () @pytest.mark.asyncio @@ -298,18 +315,15 @@ async def test_get_latest_health_checks_for_models_bounds_the_query_to_those_mod prisma_client: PrismaClient, ) -> None: """A paged caller reads health for its page; an unbounded read is the bug this exists to avoid.""" - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[]) - await prisma_client.get_latest_health_checks_for_models(["gpt-5", "claude-opus"]) - kwargs = prisma_client.db.litellm_healthchecktable.find_many.await_args.kwargs + prisma_client.db.query_raw = AsyncMock(return_value=[_raw_health_check_row(model_name="gpt-5")]) + result = await prisma_client.get_latest_health_checks_for_models(["gpt-5", "claude-opus"]) actual = { - "where": kwargs["where"], - "distinct": kwargs["distinct"], - "order": kwargs["order"], + "query": prisma_client.db.query_raw.await_args.args, + "rows": [row.model_name for row in result], } assert actual == { - "where": {"model_name": {"in": ["gpt-5", "claude-opus"]}}, - "distinct": ["model_id", "model_name"], - "order": [{"model_id": "asc"}, {"model_name": "asc"}, {"checked_at": "desc"}], + "query": (LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, ["gpt-5", "claude-opus"]), + "rows": ["gpt-5"], } @@ -317,14 +331,14 @@ async def test_get_latest_health_checks_for_models_bounds_the_query_to_those_mod async def test_get_latest_health_checks_for_models_does_not_query_for_an_empty_page( prisma_client: PrismaClient, ) -> None: - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[]) + prisma_client.db.query_raw = AsyncMock(return_value=[]) assert await prisma_client.get_latest_health_checks_for_models([]) == () - assert prisma_client.db.litellm_healthchecktable.find_many.await_count == 0 + assert prisma_client.db.query_raw.await_count == 0 @pytest.mark.asyncio -async def test_get_latest_health_checks_for_models_db_error_returns_empty_list( +async def test_get_latest_health_checks_for_models_db_error_returns_no_rows( prisma_client: PrismaClient, ) -> None: - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(side_effect=RuntimeError("oops")) + prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("oops")) assert await prisma_client.get_latest_health_checks_for_models(["gpt-5"]) == () diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py index 41b0eb3cf95..cec6dd99ce8 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py @@ -29,7 +29,7 @@ from __future__ import annotations import asyncio import time from typing import Any, Final -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, call import pytest @@ -111,6 +111,34 @@ async def test_run_reconnect_cycle_direct_path_recreates_when_probe_fails( } +@pytest.mark.asyncio +async def test_run_reconnect_cycle_direct_path_recreates_when_writer_is_read_only( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db") + prisma_client._engine_confirmed_dead = False + prisma_client._engine_pid = 0 + prisma_client._start_engine_watcher = AsyncMock() + prisma_client._cleanup_engine_watcher = MagicMock() + + writer: Final = MagicMock() + writer.query_raw = AsyncMock(side_effect=[[{"transaction_read_only": "on"}], [{"?column?": 1}]]) + writer.recreate_prisma_client = AsyncMock() + prisma_client.db = writer + + await prisma_client._run_reconnect_cycle(timeout_seconds=5) + pinned = { + "recreate_called": writer.recreate_prisma_client.await_count, + "start_watcher_called": prisma_client._start_engine_watcher.await_count, + "cleanup_called": prisma_client._cleanup_engine_watcher.call_count, + } + assert pinned == { + "recreate_called": 1, + "start_watcher_called": 1, + "cleanup_called": 1, + } + + @pytest.mark.asyncio async def test_run_reconnect_cycle_force_recreate_skips_probe_and_recreates( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch @@ -488,6 +516,93 @@ async def test_db_health_watchdog_loop_swallows_non_db_errors( assert prisma_client.attempt_db_reconnect.await_count == 0 +def _routing_db_with_writer_sessions(*transaction_read_only: str) -> tuple[RoutingPrismaWrapper, MagicMock]: + """One watchdog cycle per value, then the loop is cancelled.""" + writer: Final = MagicMock() + writer.query_raw = AsyncMock(side_effect=[[{"transaction_read_only": value}] for value in transaction_read_only]) + reader: Final = MagicMock() + reader.query_raw = AsyncMock( + side_effect=[[{"?column?": 1}] for _ in transaction_read_only] + [asyncio.CancelledError()] + ) + return RoutingPrismaWrapper(writer=writer, reader=reader), writer + + +@pytest.mark.asyncio +async def test_db_health_watchdog_loop_forces_recreate_when_writer_is_read_only( + prisma_client: PrismaClient, +) -> None: + prisma_client._db_health_watchdog_interval_seconds = 0 + prisma_client.attempt_db_reconnect = AsyncMock(side_effect=asyncio.CancelledError()) + prisma_client.db, _ = _routing_db_with_writer_sessions("on") + + await prisma_client._db_health_watchdog_loop() + assert prisma_client.attempt_db_reconnect.await_args is not None + assert prisma_client.attempt_db_reconnect.await_args.kwargs == { + "reason": "db_health_watchdog_writer_read_only", + "timeout_seconds": prisma_client._db_watchdog_reconnect_timeout_seconds, + "force_recreate": True, + } + + +@pytest.mark.asyncio +async def test_db_health_watchdog_loop_leaves_writable_writer_alone( + prisma_client: PrismaClient, +) -> None: + prisma_client._db_health_watchdog_interval_seconds = 0 + prisma_client.attempt_db_reconnect = AsyncMock() + prisma_client.db, writer = _routing_db_with_writer_sessions("off") + + await prisma_client._db_health_watchdog_loop() + assert (prisma_client.attempt_db_reconnect.await_count, writer.query_raw.await_count) == (0, 1) + + +@pytest.mark.asyncio +async def test_db_health_watchdog_loop_backs_off_while_database_stays_read_only( + prisma_client: PrismaClient, +) -> None: + prisma_client._db_health_watchdog_interval_seconds = 0 + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + prisma_client.db, writer = _routing_db_with_writer_sessions("on", "on", "on") + + await prisma_client._db_health_watchdog_loop() + assert (prisma_client.attempt_db_reconnect.await_count, writer.query_raw.await_count) == (1, 3) + + +@pytest.mark.asyncio +async def test_db_health_watchdog_loop_recreates_again_once_writer_was_writable_in_between( + prisma_client: PrismaClient, +) -> None: + prisma_client._db_health_watchdog_interval_seconds = 0 + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + prisma_client.db, _ = _routing_db_with_writer_sessions("on", "off", "on") + + await prisma_client._db_health_watchdog_loop() + assert prisma_client.attempt_db_reconnect.await_count == 2 + + +@pytest.mark.asyncio +async def test_recreate_read_only_writer_retries_after_backoff_elapses( + prisma_client: PrismaClient, +) -> None: + prisma_client._db_reconnect_cooldown_seconds = 15 + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + first: Final = await prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction") + within_backoff: Final = await prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction") + prisma_client._db_read_only_recreate_ts -= 30 + after_backoff: Final = await prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction") + prisma_client._db_read_only_recreate_ts -= 30 + still_within_doubled_backoff: Final = await prisma_client.recreate_read_only_writer( + reason="postgres_read_only_transaction" + ) + + assert (first, within_backoff, after_backoff, still_within_doubled_backoff) == (True, False, True, False) + assert prisma_client.attempt_db_reconnect.await_args_list == [ + call(reason="postgres_read_only_transaction", timeout_seconds=None, force_recreate=True), + call(reason="postgres_read_only_transaction", timeout_seconds=None, force_recreate=True), + ] + + @pytest.mark.asyncio async def test_iam_refresh_racing_reconnect_recreates_engine_only_once( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index a1eb88a7834..c8b87bd671e 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -538,10 +538,9 @@ async def test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested( """ import litellm.constants as constants_mod import litellm.proxy.utils as utils_mod - from litellm.proxy.utils import PrismaClient, request_spend_log_flush + from litellm.proxy.utils import request_spend_log_flush monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) - PrismaClient.spend_log_flush_requested.clear() mock_prisma_client.spend_log_transactions = [] mock_prisma_client.tool_usage_transactions = [] @@ -562,16 +561,107 @@ async def test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested( try: await asyncio.sleep(0.05) assert not flushed.is_set() + assert isinstance(mock_prisma_client.spend_log_flush_requested, asyncio.Event) mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r1")) - request_spend_log_flush() + request_spend_log_flush(mock_prisma_client) await asyncio.wait_for(flushed.wait(), timeout=5.0) finally: monitor.cancel() with suppress(asyncio.CancelledError): await monitor - PrismaClient.spend_log_flush_requested.clear() + + +def test_monitor_spend_logs_queue_flush_survives_an_earlier_event_loop( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A second monitor, started in a fresh event loop, is still woken by a flush request, + so a worker whose first loop is gone keeps flushing Responses rows instead of stalling. + """ + import litellm.constants as constants_mod + import litellm.proxy.utils as utils_mod + from litellm.proxy.utils import request_spend_log_flush + + monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) + mock_prisma_client.tool_usage_transactions = [] + + async def _flush_once_under_a_monitor() -> None: + flushed: Final = asyncio.Event() + + async def _fake_job(*args: Any, **kwargs: Any) -> None: + flushed.set() + + monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job) + mock_prisma_client.spend_log_transactions = [] + + monitor: Final = asyncio.create_task( + _monitor_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + ) + try: + await asyncio.sleep(0.05) + assert not flushed.is_set() + + mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r1")) + request_spend_log_flush(mock_prisma_client) + + await asyncio.wait_for(flushed.wait(), timeout=5.0) + finally: + monitor.cancel() + with suppress(asyncio.CancelledError): + await monitor + + asyncio.run(_flush_once_under_a_monitor()) + asyncio.run(_flush_once_under_a_monitor()) + + +@pytest.mark.asyncio +async def test_flush_requested_before_the_monitor_starts_costs_the_row_nothing( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A Responses row enqueued before the monitor exists still reaches the DB on its first + pass, so dropping that early request delays nothing. + """ + import litellm.constants as constants_mod + import litellm.proxy.utils as utils_mod + from litellm.proxy.utils import request_spend_log_flush + + monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) + mock_prisma_client.spend_log_flush_requested = None + mock_prisma_client.tool_usage_transactions = [] + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")] + + flushed: Final = asyncio.Event() + + async def _fake_job(*args: Any, **kwargs: Any) -> None: + flushed.set() + + monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job) + + request_spend_log_flush(mock_prisma_client) + assert mock_prisma_client.spend_log_flush_requested is None + + monitor: Final = asyncio.create_task( + _monitor_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + ) + try: + await asyncio.wait_for(flushed.wait(), timeout=5.0) + finally: + monitor.cancel() + with suppress(asyncio.CancelledError): + await monitor def test_raise_failed_update_spend_exception_emits_failure_handler() -> None: diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 2ba58bd5644..dfe106a3f52 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -10,7 +10,11 @@ Covers ``_should_use_guardrail_load_balancing``, ``_execute_guardrail_hook``, from __future__ import annotations import asyncio -from typing import Any, Dict, List +import json +from copy import deepcopy +import logging +from collections.abc import Iterator +from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -23,8 +27,13 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger -from litellm.proxy.utils import ProxyLogging -from litellm.types.guardrails import GuardrailEventHooks +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header +from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, @@ -149,9 +158,7 @@ async def test_execute_guardrail_hook_unknown_hook_type_raises(proxy_logging, ma @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_routes_through_router( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_routes_through_router(proxy_logging, make_user_api_key_auth): cb = _make_guardrail() router = MagicMock() router.get_available_guardrail = MagicMock(return_value={"callback": cb}) @@ -167,9 +174,7 @@ async def test_execute_guardrail_with_load_balancing_routes_through_router( @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_router_none_raises( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_router_none_raises(proxy_logging, make_user_api_key_auth): with patch("litellm.proxy.proxy_server.llm_router", None): with pytest.raises(ValueError, match="Router not initialized"): await proxy_logging._execute_guardrail_with_load_balancing( @@ -182,9 +187,7 @@ async def test_execute_guardrail_with_load_balancing_router_none_raises( @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_no_callback_raises( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_no_callback_raises(proxy_logging, make_user_api_key_auth): router = MagicMock() router.get_available_guardrail = MagicMock(return_value={"callback": None}) with patch("litellm.proxy.proxy_server.llm_router", router): @@ -204,9 +207,7 @@ async def test_execute_guardrail_with_load_balancing_no_callback_raises( @pytest.mark.asyncio -async def test_process_guardrail_callback_skipped_when_should_run_false( - proxy_logging, make_user_api_key_auth -): +async def test_process_guardrail_callback_skipped_when_should_run_false(proxy_logging, make_user_api_key_auth): cb = _make_guardrail() cb.should_run_guardrail = MagicMock(return_value=False) out = await proxy_logging._process_guardrail_callback( @@ -220,9 +221,7 @@ async def test_process_guardrail_callback_skipped_when_should_run_false( @pytest.mark.asyncio -async def test_process_guardrail_callback_returns_data_on_success( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_process_guardrail_callback_returns_data_on_success(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _make_guardrail() cb.should_run_guardrail = MagicMock(return_value=True) proxy_logging._should_use_guardrail_load_balancing = MagicMock(return_value=False) @@ -326,25 +325,26 @@ def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): @pytest.mark.asyncio async def test_maybe_execute_pipelines_no_pipelines_returns_data(proxy_logging, make_user_api_key_auth): data = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} - out = await proxy_logging._maybe_execute_pipelines( + out, replacement = await proxy_logging._maybe_execute_pipelines( data=data, user_api_key_dict=make_user_api_key_auth(), call_type="completion", event_hook="pre_call", ) assert out == {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} + assert replacement is None @pytest.mark.asyncio -async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_logging, make_user_api_key_auth, monkeypatch): +async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode( + proxy_logging, make_user_api_key_auth, monkeypatch +): pipeline = MagicMock() pipeline.mode = "post_call" # not pre_call data = {"metadata": {"_guardrail_pipelines": [("p1", pipeline)]}, "model": "m", "messages": []} executed = MagicMock() - monkeypatch.setattr( - "litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed - ) - out = await proxy_logging._maybe_execute_pipelines( + monkeypatch.setattr("litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed) + out, replacement = await proxy_logging._maybe_execute_pipelines( data=data, user_api_key_dict=make_user_api_key_auth(), call_type="completion", @@ -352,6 +352,7 @@ async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_log ) executed.assert_not_called() assert out is data + assert replacement is None @pytest.mark.parametrize( @@ -530,9 +531,7 @@ def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode(): litellm.callbacks = [cb] try: with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result( - result=result, data={"model": "m"}, policy_name="p" - ) + ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") finally: litellm.callbacks = saved @@ -644,9 +643,7 @@ async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch monkeypatch.setattr(litellm, "callbacks", [prom]) with pytest.raises(HTTPException): - await ProxyLogging._run_guardrail_with_metrics( - callback=cb, coro=task(), hook_type="post_call" - ) + await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call") assert detail["guardrail_name"] == "presidio" recorded = prom._record_guardrail_metrics.call_args.kwargs @@ -674,9 +671,7 @@ def _moderation_guardrail() -> MagicMock: @pytest.mark.asyncio -async def test_during_call_hook_records_latency_metric( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _moderation_guardrail() prom = _prometheus_callback() monkeypatch.setattr(litellm, "callbacks", [prom, cb]) @@ -695,9 +690,7 @@ async def test_during_call_hook_records_latency_metric( @pytest.mark.asyncio -async def test_post_call_success_hook_records_latency_metric( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _moderation_guardrail() prom = _prometheus_callback() monkeypatch.setattr(litellm, "callbacks", [prom, cb]) @@ -725,9 +718,7 @@ async def test_post_call_success_hook_records_latency_metric( async def test_process_prompt_template_no_op_when_no_prompt_spec(proxy_logging, monkeypatch): from litellm.proxy.prompts import prompt_registry - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None) data: Dict[str, Any] = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} await proxy_logging._process_prompt_template( data=data, @@ -752,9 +743,7 @@ async def test_process_prompt_template_applies_when_spec_resolves(proxy_logging, "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock( @@ -802,9 +791,7 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=RuntimeError("bad prompt")) with pytest.raises(RuntimeError): @@ -905,9 +892,7 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock( @@ -938,3 +923,1914 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p hook_kwargs = logging_obj.async_get_chat_completion_prompt.await_args.kwargs assert hook_kwargs["messages"] == [{"role": "user", "content": "Who are you?"}] assert hook_kwargs["prompt_spec"] is prompt_spec + + +# --------------------------------------------------------------------------- +# post_call pipeline execution (LIT-6410) +# --------------------------------------------------------------------------- + + +def _post_call_pipeline_data( + guardrail: str = "gr-post", step: PipelineStep | None = None, **extra: Any +) -> Dict[str, Any]: + pipeline = GuardrailPipeline( + mode="post_call", + steps=[step or PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("response-governance", pipeline)], + "_pipeline_managed_guardrails": {guardrail}, + "policy_sources": {"response-governance": "model:m"}, + }, + **extra, + } + + +@pytest.mark.asyncio +async def test_post_call_success_hook_runs_post_call_pipeline_and_reraises_block( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class OutputBlockingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["response"] = response + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + monkeypatch.setattr( + litellm, + "callbacks", + [OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = litellm.ModelResponse() + + with pytest.raises(HTTPException) as info: + await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert info.value.detail["error"] == "output blocked" + assert seen["response"] is response + + +@pytest.mark.asyncio +async def test_post_call_pipeline_pass_runs_once_and_leaves_request_data_untouched( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class RecordingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + seen["response"] = response + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [RecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = litellm.ModelResponse() + + out = await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert out is response + assert seen["response"] is response + assert seen["count"] == 1 + assert "response" not in data + assert "guardrails" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_post_call_pipeline_managed_default_on_guardrail_runs_exactly_once( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class CountingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [CountingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert seen["count"] == 1 + + +@pytest.mark.asyncio +async def test_post_call_hook_still_runs_guardrail_managed_only_by_pre_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class DualStageGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + return None + + pre_call_pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="gr-dual", on_pass="allow", on_fail="block")], + ) + monkeypatch.setattr( + litellm, + "callbacks", + [DualStageGuardrail(guardrail_name="gr-dual", event_hook=["pre_call", "post_call"], default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("request-governance", pre_call_pipeline)], + "_pipeline_managed_guardrails": {"gr-dual"}, + }, + } + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert seen["count"] == 1 + + +@pytest.mark.asyncio +async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class DualStageGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + seen["count"] += 1 + return data + + post_call_pipeline = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="gr-dual", on_pass="allow", on_fail="block")], + ) + monkeypatch.setattr( + litellm, + "callbacks", + [DualStageGuardrail(guardrail_name="gr-dual", event_hook=["pre_call", "post_call"], default_on=True)], + ) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("response-governance", post_call_pipeline)], + "_pipeline_managed_guardrails": {"gr-dual"}, + }, + } + + await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion") + + assert seen["count"] == 1 + + +@pytest.mark.asyncio +async def test_post_call_pipeline_replacement_response_reaches_caller( + proxy_logging, make_user_api_key_auth, monkeypatch +): + masked = litellm.ModelResponse() + + class MaskingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return masked + + monkeypatch.setattr( + litellm, + "callbacks", + [MaskingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + out = await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert out is masked + assert "response" not in data + + +@pytest.mark.asyncio +async def test_post_call_pipeline_replacement_chains_to_next_step_without_pass_data( + proxy_logging, make_user_api_key_auth, monkeypatch +): + masked = litellm.ModelResponse() + seen: Dict[str, Any] = {} + + class MaskingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return masked + + class RecordingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["response"] = response + return None + + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="gr-mask", on_pass="next", on_fail="block"), + PipelineStep(guardrail="gr-audit", on_pass="allow", on_fail="block"), + ], + ) + monkeypatch.setattr( + litellm, + "callbacks", + [ + MaskingGuardrail(guardrail_name="gr-mask", event_hook=GuardrailEventHooks.post_call, default_on=False), + RecordingGuardrail(guardrail_name="gr-audit", event_hook=GuardrailEventHooks.post_call, default_on=False), + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("response-governance", pipeline)], + "_pipeline_managed_guardrails": {"gr-mask", "gr-audit"}, + }, + } + + out = await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert out is masked + assert seen["response"] is masked + + +def test_handle_pipeline_result_modify_response_carries_original_response(): + result = MagicMock() + result.terminal_action = "modify_response" + result.modify_response_message = "filtered" + response = litellm.ModelResponse() + + with pytest.raises(ModifyResponseException) as info: + ProxyLogging._handle_pipeline_result( + result=result, data={"model": "m"}, policy_name="p", original_response=response + ) + + assert info.value.original_response is response + + +def test_handle_pipeline_result_allow_on_post_call_keeps_metadata_writes_only(): + data = {"a": 1, "metadata": {"guardrails": ["other"]}} + result = MagicMock() + result.terminal_action = "allow" + result.modified_data = { + "a": 2, + "metadata": {"guardrails": ["other"], "applied_guardrails": ["gr-post"]}, + "response": object(), + } + + out = ProxyLogging._handle_pipeline_result( + result=result, data=data, policy_name="p", original_response=litellm.ModelResponse() + ) + + assert out is data + assert data["a"] == 1 + assert "response" not in data + assert data["metadata"] == {"guardrails": ["other"], "applied_guardrails": ["gr-post"]} + + +@pytest.mark.asyncio +async def test_post_call_pipeline_guardrail_metadata_writes_reach_request_data( + proxy_logging, make_user_api_key_auth, monkeypatch +): + class HeaderWritingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name="gr-post") + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"verdict": "pass"}, + request_data=data, + guardrail_status="success", + ) + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [HeaderWritingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + slg_entries = data["metadata"]["standard_logging_guardrail_information"] + assert len(slg_entries) == 1 + assert slg_entries[0]["guardrail_name"] == "gr-post" + + +@pytest.mark.asyncio +async def test_post_call_pipeline_block_keeps_guardrail_metadata_writes( + proxy_logging, make_user_api_key_auth, monkeypatch +): + class BlockingWriterGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name="gr-post") + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"verdict": "fail"}, + request_data=data, + guardrail_status="guardrail_intervened", + ) + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + monkeypatch.setattr( + litellm, + "callbacks", + [BlockingWriterGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + slg_entries = data["metadata"]["standard_logging_guardrail_information"] + assert len(slg_entries) == 1 + assert slg_entries[0]["guardrail_name"] == "gr-post" + assert slg_entries[0]["guardrail_status"] == "guardrail_intervened" + + +@pytest.mark.asyncio +async def test_post_call_pipeline_managed_parallel_guardrail_runs_exactly_once( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class CountingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [ + CountingGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert seen["count"] == 1 + + +@pytest.mark.asyncio +async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class CountingGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + seen["count"] += 1 + return data + + pre_call_pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="gr-pre", on_pass="allow", on_fail="block")], + ) + monkeypatch.setattr( + litellm, + "callbacks", + [ + CountingGuardrail( + guardrail_name="gr-pre", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + ], + ) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("request-governance", pre_call_pipeline)], + "_pipeline_managed_guardrails": {"gr-pre"}, + }, + } + + await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion") + + assert seen["count"] == 1 + + +def _warnings(caplog: pytest.LogCaptureFixture) -> List[str]: + return [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] + + +@pytest.mark.asyncio +async def test_streaming_request_whose_pipeline_guardrail_is_missing_streams_verbatim( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", []) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + delivered: List[Any] = [] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + guardrails_only=True, + ) + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ): + delivered.append(item) + + assert out is not None + assert out.get("stream") is True + assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] + assert len(delivered) == 2 + assert any("response-governance" in message and "gr-post" in message for message in _warnings(caplog)) + + +def _background_response(status: str, text: str = "") -> ResponsesAPIResponse: + output = ( + [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}] if text else [] + ) + return ResponsesAPIResponse(id="resp_bg", created_at=0, output=output, status=status) + + +def _output_blocking_callbacks(seen: dict[str, object]) -> list[CustomGuardrail]: + class OutputBlockingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["response"] = response + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + return [ + OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pending_status", ["queued", "in_progress"]) +async def test_post_call_success_hook_waits_for_pending_background_response_before_running_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + pending_status: str, +) -> None: + seen: dict[str, object] = {} + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(background=True) + response = _background_response(pending_status) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert out is response + assert "response" not in seen + assert not _warnings(caplog) + assert any( + "response-governance" in record.getMessage() and pending_status in record.getMessage() + for record in caplog.records + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("final_status", ["completed", "incomplete"]) +async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_response( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + final_status: str, +) -> None: + seen: dict[str, object] = {} + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = _background_response(final_status, text="kumquat") + + with pytest.raises(HTTPException) as info: + await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert info.value.detail["error"] == "output blocked" + assert seen["response"] is response + + +def _output_passing_callbacks() -> list[CustomGuardrail]: + class OutputPassingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + return [ + OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + ] + + +def _claimed_post_call_pipeline_data( + *policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str | None = "model:m" +): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + step = {"guardrail": "gr-post", "on_pass": "allow", "on_fail": "block"} + get_policy_registry().load_policies( + { + policy_name: { + "guardrails": {"add": ["gr-post", *(extra_guardrails or {}).get(policy_name, [])]}, + "pipeline": {"mode": "post_call", "steps": [step]}, + } + for policy_name in policy_names + } + ) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(**step)]) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [(policy_name, pipeline) for policy_name in policy_names], + "_pipeline_managed_guardrails": {"gr-post"}, + "applied_policies": list(policy_names), + "applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)], + "policy_sources": {policy_name: policy_source for policy_name in policy_names if policy_source is not None}, + }, + } + + +@pytest.fixture +def clear_policy_registry() -> Iterator[None]: + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + yield + get_policy_registry().clear() + + +@pytest.mark.asyncio +async def test_pending_background_response_withdraws_the_deferred_policy_claims( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "applied_policies" not in data["metadata"] + assert "policy_sources" not in data["metadata"] + assert "applied_guardrails" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_pending_background_response_warns_when_the_deferred_policy_was_matched_through_a_tag( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance", policy_source="tag:governed+model:m") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "policy_sources" not in data["metadata"] + assert [message for message in _warnings(caplog) if "through a request tag" in message] == [ + "Policy engine: background response resp_bg matched post_call policies through a request tag at submit; " + "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " + "does not govern the completed response: response-governance" + ] + + +@pytest.mark.asyncio +async def test_pending_background_response_warns_when_the_deferred_policy_came_from_the_request_body( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("body-governance", policy_source=None) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "policy_sources" not in data["metadata"] + assert _warnings(caplog) == [ + "Policy engine: background response resp_bg matched post_call policies through the request body's policies " + "list at submit; retrieval carries no request body, so those policies do not govern the completed " + "response: body-governance" + ] + + +@pytest.mark.asyncio +async def test_pending_background_response_matched_through_its_model_does_not_warn( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await proxy_logging.post_call_success_hook( + data=_claimed_post_call_pipeline_data("response-governance"), + response=_background_response("queued"), + user_api_key_dict=make_user_api_key_auth(), + ) + + assert _warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data( + "input-and-output-governance", + "response-governance", + extra_guardrails={"input-and-output-governance": ["gr-pre"]}, + ) + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("in_progress"), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_policies"] == ["input-and-output-governance"] + assert data["metadata"]["applied_guardrails"] == ["gr-pre"] + assert data["metadata"]["policy_sources"] == {"input-and-output-governance": "model:m"} + + +@pytest.mark.asyncio +async def test_pending_background_response_keeps_the_claim_of_a_default_on_guardrail_that_ran_pre_call( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: + class DualStageGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + monkeypatch.setattr( + litellm, + "callbacks", + [DualStageGuardrail(guardrail_name="gr-post", event_hook=["pre_call", "post_call"], default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert "applied_policies" not in data["metadata"] + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + + +@pytest.mark.asyncio +async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_passing_callbacks()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("completed", text="fine"), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_policies"] == ["response-governance"] + assert data["metadata"]["policy_sources"] == {"response-governance": "model:m"} + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + + +@pytest.mark.asyncio +async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + data = _post_call_pipeline_data(background=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="aresponses", + guardrails_only=True, + ) + + assert out is not None + assert out.get("background") is True + assert not _warnings(caplog) + + +# --------------------------------------------------------------------------- +# post_call pipelines on streaming responses +# --------------------------------------------------------------------------- + + +def _unified_stream_guardrail(seen: Dict[str, Any], block: bool = False) -> CustomGuardrail: + class UnifiedStreamGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + seen["count"] = seen.get("count", 0) + 1 + seen["input_type"] = input_type + if block: + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + return inputs + + return UnifiedStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _stream_chunks() -> List[Any]: + return [ + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "hello "}, "finish_reason": None}]), + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}, "finish_reason": "stop"}]), + ] + + +async def _async_chunk_iter(chunks: List[Any]): + for chunk in chunks: + yield chunk + + +def _legacy_hook_stream_guardrail( + seen: Dict[str, Any], + rewrite: Callable[[Any], Any] | None = None, + raises: Exception | None = None, + native_lifecycle: bool = False, +) -> CustomGuardrail: + class LegacyHookGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = native_lifecycle + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] = seen.get("count", 0) + 1 + seen["data"] = data + seen["user_api_key_dict"] = user_api_key_dict + seen["response"] = deepcopy(response) + if raises is not None: + raise raises + return None if rewrite is None else rewrite(response) + + if native_lifecycle: + + class NativeLifecycleGuardrail(LegacyHookGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail") + + return NativeLifecycleGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False + ) + return LegacyHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _iterator_hook_only_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail: + class IteratorHookGuardrail(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["count"] = seen.get("count", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + return IteratorHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + +def _iterator_and_legacy_hook_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail: + class IteratorAndLegacyHookGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["success_hook_calls"] = seen.get("success_hook_calls", 0) + 1 + return None + + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["iterator_hook_calls"] = seen.get("iterator_hook_calls", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + return IteratorAndLegacyHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + +def _rewritten_model_response(response: Any) -> litellm.ModelResponse: + payload = response.model_dump() + payload["choices"][0]["message"]["content"] = "[REWRITTEN] " + payload["choices"][0]["message"]["content"] + return litellm.ModelResponse(**payload) + + +def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator_only( + make_user_api_key_auth, monkeypatch, caplog +): + supported = _unified_stream_guardrail({}) + legacy = _legacy_hook_stream_guardrail({}) + legacy.guardrail_name = "gr-legacy" + iterator_only = _iterator_hook_only_guardrail("gr-iterator", {}) + monkeypatch.setattr(litellm, "callbacks", [supported, legacy, iterator_only]) + governed = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-legacy", on_fail="block")], + ) + ungoverned = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-iterator", on_fail="block")], + ) + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-iterator", on_fail="block")]) + data = { + "metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]} + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions")) + + assert streamable == (("governed", governed),) + assert any("'ungoverned'" in message and "gr-iterator" in message for message in _warnings(caplog)) + assert not any("'governed'" in message or "gr-legacy" in message for message in _warnings(caplog)) + + +@pytest.mark.parametrize( + "request_route", + ["/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"], +) +def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response( + make_user_api_key_auth, monkeypatch, caplog, request_route +): + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail({})]) + legacy = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("legacy-governance", legacy)]}} + auth = make_user_api_key_auth(request_route=request_route) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, auth) + + assert streamable == () + assert stream_gated_guardrail_names(data, auth) == frozenset() + assert any("'legacy-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + +def test_streamable_post_call_pipelines_keeps_guardrails_with_their_own_iterator_hook_on_their_own_path( + make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", {})]) + both_hooks = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("both-hooks", both_hooks)]}} + auth = make_user_api_key_auth(request_route="/v1/chat/completions") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, auth) + + assert streamable == () + assert stream_gated_guardrail_names(data, auth) == frozenset() + assert any("'both-hooks'" in message and "gr-post" in message for message in _warnings(caplog)) + + +def test_streamable_post_call_pipelines_is_empty_on_route_without_translation( + make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail({})]) + governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("governed", governed)]}} + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/custom/stream")) + + assert streamable == () + assert any("/custom/stream" in message and "governed" in message for message in _warnings(caplog)) + + +def test_streamable_post_call_pipelines_is_empty_without_post_call_pipelines(make_user_api_key_auth, caplog): + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) + auth = make_user_api_key_auth(request_route="/custom/stream") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + assert _streamable_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", pre_call)]}}, auth) == () + assert _streamable_post_call_pipelines({"stream": True}, auth) == () + + assert _warnings(caplog) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_route", [None, "/v1/chat/completions"]) +async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_unified( + proxy_logging, make_user_api_key_auth, monkeypatch, request_route +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + data = _post_call_pipeline_data(stream=True) + + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route=request_route), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert out is not None + assert out.get("stream") is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_lifecycle", [False, True]) +async def test_streaming_iterator_hook_runs_legacy_hook_and_delivers_its_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle, caplog +): + seen: Dict[str, Any] = {} + guardrail = _legacy_hook_stream_guardrail(seen, rewrite=_rewritten_model_response, native_lifecycle=native_lifecycle) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + auth = make_user_api_key_auth(request_route="/v1/chat/completions") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.pre_call_hook( + user_api_key_dict=auth, data=data, call_type="completion", guardrails_only=True + ) + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + assert out is not None and out.get("stream") is True + assert seen["count"] == 1 + assert isinstance(seen["response"], litellm.ModelResponse) + assert seen["response"].choices[0].message.content == "hello world" + assert seen["data"]["messages"] == data["messages"] + assert seen["user_api_key_dict"] is auth + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert delivered[0].choices[0].delta.content == "[REWRITTEN] hello world" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + assert _warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert seen["count"] == 1 + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert [item.choices[0].delta.content for item in delivered] == ["hello ", "world"] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_ends_stream_with_legacy_hook_exception( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + blocked = HTTPException(status_code=400, detail={"error": "output blocked"}) + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, raises=blocked)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ): + delivered.append(item) + + with pytest.raises(HTTPException) as info: + await _drain() + + assert seen["count"] == 1 + assert delivered == [] + assert info.value is blocked + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + def rewrite(response: Any) -> Dict[str, Any]: + return {**response, "content": [{"type": "text", "text": "[REWRITTEN] " + response["content"][0]["text"]}]} + + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=rewrite)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_sse_chunks()), + request_data=data, + ) + ] + + assert seen["count"] == 1 + assert seen["response"]["content"][0]["text"] == "hello world" + assert seen["response"]["role"] == "assistant" + raw = b"".join(delivered).decode() + assert "[REWRITTEN] hello world" in raw + assert raw.count("event: content_block_delta") == 1 + for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"): + assert f"event: {expected_event}" in raw + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeline_cannot_stream( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_iterator_hook_only_guardrail("gr-post", seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ) + ] + + assert seen["count"] == 1 + assert [item.choices[0].delta.content for item in delivered] == ["[governed] hello ", "[governed] world"] + assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_runs_the_iterator_hook_of_a_guardrail_that_also_has_a_post_call_hook( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ) + ] + + assert seen == {"iterator_hook_calls": 1} + assert [item.choices[0].delta.content for item in delivered] == ["[governed] hello ", "[governed] world"] + assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rewrite_attribute, value", + [ + ("mask_response_content", True), + ("streaming_transform_mode", "incremental_diff"), + ("guardrail_config", {"streaming_transform_mode": "incremental_diff"}), + ], +) +async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_rewrites_streamed_content( + proxy_logging, make_user_api_key_auth, monkeypatch, rewrite_attribute, value +): + seen: Dict[str, Any] = {} + guardrail = _unified_stream_guardrail(seen) + setattr(guardrail, rewrite_attribute, value) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert out is not None + assert out.get("stream") is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", [ContentFilterAction.MASK, ContentFilterAction.BLOCK]) +async def test_pre_call_hook_allows_streaming_when_content_filter_step_masks_or_blocks( + proxy_logging, make_user_api_key_auth, monkeypatch, action +): + guardrail = ContentFilterGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="persimmon", action=action)], + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") + + out = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + + assert out is not None and out.get("stream") is True + + +@pytest.mark.asyncio +async def test_pre_call_hook_allows_streaming_when_content_filter_category_masks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + guardrail = ContentFilterGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + categories=[{"category": "bias_gender", "enabled": True, "action": "MASK"}], + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") + + out = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + + assert out is not None and out.get("stream") is True + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_releases_stream_when_route_has_no_guardrail_translation( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + delivered: List[Any] = [] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/custom/stream"), + data=data, + call_type="completion", + guardrails_only=True, + ) + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/custom/stream"), + response=_async_chunk_iter(chunks), + request_data=data, + ): + delivered.append(item) + + assert out is not None + assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] + assert len(delivered) == 2 + assert seen.get("count") is None + assert any("/custom/stream" in message and "response-governance" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_allow_releases_buffered_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert seen["count"] == 1 + assert seen["input_type"] == "response" + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_block_withholds_all_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ): + delivered.append(item) + + with pytest.raises(HTTPException) as info: + await _drain() + + assert delivered == [] + assert info.value.status_code == 400 + assert "output blocked" in str(info.value.detail) + + +def _rewriting_stream_guardrail(transform: Callable[[Dict[str, Any]], Dict[str, Any]]) -> CustomGuardrail: + class RewritingStreamGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, **transform(inputs)} + + return RewritingStreamGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False + ) + + +def _tool_call_stream_chunks() -> List[Any]: + tool_call = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"ssn": "123"}'}, + } + return [ + litellm.ModelResponseStream( + choices=[{"index": 0, "delta": {"tool_calls": [tool_call]}, "finish_reason": None}] + ), + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]), + ] + + +def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]: + return [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")]) +async def test_streaming_iterator_hook_pipeline_delivers_runtime_tool_call_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error, caplog +): + transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + step = PipelineStep(guardrail="gr-post", on_pass="allow", on_fail=on_fail, on_error=on_error) + data = _post_call_pipeline_data(step=step, stream=True) + delivered: List[Any] = [] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_tool_call_stream_chunks()), + request_data=data, + ): + delivered.append(item) + + assert len(delivered) == 2 + delivered_tool_call = delivered[0].choices[0].delta.tool_calls[0] + assert delivered_tool_call.function.arguments == '{"ssn": "[MASKED]"}' + assert delivered_tool_call.function.name == "lookup" + assert delivered_tool_call.id == "call_1" + assert delivered[1].choices[0].finish_reason == "tool_calls" + assert not any("discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_runtime_text_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch +): + transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert delivered[0].choices[0].delta.content == "hello [MASKED]" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_chains_text_rewrites_across_steps( + proxy_logging, make_user_api_key_auth, monkeypatch +): + second_step_saw: Dict[str, Any] = {} + + class FirstMask(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs["texts"]]} + + class SecondMask(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + second_step_saw["texts"] = list(inputs["texts"]) + return {**inputs, "texts": [text.replace("hello", "[GREETING]") for text in inputs["texts"]]} + + monkeypatch.setattr( + litellm, + "callbacks", + [ + FirstMask(guardrail_name="gr-first", event_hook=GuardrailEventHooks.post_call, default_on=False), + SecondMask(guardrail_name="gr-second", event_hook=GuardrailEventHooks.post_call, default_on=False), + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="gr-first", on_pass="next", on_fail="block"), + PipelineStep(guardrail="gr-second", on_pass="allow", on_fail="block"), + ], + ) + data = _post_call_pipeline_data(stream=True) + data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)] + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert second_step_saw["texts"] == ["hello [MASKED]"] + assert delivered[0].choices[0].delta.content == "[GREETING] [MASKED]" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "make_chunks, transform", + [ + (_stream_chunks, lambda inputs: {"texts": tuple(inputs["texts"])}), + (_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "123"}')}), + ], + ids=["texts_as_tuple", "tool_calls_as_dicts"], +) +async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_another_shape( + proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = make_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_skips_pipeline_and_warns_without_request_route( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] + assert len(delivered) == 2 + assert seen.get("count") is None + assert any("response-governance" in message and "route None" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_runs_pipeline_managed_guardrail_without_request_route( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class UnifiedRecordingGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1 + return None + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + monkeypatch.setattr( + litellm, + "callbacks", + [UnifiedRecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(), + ) + + assert result is not None + assert seen["gr-post"] == 1 + + +def _anthropic_sse_chunks() -> List[bytes]: + events = [ + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "m", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated_block( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep( + guardrail="gr-post", + on_pass="allow", + on_fail="modify_response", + modify_response_message="content policy block", + ) + ], + ) + data = _post_call_pipeline_data(stream=True) + data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)] + chunks = _anthropic_sse_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert seen["count"] == 1 + assert "content policy block" in raw + assert "hello world" not in raw + assert not any(item is chunk for item in delivered for chunk in chunks) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _anthropic_sse_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert "hello [MASKED]" in raw + assert "hello world" not in raw + assert raw.count("event: content_block_delta") == 1 + for expected_event in ( + "message_start", + "content_block_start", + "content_block_stop", + "message_delta", + "message_stop", + ): + assert f"event: {expected_event}" in raw + + +@pytest.mark.asyncio +async def test_pipeline_executor_discards_text_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation + from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor + + class NoWriteBackTranslation(BaseTranslation): + async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj): + return data + + async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj, **kwargs): + return response + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj, + user_api_key_dict=None, + request_data=None, + stream_transform_sink=None, + deliver_ended_stream_rewrites=False, + ): + assert deliver_ended_stream_rewrites is False + await guardrail_to_apply.apply_guardrail( + inputs={"texts": ["hello world"]}, + request_data=request_data or {}, + input_type="response", + ) + return responses_so_far + + transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + chunks = _stream_chunks() + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="gr-post", on_pass="allow", on_fail="block")], + mode="post_call", + data={"metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + call_type="acompletion", + policy_name="response-governance", + streaming_chunks=chunks, + endpoint_translation=NoWriteBackTranslation(), + ) + + assert result.terminal_action == "allow" + assert [chunk.choices[0].delta.content for chunk in chunks] == ["hello ", "world"] + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class RecordingGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1 + return None + + class UnifiedRecordingGuardrail(RecordingGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + managed = UnifiedRecordingGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + free = RecordingGuardrail(guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [managed, free]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + ) + + assert result is not None + assert seen.get("gr-post") is None + assert seen["gr-free"] == 1 + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_runs_guardrail_whose_pipeline_cannot_stream( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class ChunkHookGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen["count"] = seen.get("count", 0) + 1 + seen["response"] = response + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [ChunkHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + ) + + assert result is not None + assert seen["count"] == 1 + assert seen["response"] == "hello " + + +def _mask_tool_call_arguments(inputs: Dict[str, Any]) -> Dict[str, Any]: + return { + "tool_calls": [ + { + "id": stream_item_field(tool_call, "id"), + "type": "function", + "function": { + "name": stream_item_field(stream_item_field(tool_call, "function"), "name"), + "arguments": '{"fruit": "[MASKED]"}', + }, + } + for tool_call in inputs.get("tool_calls", []) + ] + } + + +def _anthropic_tool_use_sse_chunks() -> List[bytes]: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_tool_use_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert '{\\"fruit\\": \\"[MASKED]\\"}' in raw + assert "persim" not in raw + assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw + assert '"stop_reason": "tool_use"' in raw + assert raw.count("event: content_block_delta") == 2 + + +def _responses_function_call_events() -> List[Dict[str, Any]]: + def item(arguments: str, status: str) -> Dict[str, Any]: + return { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": ' "persimmon"}'}, + {"type": "response.function_call_arguments.done", "item_id": "fc_1", "output_index": 0, "arguments": '{"fruit": "persimmon"}'}, + {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": {"id": "resp_1", "created_at": 1, "model": "m", "output": [item('{"fruit": "persimmon"}', "completed")], "status": "completed"}, + }, + ] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_function_call_rewrite_on_responses_events( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"), + response=_async_chunk_iter(_responses_function_call_events()), + request_data=data, + ) + ] + + assert [event["type"] for event in delivered] == [event["type"] for event in _responses_function_call_events()] + assert [event["delta"] for event in delivered if event["type"] == "response.function_call_arguments.delta"] == ['{"fruit": "[MASKED]"}', ""] + assert delivered[3]["arguments"] == '{"fruit": "[MASKED]"}' + assert delivered[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert delivered[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' + assert "persimmon" not in json.dumps(delivered) + + +def _drop_tool_calls(inputs: Dict[str, Any]) -> Dict[str, Any]: + return {"tool_calls": []} + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_chat_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_tool_call_stream_chunks()), + request_data=data, + ) + ] + + assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}' + assert delivered[1].choices[0].finish_reason == "tool_calls" + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()), + request_data=data, + ) + ] + + assert delivered == _anthropic_tool_use_sse_chunks() + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_responses_events( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"), + response=_async_chunk_iter(_responses_function_call_events()), + request_data=data, + ) + ] + + assert delivered == _responses_function_call_events() + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 9d2a27ce9d3..af89c424f8b 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -681,7 +681,7 @@ async def test_scan_raw_request_snapshot_taken_before_pipelines( for msg in data.get("messages", []): if "SECRET" in msg.get("content", ""): msg["content"] = msg["content"].replace("SECRET", "[REDACTED]") - return data + return data, None monkeypatch.setattr(ProxyLogging, "_maybe_execute_pipelines", fake_pipelines) monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(scan_raw_request=True)]) diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py index 93049b21460..dc4f2900038 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py @@ -244,3 +244,65 @@ async def test_list_vector_stores_dashboard_session_resolves_real_teams( ), ): assert await _listed_ids(alice) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("user_team_ids", "expected_status"), + [ + (["team_a"], 200), + (["team_b"], 403), + ([], 403), + ], +) +async def test_get_vector_store_info_dashboard_session_resolves_real_teams( + user_team_ids: list[str], expected_status: int +): + """Regression for LIT-7132: /vector_store/info must grant a dashboard session the same team-owned stores + /vector_store/list shows it, instead of judging the session's reserved litellm-dashboard team id.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + get_vector_store_info, + ) + from litellm.types.vector_stores import VectorStoreInfoRequest + + alice = UserAPIKeyAuth( + team_id="litellm-dashboard", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + async def fake_get_team_object(team_id: str, **_kwargs: object) -> LiteLLM_TeamTableCachedObj: + return LiteLLM_TeamTableCachedObj(team_id=team_id) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock( + return_value=MagicMock(model_dump=lambda: dict(_TEAM_A_OWNED)) + ) + + async def outcome() -> int: + try: + response = await get_vector_store_info( + data=VectorStoreInfoRequest(vector_store_id="vs_team_a"), user_api_key_dict=alice + ) + except HTTPException as exc: + return exc.status_code + assert response["vector_store"]["vector_store_id"] == "vs_team_a" + return 200 + + with ( + patch( # test-quality-ok: the endpoint reads the store row through the module-level prisma client, no injection seam + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch( # test-quality-ok: the endpoint consults the module-level registry before the DB, no injection seam + "litellm.vector_store_registry", None + ), + patch( # test-quality-ok: team rows come from the module-level prisma client, no injection seam + "litellm.proxy.auth.auth_checks.get_team_object", new=fake_get_team_object + ), + patch( # test-quality-ok: the user row comes from the module-level prisma client, no injection seam + "litellm.proxy.vector_store_endpoints.utils.resolve_ui_session_team_ids", + new=AsyncMock(return_value=user_team_ids), + ), + ): + assert await outcome() == expected_status diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index 758d379f22c..c7f6b8ff83a 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -4,10 +4,13 @@ Tests for gateway repository layer. import json from datetime import datetime -from typing import Any, Dict, List, Optional +from types import SimpleNamespace +from typing import Any, Dict, Final, List, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from prisma import models as prisma_models +from prisma.builder import QueryBuilder from litellm.models.base import DomainModel from litellm.models.budget import LiteLLM_BudgetTable @@ -307,6 +310,23 @@ class TestModelRepository: client = MockPrismaClient() return ModelRepository(client) + @pytest.mark.asyncio + async def test_find_all_except_serializes_exclusion_for_prisma(self) -> None: + find_many: Final = AsyncMock(return_value=[]) + client: Final = SimpleNamespace( + db=SimpleNamespace(litellm_proxymodeltable=SimpleNamespace(find_many=find_many)) + ) + + await ModelRepository(client).find_all_except("current-model") + + find_many.assert_awaited_once() + query: Final = QueryBuilder( + method="find_many", + model=prisma_models.LiteLLM_ProxyModelTable, + arguments=find_many.call_args.kwargs, + ).build_query() + assert 'where: { model_id: { not: "current-model" } }' in " ".join(query.split()) + def test_table_is_wrapped_for_config_sync(self, repo): from litellm.proxy.common_utils.config_sync_pubsub import ( _PublishOnWriteActions, @@ -2387,3 +2407,60 @@ class TestCountBillableUsers: client.db.litellm_usertable = _RacyTable() repo = UserRepository(client) assert await repo.count_billable_users() == 0 + + +class TestAutoRouterSessionRepository: + ROW: Final = { + "api_key": "hashed-key", + "session_id": "s1", + "router_name": "claude-auto", + "router_type": "complexity", + "first_turn_at": datetime(2026, 9, 1, 12, 0, 0), + "last_turn_at": datetime(2026, 9, 1, 12, 5, 0), + "last_model": "anthropic/claude-sonnet-5", + "models": {"anthropic/claude-sonnet-5": {"at": 1.0, "ttl": None}}, + "turns": 3, + "spend": 0.14, + "saved_spend": 0.24, + "classifier_cost": 0.01, + "tier_turns": {"complex": 3}, + "baseline_models": {"anthropic/claude-opus-5": 3}, + } + + @staticmethod + def _repo(record: Optional[Dict[str, Any]]): + from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository + + lookups: List[Dict[str, Any]] = [] + + class _Table: + async def find_first(self, where: Dict[str, Any], order: Dict[str, str]): + lookups.append({"where": where, "order": order}) + return MockRecord(record) if record is not None else None + + client = MagicMock() + client.db.litellm_autoroutersession = _Table() + return AutoRouterSessionRepository(client), lookups + + @pytest.mark.asyncio + async def test_find_latest_for_key_reads_the_keys_own_partition_newest_router_first(self): + repo, lookups = self._repo(dict(self.ROW)) + row = await repo.find_latest_for_key("hashed-key", "s1") + assert lookups == [{"where": {"api_key": "hashed-key", "session_id": "s1"}, "order": {"last_turn_at": "desc"}}] + assert row is not None + assert (row.router_name, row.turns, row.spend, row.saved_spend) == ("claude-auto", 3, 0.14, 0.24) + assert row.baseline_models == {"anthropic/claude-opus-5": 3} + assert row.baseline_model == "anthropic/claude-opus-5" + + @pytest.mark.asyncio + async def test_find_latest_for_key_is_none_when_the_key_wrote_no_such_session(self): + repo, _ = self._repo(None) + assert await repo.find_latest_for_key("hashed-key", "unknown") is None + + def test_table_is_the_session_rollup_and_needs_a_database(self): + from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository + + client = MagicMock() + assert AutoRouterSessionRepository(client).table is client.db.litellm_autoroutersession + with pytest.raises(RuntimeError, match="No DB Connected"): + _ = AutoRouterSessionRepository(None).table diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py index 2b6cfeda2c2..0992cd9bb37 100644 --- a/tests/test_litellm/rerank_api/test_main.py +++ b/tests/test_litellm/rerank_api/test_main.py @@ -111,6 +111,72 @@ def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter): assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key" +DASHSCOPE_RERANK_BODY = { + "object": "list", + "results": [{"index": 0, "relevance_score": 0.95}], + "model": "qwen3-rerank", + "id": "rerank-mock-id", + "usage": {"total_tokens": 10}, +} + + +def test_dashscope_rerank_defaults_to_live_rerank_route(respx_mock: respx.MockRouter, monkeypatch): + """Regression for the dead default endpoint: get_llm_provider always returns the + chat base for dashscope, which used to hijack rerank onto the dead + /compatible-mode/v1/reranks route.""" + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False) + + mock_route = respx_mock.post("https://dashscope.aliyuncs.com/compatible-api/v1/reranks") + mock_route.return_value = httpx.Response(200, json=DASHSCOPE_RERANK_BODY) + + response = litellm.rerank( + model="dashscope/qwen3-rerank", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-dashscope-key", + ) + + assert mock_route.called + assert response.results[0]["relevance_score"] == 0.95 + + +def test_dashscope_rerank_chat_env_base_keeps_host_and_rerank_route(respx_mock: respx.MockRouter, monkeypatch): + """Regression: a chat-style DASHSCOPE_API_BASE must not hijack rerank onto the + chat path, while its host (the region) is preserved.""" + monkeypatch.setenv("DASHSCOPE_API_BASE", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1") + monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False) + + mock_route = respx_mock.post("https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks") + mock_route.return_value = httpx.Response(200, json=DASHSCOPE_RERANK_BODY) + + litellm.rerank( + model="dashscope/qwen3-rerank", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-dashscope-key", + ) + + assert mock_route.called + + +def test_dashscope_rerank_explicit_api_base_wins(respx_mock: respx.MockRouter, monkeypatch): + monkeypatch.setenv("DASHSCOPE_API_BASE", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1") + + mock_route = respx_mock.post("https://custom-rerank.example/v1/reranks") + mock_route.return_value = httpx.Response(200, json=DASHSCOPE_RERANK_BODY) + + litellm.rerank( + model="dashscope/qwen3-rerank", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-dashscope-key", + api_base="https://custom-rerank.example/v1", + ) + + assert mock_route.called + + DASHSCOPE_404_BODY = { "error": { "message": "The model `does-not-exist` does not exist or you do not have access to it.", diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 2068f10ea2d..46249e50572 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1,4 +1,5 @@ import json +from typing import Final import pytest @@ -1421,6 +1422,88 @@ class TestToolChoiceTransformation: ) assert result == "required" + @pytest.mark.parametrize( + "request_tool_choice,expected", + [ + ({"type": "function", "name": "run_command"}, {"type": "function", "name": "run_command"}), + ({"type": "function", "function": {"name": "run_command"}}, {"type": "function", "name": "run_command"}), + ({"type": "custom", "name": "ApplyPatch"}, {"type": "custom", "name": "ApplyPatch"}), + ({"type": "custom", "custom": {"name": "ApplyPatch"}}, {"type": "custom", "name": "ApplyPatch"}), + ({"type": "function"}, "required"), + ({"type": "tool"}, "required"), + ({"type": "auto"}, "auto"), + ("required", "required"), + ("none", "none"), + (None, "auto"), + ("any", "auto"), + ("run_command", "auto"), + ({"name": "run_command"}, "auto"), + ], + ) + def test_transform_tool_choice_for_responses_api_response( + self, request_tool_choice: object, expected: str | dict[str, str] + ) -> None: + result: Final = LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + request_tool_choice + ) + assert result == expected + + def test_non_streamed_response_echoes_named_tool_choice_in_responses_api_shape(self) -> None: + chat_completion_response: Final = ModelResponse( + id="chatcmpl-named-tool-choice", + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion", + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_pwd", + type="function", + function=Function(name="run_command", arguments='{"command":"pwd"}'), + ) + ], + ), + ) + ], + ) + + responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": {"type": "function", "name": "run_command"}}, + chat_completion_response=chat_completion_response, + ) + + assert responses_api_response.tool_choice == {"type": "function", "name": "run_command"} + + def test_non_streamed_response_with_unrecognized_tool_choice_echoes_auto(self) -> None: + chat_completion_response: Final = ModelResponse( + id="chatcmpl-unrecognized-tool-choice", + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message(role="assistant", content="/Users/dev"), + ) + ], + ) + + responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": "any"}, + chat_completion_response=chat_completion_response, + ) + + assert responses_api_response.tool_choice == "auto" + class TestContentTypeTransformation: """Test content type transformation from Responses API to Chat Completion format""" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 719d51c11e3..850ee7ba623 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -11,6 +11,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve """ import json +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -628,3 +629,79 @@ def test_streamed_anthropic_tool_call_events_correlate_on_normalized_item_id(): assert item_dones[0].item.call_id == "toolu_01AbCdEf" for evt in deltas + dones: assert evt.item_id == added[0].item.id + + +def _tool_call_chunk(finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_pwd", + "type": "function", + "function": {"name": "run_command", "arguments": '{"command":"pwd"}'}, + "index": 0, + } + ], + ), + finish_reason=finish_reason, + ) + ], + ) + + +def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape() -> None: + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), + request_input="Run the command pwd.", + responses_api_request={ + "tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}], + "tool_choice": {"type": "function", "name": "run_command"}, + }, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + events: Final = list(iterator) + + response_events: Final = [event for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] + assert [event.type for event in response_events] == [ + "response.created", + "response.in_progress", + "response.completed", + ] + assert [event.response.tool_choice for event in response_events] == [ + {"type": "function", "name": "run_command"}, + {"type": "function", "name": "run_command"}, + {"type": "function", "name": "run_command"}, + ] + assert any(getattr(event, "type", None) == "response.output_item.done" for event in events) + + +def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None: + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), + request_input="Run the command pwd.", + responses_api_request={ + "tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}], + "tool_choice": "any", + }, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + response_events: Final = [ + event for event in iterator if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES + ] + + assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"] diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index ab4c5185057..2c1845f7b92 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -1,6 +1,13 @@ +import json +import sys +import types + import pytest +import respx +from httpx import Response from unittest.mock import AsyncMock, patch +import litellm from litellm.types.utils import ModelResponse from litellm.responses.mcp import chat_completions_handler @@ -1344,3 +1351,39 @@ async def test_acompletion_with_mcp_streaming_drains_inner_stream_after_exhausti assert len(all_chunks) == 3 assert initial_stream.drained_after_exhaustion is True + + +@pytest.mark.asyncio +@respx.mock +async def test_acompletion_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + zapier_tool = {"type": "mcp", "server_label": "zapier", "server_url": "https://mcp.zapier.com/api/mcp/mcp"} + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) + monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + provider = respx.post("https://api.openai.com/v1/chat/completions").mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-zapier", + "object": "chat.completion", + "created": 0, + "model": "gpt-4.1", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + + result = await acompletion_with_mcp( + model="openai/gpt-4.1", + messages=[{"role": "user", "content": "hello"}], + tools=[zapier_tool], + api_key="sk-test", + acompletion=True, + ) + + assert isinstance(result, ModelResponse) + assert result.id == "chatcmpl-zapier" + assert json.loads(provider.calls.last.request.content)["tools"] == [zapier_tool] diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 80151d0cba8..9745a0af970 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException +from openai.types.responses.tool_param import Mcp import importlib from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing @@ -724,6 +725,178 @@ def test_extract_tool_call_details_still_prefers_openai_arguments(): assert arguments == '{"city": "Paris"}' +def _registered( + server_id: str, + name: str, + alias: str | None = None, + server_name: str | None = None, + access_groups: list[str] | None = None, +): + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id=server_id, + name=name, + alias=alias, + server_name=server_name, + transport=MCPTransport.http, + access_groups=access_groups, + ) + + +async def _no_toolset(_: str) -> bool: + return False + + +ZAPIER_TOOL: Mcp = { + "type": "mcp", + "server_label": "zapier", + "server_url": "https://mcp.zapier.com/api/mcp/mcp", + "require_approval": "never", +} +EXPLICIT_GATEWAY_TOOL = {"type": "mcp", "server_label": "github", "server_url": "litellm_proxy/mcp/github"} +FUNCTION_TOOL = {"type": "function", "name": "get_weather", "parameters": {}} + + +@pytest.mark.asyncio +async def test_gateway_served_names_matches_alias_server_name_name_access_group_and_toolset(): + from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + + servers = ( + _registered("id-1", "github-name", alias="github", server_name="github-server", access_groups=["prod-group"]), + _registered("id-2", "deepwiki"), + ) + + async def toolset_exists(name: str) -> bool: + return name == "my-toolset" + + served = await _gateway_served_names( + {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset", "mcp", "nope"}, + servers=lambda: servers, + toolset_exists=toolset_exists, + ) + + assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"} + + +@pytest.mark.asyncio +async def test_gateway_served_names_matches_server_id_short_prefix_and_alias_case_like_the_gateway(): + from litellm.proxy._experimental.mcp_server.utils import compute_short_server_prefix + from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + + server_id = "0b9ae4ca-1bd2-4faa-b183-7dd812597e3b" + short_prefix = compute_short_server_prefix(server_id) + servers = (_registered(server_id, "github-name", alias="github", access_groups=["prod-group"]),) + + served = await _gateway_served_names( + {server_id, short_prefix, "GitHub", "PROD-GROUP", "nope"}, servers=lambda: servers, toolset_exists=_no_toolset + ) + + assert served == {server_id, short_prefix, "GitHub"} + + +@pytest.mark.asyncio +async def test_routes_through_gateway_flags_explicit_and_served_tools_only(): + served_tool = {"type": "mcp", "server_label": "github", "server_url": "http://localhost:4000/mcp/github"} + + async def served_names(names): + assert names == {"github", "mcp"} + return frozenset({"github"}) + + flags = await LiteLLM_Proxy_MCP_Handler.routes_through_gateway( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, served_tool, FUNCTION_TOOL], served_names=served_names + ) + + assert flags == (False, True, True, False) + + +@pytest.mark.asyncio +async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider(): + + async def served_names(names): + assert names == {"mcp"} + return frozenset() + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [ZAPIER_TOOL, FUNCTION_TOOL] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_repoints_served_proxy_urls_at_the_gateway(): + served_tool = { + "type": "mcp", + "server_label": "toolset", + "server_url": "http://localhost:4000/mcp/my-toolset", + "require_approval": "never", + "allowed_tools": ["get_me"], + } + unserved_tool = {"type": "mcp", "server_label": "typo", "server_url": "http://localhost:4000/mcp/githb"} + + async def served_names(names): + return frozenset({"my-toolset"}) + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [served_tool, unserved_tool], served_names=served_names + ) + + assert gateway_tools == [{**served_tool, "server_url": "litellm_proxy/mcp/my-toolset"}] + assert other_tools == [unserved_tool] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_skips_resolution_when_nothing_points_at_the_proxy(): + async def served_names(names): + raise AssertionError("no lookup expected") + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [FUNCTION_TOOL] + + +def test_should_use_gateway_still_triggers_on_http_mcp_path(): + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([ZAPIER_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([EXPLICIT_GATEWAY_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([FUNCTION_TOOL]) is False + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(None) is False + + +@pytest.mark.asyncio +async def test_aresponses_api_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.responses import main as responses_main + from litellm.types.llms.openai import ResponsesAPIResponse + + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) + monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) + provider_tools: list[object] = [] + + def fake_provider(**kwargs: object) -> object: + request_params = cast(dict[str, object], kwargs["response_api_optional_request_params"]) + provider_tools.append(request_params.get("tools")) + + async def respond() -> ResponsesAPIResponse: + return ResponsesAPIResponse(id="resp_zapier", created_at=0, output=[]) + + return respond() + + monkeypatch.setattr(responses_main.base_llm_http_handler, "response_api_handler", fake_provider) + + response = await responses_main.aresponses_api_with_mcp( + input="Reply with the single word ok.", model="openai/gpt-4.1", tools=[ZAPIER_TOOL] + ) + + assert isinstance(response, ResponsesAPIResponse) + assert provider_tools == [[ZAPIER_TOOL]] + + def _response_with_reasoning_and_tool_call() -> Any: """A first-turn response as a reasoning model returns it: reasoning item, then a function call.""" return ResponsesAPIResponse( diff --git a/tests/test_litellm/responses/test_responses_api_bridge_flag.py b/tests/test_litellm/responses/test_responses_api_bridge_flag.py index 57aa2a6baa2..ed44a9f4545 100644 --- a/tests/test_litellm/responses/test_responses_api_bridge_flag.py +++ b/tests/test_litellm/responses/test_responses_api_bridge_flag.py @@ -7,10 +7,14 @@ calls so routed requests do not hit a custom api_base /v1/responses endpoint. """ from importlib import import_module +from typing import Final from unittest.mock import MagicMock, patch +import httpx +import pytest import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -18,6 +22,26 @@ from litellm.types.utils import Choices, Message, ModelResponse, Usage class TestUseResponsesApiBridgeFlag: """Test that bridge opt-in forces the chat completions path.""" + @pytest.mark.parametrize("model", ["openai/chat_completions/gpt-6-astra", "xai/test-classifier"]) + def test_encrypted_classifier_rejection_preserves_public_error(self, model: str) -> None: + respond: Final = MagicMock(side_effect=AssertionError("Incompatible classifier sent an upstream request")) + with httpx.Client(transport=httpx.MockTransport(respond)) as client: + with pytest.raises( + litellm.APIConnectionError, + match="Encrypted task classification requires a compatible native Responses deployment", + ) as error: + litellm.responses( + model=model, + input="Delegated task", + api_key="test-key", + api_base="https://classifier.test/v1", + client=HTTPHandler(client=client), + _require_encrypted_task_support=True, + num_retries=0, + ) + assert error.value.status_code == 500 + respond.assert_not_called() + @patch.object( import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/test_litellm/responses/test_responses_api_request_body.py index 3e60906ec6d..5fced458208 100644 --- a/tests/test_litellm/responses/test_responses_api_request_body.py +++ b/tests/test_litellm/responses/test_responses_api_request_body.py @@ -246,8 +246,10 @@ async def test_aresponses_keeps_include_obfuscation_in_stream_options(): @pytest.mark.asyncio +@pytest.mark.parametrize("drop_params", [True, "true"]) async def test_aresponses_request_level_drop_params_drops_bedrock_mantle_service_tier( monkeypatch, + drop_params, ): """ Request-level drop_params=True (as the proxy injects for agentic CLIs) must @@ -271,7 +273,7 @@ async def test_aresponses_request_level_drop_params_drops_bedrock_mantle_service aws_region_name="us-east-1", input="hi", service_tier="priority", - drop_params=True, + drop_params=drop_params, ) mock_post.assert_called_once() diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index cb6efa21036..9d9eefdceb3 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -440,6 +440,56 @@ class TestResponseAPILoggingUtils: assert result.completion_tokens_details.text_tokens == 20 assert result.completion_tokens_details.audio_tokens is None + def test_transform_realtime_usage_partitions_reasoning_out_of_text_tokens(self): + """Realtime nests reasoning_tokens inside text_tokens; the stored text share excludes them.""" + usage = { + "input_tokens": 237, + "output_tokens": 70, + "total_tokens": 307, + "input_token_details": {"text_tokens": 43, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52}, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.completion_tokens == 70 + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 18 + assert result.completion_tokens_details.reasoning_tokens == 52 + assert result.completion_tokens_details.audio_tokens == 0 + + def test_transform_realtime_usage_partitions_reasoning_beside_audio_output(self): + """Audio output stays as reported; only the text share sheds the nested reasoning tokens.""" + usage = { + "input_tokens": 100, + "output_tokens": 70, + "total_tokens": 170, + "input_token_details": {"text_tokens": 100, "audio_tokens": 0, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 39, "audio_tokens": 31, "reasoning_tokens": 23}, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 16 + assert result.completion_tokens_details.audio_tokens == 31 + assert result.completion_tokens_details.reasoning_tokens == 23 + + def test_transform_response_api_usage_keeps_partitioned_text_tokens(self): + """A provider already reporting text_tokens beside reasoning_tokens is stored as sent.""" + usage = { + "input_tokens": 10, + "output_tokens": 20, + "total_tokens": 30, + "output_tokens_details": {"text_tokens": 12, "reasoning_tokens": 5}, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 12 + assert result.completion_tokens_details.reasoning_tokens == 5 + def test_transform_response_api_usage_carries_extra_provider_fields(self): """Non-standard usage fields (e.g. xAI tool details) must survive chat normalization.""" details = {"web_search_calls": 2, "x_search_calls": 0} diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index e8333214ea8..fe3c4a0640d 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -17,6 +17,9 @@ from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPICon from litellm.llms.databricks.responses.transformation import ( DatabricksResponsesAPIConfig, ) +from litellm.llms.fireworks_ai.responses.transformation import ( + FireworksAIResponsesAPIConfig, +) from litellm.llms.github_copilot.responses.transformation import ( GithubCopilotResponsesAPIConfig, ) @@ -102,6 +105,12 @@ class TestResponsesAPIWebSocketSupport: def test_openai_model_in_websocket_url_default(self): assert OpenAIResponsesAPIConfig().model_in_websocket_url() is True + def test_fireworks_ai_uses_managed_websocket(self): + """Fireworks AI should use managed websocket handler""" + assert ( + FireworksAIResponsesAPIConfig().supports_native_websocket() is False + ), "Fireworks AI should use managed websocket handler" + def test_xai_uses_managed_websocket(self): """XAI should use managed websocket handler""" config = XAIResponsesAPIConfig() diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index 4b446368dbe..74d96bda336 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -4,7 +4,6 @@ import pytest from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled from litellm.rust_bridge import configuration, responses_websocket -from litellm.types.router import GenericLiteLLMParams class _FakeNativeConnection: @@ -48,22 +47,12 @@ def reset_responses_websocket(): configuration.reset_rust_configuration() -def test_rust_websocket_bridge_is_disabled_without_flag() -> None: - assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) - assert not _rust_responses_websocket_enabled("anthropic", GenericLiteLLMParams(rust=True)) - assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=True)) - - -def test_explicit_false_overrides_process_enable() -> None: +def test_rust_websocket_bridge_uses_process_enablement() -> None: + configuration.rust(False) + assert not _rust_responses_websocket_enabled("openai") configuration.rust(True) - - assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=False)) - - -def test_process_enable_applies_without_request_override() -> None: - configuration.rust(True) - - assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) + assert _rust_responses_websocket_enabled("openai") + assert not _rust_responses_websocket_enabled("anthropic") @pytest.mark.asyncio diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index c226c0b4d09..5e0e794d93e 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -18,6 +18,7 @@ from litellm.responses.streaming_iterator import ( SyncResponsesAPIStreamingIterator, ) from litellm.types.llms.openai import ( + ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, @@ -329,8 +330,6 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): def _responses_api_response_with_usage() -> ResponsesAPIResponse: - from litellm.types.llms.openai import ResponseAPIUsage - return ResponsesAPIResponse( id="resp_lit6427", created_at=int(datetime(2025, 1, 1).timestamp()), @@ -368,6 +367,53 @@ def test_stamp_responses_usage_cost_keeps_provider_reported_cost(): logging_obj._response_cost_calculator.assert_not_called() +def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse: + return ResponsesAPIResponse.model_construct( + id="resp_lit7391", + created_at=int(datetime(2025, 1, 1).timestamp()), + status="completed", + model="perplexity/deepseek-v4-flash-0731", + object="response", + output=[], + truncation="", + usage=usage, + ) + + +def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + response = _unvalidated_response_with_dict_usage( + { + "input_tokens": 29, + "output_tokens": 120, + "output_tokens_details": {"reasoning_tokens": 117}, + "total_tokens": 149, + "cost": {"currency": "USD", "input_cost": 0, "output_cost": 3e-05, "total_cost": 3e-05}, + } + ) + logging_obj = Mock(spec=LiteLLMLoggingObj) + + _stamp_responses_usage_cost(response, logging_obj) + + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.cost == pytest.approx(3e-05) + assert response.usage.output_tokens_details.reasoning_tokens == 117 + logging_obj._response_cost_calculator.assert_not_called() + + +def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149}) + logging_obj = Mock(spec=LiteLLMLoggingObj) + logging_obj._response_cost_calculator.return_value = 0.000704 + + _stamp_responses_usage_cost(response, logging_obj) + + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.cost == pytest.approx(0.000704) + logging_obj._response_cost_calculator.assert_called_once_with(result=response) + + def test_stamp_responses_usage_cost_survives_calculator_failure(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost @@ -535,5 +581,50 @@ async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched(): with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")): iterator._log_completed_response(is_async=True) - assert logged == [iterator.completed_response] + assert len(logged) == 1 + assert logged[0] is not iterator.completed_response + assert logged[0].response is not iterator.completed_response.response + assert logged[0].response._hidden_params["headers"]["apim-request-id"] == "azure-correlation-1" assert iterator.completed_response.response._hidden_params == {} + + +def _unvalidated_completed_config() -> Mock: + """Config whose completed event carries a Perplexity-style response that fails validation + (``truncation: ""``) and already holds the stamped ``ResponseAPIUsage``.""" + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + response = _unvalidated_response_with_dict_usage( + ResponseAPIUsage(input_tokens=29, output_tokens=373, total_tokens=402, cost={"total_cost": 0.0001}) + ) + return ResponseCompletedEvent(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_validation(): + """LIT-7391: the logging copy cannot round-trip a response that fails validation, and logging + rewrites the assembled response's usage to chat shape in place, so the event handed to logging + must never be the one the caller receives.""" + logging_obj = _logging_obj_stub() + logging_obj.stream = True + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator(headers={}, config=_unvalidated_completed_config(), logging_obj=logging_obj) + events = [event async for event in iterator] + + assert len(logged) == 1 + now = datetime.now() + LiteLLMLoggingObj._get_assembled_streaming_response( + logging_obj, logged[0], start_time=now, end_time=now, is_async=True, streaming_chunks=[] + ) + assert logged[0].response.usage["prompt_tokens"] == 29 + + client_usage = events[-1].response.usage + assert isinstance(client_usage, ResponseAPIUsage) + assert client_usage.input_tokens == 29 + assert client_usage.cost == pytest.approx(0.0001) diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py index cbf5635a5ae..d717c4e8c89 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py @@ -236,6 +236,38 @@ async def test_record_turn_attributes_satisfaction_to_previous_response_model(): assert smart_after.alpha == pytest.approx(smart_before.alpha) +@pytest.mark.asyncio +async def test_external_default_keeps_feedback_history_without_entering_bandit_pool(): + r = _make_router() + before = r._cells[(RequestType.GENERAL, "fast")] + await r.record_turn( + session_id="fallback", + model_name="fast", + request_type=RequestType.GENERAL, + turn=Turn(user_content="fix this retry bug", assistant_content="clear the cache"), + ) + await r.record_turn( + session_id="fallback", + model_name="external-default", + request_type=RequestType.GENERAL, + turn=Turn(user_content="the fix is still broken", assistant_content="keep cache entries"), + ) + assert r._cells[(RequestType.GENERAL, "fast")].beta > before.beta + await r.record_turn( + session_id="fallback", + model_name="smart", + request_type=RequestType.GENERAL, + turn=Turn( + user_content="the fix is still broken", + assistant_content="use the corrected entry", + tool_results=[{"is_error": True, "content": "failure"}], + ), + ) + assert r._feedback_contexts["fallback"].model_name == "smart" + assert all(model != "external-default" for _, model in r._cells) + assert r.config.available_models == ["fast", "smart"] + + @pytest.mark.asyncio async def test_record_turn_bounds_feedback_contexts_and_evicts_least_recent_session(): r = _make_router() @@ -269,7 +301,9 @@ async def test_record_turn_bounds_feedback_contexts_and_evicts_least_recent_sess @pytest.mark.asyncio -async def test_load_state_from_db_overrides_cold_start(): +async def test_load_state_from_db_adds_the_persisted_delta_to_the_cold_start_prior(): + """A row holds an accumulated delta, not a full posterior; loading must add it to the + cold-start prior, not replace the cell outright.""" r = _make_router() cold = r._cells[(RequestType.GENERAL, "fast")] @@ -284,14 +318,38 @@ async def test_load_state_from_db_overrides_cold_start(): await r.load_state_from_db(prisma) new_cell = r._cells[(RequestType.GENERAL, "fast")] - assert (new_cell.alpha, new_cell.beta) == (42.0, 13.0) - assert (new_cell.alpha, new_cell.beta) != (cold.alpha, cold.beta) + assert (new_cell.alpha, new_cell.beta) == (cold.alpha + 42.0, cold.beta + 13.0) + + +@pytest.mark.asyncio +async def test_load_state_from_db_keeps_a_one_sided_delta_row_sampleable(): + """A cell whose only DB activity is one signal type persists a one-sided row (e.g. + beta=0.0); loading it must not zero out a Beta shape parameter and crash thompson_sample().""" + from litellm.router_strategy.adaptive_router.bandit import thompson_sample + + r = _make_router() + + one_sided_row = MagicMock() + one_sided_row.request_type = "general" + one_sided_row.model_name = "fast" + one_sided_row.alpha = 1.0 + one_sided_row.beta = 0.0 + + prisma = MagicMock() + prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[one_sided_row]) + await r.load_state_from_db(prisma) + + loaded_cell = r._cells[(RequestType.GENERAL, "fast")] + assert loaded_cell.alpha > 0.0 + assert loaded_cell.beta > 0.0 + thompson_sample(loaded_cell) # must not raise @pytest.mark.asyncio async def test_load_state_from_db_handles_unknown_request_type(): r = _make_router() - cold = r._cells[(RequestType.GENERAL, "fast")] + cold_general = r._cells[(RequestType.GENERAL, "fast")] + cold_writing = r._cells[(RequestType.WRITING, "fast")] bad_row = MagicMock() bad_row.request_type = "nonexistent_type_v999" @@ -309,10 +367,11 @@ async def test_load_state_from_db_handles_unknown_request_type(): prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[bad_row, good_row]) await r.load_state_from_db(prisma) - # Unknown skipped; good applied. - assert r._cells[(RequestType.GENERAL, "fast")].alpha == 7.0 - # Other request types kept their cold-start values. - assert r._cells[(RequestType.WRITING, "fast")] == cold or True + # Unknown skipped; good added to the cold-start prior. + new_general = r._cells[(RequestType.GENERAL, "fast")] + assert new_general.alpha == cold_general.alpha + 7.0 + # Other request types kept their own cold-start values. + assert r._cells[(RequestType.WRITING, "fast")] == cold_writing # ---- Session state eviction --------------------------------------------- diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py index 3071f916ef1..23fc859d4a6 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py @@ -186,8 +186,12 @@ async def test_failure_signal_increments_beta_after_flush(): @pytest.mark.asyncio -async def test_load_state_from_db_overrides_cold_start(): +async def test_load_state_from_db_adds_persisted_delta_to_cold_start(): + """A row holds an accumulated delta, not a full posterior; loading must add it to the + cold-start prior, not replace the cell outright.""" router = _make_router() + cold = router._cells[(RequestType.GENERAL, "gpt-4o")] + fake_row = MagicMock() fake_row.request_type = RequestType.GENERAL.value fake_row.model_name = "gpt-4o" @@ -200,8 +204,8 @@ async def test_load_state_from_db_overrides_cold_start(): await router.load_state_from_db(prisma) cell = router._cells[(RequestType.GENERAL, "gpt-4o")] - assert cell.alpha == 90.0 - assert cell.beta == 10.0 + assert cell.alpha == cold.alpha + 90.0 + assert cell.beta == cold.beta + 10.0 @pytest.mark.asyncio diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py b/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py index a4e803f59ad..6fe39599b0f 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py @@ -133,6 +133,102 @@ def test_init_adaptive_router_reads_cost_from_litellm_params(): } +def test_init_adaptive_router_falls_back_to_model_info_cost(): + """Custom pricing declared under model_info (the conventional location everywhere else in + LiteLLM: cost_calculator.py, add_deployment's litellm.model_cost registration) must still + feed cost-weighted routing, not silently zero it out.""" + r = Router( + model_list=[ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": { + "available_models": ["fast", "smart"], + }, + }, + }, + { + "model_name": "fast", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"input_cost_per_token": 0.00000015}, + }, + { + "model_name": "smart", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"input_cost_per_token": 0.0000050}, + }, + ] + ) + assert _adaptive(r, "smart-cheap-router").model_to_cost == { + "fast": 0.00000015, + "smart": 0.0000050, + } + + +@pytest.mark.asyncio +async def test_pick_model_favors_the_cheaper_model_info_priced_deployment(): + """Same fix, exercised through pick_model's actual scoring rather than the model_to_cost + dict alone: with cost as the only weight and equal quality priors, the cheaper deployment + must win every draw. `smart` (expensive) is listed first deliberately: before the fix both + models silently cost 0.0, tying every score, and pick_best's insertion-order tie-break would + hand every request to the first-listed (expensive) model instead.""" + r = Router( + model_list=[ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": { + "available_models": ["smart", "fast"], + "weights": {"quality": 0.0, "cost": 1.0}, + }, + }, + }, + { + "model_name": "smart", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"input_cost_per_token": 0.0000050}, + }, + { + "model_name": "fast", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"input_cost_per_token": 0.00000015}, + }, + ] + ) + adaptive = _adaptive(r, "smart-cheap-router") + + picks = [await adaptive.pick_model(RequestType.GENERAL) for _ in range(10)] + + assert picks == ["fast"] * 10 + + +def test_init_adaptive_router_prefers_litellm_params_cost_over_model_info(): + r = Router( + model_list=[ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": { + "available_models": ["fast"], + }, + }, + }, + { + "model_name": "fast", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.00000015, + }, + "model_info": {"input_cost_per_token": 0.0000050}, + }, + ] + ) + assert _adaptive(r, "smart-cheap-router").model_to_cost == {"fast": 0.00000015} + + # ---- Fix 4: pre-routing dispatch --------------------------------------- diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py index 123ada83ca4..01bf0c2a5ad 100644 --- a/tests/test_litellm/router_strategy/test_auto_router.py +++ b/tests/test_litellm/router_strategy/test_auto_router.py @@ -316,6 +316,11 @@ class TestAutoRouter: semantic_router = pytest.importorskip("semantic_router", reason="auto-router needs the semantic-router extra") +# SemanticRouter(auto_sync="local") calls the encoder's sync embedding path twice per build: +# once to probe the encoder's output dimension, once to embed ROUTER_CONFIG's one route's +# utterances. +_EMBEDDING_CALLS_PER_ROUTELAYER_BUILD: Final = 2 + ROUTER_CONFIG: Final = json.dumps( { "routes": [ @@ -604,3 +609,164 @@ class TestAutoRouterAttributesItsEmbeddingSpend: assert router.aembedding_kwargs["proxy_server_request"] == { "body": {"model": "text-embedding-3-small", "input": ["fix this stack trace"]} } + + +class ThreadTrackingEmbeddingRouter(StubEmbeddingRouter): + """Records which OS thread and how many times `embedding()` was called during a build.""" + + def __init__(self) -> None: + super().__init__() + self.embedding_call_threads: list[int] = [] + + def embedding(self, input: list[str], model: str, **kwargs: Any) -> Any: + import threading + + self.embedding_call_threads.append(threading.get_ident()) + return super().embedding(input, model, **kwargs) + + +class TestAutoRouterColdStartDoesNotBlockTheEventLoop: + """The first request through a fresh alias builds the route layer off the event loop thread, + and concurrent first requests build it exactly once.""" + + @pytest.mark.asyncio + async def test_should_build_the_routelayer_on_a_worker_thread_not_the_event_loop_thread(self): + import threading + + embedding_router: Final = ThreadTrackingEmbeddingRouter() + auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router) + event_loop_thread: Final = threading.get_ident() + + result: Final = await auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs={}, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + + assert result is not None + assert len(embedding_router.embedding_call_threads) == _EMBEDDING_CALLS_PER_ROUTELAYER_BUILD + assert set(embedding_router.embedding_call_threads) == {embedding_router.embedding_call_threads[0]} + assert embedding_router.embedding_call_threads[0] != event_loop_thread + + @pytest.mark.asyncio + async def test_should_build_the_routelayer_exactly_once_under_concurrent_cold_start_requests(self): + embedding_router: Final = ThreadTrackingEmbeddingRouter() + auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router) + + results: Final = await asyncio.gather( + *( + auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs={}, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + for _ in range(10) + ) + ) + + assert all(result is not None for result in results) + assert len(embedding_router.embedding_call_threads) == _EMBEDDING_CALLS_PER_ROUTELAYER_BUILD + + @pytest.mark.asyncio + async def test_should_not_duplicate_the_build_when_a_caller_is_cancelled_mid_build(self): + """A caller arriving while the first is cancelled mid-build must reuse it, not duplicate it.""" + import threading + + class BlockingEmbeddingRouter(ThreadTrackingEmbeddingRouter): + def __init__(self) -> None: + super().__init__() + self.started = threading.Event() + self.release = threading.Event() + + def embedding(self, input: list[str], model: str, **kwargs: Any) -> Any: + self.started.set() + self.release.wait(timeout=5) + return super().embedding(input, model, **kwargs) + + embedding_router: Final = BlockingEmbeddingRouter() + auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router) + + first_call: Final = asyncio.ensure_future( + auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs={}, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + ) + while not embedding_router.started.is_set(): + await asyncio.sleep(0.01) + + first_call.cancel() + with pytest.raises(asyncio.CancelledError): + await first_call + + second_call: Final = asyncio.ensure_future( + auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs={}, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + ) + await asyncio.sleep(0.01) # let the second call observe the still-running build + embedding_router.release.set() + result: Final = await second_call + + assert result is not None + assert len(embedding_router.embedding_call_threads) == _EMBEDDING_CALLS_PER_ROUTELAYER_BUILD + + @pytest.mark.asyncio + async def test_should_clear_a_failed_build_even_with_no_caller_left_to_observe_it(self): + """A build failing after its only caller was cancelled must still clear, not stay cached.""" + import threading + + class FailsOnFirstAttemptEmbeddingRouter(ThreadTrackingEmbeddingRouter): + def __init__(self) -> None: + super().__init__() + self.started = threading.Event() + self.release = threading.Event() + self.attempts = 0 + + def embedding(self, input: list[str], model: str, **kwargs: Any) -> Any: + self.attempts += 1 + attempt = self.attempts + self.started.set() + self.release.wait(timeout=5) + if attempt == 1: + raise ValueError("boom") + return super().embedding(input, model, **kwargs) + + embedding_router: Final = FailsOnFirstAttemptEmbeddingRouter() + auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router) + + first_call: Final = asyncio.ensure_future( + auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs={}, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + ) + while not embedding_router.started.is_set(): + await asyncio.sleep(0.01) + first_call.cancel() + with pytest.raises(asyncio.CancelledError): + await first_call + + # Nobody awaits the build now. Let the first attempt fail on its own. + embedding_router.release.set() + build_task = auto_router._routelayer_build_task + assert build_task is not None + while not build_task.done(): + await asyncio.sleep(0.01) + await asyncio.sleep(0.01) # let the done-callback (scheduled via call_soon) run + + assert auto_router._routelayer_build_task is None + + embedding_router.started.clear() + embedding_router.release.clear() + result: Final = await auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs={}, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + + assert result is not None diff --git a/tests/test_litellm/router_strategy/test_base_routing_strategy.py b/tests/test_litellm/router_strategy/test_base_routing_strategy.py index 154042692d0..dc75c3d2916 100644 --- a/tests/test_litellm/router_strategy/test_base_routing_strategy.py +++ b/tests/test_litellm/router_strategy/test_base_routing_strategy.py @@ -1,4 +1,5 @@ import json +import logging from typing import Any, Dict, List, Optional, Set, Union import pytest @@ -9,7 +10,7 @@ from unittest.mock import MagicMock, patch from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, RedisPipelineIncrementOperation from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy @@ -146,3 +147,18 @@ async def test_cache_keys_management(base_strategy): # Test resetting cache keys base_strategy.reset_in_memory_keys_to_update() assert len(base_strategy.get_in_memory_keys_to_update()) == 0 + + +@pytest.mark.asyncio +async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_error(base_strategy, mock_dual_cache, caplog): + """The sync loop pushes every 100 ms under usage-based routing, so an open breaker must not add an error line per cycle.""" + mock_dual_cache.redis_cache.async_increment_pipeline.side_effect = RedisCircuitBreakerOpenError( + "Redis circuit breaker is open - skipping async_increment_pipeline" + ) + base_strategy.redis_increment_operation_queue = [{"key": "k", "increment_value": 1.0, "ttl": 60}] + + with caplog.at_level(logging.ERROR): + await base_strategy._push_in_memory_increments_to_redis() + + assert caplog.records == [] + assert base_strategy.redis_increment_operation_queue == [] diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py index 36fa38bacb5..4cc8fe78811 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -1,7 +1,13 @@ +import asyncio +import gc +import logging +from unittest.mock import AsyncMock, MagicMock + import pytest import litellm from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.router import LiteLLM_Params from litellm.types.utils import BudgetConfig @@ -303,3 +309,90 @@ def test_router_add_deployment_registers_deployment_budget( ) assert config is not None assert config.max_budget == 0.000000000001 + + +@pytest.mark.asyncio +async def test_sync_refused_by_the_open_circuit_breaker_is_quiet_and_leaks_no_task(disable_budget_sync, caplog): + """The budget sync runs every second, so an open breaker must not add an error line or an unretrieved task exception per cycle.""" + refused = RedisCircuitBreakerOpenError("Redis circuit breaker is open - skipping async_increment_pipeline") + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment_pipeline = AsyncMock(side_effect=refused) + redis_cache.async_batch_get_cache = AsyncMock(side_effect=refused) + limiter = RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=redis_cache), + provider_budget_config={"openai": BudgetConfig(max_budget=1.0, budget_duration="1d")}, + ) + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + limiter.redis_increment_operation_queue = [{"key": "provider_spend:openai:1d", "increment_value": 0.5, "ttl": 60}] + loop = asyncio.get_running_loop() + unretrieved = MagicMock() + loop.set_exception_handler(unretrieved) + + try: + with caplog.at_level(logging.ERROR): + await limiter._sync_in_memory_spend_with_redis() + await asyncio.sleep(0) + gc.collect() + finally: + loop.set_exception_handler(None) + + assert caplog.records == [] + unretrieved.assert_not_called() + assert limiter.redis_increment_operation_queue == [] + assert redis_cache.async_increment_pipeline.await_count == 1 + + +async def _limiter_with_redis(redis_cache: MagicMock) -> RouterBudgetLimiting: + limiter = RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=redis_cache), + provider_budget_config={"openai": BudgetConfig(max_budget=1.0, budget_duration="1d")}, + ) + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + limiter.redis_increment_operation_queue = [{"key": "provider_spend:openai:1d", "increment_value": 0.5, "ttl": 60}] + return limiter + + +@pytest.mark.asyncio +async def test_push_returns_before_redis_answers(disable_budget_sync): + """The push runs inside the request success callback, so it must hand the Redis round trip to a task instead of waiting on it.""" + redis_answered = asyncio.Event() + + async def wait_for_redis(**_: object) -> None: + await redis_answered.wait() + + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment_pipeline = AsyncMock(side_effect=wait_for_redis) + limiter = await _limiter_with_redis(redis_cache) + + await asyncio.wait_for(limiter._push_in_memory_increments_to_redis(), timeout=1) + await asyncio.sleep(0) + + assert not redis_answered.is_set() + assert redis_cache.async_increment_pipeline.await_count == 1 + assert limiter.redis_increment_operation_queue == [] + redis_answered.set() + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + + +@pytest.mark.asyncio +async def test_push_task_failure_is_logged_once_and_not_leaked(disable_budget_sync, caplog): + """A real Redis failure on the background push must surface as one error line, never as an unretrieved task exception.""" + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("Error 61 connecting to 127.0.0.1:6379")) + limiter = await _limiter_with_redis(redis_cache) + loop = asyncio.get_running_loop() + unretrieved = MagicMock() + loop.set_exception_handler(unretrieved) + + try: + with caplog.at_level(logging.ERROR): + await limiter._push_in_memory_increments_to_redis() + await asyncio.sleep(0) + gc.collect() + finally: + loop.set_exception_handler(None) + + assert [record.getMessage() for record in caplog.records] == [ + "Error syncing in-memory cache with Redis: Error 61 connecting to 127.0.0.1:6379" + ] + unretrieved.assert_not_called() diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 918ec7bc100..dba44d1e2e8 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -5,16 +5,24 @@ Tests the rule-based complexity scoring and tier assignment logic. """ import asyncio +import json import logging import sys -from typing import Dict, List +import time +from collections.abc import AsyncIterator, Mapping +from copy import deepcopy +from functools import partial +from typing import Dict, Final, List, Literal from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from pydantic import ValidationError import litellm from litellm import Router +from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.auto_router_model_naming import ( CUSTOMIZATION_CAPABILITY, GATED_AUTO_ROUTER_CAPABILITIES, @@ -23,7 +31,12 @@ from litellm.router_utils.auto_router_model_naming import ( ) from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import ( + OUTPUT_TOKEN_CEILING_PARAMS, + RETURN_RAW_MODEL_NAME_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) +from litellm.router import as_output_cap from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -43,10 +56,12 @@ from litellm.router_strategy.complexity_router.config import ( DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, DEFAULT_COMPLEXITY_CONFIG, DEFAULT_TECHNICAL_KEYWORDS, + TIER_SEVERITY_ORDER, ClassificationRubric, ClassifierLLMConfig, ComplexityRouterConfig, ComplexityTier, + custom_pattern_work, ) from litellm.router_strategy.complexity_router.tier_predictor import ( TierGlobalStatistic, @@ -55,8 +70,11 @@ from litellm.router_strategy.complexity_router.tier_predictor import ( from litellm.types.router import ( Deployment, LiteLLM_Params, + RouterErrors, TaggedPreRoutingStrategy, ) +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler requires_semantic_router = pytest.mark.skipif( @@ -185,7 +203,14 @@ class TestComplexityRouterInit: complexity_router_config=basic_config, ) - assert router._reminder_markers == (("", ""),) + from litellm.router_strategy.complexity_router.complexity_router import _extract_current_ask_and_system_prompt + + assert ( + _extract_current_ask_and_system_prompt( + [{"role": "user", "content": "noisehello"}], router._reminder_markers + )[0] + == "hello" + ) def test_init_without_config(self, mock_router_instance): """Test initialization without configuration uses defaults.""" @@ -756,6 +781,289 @@ class TestCustomTechnicalKeywords: assert custom_score > baseline_score +class TestCustomDimensions: + @pytest.mark.parametrize( + "matchers,prompt", + [ + pytest.param( + {"keywords": ["orbitmesh", "fluxgate"]}, + "Connect ORBITMESH and fluxgate for the requested change", + id="keywords", + ), + pytest.param( + {"patterns": [r"\bCREATE\s{1,4}TABLE\b", r"\bALTER\s{1,4}TABLE\b"]}, + "create table widgets (id integer); ALTER TABLE widgets ADD label text;", + id="regex", + ), + ], + ) + def test_custom_dimension_changes_only_matching_requests( + self, mock_router_instance: MagicMock, matchers: dict[str, object], prompt: str + ) -> None: + baseline: Final = ComplexityRouter("test-router", mock_router_instance) + configured: Final = ComplexityRouter( + "test-router", + mock_router_instance, + {"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, **matchers}]}, + ) + baseline_tier, baseline_score, baseline_signals = baseline.classify(prompt) + tier, score, signals = configured.classify(prompt) + assert baseline_tier == ComplexityTier.SIMPLE + assert tier != ComplexityTier.SIMPLE + assert score == pytest.approx(baseline_score + 0.7) + assert signals == [*baseline_signals, "custom (internalFrameworks)"] + plain: Final = "Hello!" + assert configured.classify(plain) == baseline.classify(plain) + assert configured.classify(plain)[0] == ComplexityTier.SIMPLE + + @pytest.mark.parametrize( + "dimension_overrides,config_overrides", + [ + pytest.param({"keywords": []}, {}, id="missing-matchers"), + pytest.param({"keywords": [" "]}, {}, id="blank-keyword"), + pytest.param({"patterns": ["\t"]}, {}, id="blank-pattern"), + pytest.param({"patterns": ["("]}, {}, id="invalid-regex"), + pytest.param({"patterns": [r"a*b"]}, {}, id="unbounded-star"), + pytest.param({"patterns": [r"a{2,}b"]}, {}, id="unbounded-brace"), + pytest.param({"patterns": [r"a{0,65}b"]}, {}, id="repeat-over-64"), + pytest.param({"patterns": [r"(a{0,8}){0,8}b"]}, {}, id="nested-repeat"), + pytest.param({"patterns": [r"(a|aa){0,12}b"]}, {}, id="alternation-in-repeat"), + pytest.param({"patterns": [r"(?:ab){0,64}c"]}, {}, id="group-repeat"), + pytest.param({"patterns": ["a?" * 9 + "b"]}, {}, id="pattern-work-over-budget"), + pytest.param({"patterns": ["(?:a|aa)" * 9 + "z"]}, {}, id="ambiguous-alternation-chain"), + pytest.param({"patterns": ["a?" * 8 + "a{64}" * 10 + "z"]}, {}, id="cheap-prefix-expensive-tail"), + pytest.param({"patterns": [r"(a)\1"]}, {}, id="backreference"), + pytest.param({"patterns": [r"(?=x)y"]}, {}, id="lookahead"), + pytest.param({"patterns": [r"(?>ab)"]}, {}, id="atomic-group"), + pytest.param({"patterns": [r"a*+b"]}, {}, id="possessive"), + pytest.param({"name": "CODEPRESENCE"}, {"dimension_weights": {"tokenCount": 0.1}}, id="reserved-name"), + pytest.param({}, {"dimension_weights": {"INTERNALFRAMEWORKS": 0.7}}, id="weight-in-map"), + pytest.param({"weight": 0}, {}, id="zero-weight"), + pytest.param({"weight": 1.1}, {}, id="excess-weight"), + pytest.param({"weight": float("nan")}, {}, id="nan-weight"), + pytest.param({"weight": float("inf")}, {}, id="infinite-weight"), + pytest.param({"name": "bad-name"}, {}, id="invalid-name"), + pytest.param({"name": "x" * 65}, {}, id="long-name"), + pytest.param({"keywords": [""]}, {}, id="empty-matcher"), + pytest.param({"keywords": ["x" * 257]}, {}, id="long-matcher"), + pytest.param({"keywords": ["x"] * 32, "patterns": ["y"]}, {}, id="combined-matcher-count"), + pytest.param({"keywords": ["x" * 256] * 17}, {}, id="matcher-character-budget"), + pytest.param({"unknown": True}, {}, id="extra-field"), + pytest.param({"scoring_mode": "graded"}, {}, id="unknown-scoring-mode"), + pytest.param({"scoring_mode": None}, {}, id="null-scoring-mode"), + ], + ) + def test_custom_dimension_invalid_configuration_rejected( + self, dimension_overrides: dict[str, object], config_overrides: dict[str, object] + ) -> None: + with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"): + ComplexityRouterConfig.model_validate( + { + "custom_dimensions": [ + { + "name": "internalFrameworks", + "weight": 0.7, + "keywords": ["orbitmesh"], + **dimension_overrides, + } + ], + **config_overrides, + } + ) + + @pytest.mark.parametrize( + "names", + [ + pytest.param(("internalFrameworks", "INTERNALFRAMEWORKS"), id="duplicate-casefolded-name"), + pytest.param(tuple(f"dimension{i}" for i in range(17)), id="dimension-count"), + ], + ) + def test_custom_dimension_names_and_count_are_bounded(self, names: tuple[str, ...]) -> None: + with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"): + ComplexityRouterConfig.model_validate( + {"custom_dimensions": [{"name": name, "weight": 0.7, "keywords": ["orbitmesh"]} for name in names]} + ) + + @pytest.mark.parametrize("classifier_type", ("heuristic_v2", "llm", "custom")) + def test_custom_dimensions_reject_classifiers_outside_the_tuning_gate(self, classifier_type: str) -> None: + classifier_config: Final = ( + {"classifier_plugin": _FixedTierClassifier("SIMPLE")} + if classifier_type == "custom" + else {"classifier_llm_config": {"model": "judge"}} + if classifier_type == "llm" + else {} + ) + with pytest.raises(ValidationError, match="custom_dimensions requires classifier_type"): + ComplexityRouterConfig.model_validate( + { + "classifier_type": classifier_type, + "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}], + **classifier_config, + } + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("scoring_mode", ("binary", "match_count")) + @pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh", "orbitmesh fluxgate")) + async def test_custom_dimensions_public_hook_scores_only_current_ask( + self, mock_router_instance: MagicMock, current_ask: str, scoring_mode: str + ) -> None: + router: Final = ComplexityRouter( + "test-router", + mock_router_instance, + { + "tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}, + "dimension_weights": {}, + "custom_dimensions": [ + { + "name": "internalFrameworks", + "weight": 0.8, + "keywords": ["orbitmesh", "fluxgate"], + "scoring_mode": scoring_mode, + } + ], + }, + ) + result: Final = await router.async_pre_routing_hook( + model="test-router", + request_kwargs={}, + messages=[ + {"role": "system", "content": "orbitmesh fluxgate"}, + {"role": "user", "content": "orbitmesh fluxgate"}, + {"role": "assistant", "content": "orbitmesh fluxgate is ready"}, + {"role": "user", "content": current_ask}, + {"role": "tool", "tool_call_id": "previous", "content": "orbitmesh fluxgate"}, + ], + ) + assert result is not None + assert result.routing_decision is not None + expected_score: Final = ( + 0.0 + if current_ask == "Hello!" + else 0.4 + if scoring_mode == "match_count" and current_ask == "orbitmesh" + else 0.8 + ) + assert result.routing_decision["score"] == expected_score + assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (expected_score > 0) + assert result.model == ("cheap" if expected_score == 0 else "strong" if expected_score == 0.4 else "top") + assert "orbitmesh" not in " ".join(result.routing_decision["signals"]) + + @pytest.mark.parametrize("scoring_mode", ("binary", "match_count")) + def test_custom_patterns_scan_only_the_first_2048_characters( + self, mock_router_instance: MagicMock, scoring_mode: str + ) -> None: + router: Final = ComplexityRouter( + "test-router", + mock_router_instance, + { + "custom_dimensions": [ + { + "name": "late", + "weight": 0.7, + "patterns": [r"zzz{1,3}", r"yyy{1,3}"], + "scoring_mode": scoring_mode, + } + ] + }, + ) + baseline: Final = ComplexityRouter("test-router", mock_router_instance) + assert "custom (late)" in router.classify("a" * 2040 + " zzz")[2] + assert "custom (late)" not in router.classify("a" * 2048 + " zzz")[2] + second_hit_past_the_bound: Final = "yyy " + "a" * 2044 + " zzz" + contribution: Final = ( + router.classify(second_hit_past_the_bound)[1] - baseline.classify(second_hit_past_the_bound)[1] + ) + assert contribution == pytest.approx(0.7 if scoring_mode == "binary" else 0.35) + + @pytest.mark.parametrize( + "prompt,expected_score", + [ + pytest.param("Hello!", 0.0, id="no-hit"), + pytest.param("orbitmesh orbitmesh ORBITMESH again", 0.5, id="one-keyword-repeated"), + pytest.param("create table a; CREATE TABLE b; create table c", 0.5, id="one-pattern-repeated"), + pytest.param("orbitmesh and fluxgate", 1.0, id="two-keywords"), + pytest.param("orbitmesh then create table t", 1.0, id="keyword-plus-pattern"), + pytest.param("create table a; alter table b", 1.0, id="two-patterns"), + pytest.param("orbitmesh fluxgate create table a alter table b", 1.0, id="all-matchers"), + ], + ) + def test_match_count_grades_distinct_matchers( + self, mock_router_instance: MagicMock, prompt: str, expected_score: float + ) -> None: + dimension: Final = { + "name": "graded", + "weight": 0.6, + "keywords": ["orbitmesh", "ORBITMESH", "fluxgate"], + "patterns": [r"\bcreate\s{1,4}table\b", r"\bcreate\s{1,4}table\b", r"\balter\s{1,4}table\b"], + } + baseline: Final = ComplexityRouter("test-router", mock_router_instance) + binary: Final = ComplexityRouter("test-router", mock_router_instance, {"custom_dimensions": [dimension]}) + graded: Final = ComplexityRouter( + "test-router", + mock_router_instance, + {"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]}, + ) + _, baseline_score, baseline_signals = baseline.classify(prompt) + _, binary_score, binary_signals = binary.classify(prompt) + _, graded_score, graded_signals = graded.classify(prompt) + assert graded_score == pytest.approx(baseline_score + 0.6 * expected_score) + assert binary_score == pytest.approx(baseline_score + (0.6 if expected_score else 0.0)) + expected_signals: Final = [*baseline_signals, *(["custom (graded)"] if expected_score else [])] + assert graded_signals == expected_signals + assert binary_signals == expected_signals + + def test_scoring_mode_round_trips_and_defaults_to_binary(self) -> None: + dimension: Final = {"name": "graded", "weight": 0.6, "keywords": ["orbitmesh"]} + legacy: Final = ComplexityRouterConfig.model_validate({"custom_dimensions": [dimension]}) + graded: Final = ComplexityRouterConfig.model_validate( + {"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]} + ) + assert legacy.custom_dimensions[0].scoring_mode == "binary" + assert graded.model_dump(mode="json")["custom_dimensions"][0]["scoring_mode"] == "match_count" + assert ComplexityRouterConfig.model_validate(graded.model_dump(mode="json")) == graded + + def test_custom_dimensions_router_wide_regex_work_is_capped(self) -> None: + heavy: Final = {"weight": 0.5, "patterns": ["a?" * 8 + "z"]} + ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(6)]}) + with pytest.raises(ValidationError, match="regex work estimate is 8939"): + ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(7)]}) + + @pytest.mark.parametrize( + "pattern,work", + [ + pytest.param(r"\b(create|alter|drop)\s{1,4}table\b", 135, id="sql-ddl"), + pytest.param("a?" * 8 + "z", 1277, id="optional-chain-near-cap"), + pytest.param(r"a{0,15}a{0,15}z", 801, id="adjacent-bounded-near-cap"), + pytest.param(r"[a-z0-9_]{3,63}\.(com|net|io)", 1291, id="class-repeat-plus-alternation"), + pytest.param("(?:a|aa)" * 8 + "z", 1787, id="ambiguous-alternation-near-cap"), + pytest.param("a{64}" * 10 + "z", 662, id="long-deterministic-tail"), + ], + ) + def test_custom_pattern_work_stays_cheap_on_adversarial_text( + self, mock_router_instance: MagicMock, pattern: str, work: int + ) -> None: + assert custom_pattern_work(pattern) == work + router: Final = ComplexityRouter( + "test-router", + mock_router_instance, + { + "custom_dimensions": [ + {"name": "bounded", "weight": 0.7, "patterns": [pattern]}, + {"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}, + ] + }, + ) + adversarial: Final = "orbitmesh " + "a" * 4000 + started: Final = time.perf_counter() + tier, score, signals = router.classify(adversarial) + elapsed: Final = time.perf_counter() - started + assert signals == ["long (1002 tokens)", "custom (internalFrameworks)"] + assert score == pytest.approx(0.8) + assert tier == ComplexityTier.REASONING + assert elapsed < 0.1 + + class TestAsyncPreRoutingHookEdgeCases: """Test edge cases for async_pre_routing_hook method.""" @@ -1310,6 +1618,7 @@ class TestRouterComplexityDeploymentMethods: def test_the_shipped_rubric_and_default_prompt_stay_free(self) -> None: """Only an operator-written prompt is gated: picking a shipped rubric preset, or writing no prompt at all, leaves a router unmetered, so several of them register under a ceiling of one.""" + def rubric(model_name: str, model_id: str, preset: str | None) -> dict[str, object]: llm_config: dict[str, object] = {"model": "gpt-4o-mini"} if preset is not None: @@ -1446,6 +1755,7 @@ class TestRouterComplexityDeploymentMethods: def test_renaming_built_in_tiers_is_not_a_custom_tier_set(self) -> None: """tier_labels renames the built-in ladder without defining one, so it stays ungated: two such routers register under a ceiling of one.""" + def labeled(model_name: str, model_id: str) -> dict[str, object]: row = self._router_row(model_name, model_id, "heuristic") row["litellm_params"]["complexity_router_config"]["tier_labels"] = {"SIMPLE": "Cheap", "MEDIUM": "Standard"} @@ -1512,6 +1822,110 @@ class TestRouterComplexityDeploymentMethods: assert adaptive.model_to_prefs["cheap"].quality_tier == 1 assert adaptive.model_to_prefs["premium"].quality_tier == 3 + def test_hybrid_adaptive_router_falls_back_to_model_info_cost(self): + """Custom pricing declared under model_info (the conventional location everywhere else + in LiteLLM) must still feed the hybrid adaptive router's cost-weighted scoring, not + silently cost the deployment at 0.0.""" + router = Router( + model_list=[ + { + "model_name": "hybrid", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "cheap", + "complexity_router_config": { + "adaptive": True, + "tiers": {"SIMPLE": ["cheap"], "MEDIUM": ["cheap", "premium"]}, + }, + }, + }, + { + "model_name": "cheap", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"input_cost_per_token": 0.00000015}, + }, + { + "model_name": "premium", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"input_cost_per_token": 0.000005}, + }, + ] + ) + + adaptive = router.adaptive_routers["hybrid"][0].strategy + assert adaptive.model_to_cost == { + "cheap": pytest.approx(0.00000015), + "premium": pytest.approx(0.000005), + } + + @pytest.mark.asyncio + async def test_hybrid_adaptive_router_pick_model_favors_the_cheaper_model_info_priced_deployment(self): + """Same fix, exercised through pick_model's actual scoring rather than the model_to_cost + dict alone. `premium` is listed first (SIMPLE tier) deliberately: before the fix both + models silently cost 0.0, tying every score, and pick_best's insertion-order tie-break + would hand every request to the first-listed (expensive) model instead.""" + from litellm.types.router import RequestType + + router = Router( + model_list=[ + { + "model_name": "hybrid", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "cheap", + "complexity_router_config": { + "adaptive": True, + "adaptive_weights": {"quality": 0.0, "cost": 1.0}, + "tiers": {"SIMPLE": ["premium"], "MEDIUM": ["premium", "cheap"]}, + }, + }, + }, + { + "model_name": "premium", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"input_cost_per_token": 0.000005}, + }, + { + "model_name": "cheap", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"input_cost_per_token": 0.00000015}, + }, + ] + ) + adaptive = router.adaptive_routers["hybrid"][0].strategy + + picks = [await adaptive.pick_model(RequestType.GENERAL) for _ in range(10)] + + assert picks == ["cheap"] * 10 + + def test_hybrid_adaptive_router_prefers_litellm_params_cost_over_model_info(self): + router = Router( + model_list=[ + { + "model_name": "hybrid", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "cheap", + "complexity_router_config": { + "adaptive": True, + "tiers": {"SIMPLE": ["cheap"]}, + }, + }, + }, + { + "model_name": "cheap", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.00000015, + }, + "model_info": {"input_cost_per_token": 0.000005}, + }, + ] + ) + + adaptive = router.adaptive_routers["hybrid"][0].strategy + assert adaptive.model_to_cost == {"cheap": pytest.approx(0.00000015)} + class TestComplexityRouterTagBasedRouting: """Regression tests for https://github.com/BerriAI/litellm/issues/33655. @@ -2084,6 +2498,367 @@ class TestTierLabels: assert set(config.tier_boundaries) == {"simple_medium", "medium_complex", "complex_reasoning"} +def _encrypted_agent_task() -> dict[str, object]: + return { + "type": "agent_message", + "author": "/root", + "recipient": "/root/child", + "content": [ + {"type": "input_text", "text": "Message Type: NEW_TASK\nTask name: /root/child\nPayload:\nHello"}, + {"type": "encrypted_content", "encrypted_content": "opaque-provider-task"}, + ], + } + + +def _native_classifier_response(content: str) -> ResponsesAPIResponse: + response: Final = ResponsesAPIResponse( + id="resp_classifier", + created_at=0, + status="completed", + output=[{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": content}]}], + ) + response._hidden_params = {"response_cost": 0.0001} + return response + + +def _native_classifier_router( + output: str = '{"tier":"REASONING"}', + classifier_type: str = "llm", + deployment_model: str = "openai/gpt-6-astra", + failure: Exception | None = None, + native_router: Router | None = None, + http_handler: AsyncHTTPHandler | None = None, +) -> tuple[ComplexityRouter, MagicMock]: + dependency: Final = MagicMock( + aresponses=( + native_router.factory_function(partial(litellm.aresponses, client=http_handler), call_type="aresponses") + if native_router is not None + else AsyncMock(return_value=_native_classifier_response(output), side_effect=failure) + ), + acompletion=AsyncMock(return_value=_llm_response('{"tier":"SIMPLE"}')), + get_model_list=( + native_router.get_model_list + if native_router is not None + else MagicMock(return_value=[{"litellm_params": {"model": deployment_model}}]) + ), + ) + return ( + ComplexityRouter( + model_name="encrypted-router", + litellm_router_instance=dependency, + complexity_router_config={ + "tiers": {"SIMPLE": "cheap-model", "REASONING": "deep-model"}, + "classifier_type": classifier_type, + "classifier_llm_config": { + "model": "classifier", + "timeout_ms": 5000 if native_router is not None else 100, + "reasoning_effort": "low", + }, + "heuristic_first_max_tier": "SIMPLE" if classifier_type == "heuristic_first" else None, + "hybrid_boundary_margin": 0.01 if classifier_type == "hybrid" else None, + "classifier_fallback": "default_model", + "default_model": "deep-model", + "session_affinity": False, + "deployment_affinity": False, + }, + ), + dependency, + ) + + +@pytest.fixture +async def native_classifier_http() -> AsyncIterator[tuple[AsyncHTTPHandler, MagicMock]]: + respond: Final = MagicMock( + return_value=httpx.Response(200, json=_native_classifier_response('{"tier":"REASONING"}').model_dump()) + ) + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + handler.client = client + yield handler, respond + + +class TestEncryptedTaskClassifier: + @pytest.mark.asyncio + @pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"]) + @pytest.mark.parametrize("codex", [True, False]) + @pytest.mark.parametrize( + "reminder", + [ + "cwd=/repo", + "Keep answers concise", + ], + ) + async def test_encrypted_task_detection_uses_request_reminder_markers( + self, classifier_type: str, codex: bool, reminder: str + ): + router, dependency = _native_classifier_router(classifier_type=classifier_type) + task: Final = _encrypted_agent_task() + request: Final = { + "input": [task, {"role": "user", "content": reminder}], + "metadata": {"user_agent": "codex-tui" if codex else "curl/8.7.1"}, + } + original: Final = deepcopy(request) + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs=request) + + assert request == original + assert result.model == ("deep-model" if codex else "cheap-model") + if codex: + assert result.routing_decision["cause"] == "llm_classifier" + assert result.routing_decision["tier"] == "REASONING" + dependency.aresponses.assert_awaited_once() + assert dependency.aresponses.call_args.kwargs["input"][-1] == task + dependency.acompletion.assert_not_called() + else: + dependency.aresponses.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"]) + @pytest.mark.parametrize("tier,model", [("SIMPLE", "cheap-model"), ("REASONING", "deep-model")]) + async def test_encrypted_task_routes_by_native_verdict(self, classifier_type: str, tier: str, model: str): + router, dependency = _native_classifier_router(json.dumps({"tier": tier}), classifier_type) + task: Final = _encrypted_agent_task() + request: Final = { + "input": [ + {"role": "user", "content": "Prior task context"}, + task, + {"type": "function_call_output", "call_id": "call_1", "output": "Tool output"}, + {"role": "user", "content": "Injected reminder"}, + ], + "instructions": "Caller constraints", + "proxy_server_request": {"body": {"input": [task], "metadata": {"authorization": "source-secret"}}}, + "tools": [{"type": "function", "name": "execute"}], + "previous_response_id": "resp_parent", + "litellm_session_id": "parent-session", + "litellm_trace_id": "parent-trace", + "turn_off_message_logging": True, + "litellm_metadata": {"user_api_key_hash": "caller-key-hash"}, + } + original: Final = deepcopy(request) + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs=request) + + assert result.model == model + assert result.routing_decision["tier"] == tier + assert result.routing_decision["cause"] == "llm_classifier" + assert result.routing_decision["classifier_cost"] == 0.0001 + assert result.messages is None + assert request == original + dependency.acompletion.assert_not_called() + call: Final = dependency.aresponses.call_args.kwargs + assert call["input"][-1] == task + assert "opaque-provider-task" not in json.dumps(call["input"][:-1]) + assert "Prior task context" in json.dumps(call["input"][:-1]) + assert "Caller constraints" in json.dumps(call["input"][:-1]) + assert "Caller constraints" not in call["instructions"] + assert "SIMPLE" in call["instructions"] and "REASONING" in call["instructions"] + assert call["text"]["format"]["schema"]["properties"]["tier"]["enum"] == [ + "SIMPLE", + "MEDIUM", + "COMPLEX", + "REASONING", + ] + assert call["text"]["format"]["strict"] is True + assert call["reasoning"] == {"effort": "low"} + assert call["store"] is False + assert call["_require_encrypted_task_support"] is True + assert call["stream"] is False + assert "tools" not in call and "previous_response_id" not in call + assert "messages" not in call and "response_format" not in call + assert call["timeout"] == 0.1 and call["num_retries"] == 0 and call["disable_fallbacks"] is True + assert call["litellm_session_id"] == "parent-session" + assert call["litellm_trace_id"] == "parent-trace" + assert call["turn_off_message_logging"] is True + assert call["metadata"]["user_api_key_hash"] == "caller-key-hash" + assert call["proxy_server_request"]["body"]["input"] == call["input"] + assert call["proxy_server_request"]["originating_request_masked"] == { + "input": [task], + "metadata": {"authorization": "REDACTED"}, + } + assert "source-secret" not in json.dumps(call) + assert "originating_request_masked" not in call["proxy_server_request"]["body"] + + @pytest.mark.asyncio + async def test_claude_code_encrypted_task_omits_caller_instructions(self): + router, dependency = _native_classifier_router() + task: Final = _encrypted_agent_task() + request: Final = { + "input": [task], + "instructions": "CLAUDE_CODE_SYSTEM", + "litellm_metadata": {"user_agent": "claude-cli/2.1.233"}, + } + original: Final = deepcopy(request) + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs=request) + + assert result.routing_decision["cause"] == "llm_classifier" + assert request == original + call: Final = dependency.aresponses.call_args.kwargs + assert call["instructions"] == classification_system_prompt(router.config.classifier_context_window_size) + assert "CLAUDE_CODE_SYSTEM" not in json.dumps(call["input"][:-1]) + assert call["input"][-1] == task + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "items", + [ + [ + {"type": "reasoning", "encrypted_content": "opaque-history", "summary": []}, + {"role": "user", "content": "hi"}, + ], + [_encrypted_agent_task(), {"role": "user", "content": "hi"}], + [{**_encrypted_agent_task(), "content": [{"type": "input_text", "text": "hi"}]}], + [{"role": "user", "content": "gAAAA is plain text"}], + [ + {"role": "user", "content": "hi"}, + {"type": "function_call_output", "call_id": "call_1", "output": "opaque-provider-task"}, + ], + ], + ids=[ + "historical-reasoning", + "older-encrypted-task", + "plaintext-agent", + "ciphertext-looking-text", + "tool-output", + ], + ) + async def test_other_asks_keep_chat_classifier(self, items: list[dict[str, object]]): + router, dependency = _native_classifier_router() + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs={"input": items}) + + assert result.model == "cheap-model" + assert result.routing_decision["cause"] == "llm_classifier" + dependency.aresponses.assert_not_called() + dependency.acompletion.assert_awaited_once() + + @pytest.mark.asyncio + @pytest.mark.parametrize("output", ["", "not-json", '{"tier":"UNKNOWN"}']) + async def test_invalid_native_verdict_uses_existing_fallback(self, output: str): + router, dependency = _native_classifier_router(output=output) + + result: Final = await router.async_pre_routing_hook( + model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]} + ) + + assert result.model == "deep-model" + assert result.routing_decision["cause"] == "default_model_fallback" + dependency.aresponses.assert_awaited_once() + dependency.acompletion.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "deployment_model", + ["anthropic/test-classifier", "openai/chat_completions/gpt-6-astra", "xai/test-classifier"], + ) + async def test_incompatible_classifier_does_not_flatten_encryption( + self, deployment_model: str, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock] + ): + handler, respond = native_classifier_http + native: Final = Router( + model_list=[ + { + "model_name": "classifier", + "litellm_params": { + "model": deployment_model, + "api_key": "test-key", + "api_base": "https://classifier.test/v1", + }, + } + ], + num_retries=0, + ) + router, _ = _native_classifier_router(native_router=native, http_handler=handler) + + result: Final = await router.async_pre_routing_hook( + model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]} + ) + + assert result.model == "deep-model" + assert result.routing_decision["cause"] == "default_model_fallback" + respond.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("blocked", [True, False]) + async def test_native_classifier_validates_selected_deployment( + self, blocked: bool, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock] + ): + handler, respond = native_classifier_http + native: Final = Router( + model_list=[ + { + "model_name": "classifier", + "litellm_params": {"model": "anthropic/test-classifier", "api_key": "test-key", "order": 0}, + "model_info": {"id": "incompatible", "blocked": blocked}, + }, + { + "model_name": "classifier", + "litellm_params": { + "model": "openai/gpt-6-astra", + "api_key": "test-key", + "order": 1, + "api_base": "https://classifier.test/v1", + }, + "model_info": {"id": "compatible"}, + }, + ], + num_retries=0, + ) + router, _ = _native_classifier_router(native_router=native, http_handler=handler) + task: Final = _encrypted_agent_task() + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs={"input": [task]}) + + assert result.model == "deep-model" + assert result.routing_decision["cause"] == ("llm_classifier" if blocked else "default_model_fallback") + if blocked: + respond.assert_called_once() + request: Final = respond.call_args.args[0] + assert request.url.path == "/v1/responses" + body: Final = json.loads(request.content) + assert body["input"][-1] == task + assert "_require_encrypted_task_support" not in body + else: + respond.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("classifier_type", ["llm", "heuristic_first", "hybrid"]) + @pytest.mark.parametrize( + "input_items", + [ + ["unsupported-input-item"], + [{**_encrypted_agent_task(), "content": [{"type": "input_text", "text": "hi"}, None]}], + ], + ) + async def test_encrypted_detection_does_not_reject_other_input_shapes( + self, classifier_type: str, input_items: list[object] + ): + router, dependency = _native_classifier_router(classifier_type=classifier_type) + + result: Final = await router.aclassify("hi", request_kwargs={"input": input_items}) + + assert result.cause != "default_model_fallback" + assert result.tier == ComplexityTier.SIMPLE + dependency.aresponses.assert_not_called() + if classifier_type == "llm": + dependency.acompletion.assert_awaited_once() + + @pytest.mark.asyncio + @pytest.mark.parametrize("failure", [ValueError("invalid_encrypted_content"), TimeoutError("classifier timed out")]) + async def test_native_provider_failure_uses_existing_fallback(self, failure: Exception): + router, dependency = _native_classifier_router(failure=failure) + + result: Final = await router.async_pre_routing_hook( + model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]} + ) + + assert result.model == "deep-model" + assert result.routing_decision["cause"] == "default_model_fallback" + dependency.aresponses.assert_awaited_once() + dependency.acompletion.assert_not_called() + + class TestLLMClassifier: """Test the LLM-based classifier path (aclassify) and its fallback behavior.""" @@ -2215,9 +2990,7 @@ class TestLLMClassifier: assert outcome.classifier_cost == pytest.approx(1.35e-05) @pytest.mark.asyncio - async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks( - self, llm_classifier_config - ): + async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks(self, llm_classifier_config): real_router = Router( model_list=[ { @@ -2260,9 +3033,7 @@ class TestLLMClassifier: assert real_router.total_calls["openai/mock-backup-classifier"] == 0 @pytest.mark.asyncio - async def test_aclassify_enforces_total_classifier_deadline( - self, mock_router_instance, llm_classifier_config - ): + async def test_aclassify_enforces_total_classifier_deadline(self, mock_router_instance, llm_classifier_config): cancelled = asyncio.Event() async def slow_classifier(**_kwargs: object) -> None: @@ -2546,6 +3317,33 @@ class TestLLMClassifier: "REASONING", ] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "source_body", + [ + {"model": "router", "messages": [{"role": "user", "content": "source-only"}]}, + {"model": "router", "system": "source-only", "messages": [{"role": "user", "content": "ask"}]}, + {"model": "router", "instructions": "source-only", "input": "ask"}, + ], + ) + async def test_classifier_source_is_masked_and_separate_from_provider_input( + self, llm_complexity_router, mock_router_instance, source_body + ): + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) + outcome = await llm_complexity_router.aclassify( + "classify-this-ask", + request_kwargs={ + "proxy_server_request": {"body": {**source_body, "metadata": {"authorization": "source-secret"}}} + }, + ) + assert outcome.cause == "llm_classifier" + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + source = call_kwargs["proxy_server_request"]["originating_request_masked"] + assert source == {**source_body, "metadata": {"authorization": "REDACTED"}} + assert "source-only" not in str(call_kwargs["messages"]) + assert "source-only" not in str(call_kwargs["proxy_server_request"]["body"]) + assert "classify-this-ask" in str(call_kwargs["messages"]) + @pytest.mark.asyncio @pytest.mark.parametrize("reasoning_effort", [None, "none", "low"], ids=["omitted", "none", "low"]) async def test_classifier_reasoning_effort_reaches_only_classifier_call( @@ -3028,11 +3826,11 @@ class TestRouterPreRoutingAliasOverrides: def test_drop_client_effort_carriers_helper_edge_shapes(self): no_pin: Dict = {"thinking": {"type": "adaptive"}} - Router._drop_client_effort_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1}) + Router._drop_client_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1}) assert no_pin == {"thinking": {"type": "adaptive"}} non_dict_carriers: Dict = {"output_config": "max", "reasoning": 3} - Router._drop_client_effort_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"}) + Router._drop_client_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"}) assert non_dict_carriers == {"output_config": "max", "reasoning": 3} effort_only: Dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}} @@ -3452,11 +4250,11 @@ class TestRouterPreRoutingSharedAliasName: def test_forwardable_alias_marker_params_reads_the_marker_entry_only(self): router = Router(model_list=[self._plain_entry(), self._marker_entry(), self._tier_entry()]) - forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=())) + forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=(), request_kwargs={})) assert forwarded["drop_params"] is True assert "api_key" not in forwarded and "api_base" not in forwarded - assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=()) == () + assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=(), request_kwargs={}) == () @staticmethod def _region_marker_entry() -> dict: @@ -3486,11 +4284,11 @@ class TestRouterPreRoutingSharedAliasName: } @staticmethod - async def _routed_call_kwargs(router: Router, **request_params) -> dict: + async def _routed_call_kwargs(router: Router, prompt: str = "hi", **request_params) -> dict: mock_acompletion = AsyncMock(return_value=litellm.ModelResponse(choices=[{"message": {"content": "hi"}}])) with patch.object(litellm, "acompletion", mock_acompletion): await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "hi"}], **request_params + model="smart-router", messages=[{"role": "user", "content": prompt}], **request_params ) return mock_acompletion.call_args.kwargs @@ -7207,11 +8005,384 @@ _ASKED = {"role": "user", "content": _ASK} _ANSWERED = {"role": "assistant", "content": "Working on it."} _TOOL_RESULT = {"type": "tool_result", "tool_use_id": "x", "content": "out"} _REMINDER = "Budget: 42 tokens remaining. Do not mention this." +_CODEX_NEW_TASK: Final = ( + "Message Type: NEW_TASK\nTask name: /root/cache_worker\nSender: /root\nPayload:\n" + "Implement and test a thread-safe bounded LRU cache." +) +_CODEX_ENVELOPES: Final = ( + "LITELLM ESCALATE cwd=/repo", + "LITELLM ESCALATE plugin list", + "LITELLM ESCALATE preferences", + "LITELLM ESCALATE environment", + "# AGENTS.md instructions for /repo with spaces/中文\nLITELLM ESCALATE instructions", +) class TestContextAwareClassifier: """Test the new classifier context window and trajectory signals.""" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "request_metadata,forwards_system", + [ + ({"metadata": {"user_agent": "claude-cli/2.1.233"}}, False), + ({"litellm_metadata": {"user_agent": "claude-code/2.1.233"}}, False), + ({"metadata": {"user_agent": "curl/8.7.1"}}, True), + ({"litellm_metadata": {}}, True), + ( + {"metadata": {"user_agent": "claude-cli/2.1.233"}, "litellm_metadata": {"user_agent": "curl/8.7.1"}}, + False, + ), + ({"metadata": {"user_agent": "Claude-Code/2.1.233"}}, True), + ], + ) + async def test_claude_code_classifier_omits_harness_system_prompt( + self, + llm_classifier_config: dict[str, object], + request_metadata: dict[str, object], + forwards_system: bool, + ) -> None: + dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}'))) + router: Final = ComplexityRouter( + "test-complexity-router", + dependency, + { + **llm_classifier_config, + "classifier_context_include_assistant_turns": True, + }, + ) + messages: Final = [ + {"role": "user", "content": "Design the retry state machine"}, + {"role": "assistant", "content": "The design needs a lease and fencing token"}, + {"role": "user", "content": "Now prove it cannot livelock"}, + { + "role": "system", + "content": [{"type": "text", "text": "ENVIRONMENT_CATALOG\nAGENT_CATALOG\nSKILL_CATALOG"}], + }, + ] + top_level_system: Final = [{"type": "text", "text": "TOP_LEVEL_HARNESS_SYSTEM"}] + claude_kwargs: Final = { + "metadata": {"user_agent": "claude-cli/2.1.233"}, + "system": top_level_system, + "proxy_server_request": {"body": {"system": top_level_system}}, + } + compared_kwargs: Final = { + **request_metadata, + "system": top_level_system, + "proxy_server_request": {"body": {"system": top_level_system}}, + } + original_messages: Final = deepcopy(messages) + original_kwargs: Final = deepcopy((claude_kwargs, compared_kwargs)) + results: Final = ( + await router.async_pre_routing_hook("test-complexity-router", claude_kwargs, messages), + await router.async_pre_routing_hook("test-complexity-router", compared_kwargs, messages), + ) + + assert all(result is not None and result.routing_decision["cause"] == "llm_classifier" for result in results) + assert all(result is not None and result.messages == original_messages for result in results) + assert messages == original_messages + assert (claude_kwargs, compared_kwargs) == original_kwargs + calls: Final = tuple(call.kwargs["messages"] for call in dependency.acompletion.await_args_list) + assert calls[0][0]["content"] == calls[1][0]["content"] == classification_system_prompt( + router.config.classifier_context_window_size + ) + payloads: Final = (calls[0][1]["content"], calls[1][1]["content"]) + for payload, expected_system in zip(payloads, (False, forwards_system)): + assert payload.endswith("Classify this message:\nNow prove it cannot livelock") + assert ("ENVIRONMENT_CATALOG" in payload) is expected_system + assert ("AGENT_CATALOG" in payload) is expected_system + assert ("SKILL_CATALOG" in payload) is expected_system + assert "Design the retry state machine" in payload + assert "lease and fencing token" in payload + assert "TOP_LEVEL_HARNESS_SYSTEM" not in payload + assert "Conversation so far: ~35 tokens across the request" in payload + + @pytest.mark.asyncio + async def test_claude_code_first_turn_without_context_omits_harness_system_prompt( + self, llm_classifier_config: dict[str, object] + ) -> None: + dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}'))) + router: Final = ComplexityRouter( + "test-complexity-router", + dependency, + {**llm_classifier_config, "classifier_context_window_size": 0}, + ) + messages: Final = [ + {"role": "user", "content": "What is two plus two?"}, + { + "role": "system", + "content": [{"type": "text", "text": "ENVIRONMENT_CATALOG\nAGENT_CATALOG\nSKILL_CATALOG"}], + }, + ] + request_kwargs: Final = {"litellm_metadata": {"user_agent": "claude-code/2.1.233"}} + original: Final = deepcopy((messages, request_kwargs)) + + result: Final = await router.async_pre_routing_hook("test-complexity-router", request_kwargs, messages) + + assert result is not None and result.routing_decision["cause"] == "llm_classifier" + assert result.messages == messages == original[0] + assert request_kwargs == original[1] + classifier_messages: Final = dependency.acompletion.call_args.kwargs["messages"] + assert classifier_messages[0]["content"] == classification_system_prompt( + router.config.classifier_context_window_size + ) + assert classifier_messages[1]["content"].strip() == "Classify this message:\nWhat is two plus two?" + + @pytest.mark.parametrize( + "tail,expected", + ( + ([{"role": "user", "content": [{"type": "text", "text": _CODEX_ENVELOPES[0]}]}], True), + ([{"role": "assistant", "content": _CODEX_ENVELOPES[0]}], False), + ([{"role": "tool", "content": _CODEX_ENVELOPES[0]}], False), + ([{"role": "user", "content": " "}], False), + ( + [{"role": "user", "content": [_TOOL_RESULT, {"type": "text", "text": _CODEX_ENVELOPES[0]}]}], + False, + ), + ( + [{"role": "user", "content": [{"type": "image_url"}, {"type": "text", "text": _CODEX_ENVELOPES[0]}]}], + False, + ), + ), + ) + def test_only_text_reminder_tails_are_ignored_for_new_asks( + self, tail: list[dict[str, object]], expected: bool + ) -> None: + from litellm.router_strategy.complexity_router.complexity_router import ( + _CODEX_REMINDER_MARKERS, + _newest_turn_is_human_ask, + ) + + assert _newest_turn_is_human_ask([_ASKED, *tail], _CODEX_REMINDER_MARKERS) is expected + assert _newest_turn_is_human_ask(tail, _CODEX_REMINDER_MARKERS) is False + + @pytest.mark.asyncio + @pytest.mark.parametrize("new_ask", (_CODEX_NEW_TASK, "Now design cache invalidation")) + @pytest.mark.parametrize("responses_api", (False, True)) + @pytest.mark.parametrize("session_affinity", (False, True)) + async def test_codex_tail_preserves_new_ask_and_tool_continuation_boundaries( + self, new_ask: str, responses_api: bool, session_affinity: bool + ) -> None: + completion: Final = AsyncMock( + side_effect=[_llm_response('{"tier":"SIMPLE"}'), _llm_response('{"tier":"COMPLEX"}')] + ) + router: Final = ComplexityRouter( + model_name="router", + litellm_router_instance=MagicMock(acompletion=completion, cache=DualCache()), + complexity_router_config={ + "tiers": {"SIMPLE": "simple-model", "COMPLEX": "task-model"}, + "classifier_type": "llm", + "classifier_llm_config": {"model": "classifier-model"}, + "classification_mode": "user_turn", + "session_affinity": session_affinity, + "escalation_keywords": [], + }, + ) + metadata: Final = {"user_agent": "codex-tui", "session_id": "codex-tail-session"} + first_messages: Final = [{"role": "user", "content": "Hello"}] + tail: Final = [{"role": "user", "content": envelope} for envelope in _CODEX_ENVELOPES] + new_messages: Final = [ + *first_messages, + {"role": "assistant", "content": "Hello"}, + {"role": "user", "content": new_ask}, + *tail, + ] + continuation: Final = [ + *new_messages, + {"role": "assistant", "content": "Working on it"}, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "read-cache", "content": "cache source"}, + {"type": "text", "text": _CODEX_ENVELOPES[0]}, + ], + }, + *tail, + ] + results: Final = [ + await router.async_pre_routing_hook( + model="router", + request_kwargs=( + {"input": messages, "litellm_metadata": {**metadata, "user_api_key_request_route": "/v1/responses"}} + if responses_api + else {"metadata": metadata} + ), + messages=None if responses_api else messages, + input=messages if responses_api else None, + ) + for messages in (first_messages, new_messages, continuation) + ] + + assert [result.model for result in results] == ( + ["simple-model", "simple-model", "simple-model"] + if session_affinity + else ["simple-model", "task-model", "task-model"] + ) + assert completion.await_count == (1 if session_affinity else 2) + assert results[-1].routing_decision["cause"] == ( + "session_affinity_pin" if session_affinity else "user_turn_continuation" + ) + if not session_affinity: + assert completion.call_args.kwargs["messages"][1]["content"].endswith(f"Classify this message:\n{new_ask}") + assert results[1].messages == (None if responses_api else new_messages) + + @pytest.mark.asyncio + @pytest.mark.parametrize("envelope", _CODEX_ENVELOPES) + @pytest.mark.parametrize("user_agent", (None, "curl/8.7.1", "codexify/1.0")) + async def test_non_codex_requests_preserve_tagged_asks(self, envelope: str, user_agent: str | None) -> None: + completion: Final = AsyncMock(return_value=_llm_response('{"tier":"COMPLEX"}')) + router: Final = ComplexityRouter( + model_name="router", + litellm_router_instance=MagicMock(acompletion=completion), + complexity_router_config={ + "tiers": {"COMPLEX": "task-model"}, + "default_model": "fallback-model", + "classifier_type": "llm", + "classifier_llm_config": {"model": "classifier-model"}, + "escalation_keywords": [], + }, + ) + + response: Final = await router.async_pre_routing_hook( + model="router", + request_kwargs={"metadata": {"user_agent": user_agent}} if user_agent is not None else {}, + messages=[{"role": "user", "content": envelope}], + ) + + assert response is not None + assert response.model == "task-model" + completion.assert_awaited_once() + assert completion.call_args.kwargs["messages"][1]["content"].strip() == f"Classify this message:\n{envelope}" + + @pytest.mark.parametrize("envelope", _CODEX_ENVELOPES) + def test_codex_envelopes_preserve_delegated_task_and_prior_context(self, envelope: str) -> None: + from litellm.router_strategy.complexity_router.complexity_router import ( + _CODEX_REMINDER_MARKERS, + _extract_current_ask_and_system_prompt, + _extract_prior_turns, + _newest_turn_ask, + _newest_turn_is_human_ask, + ) + + messages: Final = [ + {"role": "user", "content": f"{envelope}\nDesign cache invalidation"}, + { + "role": "user", + "content": [{"type": "text", "text": envelope}, {"type": "text", "text": _CODEX_NEW_TASK}], + }, + {"role": "developer", "content": "developer scope"}, + {"role": "user", "content": envelope}, + ] + + assert _extract_current_ask_and_system_prompt(messages, _CODEX_REMINDER_MARKERS)[0] == _CODEX_NEW_TASK + assert _extract_prior_turns(messages, _CODEX_NEW_TASK, 1, 100, None, False, _CODEX_REMINDER_MARKERS) == ( + ("user", "Design cache invalidation"), + ) + assert _newest_turn_ask(messages, _CODEX_REMINDER_MARKERS) is None + assert _newest_turn_is_human_ask(messages, _CODEX_REMINDER_MARKERS) is False + assert _extract_current_ask_and_system_prompt([messages[-1]], _CODEX_REMINDER_MARKERS)[0] is None + + @pytest.mark.parametrize("envelope", _CODEX_ENVELOPES) + def test_codex_marker_override_and_incomplete_blocks_preserve_text(self, envelope: str) -> None: + from litellm.router_strategy.complexity_router.complexity_router import ( + _CODEX_REMINDER_MARKERS, + _strip_reminder_blocks, + ) + + incomplete: Final = envelope.rsplit("noise{envelope}", (("", ""),)) == envelope + + @pytest.mark.asyncio + @pytest.mark.parametrize("responses_api", (False, True)) + async def test_codex_routing_preserves_original_request(self, responses_api: bool) -> None: + completion: Final = AsyncMock(return_value=_llm_response('{"tier":"COMPLEX"}')) + router: Final = ComplexityRouter( + model_name="codex-router", + litellm_router_instance=MagicMock(acompletion=completion), + complexity_router_config={ + "tiers": {"COMPLEX": "task-model", "REASONING": "escalated-model"}, + "classifier_type": "llm", + "classifier_llm_config": {"model": "classifier-model"}, + "keyword_tier_rules": [{"keywords": ["LITELLM ESCALATE"], "tier": "REASONING"}], + }, + ) + messages: Final = [ + {"role": "user", "content": _CODEX_NEW_TASK}, + {"role": "user", "content": "\n".join(_CODEX_ENVELOPES)}, + ] + original: Final = deepcopy(messages) + request_kwargs: Final = ( + { + "input": messages, + "litellm_metadata": {"user_api_key_request_route": "/v1/responses", "user_agent": "codex-tui"}, + } + if responses_api + else {"metadata": {"user_agent": "codex-tui"}} + ) + + response: Final = await router.async_pre_routing_hook( + model="codex-router", + request_kwargs=request_kwargs, + messages=None if responses_api else messages, + input=messages if responses_api else None, + ) + + assert response is not None + assert response.model == "task-model" + completion.assert_awaited_once() + assert completion.call_args.kwargs["messages"][1]["content"].strip() == ( + f"Classify this message:\n{_CODEX_NEW_TASK}" + ) + assert messages == original + if responses_api: + assert response.messages is None + assert request_kwargs["input"] == original + else: + assert response.messages == original + + @pytest.mark.asyncio + @pytest.mark.parametrize("custom_markers", (False, True)) + async def test_codex_markers_are_request_scoped_and_respect_overrides(self, custom_markers: bool) -> None: + completion: Final = AsyncMock(return_value=_llm_response('{"tier":"COMPLEX"}')) + router: Final = ComplexityRouter( + model_name="router", + litellm_router_instance=MagicMock(acompletion=completion), + complexity_router_config={ + "tiers": {"COMPLEX": "task-model"}, + "classifier_type": "llm", + "classifier_llm_config": {"model": "classifier-model"}, + "classifier_context_window_size": 2, + "escalation_keywords": [], + **({"reminder_markers": [{"open": "", "close": ""}]} if custom_markers else {}), + }, + ) + envelope: Final = "\n".join(_CODEX_ENVELOPES) + prior: Final = f"{envelope}\nDesign cache invalidation" + messages: Final = [ + {"role": "user", "content": prior}, + {"role": "user", "content": _CODEX_NEW_TASK}, + {"role": "user", "content": envelope}, + ] + for user_agent in ("codex-tui", "curl/8.7.1", "codex_cli_rs/0.62.0"): + response: Final = await router.async_pre_routing_hook( + model="router", request_kwargs={"metadata": {"user_agent": user_agent}}, messages=messages + ) + assert response is not None + assert response.model == "task-model" + payload: Final = completion.call_args.kwargs["messages"][1]["content"] + if user_agent.startswith("codex") and not custom_markers: + assert payload.endswith(f"Classify this message:\n{_CODEX_NEW_TASK}") + assert "Design cache invalidation" in payload + assert "LITELLM ESCALATE" not in payload + else: + assert payload.endswith(f"Classify this message:\n{envelope}") + assert prior in payload + assert response.messages == messages + assert completion.await_count == 3 + @pytest.mark.parametrize( "messages,expected_ask", [ @@ -10739,7 +11910,7 @@ async def test_tier_model_params_reach_the_hook_response_and_override_client_val async def test_tier_params_mask_credentials_in_routing_decision(route, mock_router_instance): params = {"reasoning_effort": "xhigh", "api_key": "secret-tier-key"} config = { - "tiers": {tier.value: {"model_name": "opus", "litellm_params": params} for tier in ComplexityTier}, + "tiers": {tier.value: {"model_name": "opus", "litellm_params": params} for tier in TIER_SEVERITY_ORDER}, "keyword_tier_rules": [{"keywords": ["reason carefully"], "tier": "REASONING"}] if route == "keyword" else None, "session_affinity": route == "session", } @@ -12060,6 +13231,528 @@ class TestModalityRouting: assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"} +@pytest.mark.usefixtures("local_model_cost_map") +class TestHealthFallbackDispatch: + @pytest.fixture(autouse=True) + def httpx_transport(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + @staticmethod + def _router( + surface: str = "chat", + *, + peer: bool = False, + session: bool = False, + tagged: bool = False, + budgeted: bool = False, + config: Mapping[str, object] | None = None, + ) -> Router: + provider: Final = "anthropic/claude-sonnet-5" if surface == "messages" else "openai/gpt-5.6" + base_suffix: Final = "" if surface == "messages" else "/v1" + return Router( + model_list=[ + { + "model_name": "health-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": (config or {}).get("default_model", "fallback"), + "complexity_router_config": { + "tiers": {"SIMPLE": ["primary", "peer"] if peer else "primary", "MEDIUM": "primary"}, + "session_affinity": session, + "deployment_affinity": False, + "max_tokens_from_tier_model": False, + **(config or {}), + }, + }, + }, + *[ + { + "model_name": name, + "litellm_params": { + "model": provider, + "api_key": "test-only", + "api_base": f"https://{name}.test{base_suffix}", + **({"tags": [name]} if tagged else {}), + **( + {"max_budget": 1.0, "budget_duration": "1d"} + if budgeted and name == "primary" + else {} + ), + }, + "model_info": {"id": f"{name}-id"}, + } + for name in ("primary", "peer", "fallback") + ], + ], + num_retries=0, + enable_health_check_routing=True, + enable_tag_filtering=tagged, + ) + + @staticmethod + def _unavailable(router: Router, model_id: str, source: Literal["health", "cooldown"]) -> None: + if source == "health": + router.health_state_cache.set_deployment_health_states( + {model_id: {"is_healthy": False, "timestamp": time.time()}} + ) + else: + router.cooldown_cache.add_deployment_to_cooldown( + model_id=model_id, + original_exception=RuntimeError("unavailable"), + exception_status=503, + cooldown_time=60, + ) + + @staticmethod + def _http_response(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + text: Final = request.url.host.split(".")[0] + payload: Final[Mapping[str, object]] + events: Final[tuple[Mapping[str, object], ...]] + if request.url.path.endswith("/responses"): + from litellm.responses.main import mock_responses_api_response + + payload = mock_responses_api_response(text).model_dump() + events = ( + {"type": "response.created", "response": {**payload, "status": "in_progress"}, "sequence_number": 0}, + { + "type": "response.output_text.delta", + "delta": text, + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + "sequence_number": 1, + }, + {"type": "response.completed", "response": payload, "sequence_number": 2}, + ) + elif request.url.path.endswith("/messages"): + payload = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": body["model"], + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + } + events = ( + {"type": "message_start", "message": {**payload, "content": [], "stop_reason": None}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}, + {"type": "message_stop"}, + ) + else: + payload = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": body["model"], + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, + } + events = ( + { + **payload, + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}], + }, + { + **payload, + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + ) + if not body.get("stream"): + return httpx.Response(200, json=payload) + wire: Final = "".join( + (f"event: {event['type']}\n" if "type" in event else "") + f"data: {json.dumps(event)}\n\n" + for event in events + ) + return httpx.Response( + 200, + text=wire + ("data: [DONE]\n\n" if "type" not in events[0] else ""), + headers={"content-type": "text/event-stream"}, + ) + + @staticmethod + async def _request(router: Router, surface: str, stream: bool, metadata: dict[str, object]) -> str: + if surface == "responses": + result = await router.aresponses( + model="health-router", input="Hello!", stream=stream, litellm_metadata=metadata + ) + elif surface == "messages": + result = await router.aanthropic_messages( + model="health-router", + messages=[{"role": "user", "content": "Hello!"}], + max_tokens=32, + stream=stream, + litellm_metadata=metadata, + ) + else: + result = await router.acompletion( + model="health-router", + messages=[{"role": "user", "content": "Hello!"}], + stream=stream, + metadata=metadata, + ) + if not stream: + payload = result if isinstance(result, dict) else result.model_dump() + if surface == "responses": + return payload["output"][0]["content"][0]["text"] + if surface == "messages": + return payload["content"][0]["text"] + return payload["choices"][0]["message"]["content"] + if surface == "messages": + wire: Final = b"".join([chunk async for chunk in result]).decode() + events = tuple(json.loads(line[6:]) for line in wire.splitlines() if line.startswith("data: ")) + assert events[-1]["type"] == "message_stop" + return "".join(c["delta"]["text"] for c in events if c["type"] == "content_block_delta") + chunks: Final = [chunk.model_dump() async for chunk in result] + if surface == "responses": + assert chunks[-1]["type"] == "response.completed" + return "".join(c["delta"] for c in chunks if c["type"] == "response.output_text.delta") + assert chunks[-1]["choices"][0]["finish_reason"] == "stop" + return "".join(c["choices"][0]["delta"].get("content") or "" for c in chunks if c["choices"]) + + @pytest.mark.asyncio + @pytest.mark.parametrize("surface", ["chat", "responses", "messages"]) + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("source", ["health", "cooldown"]) + async def test_public_call_falls_back_and_recovers( + self, surface: str, stream: bool, source: Literal["health", "cooldown"] + ) -> None: + router: Final = self._router(surface, session=True) + self._unavailable(router, "primary-id", source) + metadata: Final[dict[str, object]] = {"session_id": "outage"} + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host__regex=r"^(primary|peer|fallback)\.test$").mock(side_effect=self._http_response) + assert await self._request(router, surface, stream, metadata) == "fallback" + assert metadata["routing_decision"]["cause"] == "health_default_fallback" + assert "tier" not in metadata["routing_decision"] + assert "health_displaced:primary" in metadata["routing_decision"]["signals"] + assert [c.request.url.host for c in upstream.calls] == ["fallback.test"] + strategy: Final = router.complexity_routers["health-router"][0].strategy + key: Final = strategy._get_session_affinity_cache_key("outage", {}) + assert await router.cache.async_get_cache(key=key) is None + if source == "health": + router.health_state_cache.set_deployment_health_states( + {"primary-id": {"is_healthy": True, "timestamp": time.time()}} + ) + else: + router.cooldown_cache.cooldown_store.delete_cache( + router.cooldown_cache.get_cooldown_cache_key("primary-id") + ) + recovered: Final[dict[str, object]] = {"session_id": "outage"} + assert await self._request(router, surface, stream, recovered) == "primary" + assert recovered["routing_decision"]["routed_model"] == "primary" + assert [c.request.url.host for c in upstream.calls] == ["fallback.test", "primary.test"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("source", ["health", "cooldown"]) + async def test_partial_group_then_peer_then_default(self, source: Literal["health", "cooldown"]) -> None: + router: Final = self._router(peer=True, session=True) + router.add_deployment( + Deployment( + model_name="primary", + litellm_params=LiteLLM_Params( + model="openai/gpt-5.6", api_key="test-only", api_base="https://primary.test/v1" + ), + model_info={"id": "primary-sibling-id"}, + ) + ) + strategy: Final = router.complexity_routers["health-router"][0].strategy + key: Final = strategy._get_session_affinity_cache_key("precedence", {}) + await router.cache.async_set_cache(key=key, value={"model": "primary", "tier": "SIMPLE"}, ttl=600) + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host__regex=r"^(primary|peer|fallback)\.test$").mock(side_effect=self._http_response) + for model_id, expected, cause in ( + ("primary-id", "primary", "session_affinity_pin"), + ("primary-sibling-id", "peer", "health_failover"), + ("peer-id", "fallback", "health_default_fallback"), + ): + self._unavailable(router, model_id, source) + metadata: Final[dict[str, object]] = {"session_id": "precedence"} + assert await self._request(router, "chat", False, metadata) == expected + assert metadata["routing_decision"]["cause"] == cause + assert await router.cache.async_get_cache(key=key) == {"model": "primary", "tier": "SIMPLE"} + assert [c.request.url.host for c in upstream.calls] == ["primary.test", "peer.test", "fallback.test"] + + @pytest.mark.asyncio + async def test_spent_deployment_budget_falls_back_to_the_default(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A spent budget leaves the tier with nothing that may serve the request, and the budget + filter reports that as a bare ValueError instead of a typed router error. Reading it as + capacity skips the recovery and fails the request the recovery exists for.""" + + async def _no_sync(*args: object, **kwargs: object) -> None: + return None + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis", + _no_sync, + ) + monkeypatch.setattr(litellm, "callbacks", []) + router: Final = self._router(budgeted=True) + limiter: Final = router.router_budget_logger + assert limiter is not None, "a deployment max_budget must install the budget limiter" + await router.cache.async_set_cache(key="deployment_spend:primary-id:1d", value=2.0) + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host__regex=r"^(primary|fallback)\.test$").mock(side_effect=self._http_response) + metadata: Final[dict[str, object]] = {} + assert await self._request(router, "chat", False, metadata) == "fallback" + assert metadata["routing_decision"]["cause"] == "health_default_fallback" + assert [c.request.url.host for c in upstream.calls] == ["fallback.test"] + + @pytest.mark.asyncio + async def test_concurrent_tag_scopes_keep_fallbacks_request_local(self) -> None: + router: Final = self._router(tagged=True) + router.add_deployment( + Deployment( + model_name="fallback", + litellm_params=LiteLLM_Params( + model="openai/gpt-5.6", api_key="test-only", api_base="https://peer.test/v1", tags=["peer"] + ), + model_info={"id": "fallback-peer-id"}, + ) + ) + self._unavailable(router, "primary-id", "cooldown") + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host__regex=r"^(peer|fallback)\.test$").mock(side_effect=self._http_response) + scopes: Final = tuple({"tags": [name], "session_id": name} for name in ("peer", "fallback")) + results: Final = await asyncio.gather( + *(self._request(router, "chat", False, metadata) for metadata in scopes) + ) + assert results == ["peer", "fallback"] + assert [m["tags"] for m in scopes] == [["peer"], ["fallback"]] + assert [m["routing_decision"]["routed_model"] for m in scopes] == ["fallback", "fallback"] + assert sorted(c.request.url.host for c in upstream.calls) == ["fallback.test", "peer.test"] + + @pytest.mark.asyncio + async def test_probe_preserves_consumed_request_exclusions(self) -> None: + router: Final = self._router() + self._unavailable(router, "primary-id", "cooldown") + kwargs: Final = {"_excluded_deployment_ids": ["fallback-id"], "_target_order": 1} + strategy: Final = router.complexity_routers["health-router"][0].strategy + response: Final = await strategy.async_pre_routing_hook( + model="health-router", messages=[{"role": "user", "content": "Hello!"}], request_kwargs=kwargs + ) + assert response.model == "primary" + assert kwargs == {"_excluded_deployment_ids": ["fallback-id"], "_target_order": 1} + + @pytest.mark.asyncio + @pytest.mark.parametrize("default_state", ["cooldown", "unconfigured", "same-model"]) + async def test_unavailable_default_preserves_no_deployment_error(self, default_state: str) -> None: + from litellm.types.router import RouterRateLimitError + + router: Final = self._router(config={"default_model": "primary"} if default_state == "same-model" else None) + self._unavailable(router, "primary-id", "cooldown") + if default_state == "unconfigured": + router.delete_deployment(id="fallback-id") + elif default_state == "cooldown": + self._unavailable(router, "fallback-id", "cooldown") + with respx.mock(assert_all_mocked=True) as upstream: + with pytest.raises(RouterRateLimitError, match="No deployments available"): + await self._request(router, "chat", False, {}) + assert not upstream.calls + + @pytest.mark.asyncio + @pytest.mark.parametrize("plan_active", [False, True]) + async def test_plan_floor_outage_cannot_use_untiered_default(self, plan_active: bool) -> None: + from litellm.types.router import RouterRateLimitError + + router: Final = self._router( + config={"tiers": {"SIMPLE": "primary", "MEDIUM": "peer"}, "plan_mode_min_tier": "MEDIUM"} + ) + self._unavailable(router, "primary-id", "cooldown") + self._unavailable(router, "peer-id", "cooldown") + metadata: Final = {} + with respx.mock(assert_all_mocked=True, assert_all_called=False) as upstream: + upstream.post(host="fallback.test").mock(side_effect=self._http_response) + if plan_active: + with pytest.raises(RouterRateLimitError, match="No deployments available"): + await router.acompletion( + model="health-router", + messages=[ + {"role": "system", "content": "Plan mode is active"}, + {"role": "user", "content": "Hello!"}, + ], + metadata=metadata, + ) + assert not upstream.calls + assert metadata["routing_decision"]["routed_model"] == "peer" + assert metadata["routing_decision"]["tier"] == "MEDIUM" + else: + assert await self._request(router, "chat", False, metadata) == "fallback" + + @pytest.mark.asyncio + async def test_default_dispatch_drops_displaced_tier_params(self) -> None: + router: Final = self._router( + config={"tiers": {"SIMPLE": {"model_name": "primary", "litellm_params": {"max_tokens": 9}}}} + ) + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host__regex=r"^(primary|fallback)\.test$").mock(side_effect=self._http_response) + await router.acompletion( + model="health-router", messages=[{"role": "user", "content": "Hello!"}], max_tokens=32 + ) + assert json.loads(upstream.calls[-1].request.content)["max_completion_tokens"] == 9 + self._unavailable(router, "primary-id", "cooldown") + await router.acompletion( + model="health-router", messages=[{"role": "user", "content": "Hello!"}], max_tokens=32 + ) + assert json.loads(upstream.calls[-1].request.content)["max_completion_tokens"] == 32 + assert upstream.calls[-1].request.url.host == "fallback.test" + + @pytest.mark.asyncio + @pytest.mark.parametrize("source", ["health", "cooldown"]) + async def test_pinned_session_returns_to_primary_after_outage(self, source: Literal["health", "cooldown"]) -> None: + router: Final = self._router(session=True) + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host__regex=r"^(primary|fallback)\.test$").mock(side_effect=self._http_response) + assert await self._request(router, "chat", False, {"session_id": "pinned"}) == "primary" + self._unavailable(router, "primary-id", source) + outage: Final[dict[str, object]] = {"session_id": "pinned"} + assert await self._request(router, "chat", False, outage) == "fallback" + assert outage["routing_decision"]["cause"] == "health_default_fallback" + if source == "health": + router.health_state_cache.set_deployment_health_states( + {"primary-id": {"is_healthy": True, "timestamp": time.time()}} + ) + else: + router.cooldown_cache.cooldown_store.delete_cache( + router.cooldown_cache.get_cooldown_cache_key("primary-id") + ) + recovered: Final[dict[str, object]] = {"session_id": "pinned"} + assert await self._request(router, "chat", False, recovered) == "primary" + assert recovered["routing_decision"]["cause"] == "session_affinity_pin" + assert [c.request.url.host for c in upstream.calls] == ["primary.test", "fallback.test", "primary.test"] + + @pytest.mark.asyncio + async def test_policy_plugin_does_not_escape_to_live_default(self) -> None: + from litellm.types.router import RouterRateLimitError, RoutingContext + + class PrimaryOnly: + async def run(self, context: RoutingContext) -> RoutingContext: + context.candidate_models = [name for name in context.candidate_models if name == "primary"] + return context + + router: Final = self._router(peer=True, config={"plugins": [PrimaryOnly()]}) + self._unavailable(router, "primary-id", "cooldown") + with respx.mock(assert_all_mocked=True) as upstream: + with pytest.raises(RouterRateLimitError, match="No deployments available"): + await self._request(router, "chat", False, {}) + assert not upstream.calls + + @pytest.mark.asyncio + @pytest.mark.parametrize("live_tier", [True, False]) + @pytest.mark.parametrize("default_fits", [True, False]) + async def test_context_recovery_precedes_default_with_prechecks_off( + self, live_tier: bool, default_fits: bool + ) -> None: + from litellm.types.router import RouterRateLimitError + + router: Final = self._router(config={"tiers": {"SIMPLE": "primary", "MEDIUM": "peer", "COMPLEX": "large"}}) + router.add_deployment( + Deployment( + model_name="large", + litellm_params=LiteLLM_Params( + model="openai/gpt-5.6", api_key="test-only", api_base="https://large.test/v1" + ), + model_info={"id": "large-id", "max_input_tokens": 10000}, + ) + ) + for deployment in router.model_list: + deployment["model_info"]["max_input_tokens"] = ( + 10 + if deployment["model_name"] == "primary" + or (deployment["model_name"] == "fallback" and not default_fits) + else 10000 + ) + self._unavailable(router, "peer-id", "cooldown") + if not live_tier: + self._unavailable(router, "large-id", "cooldown") + assert router.enable_pre_call_checks is False + metadata: Final = {} + messages: Final = [{"role": "user", "content": "hello " * 100}] + with respx.mock(assert_all_mocked=True, assert_all_called=False) as upstream: + upstream.post(host__regex=r"^(large|fallback)\.test$").mock(side_effect=self._http_response) + if not live_tier and not default_fits: + with pytest.raises(RouterRateLimitError, match="No deployments available"): + await router.acompletion(model="health-router", messages=messages, metadata=metadata) + assert not upstream.calls + else: + result: Final = await router.acompletion(model="health-router", messages=messages, metadata=metadata) + expected: Final = "large" if live_tier else "fallback" + assert result.choices[0].message.content == expected + assert upstream.calls[-1].request.url.host == f"{expected}.test" + assert metadata["routing_decision"].get("tier") == ("COMPLEX" if live_tier else None) + + @pytest.mark.asyncio + @pytest.mark.parametrize("live_tier", [True, False]) + async def test_modality_recovery_precedes_default(self, live_tier: bool) -> None: + router: Final = self._router( + config={"modality_routing": True, "tiers": {"SIMPLE": "primary", "MEDIUM": "peer", "COMPLEX": "vision"}} + ) + router.add_deployment( + Deployment( + model_name="vision", + litellm_params=LiteLLM_Params( + model="openai/gpt-5.6", api_key="test-only", api_base="https://vision.test/v1" + ), + model_info={"id": "vision-id", "supports_vision": True}, + ) + ) + for deployment in router.model_list: + deployment["model_info"]["supports_vision"] = deployment["model_name"] != "primary" + self._unavailable(router, "peer-id", "cooldown") + if not live_tier: + self._unavailable(router, "vision-id", "cooldown") + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host__regex=r"^(vision|fallback)\.test$").mock(side_effect=self._http_response) + result: Final = await router.acompletion( + model="health-router", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello!"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGk="}}, + ], + } + ], + ) + expected: Final = "vision" if live_tier else "fallback" + assert result.choices[0].message.content == expected + assert upstream.calls[-1].request.url.host == f"{expected}.test" + + @pytest.mark.asyncio + @pytest.mark.parametrize("default_fits", [True, False]) + async def test_modality_default_must_also_fit_context(self, default_fits: bool) -> None: + router: Final = self._router(config={"modality_routing": True, "tiers": {"SIMPLE": "primary"}}) + for deployment in router.model_list: + deployment["model_info"]["supports_vision"] = deployment["model_name"] == "fallback" + deployment["model_info"]["max_input_tokens"] = 10000 if default_fits else 10 + with respx.mock(assert_all_mocked=True, assert_all_called=False) as upstream: + upstream.post(host="fallback.test").mock(side_effect=self._http_response) + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello " * 100}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGk="}}, + ], + } + ] + if default_fits: + result: Final = await router.acompletion(model="health-router", messages=messages) + assert result.choices[0].message.content == "fallback" + else: + with pytest.raises(litellm.BadRequestError, match="modality_routing is enabled"): + await router.acompletion(model="health-router", messages=messages) + assert not upstream.calls + + class TestTierHealthFailover: """A tier whose decided model group is entirely in cooldown falls back to a live peer.""" @@ -12094,7 +13787,7 @@ class TestTierHealthFailover: probed_prompts = [] async def get_healthy_deployments( - model, request_kwargs, messages=None, input=None, parent_otel_span=None, **kwargs + model, request_kwargs, messages=None, input=None, parent_otel_span=None, health_check_probe=False ): probed_kwargs.append(request_kwargs) probed_prompts.append((messages, input)) @@ -12109,9 +13802,7 @@ class TestTierHealthFailover: llm_provider="", ) filtered = (*cooling, *blocked, *excluded) - healthy = [ - {"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered - ] + healthy = [{"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered] if not healthy: raise RouterRateLimitError( model=model, cooldown_time=60.0, enable_pre_call_checks=False, cooldown_list=[] @@ -12540,9 +14231,7 @@ class TestTierHealthFailover: assert all(probed is not request_kwargs for probed in router.litellm_router_instance.probed_kwargs) @pytest.mark.asyncio - async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target( - self, mock_router_instance - ): + async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target(self, mock_router_instance): """RPM exhaustion is its own verdict from the owner (RouterRateLimitErrorBasic). A peer in that state would be rejected downstream, so it cannot be the substitute.""" from litellm.types.router import RouterRateLimitErrorBasic @@ -12575,9 +14264,7 @@ class TestTierHealthFailover: assert {r.model for r in results} == {"live-c"} @pytest.mark.asyncio - async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces( - self, mock_router_instance - ): + async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces(self, mock_router_instance): """The Responses API carries its prompt as `input`, never as messages. The owner only runs its context-window pre-call check when one of them is present, so dropping `input` would silently skip window filtering on that whole surface.""" @@ -12603,9 +14290,52 @@ class TestTierHealthFailover: ), "the eligibility probe must forward `input` to the owner" @pytest.mark.asyncio - async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target( - self, mock_router_instance + @pytest.mark.parametrize( + "raised, expected", + [ + (ValueError(f"{RouterErrors.no_deployments_with_tag_routing.value}. Passed model=b"), {"live-c"}), + ( + ValueError(f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: b over budget"), + {"live-c"}, + ), + (ValueError("cannot unpack non-sequence"), {"exhausted-b", "live-c"}), + ], + ) + async def test_a_marked_exhaustion_value_error_is_a_verdict_and_an_unmarked_one_is_not( + self, mock_router_instance, raised, expected ): + """Budget and tag filters exhaust a group without a typed error, signalling it only by a + RouterErrors marker on a bare ValueError. Those are verdicts; any other ValueError is a + fault, and a fault must still read as capacity rather than silently rerouting.""" + router = self._router( + mock_router_instance, + { + "tiers": { + "SIMPLE": ["dead-a", "exhausted-b", "live-c"], + "MEDIUM": "mid", + "COMPLEX": "big", + "REASONING": "top", + }, + "session_affinity": True, + }, + {"dead-a": ["id-a1"], "exhausted-b": ["id-b1"], "live-c": ["id-c1"]}, + cooling=("id-a1",), + raises_for={"exhausted-b": raised}, + ) + key = router._get_session_affinity_cache_key("sess-exhausted", {}) + await router.litellm_router_instance.cache.async_set_cache( + key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 + ) + results = [ + await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": "sess-exhausted"}}, messages=self.SIMPLE_MESSAGE + ) + for _ in range(20) + ] + assert {r.model for r in results} == expected + + @pytest.mark.asyncio + async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(self, mock_router_instance): """The owner answers an unconfigured group with BadRequestError. Reading that as live would both skip failover off it and let it be chosen as a substitute.""" router = self._router( @@ -12794,9 +14524,7 @@ class TestClassifierVision: routed as default_fallback on text the request never contained. """ router = self._router(mock_router_instance, vision={"enabled": True}) - response = await router.async_pre_routing_hook( - model="m", request_kwargs={}, messages=self._turn(IMG_PART) - ) + response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART)) assert response.routing_decision["cause"] == "llm_classifier" assert response.model == "t-complex" assert [block["type"] for block in self._classifier_user_content(mock_router_instance)] == [ @@ -12807,9 +14535,7 @@ class TestClassifierVision: @pytest.mark.asyncio async def test_image_only_turn_still_falls_back_when_vision_is_off(self, mock_router_instance): router = self._router(mock_router_instance, vision={"enabled": False}) - response = await router.async_pre_routing_hook( - model="m", request_kwargs={}, messages=self._turn(IMG_PART) - ) + response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART)) assert response.routing_decision["cause"] == "default_fallback" mock_router_instance.acompletion.assert_not_awaited() @@ -12879,9 +14605,7 @@ class TestClassifierVision: makes the image the only variable; a margin loose enough to leave the score undecided would pass whether or not the guard exists. """ - router = self._router( - mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra - ) + router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra) response = await router.async_pre_routing_hook( model="m", request_kwargs={}, messages=self._turn({"type": "text", "text": "what is this"}, IMG_PART) ) @@ -12895,9 +14619,7 @@ class TestClassifierVision: self, mock_router_instance, classifier_type, extra, short_circuit_cause ): """The negative class: same router, same text, no image, and the scorer still decides.""" - router = self._router( - mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra - ) + router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra) response = await router.async_pre_routing_hook( model="m", request_kwargs={}, messages=[{"role": "user", "content": "what is this"}] ) @@ -12907,3 +14629,558 @@ class TestClassifierVision: def test_max_images_must_be_positive(self): with pytest.raises(ValidationError): ClassifierLLMConfig(model="clf", vision={"enabled": True, "max_images": 0}) + + +class TestMaxTokensFromTierModel: + """The auto-router replaces the caller's output ceiling with the tier model's own, so one + client-side value no longer starves a bigger tier or gets rejected by a smaller one.""" + + COMPLEX_PROMPT: Final = ( + "Design a distributed rate limiter with Redis, sharding and failover. Analyze the consistency " + "tradeoffs and implement the algorithm step by step with tests." + ) + SMALL: Final = { + "model_name": "small", + "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"}, + "model_info": {"max_output_tokens": 8192}, + } + + @staticmethod + def _router( + tier_litellm_params: dict | None = None, + max_tokens_from_tier_model: bool | None = None, + simple_deployments: list[dict] | None = None, + extra_config: dict | None = None, + ) -> Router: + simple_tier: dict = {"model_name": "small"} + if tier_litellm_params: + simple_tier["litellm_params"] = tier_litellm_params + config: dict = { + "tiers": {"SIMPLE": simple_tier, "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"}, + **(extra_config or {}), + } + if max_tokens_from_tier_model is not None: + config["max_tokens_from_tier_model"] = max_tokens_from_tier_model + return Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": config}, + }, + *(simple_deployments or [TestMaxTokensFromTierModel.SMALL]), + { + "model_name": "big", + "litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "k"}, + "model_info": {"max_output_tokens": 64000}, + }, + ] + ) + + @staticmethod + async def _routed(router: Router, prompt: str = "hi", **request_kwargs) -> dict: + """Drive the real routing entry point and return the request kwargs it leaves behind.""" + deployment = await router.async_get_available_deployment( + model="smart-router", request_kwargs=request_kwargs, messages=[{"role": "user", "content": prompt}] + ) + return {"model": deployment["litellm_params"]["model"], **request_kwargs} + + @staticmethod + async def _routed_responses(router: Router, prompt: str = "hi", **request_kwargs) -> dict: + """The Responses surface hands the router `input` both as the prompt argument and inside the + request kwargs, so the hook sees the same shape the real call carries.""" + routed: dict = {"input": prompt, **request_kwargs} + deployment = await router.async_get_available_deployment( + model="smart-router", request_kwargs=routed, input=prompt + ) + return {"model": deployment["litellm_params"]["model"], **routed} + + @pytest.mark.asyncio + async def test_client_ceiling_is_replaced_by_the_routed_tier_models_ceiling(self): + router = self._router() + + simple = await self._routed(router, max_tokens=8192) + complex_ = await self._routed(router, self.COMPLEX_PROMPT, max_tokens=8192) + + assert (simple["model"], simple["max_tokens"]) == ("anthropic/claude-haiku-4-5", 8192) + assert (complex_["model"], complex_["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + assert "max_output_tokens" not in complex_ + + @pytest.mark.asyncio + async def test_every_client_carrier_of_the_ceiling_is_replaced(self): + sent = await self._routed(self._router(), self.COMPLEX_PROMPT, max_completion_tokens=8192) + + assert sent["max_tokens"] == 64000 + assert "max_completion_tokens" not in sent + + @pytest.mark.asyncio + async def test_responses_surface_gets_the_ceiling_under_its_own_name(self): + sent = await self._routed_responses(self._router(), self.COMPLEX_PROMPT, max_output_tokens=8192) + + assert (sent["model"], sent["max_output_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + assert "max_tokens" not in sent + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "tier_params, responses_call", + [ + ({"max_tokens": 4321}, False), + ({"max_tokens": 4321}, True), + ({"max_completion_tokens": 4321}, False), + ({"max_completion_tokens": 4321}, True), + ({"max_output_tokens": 4321}, False), + ], + ) + async def test_operators_own_tier_ceiling_wins_under_the_surface_name(self, tier_params, responses_call): + router = self._router(tier_litellm_params=tier_params) + if responses_call: + sent = await self._routed_responses(router, max_output_tokens=8192) + else: + sent = await self._routed(router, max_tokens=8192) + + surface_key = "max_output_tokens" if responses_call else "max_tokens" + assert sent[surface_key] == 4321 + assert not (OUTPUT_TOKEN_CEILING_PARAMS - {surface_key}) & sent.keys() + + @pytest.mark.asyncio + async def test_opting_out_forwards_the_client_value_unchanged(self): + sent = await self._routed(self._router(max_tokens_from_tier_model=False), self.COMPLEX_PROMPT, max_tokens=8192) + + assert sent["max_tokens"] == 8192 + + @pytest.mark.asyncio + async def test_a_tier_model_with_an_unknown_ceiling_keeps_the_client_value(self): + unmapped: dict = {"model_name": "small", "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}} + + sent = await self._routed(self._router(simple_deployments=[self.SMALL, unmapped]), max_tokens=4000) + + assert sent["max_tokens"] == 4000 + + @pytest.mark.asyncio + async def test_a_multi_deployment_tier_model_uses_its_smallest_ceiling(self): + smaller: dict = { + **self.SMALL, + "litellm_params": {**self.SMALL["litellm_params"], "api_key": "k2"}, + "model_info": {"max_output_tokens": 4096}, + } + + sent = await self._routed(self._router(simple_deployments=[self.SMALL, smaller]), max_tokens=100000) + + assert sent["max_tokens"] == 4096 + + @pytest.mark.asyncio + async def test_ceiling_falls_back_to_the_cost_map(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + "auto-cap-probe-model", + {"litellm_provider": "openai", "mode": "chat", "max_output_tokens": 4242, "max_input_tokens": 100000}, + ) + mapped_only: dict = { + "model_name": "small", + "litellm_params": {"model": "openai/auto-cap-probe-model", "api_key": "k"}, + } + + sent = await self._routed(self._router(simple_deployments=[mapped_only]), max_tokens=8192) + + assert sent["max_tokens"] == 4242 + + @pytest.mark.asyncio + @pytest.mark.parametrize("client_kwargs", [{}, {"max_tokens": 0}], ids=["omitted", "zero"]) + async def test_omitted_and_zero_are_replaced_like_any_other_value(self, client_kwargs): + sent = await self._routed(self._router(), self.COMPLEX_PROMPT, **client_kwargs) + + assert sent["max_tokens"] == 64000 + + @pytest.mark.parametrize( + "tier_params, responses_call, expected", + [ + ({"max_tokens": 1, "temperature": 0.2}, False, {"max_tokens": 1, "temperature": 0.2}), + ({"max_tokens": 1}, True, {"max_output_tokens": 1}), + ({"max_completion_tokens": 2}, False, {"max_tokens": 2}), + ({"max_completion_tokens": 2}, True, {"max_output_tokens": 2}), + ({"max_output_tokens": 3}, False, {"max_tokens": 3}), + ({"max_output_tokens": 3}, True, {"max_output_tokens": 3}), + ({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 1}), + ({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, True, {"max_output_tokens": 3}), + ({"max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 2}), + ({"reasoning_effort": "low"}, True, {"reasoning_effort": "low"}), + ], + ) + def test_every_tier_alias_collapses_onto_the_surface_key(self, tier_params, responses_call, expected): + assert dict(Router._tier_ceiling_under_the_surface_name(tier_params, responses_call=responses_call)) == expected + + @pytest.mark.asyncio + async def test_the_default_fallback_exit_carries_the_ceiling(self): + routed: dict = {"max_tokens": 8192} + deployment = await self._router().async_get_available_deployment( + model="smart-router", request_kwargs=routed, messages=[{"role": "system", "content": "be nice"}] + ) + + assert routed["metadata"]["routing_decision"]["cause"] == "default_fallback" + assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + + @pytest.mark.asyncio + async def test_the_plan_mode_exit_carries_the_ceiling(self): + routed: dict = {"max_tokens": 8192} + deployment = await self._router( + extra_config={"plan_mode_min_tier": "REASONING"} + ).async_get_available_deployment( + model="smart-router", + request_kwargs=routed, + messages=[ + {"role": "user", "content": "plan the refactor"}, + {"role": "system", "content": "Plan mode is active"}, + ], + ) + + assert routed["metadata"]["routing_decision"]["cause"] == "plan_mode" + assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + + @pytest.mark.asyncio + async def test_a_default_model_landing_with_no_tier_still_gets_its_ceiling(self): + strategy = ComplexityRouter( + model_name="smart-router", + litellm_router_instance=self._router(), + complexity_router_config={"tiers": {"SIMPLE": "small"}, "default_model": "big"}, + ) + + assert dict(strategy._litellm_params_for_model(None, "big")) == {"max_tokens": 64000} + + @pytest.mark.asyncio + async def test_a_fallback_into_a_plain_group_gets_the_callers_ceiling_back(self): + """A model-group fallback re-enters routing with the same kwargs; a Sonnet-sized ceiling + must not ride onto the plain group the caller configured as the fallback.""" + big: dict = { + "model_name": "big", + "litellm_params": { + "model": "anthropic/claude-sonnet-5", + "api_key": "k", + "mock_response": "litellm.InternalServerError", + }, + "model_info": {"max_output_tokens": 64000}, + } + plain: dict = { + "model_name": "plain", + "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k", "mock_response": "ok"}, + } + router = Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "big", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"} + }, + }, + }, + big, + plain, + ], + fallbacks=[{"smart-router": ["plain"]}], + num_retries=0, + ) + recorder = _OutputCeilingRecorder() + litellm.callbacks.append(recorder) + try: + await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": self.COMPLEX_PROMPT}], max_tokens=8192 + ) + finally: + litellm.callbacks.remove(recorder) + + assert recorder.seen == [("claude-sonnet-5", 64000), ("claude-haiku-4-5", 8192)] + + @pytest.mark.asyncio + async def test_a_caller_seeded_stamp_cannot_inject_kwargs_on_a_plain_group(self): + """The stamp sits in a metadata bucket a caller can write; a planted one must yield + nothing but integer ceiling carriers, never a redirected api_base or credential.""" + planted: dict = { + "api_base": "https://attacker.example", + "api_key": "stolen", + "max_tokens": "not-an-int", + "max_completion_tokens": True, + "max_output_tokens": 321, + } + routed: dict = {"max_tokens": 8192, "metadata": {"_client_output_ceiling": planted}} + + await self._router().async_get_available_deployment( + model="big", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}] + ) + + assert {k: v for k, v in routed.items() if k not in ("metadata", "model_info")} == {"max_output_tokens": 321} + + @pytest.mark.asyncio + async def test_the_pass_through_routing_entry_point_pins_and_restores_the_same_way(self): + pass_through: dict = {**self.SMALL["litellm_params"], "use_in_pass_through": True} + small: dict = {**self.SMALL, "litellm_params": pass_through} + plain: dict = {**small, "model_name": "plain"} + router = self._router(simple_deployments=[small, plain]) + for deployment in router.model_list: + deployment["litellm_params"]["use_in_pass_through"] = True + routed: dict = {"max_tokens": 8192} + + deployment = await router.async_get_available_deployment_for_pass_through( + model="smart-router", request_kwargs=routed, messages=[{"role": "user", "content": self.COMPLEX_PROMPT}] + ) + pinned = routed["max_tokens"] + await router.async_get_available_deployment_for_pass_through( + model="plain", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}] + ) + + assert (deployment["litellm_params"]["model"], pinned, routed["max_tokens"]) == ( + "anthropic/claude-sonnet-5", + 64000, + 8192, + ) + + @pytest.mark.asyncio + async def test_the_classifier_fallback_exit_carries_the_ceiling(self): + router = self._router( + extra_config={ + "classifier_type": "llm", + "classifier_llm_config": {"model": "no-such-classifier", "timeout_ms": 400}, + "classifier_fallback": "default_model", + "default_model": "big", + } + ) + routed: dict = {"max_tokens": 8192} + + deployment = await router.async_get_available_deployment( + model="smart-router", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}] + ) + + assert routed["metadata"]["routing_decision"]["cause"] == "default_model_fallback" + assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + + @pytest.mark.parametrize( + "value, expected", + [(8192, 8192), ("8192", 8192), (100.9, 100), (0, 0), (-1, None), (True, None), ("x", None), (None, None)], + ) + def test_a_client_cap_is_read_as_an_integer_or_ignored(self, value, expected): + assert as_output_cap(value) == expected + + def test_restoring_the_callers_ceiling_reads_the_stamp_and_replaces_every_carrier(self): + stamped: dict = {"max_output_tokens": 500, "metadata": {"_client_output_ceiling": {"max_tokens": 8192}}} + Router._restore_client_ceiling_no_tier_pins(stamped) + assert {k: v for k, v in stamped.items() if k != "metadata"} == {"max_tokens": 8192} + + coerced: dict = { + "max_tokens": 64000, + "metadata": {"_client_output_ceiling": {"max_tokens": "8192", "max_completion_tokens": 100.0}}, + } + Router._restore_client_ceiling_no_tier_pins(coerced) + assert {k: v for k, v in coerced.items() if k != "metadata"} == { + "max_tokens": 8192, + "max_completion_tokens": 100, + } + + unstamped: dict = {"max_tokens": 64000, "metadata": {}} + Router._restore_client_ceiling_no_tier_pins(unstamped) + assert unstamped["max_tokens"] == 64000 + + @pytest.mark.asyncio + async def test_pinning_stamps_the_callers_carriers_once(self): + router = self._router() + request_kwargs: dict = {"max_completion_tokens": 8192} + + first = router._pin_tier_params_onto_request( + model="big", tier_litellm_params={"max_tokens": 64000}, request_kwargs=request_kwargs, responses_call=False + ) + second = router._pin_tier_params_onto_request( + model="big", tier_litellm_params={"max_tokens": 32000}, request_kwargs=request_kwargs, responses_call=False + ) + none = router._pin_tier_params_onto_request( + model="big", tier_litellm_params=None, request_kwargs=request_kwargs, responses_call=False + ) + + assert (first, second, none) == (True, True, False) + assert request_kwargs["max_tokens"] == 32000 + assert request_kwargs["metadata"]["_client_output_ceiling"] == {"max_completion_tokens": 8192} + + +class _OutputCeilingRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen: list[tuple[str, int | None]] = [] + + def log_pre_api_call(self, model, messages, kwargs): + self.seen.append((model, kwargs.get("optional_params", {}).get("max_tokens"))) + + +NON_REASONING_TIERS: Final = { + "NON_REASONING": "gpt-4o-mini", + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + "COMPLEX": "claude-sonnet-4-20250514", + "REASONING": "o1-preview", +} + + +class TestNonReasoningTier: + """The opt-in fifth built-in tier below SIMPLE: inert unless enabled, reachable when it is.""" + + @staticmethod + def _router(mock_router_instance, **overrides) -> ComplexityRouter: + config: Final = { + "tiers": dict(NON_REASONING_TIERS), + "enable_non_reasoning_tier": True, + "classifier_type": "llm", + "classifier_llm_config": {"model": "haiku-classifier"}, + **overrides, + } + return ComplexityRouter( + model_name="test-non-reasoning-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + + def test_ladder_gains_a_rung_below_simple_only_when_enabled(self): + """Tier 0 sits at the bottom; anywhere else and escalation and the baseline shift.""" + enabled: Final = ComplexityRouterConfig( + tiers=dict(NON_REASONING_TIERS), + enable_non_reasoning_tier=True, + classifier_type="llm", + classifier_llm_config={"model": "clf"}, + ) + assert enabled.tier_names() == ("NON_REASONING", "SIMPLE", "MEDIUM", "COMPLEX", "REASONING") + assert ComplexityRouterConfig().tier_names() == ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING") + + def test_default_router_is_unchanged_by_the_tier_existing(self): + """The enum grew a member, and nothing a four-tier router sends or resolves may change.""" + default: Final = ComplexityRouterConfig() + assert default.enable_non_reasoning_tier is False + assert "NON_REASONING" not in DEFAULT_COMPLEXITY_CONFIG.tiers + assert default.classifier_wire_labels() == ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING") + assert default.labeled_tiers() == TIER_SEVERITY_ORDER_LABELED + assert default.resolve_classified_tier("NON_REASONING") is None + + @pytest.mark.parametrize("preset", tuple(ClassificationRubric)) + def test_rubric_gains_the_bullet_only_when_enabled(self, preset): + """An unset toggle leaves every shipped rubric byte-identical; an enabled one adds a bullet.""" + enabled: Final = ComplexityRouterConfig( + tiers=dict(NON_REASONING_TIERS), + enable_non_reasoning_tier=True, + classifier_type="llm", + classifier_llm_config={"model": "clf"}, + ) + on: Final = classification_system_prompt(3, None, enabled.labeled_tiers(), preset) + off: Final = classification_system_prompt(3, None, ComplexityRouterConfig().labeled_tiers(), preset) + assert "- NON_REASONING:" in on + assert "- NON_REASONING" not in off + + def test_enabled_router_puts_the_tier_on_the_classifier_wire(self, mock_router_instance): + """The schema enum bounds what the classifier may return, whatever the rubric says.""" + router: Final = self._router(mock_router_instance) + enum: Final = router._classifier_response_format["json_schema"]["schema"]["properties"]["tier"]["enum"] + assert enum == ["NON_REASONING", "SIMPLE", "MEDIUM", "COMPLEX", "REASONING"] + + @pytest.mark.asyncio + async def test_classifier_verdict_routes_to_the_tier_model(self, mock_router_instance): + """The classifier names the tier and the request lands on that tier's model.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "NON_REASONING"}')) + router: Final = self._router( + mock_router_instance, tiers={**NON_REASONING_TIERS, "NON_REASONING": "cheap-relay"} + ) + response = await router.async_pre_routing_hook( + model="test-non-reasoning-router", + request_kwargs={}, + messages=[{"role": "user", "content": "here is the file, pass it along"}], + ) + assert response.model == "cheap-relay" + assert response.routing_decision["tier"] == "NON_REASONING" + assert response.routing_decision["cause"] == "llm_classifier" + + @pytest.mark.asyncio + async def test_a_four_tier_router_ignores_a_non_reasoning_verdict( + self, llm_complexity_router, mock_router_instance + ): + """Naming the tier at a router that never opted in falls back instead of routing there.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "NON_REASONING"}')) + outcome = await llm_complexity_router.aclassify("relay this") + assert outcome.tier != ComplexityTier.NON_REASONING + assert outcome.cause != "llm_classifier" + + def test_escalation_walks_up_off_the_tier(self, mock_router_instance): + """Escalation is a built-in-ladder feature and the issue asks for it from the new tier.""" + router: Final = self._router(mock_router_instance) + assert router._escalate_tier(ComplexityTier.NON_REASONING) == ComplexityTier.SIMPLE + assert router._escalate_tier(ComplexityTier.REASONING) == ComplexityTier.REASONING + + def test_escalation_skips_the_tier_when_unconfigured(self, mock_router_instance): + """SIMPLE still escalates to MEDIUM, so escalation never routes below the caller's model.""" + router: Final = self._router( + mock_router_instance, + tiers={"NON_REASONING": "cheap-relay", "SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + ) + assert router._escalate_tier(ComplexityTier.SIMPLE) == ComplexityTier.MEDIUM + + def test_tier_zero_is_never_the_savings_baseline(self, mock_router_instance): + """Savings use the hardest configured tier; tier 0 winning would invert every figure.""" + assert self._router(mock_router_instance)._hardest_tier_models() == ("o1-preview",) + cheap_only: Final = self._router( + mock_router_instance, tiers={"NON_REASONING": "cheap-relay", "SIMPLE": "gpt-4o-mini"} + ) + assert cheap_only._hardest_tier_models() == ("gpt-4o-mini",) + + def test_the_tier_gets_its_own_display_label(self, mock_router_instance): + """tier_labels covers the built-in tiers, so the new rung must be renameable like the rest.""" + router: Final = self._router(mock_router_instance, tier_labels={"NON_REASONING": "Relay"}) + assert router.config.classifier_wire_labels()[0] == "Relay" + assert router.config.resolve_classified_tier("relay") == ComplexityTier.NON_REASONING + + @pytest.mark.parametrize( + "overrides, expected", + ( + ({"classifier_type": "heuristic", "classifier_llm_config": None}, "requires classifier_type"), + ({"classifier_type": "heuristic_v2", "classifier_llm_config": None}, "requires classifier_type"), + ({"tiers": {"SIMPLE": "a", "MEDIUM": "b"}}, "at least one model"), + ), + ids=["heuristic", "heuristic_v2", "no_model"], + ) + def test_unreachable_or_unroutable_configs_are_rejected(self, overrides, expected): + """Refused where it could do nothing: no scorer emits the tier, no pool routes it.""" + config: Final = { + "tiers": dict(NON_REASONING_TIERS), + "enable_non_reasoning_tier": True, + "classifier_type": "llm", + "classifier_llm_config": {"model": "clf"}, + **overrides, + } + with pytest.raises(ValidationError, match=expected): + ComplexityRouterConfig.model_validate(config) + + def test_the_tier_cannot_be_configured_without_the_toggle(self): + """Silently ignoring the key would leave an operator paying for a pool nothing routes to.""" + with pytest.raises(ValidationError, match="no request can route there"): + ComplexityRouterConfig(tiers={"NON_REASONING": "cheap", "SIMPLE": "a"}) + + def test_the_toggle_is_refused_alongside_a_custom_tier_set(self): + """A custom tier set replaces the built-in ladder, so both at once has no meaning.""" + with pytest.raises(ValidationError, match="cannot be combined with tier_definitions"): + ComplexityRouterConfig( + enable_non_reasoning_tier=True, + classifier_type="llm", + classifier_llm_config={"model": "clf"}, + tier_definitions=({"name": "lo", "description": "d"}, {"name": "hi", "description": "d"}), + tiers={"lo": "a", "hi": "b"}, + fallback_tier="lo", + ) + + def test_heuristic_v2_predictions_never_reach_the_new_tier(self, mock_router_instance): + """The four-class artifact's 1-based index must keep mapping onto SIMPLE..REASONING.""" + router: Final = ComplexityRouter( + model_name="v2-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {k: v for k, v in NON_REASONING_TIERS.items() if k != "NON_REASONING"}, + "classifier_type": "heuristic_v2", + }, + ) + outcome = router._classify_with_heuristic_v2("implement a distributed rate limiter under concurrency") + assert outcome.tier in TIER_SEVERITY_ORDER + assert tuple(signal.split(":")[1].split("=")[0] for signal in outcome.signals[1:]) == ( + "simple", + "medium", + "complex", + "reasoning", + ) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py new file mode 100644 index 00000000000..c2fa41f4ca8 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -0,0 +1,210 @@ +import logging +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler + +GROUP: Final = "least-busy-group" +DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} +DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} +HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] + + +def _call_kwargs(deployment_id: str) -> dict[str, object]: + return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} + + +class SharedRedisCounters: + """Mirrors what Redis gives the handler: increments clamped at zero, a TTL set once when + the key is created, and ordered reads that raise rather than invent a value.""" + + def __init__(self) -> None: + self.counts: dict[str, int] = {} + self.ttls: dict[str, int] = {} + + def count(self, key: str) -> int | None: + return self.counts.get(key) + + def expire(self, key: str) -> None: + self.counts.pop(key, None) + self.ttls.pop(key, None) + + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + incremented: Final = max(0, self.counts.get(key, 0) + value) + self.counts[key] = incremented + self.ttls.setdefault(key, ttl) + return incremented + + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + return self.increment_with_floor(key, value, ttl) + + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return tuple(self.counts.get(key) for key in key_list) + + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return self.batch_get_counts(key_list) + + +def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: + cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double + return LeastBusyLoggingHandler(router_cache=cache) + + +@pytest.mark.asyncio +async def test_worker_routes_around_a_request_another_worker_started() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + await picking_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await streaming_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_sync_pick_reads_the_shared_counts() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + picking_worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + streaming_worker.log_failure_event(_call_kwargs("dep-a"), None, None, None) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: + """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that + stuck count an hour after the key was created, which only works while nothing writes the TTL + again: a handler that refreshed it on every touch would keep the count alive for as long as + the group takes traffic, and the deployment would read busier than it is forever.""" + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + key: Final = f"{GROUP}_request_count:dep-a" + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + + shared.ttls[key] = 5 + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(key) == 1 + assert shared.ttls == {key: 5} + + +@pytest.mark.asyncio +async def test_counts_stay_in_memory_without_redis() -> None: + worker: Final = _worker(None) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + +class UnavailableRedis(SharedRedisCounters): + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + raise ConnectionError("redis is down") + + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + raise ConnectionError("redis is down") + + +@pytest.mark.asyncio +async def test_a_redis_outage_falls_back_to_this_workers_own_counts() -> None: + worker: Final = _worker(UnavailableRedis()) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + shared.expire(f"{GROUP}_request_count:dep-a") + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 1 + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +@pytest.mark.asyncio +async def test_a_local_counter_that_expired_mid_request_cannot_go_negative() -> None: + worker: Final = _worker(None) + in_memory: Final = worker.router_cache.in_memory_cache + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + in_memory.delete_cache(f"{GROUP}_request_count:dep-a") + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +def test_calls_without_a_deployment_are_ignored() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}}) + worker.log_pre_api_call(model="m", messages=[], kwargs={}) + + assert shared.counts == {} + + +class OpenBreakerRedis(SharedRedisCounters): + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") + + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") + + +@pytest.mark.asyncio +async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(caplog: pytest.LogCaptureFixture) -> None: + worker: Final = _worker(OpenBreakerRedis()) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + picked: Final = worker.get_available_deployments(GROUP, HEALTHY) + + assert picked is DEPLOYMENT_B + assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] + assert sum("circuit breaker is open" in record.getMessage() for record in caplog.records) == 2 diff --git a/tests/test_litellm/router_strategy/test_lowest_cost.py b/tests/test_litellm/router_strategy/test_lowest_cost.py new file mode 100644 index 00000000000..ab3ef099410 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_lowest_cost.py @@ -0,0 +1,103 @@ +import copy +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" +COST_KEY = "cost_map:gpt-5.5-pool" +LATENCY_KEYS = ("gpt-5.5-pool_map", "gpt-5.5-pool_cost_map") +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=COST_KEY) 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 +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +async def test_log_success_event_keeps_cost_bookkeeping_out_of_the_latency_routing_entry(use_async: bool): + cache = DualCache() + latency_entry = {DEPLOYMENT_ID: {"latency": [0.5], "time_to_first_token": [0.1]}} + for latency_key in LATENCY_KEYS: + cache.set_cache(key=latency_key, value=copy.deepcopy(latency_entry)) + handler = LowestCostLoggingHandler(router_cache=cache) + call_args = { + "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), + } + + if use_async: + await handler.async_log_success_event(**call_args) + else: + handler.log_success_event(**call_args) + + assert [cache.get_cache(key=latency_key) for latency_key in LATENCY_KEYS] == [latency_entry, latency_entry] + assert _recorded_minute_counters(cache) == {"tpm": 12, "rpm": 1} + + +@pytest.mark.asyncio +async def test_async_get_available_deployments_applies_rpm_limit_from_the_cost_entry(): + cache = DualCache() + handler = LowestCostLoggingHandler(router_cache=cache) + precise_minute = datetime.now().strftime("%Y-%m-%d-%H-%M") + cache.set_cache(key=COST_KEY, value={DEPLOYMENT_ID: {precise_minute: {"tpm": 12, "rpm": 1}}}) + healthy_deployments = [{"model_info": {"id": DEPLOYMENT_ID}, "litellm_params": {"model": "gpt-5.5", "rpm": 1}}] + + picked = await handler.async_get_available_deployments( + model_group="gpt-5.5-pool", + healthy_deployments=healthy_deployments, + messages=[{"role": "user", "content": "hi"}], + ) + + assert picked is None + + +@pytest.mark.asyncio +async def test_async_log_success_event_counts_a_response_with_no_completion_tokens(): + cache = DualCache() + handler = LowestCostLoggingHandler(router_cache=cache) + + await handler.async_log_success_event( + kwargs=KWARGS, + response_obj=_chat_response_with_no_completion_tokens(), + start_time=datetime(2026, 1, 1, 12, 0, 0), + end_time=datetime(2026, 1, 1, 12, 0, 2), + ) + + assert _recorded_minute_counters(cache) == {"tpm": 12, "rpm": 1} diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/test_litellm/router_strategy/test_lowest_latency.py index 6701f4a7aa2..1a8614e3fca 100644 --- a/tests/test_litellm/router_strategy/test_lowest_latency.py +++ b/tests/test_litellm/router_strategy/test_lowest_latency.py @@ -8,11 +8,12 @@ import json from datetime import datetime, timedelta import pytest - +from pydantic import ValidationError import litellm from litellm.caching.caching import DualCache -from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler +from litellm.router import Router +from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler, RoutingArgs DEPLOYMENT_ID = "9876" KWARGS = { @@ -58,9 +59,9 @@ def test_sync_embedding_latency_is_json_serializable(): latencies = _recorded_latencies(cache) assert latencies, "expected a latency entry to be recorded" - assert all( - not isinstance(value, timedelta) for value in latencies - ), f"raw timedelta leaked into latency list: {latencies}" + assert all(not isinstance(value, timedelta) for value in latencies), ( + f"raw timedelta leaked into latency list: {latencies}" + ) assert latencies[-1] == pytest.approx(2.0) # the exact failure mode from production: redis cache sync json.dumps json.dumps({"latency": latencies}) @@ -84,9 +85,9 @@ async def test_async_embedding_latency_is_json_serializable(): latencies = _recorded_latencies(cache) assert latencies, "expected a latency entry to be recorded" - assert all( - not isinstance(value, timedelta) for value in latencies - ), f"raw timedelta leaked into latency list: {latencies}" + assert all(not isinstance(value, timedelta) for value in latencies), ( + f"raw timedelta leaked into latency list: {latencies}" + ) assert latencies[-1] == pytest.approx(3.0) json.dumps({"latency": latencies}) @@ -163,3 +164,315 @@ 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}) + + +MODEL_GROUP = "gpt-4o-mini" +FAST_TTFT_ID = "fast-ttft-short-output" +SLOW_TTFT_ID = "slow-ttft-long-output" +STREAMING_DEPLOYMENTS = [ + {"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {}}, + {"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {}}, +] + + +def _streaming_kwargs(deployment_id: str, start_time: datetime, ttft_seconds: float): + return { + "litellm_params": { + "metadata": {"model_group": MODEL_GROUP}, + "model_info": {"id": deployment_id}, + }, + "stream": True, + "completion_start_time": start_time + timedelta(seconds=ttft_seconds), + } + + +def _recorded_ttft(cache: DualCache, deployment_id: str): + cached = cache.get_cache(key=f"{MODEL_GROUP}_map") or {} + return cached.get(deployment_id, {}).get("time_to_first_token_seconds", []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_streaming_ttft_ranking_ignores_completion_length(sync_mode: bool): + """Deployment A: TTFT 1s, 50 completion tokens. Deployment B: TTFT 3s, 500 + completion tokens. Dividing TTFT by completion tokens made B look faster + (3/500 = 0.006 beats 1/50 = 0.02); actual TTFT must win.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + start_time = datetime(2026, 1, 1, 12, 0, 0) + end_time = start_time + timedelta(seconds=10) + + samples = ( + (FAST_TTFT_ID, 1.0, 50), + (SLOW_TTFT_ID, 3.0, 500), + ) + for deployment_id, ttft, completion_tokens in samples: + kwargs = _streaming_kwargs(deployment_id, start_time, ttft) + response_obj = _chat_response(completion_tokens=completion_tokens) + if sync_mode: + handler.log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time + ) + else: + await handler.async_log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time + ) + + assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(1.0)] + assert _recorded_ttft(cache, SLOW_TTFT_ID) == [pytest.approx(3.0)] + + request_kwargs = {"stream": True, "metadata": {}} + if sync_mode: + picked = handler.get_available_deployments( + model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs + ) + else: + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_ttft_window_keeps_newest_samples_when_full(sync_mode: bool): + """Float timestamps, as the SDK passes them. Once max_latency_list_size + samples exist the oldest TTFT is dropped so the window slides.""" + max_size = 3 + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": max_size}) + start_time = 1_700_000_000.0 + ttfts = (0.1, 0.2, 0.3, 0.4) + + for ttft in ttfts: + kwargs = { + "litellm_params": { + "metadata": {"model_group": MODEL_GROUP}, + "model_info": {"id": FAST_TTFT_ID}, + }, + "stream": True, + "completion_start_time": start_time + ttft, + } + response_obj = _chat_response(completion_tokens=1) + if sync_mode: + handler.log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0 + ) + else: + await handler.async_log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0 + ) + + assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(ttft) for ttft in ttfts[-max_size:]] + + +@pytest.mark.asyncio +async def test_streaming_routing_ignores_per_token_ttft_samples_from_older_workers(): + """Workers on the previous release share the Redis map and keep writing + seconds-per-token under the old "time_to_first_token" key during a rolling + deploy. Those samples favor SLOW; routing must only read the seconds key.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token": [0.02], "time_to_first_token_seconds": [1.0]}, + SLOW_TTFT_ID: {"time_to_first_token": [0.006], "time_to_first_token_seconds": [3.0]}, + }, + ) + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("ttft_percentile", "first_samples", "second_samples", "expected_id"), + [ + (None, [0.1, 0.1, 1.0], [0.3, 0.3, 0.3], SLOW_TTFT_ID), + (0.5, [0.1, 0.1, 1.0], [0.3, 0.3, 0.3], FAST_TTFT_ID), + (0.9, [0.1, 0.1, 0.1, 0.1, 1.5], [0.3, 0.3, 0.3, 0.3, 0.3], SLOW_TTFT_ID), + ], + ids=["default_average", "p50", "p90"], +) +async def test_streaming_ttft_ranking_percentile( + sync_mode: bool, + ttft_percentile: float | None, + first_samples: list[float], + second_samples: list[float], + expected_id: str, +): + cache = DualCache() + routing_args = {} if ttft_percentile is None else {"ttft_percentile": ttft_percentile} + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args=routing_args) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token_seconds": first_samples}, + SLOW_TTFT_ID: {"time_to_first_token_seconds": second_samples}, + }, + ) + + if sync_mode: + picked = handler.get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + else: + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == expected_id + + +@pytest.mark.parametrize("ttft_percentile", [0, -0.1, 1.1]) +def test_ttft_percentile_validation(ttft_percentile: float): + with pytest.raises(ValidationError): + RoutingArgs(ttft_percentile=ttft_percentile) + + +@pytest.mark.parametrize("ttft_percentile", [0.5, 0.9, 0.95, 1.0]) +def test_ttft_percentile_accepts_valid_values(ttft_percentile: float): + assert RoutingArgs(ttft_percentile=ttft_percentile).ttft_percentile == ttft_percentile + + +@pytest.mark.asyncio +async def test_ttft_percentile_does_not_change_non_streaming_routing(): + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"ttft_percentile": 0.9}) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"latency": [1.0], "time_to_first_token_seconds": [0.1]}, + SLOW_TTFT_ID: {"latency": [0.2], "time_to_first_token_seconds": [1.5]}, + }, + ) + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": False, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == SLOW_TTFT_ID + + +@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 + + +def _latency_router(routing_strategy_args: dict) -> Router: + return Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": {"model": f"openai/{MODEL_GROUP}", "api_key": "sk-fake"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in (FAST_TTFT_ID, SLOW_TTFT_ID) + ], + routing_strategy="latency-based-routing", + routing_strategy_args=routing_strategy_args, + ) + + +def _seed_streaming_ttft(router: Router) -> None: + router.cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token_seconds": [0.1, 0.1, 1.0]}, + SLOW_TTFT_ID: {"time_to_first_token_seconds": [0.3, 0.3, 0.3]}, + }, + ) + + +async def _pick_streaming(router: Router) -> str: + picked = await router.async_get_available_deployment( + model=MODEL_GROUP, + request_kwargs={"stream": True, "metadata": {}}, + ) + return picked["model_info"]["id"] + + +@pytest.mark.asyncio +async def test_runtime_routing_strategy_args_update_applies_ttft_percentile(): + """A config reload that adds ttft_percentile must reach the live selector, + not sit unused until the proxy restarts.""" + router = _latency_router({"max_latency_list_size": 50}) + _seed_streaming_ttft(router) + + assert await _pick_streaming(router) == SLOW_TTFT_ID + + router.update_settings(routing_strategy_args={"max_latency_list_size": 50, "ttft_percentile": 0.5}) + + assert await _pick_streaming(router) == FAST_TTFT_ID + + +@pytest.mark.asyncio +async def test_runtime_routing_strategy_args_update_keeps_previous_args_when_invalid(): + router = _latency_router({"ttft_percentile": 0.5}) + _seed_streaming_ttft(router) + + router.update_settings(routing_strategy_args={"ttft_percentile": 5}) + + assert await _pick_streaming(router) == FAST_TTFT_ID + + +@pytest.mark.asyncio +async def test_runtime_routing_strategy_args_update_is_a_noop_without_a_selector(): + """simple-shuffle has no selector to re-link, so an args update must leave + the router alone instead of blowing up on a missing selector attribute.""" + router = Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": {"model": f"openai/{MODEL_GROUP}", "api_key": "sk-fake"}, + "model_info": {"id": FAST_TTFT_ID}, + } + ], + routing_strategy="simple-shuffle", + ) + + router.update_settings(routing_strategy_args={"ttl": 5}) + + assert router.routing_strategy_args == {"ttl": 5} + assert await _pick_streaming(router) == FAST_TTFT_ID diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 5599c5aad63..5f37842305d 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm import Router +from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import RoutingGroup, RoutingStrategy @@ -435,6 +436,81 @@ def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeyp assert router._group_selectors == {} +def test_two_least_busy_groups_count_a_request_once(monkeypatch): + """ + Least-busy counts a request up from the pre-call hooks on `litellm.input_callback` and + back down from the success hooks on `litellm.callbacks`. The success list drops a second + selector of the same class, so a pre-call list that kept both counted every request twice + and released it once, and the deployment's in-flight count climbed until it looked pinned. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + router = _build_router( + routing_strategy="least-busy", + routing_groups=[ + { + "group_name": "fast", + "models": ["filtered-model"], + "routing_strategy": "least-busy", + } + ], + ) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + +def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): + """ + Least-busy hangs its counting off litellm's global callback lists, and those lists keep one + logger per class unless the instances differ in a plain attribute. Two routers in one process + (a second Router, or a per-request `user_config` one) therefore have to register separately: + a second router whose selector is dropped counts nothing, reads zero for every deployment, + and sends every request to whichever one is listed first. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + first = _build_router(routing_strategy="least-busy") + second = _build_router(routing_strategy="least-busy") + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 1 + assert ( + second.get_available_deployment(model="filtered-model", messages=[])["model_info"]["id"] + == "deploy-2" + ) + + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert first.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + # --------------------------------------------------------------------------- # Direct helper coverage # --------------------------------------------------------------------------- @@ -716,15 +792,168 @@ def test_strategy_reinit_unregisters_override_selectors(): router = _build_router(routing_strategy="least-busy") override_selector = router._get_override_strategy_selector("latency-based-routing") assert override_selector is not None - assert any(id(cb) == id(override_selector) for cb in litellm.callbacks) + assert not any(cb is override_selector for cb in litellm.callbacks) router.update_settings(routing_strategy="latency-based-routing") assert router._override_selectors == {} - assert not any(id(cb) == id(override_selector) for cb in litellm.callbacks) + assert not any(cb is override_selector for cb in litellm.callbacks) assert router._get_override_strategy_selector("latency-based-routing") is router.lowestlatency_logger +def test_override_selectors_are_not_registered_process_wide(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router = _build_router(routing_strategy="simple-shuffle") + v2_selector = router._get_override_strategy_selector("usage-based-routing-v2") + least_busy_selector = router._get_override_strategy_selector("least-busy") + assert v2_selector is not None and least_busy_selector is not None + + assert litellm.callbacks == [] + assert litellm.input_callback == [] + + +def _rpm_limited_model_list(): + return [ + { + "model_name": "other-model", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-test-3", + "api_base": "https://example.invalid", + "rpm": 1, + }, + "model_info": {"id": "deploy-3"}, + }, + ] + + +async def _mock_completion(router, **override): + return await router.acompletion( + model="other-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok", **override + ) + + +@pytest.mark.asyncio +async def test_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(): + router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) + + await _mock_completion(router, routing_strategy="usage-based-routing-v2") + override_selector = router._override_selectors["usage-based-routing-v2"] + assert not any(cb is override_selector for cb in litellm.callbacks) + + with patch.object( + override_selector, "async_pre_call_check", wraps=override_selector.async_pre_call_check + ) as pre_call_spy: + for _ in range(2): + plain = await _mock_completion(router) + assert plain.choices[0].message.content == "ok" + assert not pre_call_spy.called + + with pytest.raises(litellm.RateLimitError): + await _mock_completion(router, routing_strategy="usage-based-routing-v2") + + +def test_sync_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(): + router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) + messages = [{"role": "user", "content": "hi"}] + + router.completion( + model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2" + ) + override_selector = router._override_selectors["usage-based-routing-v2"] + assert not any(cb is override_selector for cb in litellm.callbacks) + + for _ in range(2): + plain = router.completion(model="other-model", messages=messages, mock_response="ok") + assert plain.choices[0].message.content == "ok" + + with pytest.raises(ValueError, match="No deployments available"): + router.completion( + model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2" + ) + + +@pytest.mark.asyncio +async def test_override_selector_pre_call_check_only_runs_for_override_selectors(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + deployment = _rpm_limited_model_list()[0] + + override_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle") + override_selector = override_router._get_override_strategy_selector("usage-based-routing-v2") + await override_router._async_override_selector_pre_call_check( + "usage-based-routing-v2", override_selector, deployment, None + ) + with pytest.raises(litellm.RateLimitError): + override_router._override_selector_pre_call_check("usage-based-routing-v2", override_selector, deployment) + + default_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="usage-based-routing-v2") + for _ in range(2): + await default_router._async_override_selector_pre_call_check( + "usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment, None + ) + default_router._override_selector_pre_call_check( + "usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment + ) + await default_router._async_override_selector_pre_call_check(None, None, deployment, None) + default_router._override_selector_pre_call_check(None, None, deployment) + + +@pytest.mark.asyncio +async def test_usage_based_v2_override_enforces_rpm_when_a_specific_deployment_is_requested(): + router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) + kwargs = {"model": "deploy-3", "messages": [{"role": "user", "content": "hi"}], "mock_response": "ok"} + + first = await router.acompletion(**kwargs, routing_strategy="usage-based-routing-v2") + assert first.choices[0].message.content == "ok" + with pytest.raises(litellm.RateLimitError): + await router.acompletion(**kwargs, routing_strategy="usage-based-routing-v2") + assert (await router.acompletion(**kwargs)).choices[0].message.content == "ok" + + +def test_sync_usage_based_v2_override_enforces_rpm_when_a_specific_deployment_is_requested(): + router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) + kwargs = {"model": "deploy-3", "messages": [{"role": "user", "content": "hi"}], "mock_response": "ok"} + + first = router.completion(**kwargs, routing_strategy="usage-based-routing-v2") + assert first.choices[0].message.content == "ok" + with pytest.raises(litellm.RateLimitError): + router.completion(**kwargs, routing_strategy="usage-based-routing-v2") + assert router.completion(**kwargs).choices[0].message.content == "ok" + + +def _pass_through_rpm_limited_model_list(): + deployment = _rpm_limited_model_list()[0] + return [{**deployment, "litellm_params": {**deployment["litellm_params"], "use_in_pass_through": True}}] + + +@pytest.mark.asyncio +async def test_async_early_return_paths_run_the_override_pre_call_check(): + router = Router(model_list=_pass_through_rpm_limited_model_list(), routing_strategy="simple-shuffle") + override = {"routing_strategy": "usage-based-routing-v2"} + + pinned = await router.async_get_available_deployment( + model="other-model", request_kwargs={**override, "_encrypted_content_affinity_pinned": True} + ) + assert pinned["model_info"]["id"] == "deploy-3" + with pytest.raises(litellm.RateLimitError): + await router.async_get_available_deployment_for_pass_through(model="deploy-3", request_kwargs=override) + plain = await router.async_get_available_deployment_for_pass_through(model="deploy-3", request_kwargs={}) + assert plain["model_info"]["id"] == "deploy-3" + + +def test_sync_pass_through_specific_deployment_runs_the_override_pre_call_check(): + router = Router(model_list=_pass_through_rpm_limited_model_list(), routing_strategy="simple-shuffle") + override = {"routing_strategy": "usage-based-routing-v2"} + + first = router.get_available_deployment_for_pass_through(model="deploy-3", request_kwargs=override) + assert first["model_info"]["id"] == "deploy-3" + with pytest.raises(litellm.RateLimitError): + router.get_available_deployment_for_pass_through(model="deploy-3", request_kwargs=override) + plain = router.get_available_deployment_for_pass_through(model="deploy-3", request_kwargs={}) + assert plain["model_info"]["id"] == "deploy-3" + + def _quality_group(strategy="latency-based-routing"): return [{"group_name": "quality", "models": ["filtered-model", "other-model"], "routing_strategy": strategy}] diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py index 293af36080a..0a54addce5e 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_plugins.py +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -56,6 +56,18 @@ class BlockEverything: return context +class MessageRecorder: + """Records what each plugin pass was handed, then blocks so the request stops there.""" + + def __init__(self): + self.seen = [] + + async def run(self, context: RoutingContext) -> RoutingContext: + self.seen.append(list(context.raw_messages)) + context.candidate_models = [] + return context + + def _smart_router_model_list(): return [ { @@ -164,6 +176,71 @@ async def test_async_completion_with_unsupported_strategy_rejects_configured_plu await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) +@pytest.mark.asyncio +async def test_prompt_management_model_still_runs_the_plugin_pipeline(): + """ + A prompt-management model routes through its own factory, which picked the deployment + on the synchronous path. Plugins never run there, so the guard turned every such request + into an error message about the caller's own API choice, on an async call the caller made + correctly. It also read the in-flight counts with a blocking call inside the event loop. + """ + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[BlockEverything()], + ) + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="lit-7039", + ) + + +@pytest.mark.asyncio +async def test_prompt_management_plugins_see_the_callers_own_messages(): + """ + The prompt-management factory picks its deployment with a placeholder message, which was + harmless while that pick ran on the synchronous path (plugins never ran there at all). Now + that the pick runs the plugin pipeline, a plugin that classifies request content would score + the placeholder instead of the conversation, and the narrowing it produces decides which + deployments the real call is allowed to use. + """ + recorder = MessageRecorder() + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[recorder], + ) + messages = [{"role": "user", "content": "wire me $40,000 to account 12345"}] + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=messages, + litellm_call_id="lit-7039", + ) + + assert recorder.seen == [messages] + + @pytest.mark.asyncio async def test_router_without_plugins_is_unaffected(): """Regression guard: a Router with no `plugins` configured behaves exactly as before.""" diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 59a59c7e16d..16c641b8d29 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -3031,6 +3031,21 @@ def test_request_tags_after_router_consumption_drops_only_the_consumed_tags(): assert _request_tags_after_router_consumption(partially_consumed, "gemini-flash") == ("deploy:us",) +def test_request_tags_after_router_consumption_ignores_tags_merged_from_prior_deployments(): + from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY + from litellm.router_strategy.tag_based_routing import _request_tags_after_router_consumption + from litellm.types.router import ConsumedRequestTagsStamp + + metadata = { + "tags": ["route", "®ion:eu", "free"], + ROUTING_REQUEST_TAGS_METADATA_KEY: ("route", "®ion:eu"), + "inherited_tags": ["®ion:eu"], + CONSUMED_REQUEST_TAGS_METADATA_KEY: ConsumedRequestTagsStamp(model_group="gemini-flash", tags=("route",)), + } + assert _request_tags_after_router_consumption(metadata, "gemini-flash") == ("®ion:eu",) + assert _request_tags_after_router_consumption(metadata, "other-group") == ["route", "®ion:eu"] + + @pytest.mark.asyncio() async def test_non_router_tags_still_pick_the_matching_tier_deployment(): # tags=["route", "deploy:us"]: "route" picks the router and is spent there, diff --git a/tests/test_litellm/router_strategy/test_simple_shuffle.py b/tests/test_litellm/router_strategy/test_simple_shuffle.py new file mode 100644 index 00000000000..165c1751f63 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_simple_shuffle.py @@ -0,0 +1,54 @@ +from collections import Counter + +import pytest + +from litellm import Router +from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict + +DRAWS = 200 + + +def _deployment(dep_id: str, metric: LiteLLMParamsTypedDict | None = None) -> DeploymentTypedDict: + params: LiteLLMParamsTypedDict = {"model": "gpt-4o", "api_key": "key", "mock_response": f"from {dep_id}"} + return { + "model_name": "test-model", + "litellm_params": {**params, **(metric or {})}, + "model_info": {"id": dep_id}, + } + + +async def _draw_model_ids(router: Router) -> Counter[str]: + counts: Counter[str] = Counter() + for _ in range(DRAWS): + response = await router.acompletion(model="test-model", messages=[{"role": "user", "content": "hi"}]) + counts[response._hidden_params["model_id"]] += 1 + return counts + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metric", [{"weight": 5}, {"rpm": 5}, {"tpm": 5}], ids=["weight", "rpm", "tpm"]) +async def test_weighted_pick_when_only_a_later_deployment_carries_the_metric(metric: LiteLLMParamsTypedDict): + router = Router( + model_list=[_deployment("unweighted"), _deployment("weighted", metric)], + routing_strategy="simple-shuffle", + num_retries=0, + ) + + counts = await _draw_model_ids(router) + + assert counts["weighted"] == DRAWS + assert counts["unweighted"] == 0 + + +@pytest.mark.asyncio +async def test_uniform_pick_when_every_configured_weight_is_zero(): + router = Router( + model_list=[_deployment("unweighted"), _deployment("standby", {"weight": 0})], + routing_strategy="simple-shuffle", + num_retries=0, + ) + + counts = await _draw_model_ids(router) + + assert counts["unweighted"] > 0 + assert counts["standby"] > 0 diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index dac991a41c4..ea8e2eacaa6 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -16,12 +16,10 @@ The mechanism works without any cache and supports two encoding strategies: """ import time -from typing import List, Optional from unittest.mock import AsyncMock, patch import pytest - import litellm from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse @@ -68,9 +66,7 @@ class TestEncryptedItemIdCodec: def test_roundtrip(self): model_id = "deployment-1" original_item_id = "rs_abc123def456" - encoded = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_item_id - ) + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) assert encoded.startswith("encitem_") decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded) assert decoded is not None @@ -81,9 +77,7 @@ class TestEncryptedItemIdCodec: """Decoding must succeed even if base64 padding (=) was stripped in transit.""" model_id = "gpt-5.1-codex-openai-2" original_item_id = "rs_0efb96cb222403210069a01d5d52588196a9dc394ffdb89d00" - encoded = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_item_id - ) + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) # Strip any trailing '=' to simulate what happens in transit stripped = encoded.rstrip("=") decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(stripped) @@ -100,9 +94,7 @@ class TestEncryptedItemIdCodec: """item_id values containing ';' must survive the roundtrip.""" model_id = "deployment-1" original_item_id = "rs_part1;part2;part3" - encoded = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_item_id - ) + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded) assert decoded is not None assert decoded["item_id"] == original_item_id @@ -118,11 +110,7 @@ class TestUpdateEncryptedContentItemIds: {"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"}, ], } - result = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response, model_id - ) - ) + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, model_id) # Plain message item untouched assert result["output"][0]["id"] == "msg_abc" # Reasoning item with encrypted_content gets encoded @@ -133,16 +121,8 @@ class TestUpdateEncryptedContentItemIds: assert decoded["item_id"] == "rs_xyz" def test_no_op_when_model_id_is_none(self): - response = { - "output": [ - {"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"} - ] - } - result = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response, None - ) - ) + response = {"output": [{"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"}]} + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, None) assert result["output"][0]["id"] == "rs_xyz" @@ -151,9 +131,7 @@ class TestEncryptedContentWrapping: """Test wrapping encrypted_content with model_id metadata.""" model_id = "deployment-1" original_content = "gAAAAABpnW_yEYmSNEyOG_original_encrypted_data" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) assert wrapped.startswith("litellm_enc:") assert wrapped != original_content @@ -170,9 +148,7 @@ class TestEncryptedContentWrapping: ( model_id, content, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( - plain_content - ) + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(plain_content) assert model_id is None assert content == plain_content @@ -189,11 +165,7 @@ class TestEncryptedContentWrapping: }, ], } - result = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response, model_id - ) - ) + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, model_id) assert result["output"][0].get("encrypted_content") is None wrapped = result["output"][1]["encrypted_content"] assert wrapped.startswith("litellm_enc:") @@ -210,19 +182,13 @@ class TestRestoreEncryptedContentItemIds: def test_restores_encoded_ids(self): model_id = "deployment-1" original_id = "rs_encrypted_item_456" - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_id - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_id) request_input = [ {"type": "message", "id": "msg_abc123", "role": "assistant"}, {"type": "reasoning", "id": encoded_id, "encrypted_content": "secret"}, ] - restored = ( - ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - request_input - ) - ) + restored = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(request_input) assert restored[0]["id"] == "msg_abc123" assert restored[1]["id"] == original_id @@ -230,33 +196,21 @@ class TestRestoreEncryptedContentItemIds: """Test that wrapped encrypted_content is unwrapped before forwarding.""" model_id = "deployment-1" original_content = "gAAAAABpnW_yEYmSNEyOG_original" - wrapped_content = ( - ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) - ) + wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) request_input = [ {"type": "reasoning", "encrypted_content": wrapped_content}, ] - restored = ( - ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - request_input - ) - ) + restored = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(request_input) assert restored[0]["encrypted_content"] == original_content def test_no_op_for_plain_string_input(self): - result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - "Hello world" - ) + result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input("Hello world") assert result == "Hello world" def test_no_op_for_unencoded_ids(self): request_input = [{"type": "message", "id": "msg_plain"}] - result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - request_input - ) + result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(request_input) assert result[0]["id"] == "msg_plain" @@ -283,9 +237,7 @@ async def test_encrypted_content_affinity_tracks_and_routes(): "id": "msg_abc123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], }, { "type": "reasoning", @@ -347,9 +299,9 @@ async def test_encrypted_content_affinity_tracks_and_routes(): # The response must have rewritten the encrypted item's ID to encoded form encoded_item_id = _extract_encoded_item_id(first_response) - assert encoded_item_id.startswith( - "encitem_" - ), f"Expected output item ID to be rewritten to encitem_... but got {encoded_item_id!r}" + assert encoded_item_id.startswith("encitem_"), ( + f"Expected output item ID to be rewritten to encitem_... but got {encoded_item_id!r}" + ) # Verify the encoded ID decodes back to the correct deployment + original ID decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded_item_id) @@ -371,9 +323,9 @@ async def test_encrypted_content_affinity_tracks_and_routes(): ) second_model_id = second_response._hidden_params["model_id"] - assert ( - second_model_id == first_model_id - ), f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + assert second_model_id == first_model_id, ( + f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + ) @pytest.mark.asyncio @@ -478,9 +430,7 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits(): # Extract encoded item ID from the first response output encoded_item_id = _extract_encoded_item_id(first_response) - assert encoded_item_id.startswith( - "encitem_" - ), f"Expected encitem_... but got {encoded_item_id!r}" + assert encoded_item_id.startswith("encitem_"), f"Expected encitem_... but got {encoded_item_id!r}" # Follow-up with the encoded item ID — should pin to same deployment second_response = await router.aresponses( @@ -628,17 +578,13 @@ async def test_encrypted_content_affinity_with_wrapped_content_no_id(): if hasattr(first_item, "encrypted_content") else first_item.get("encrypted_content") ) - assert wrapped_content.startswith( - "litellm_enc:" - ), f"Expected wrapped content but got {wrapped_content[:50]}..." + assert wrapped_content.startswith("litellm_enc:"), f"Expected wrapped content but got {wrapped_content[:50]}..." # Verify we can extract model_id from wrapped content ( extracted_model_id, _, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( - wrapped_content - ) + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(wrapped_content) assert extracted_model_id == first_model_id # Second request: use wrapped encrypted_content WITHOUT an ID (Codex behavior) @@ -653,9 +599,9 @@ async def test_encrypted_content_affinity_with_wrapped_content_no_id(): ) second_model_id = second_response._hidden_params["model_id"] - assert ( - second_model_id == first_model_id - ), f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + assert second_model_id == first_model_id, ( + f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + ) def test_encrypted_content_wrapping_preserves_original_content(): @@ -664,13 +610,9 @@ def test_encrypted_content_wrapping_preserves_original_content(): This is critical for streaming responses where content must round-trip correctly. """ model_id = "test-deployment-1" - original_encrypted_content = ( - "gAAAAABpnW_yEYmSNEyOG_streaming_test_content_with_special_chars==+/" - ) + original_encrypted_content = "gAAAAABpnW_yEYmSNEyOG_streaming_test_content_with_special_chars==+/" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_encrypted_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_encrypted_content, model_id) assert wrapped.startswith("litellm_enc:") assert wrapped != original_encrypted_content @@ -691,9 +633,7 @@ def test_encrypted_content_wrapping_with_multiple_semicolons(): model_id = "deployment-with-semicolons" original_content = "gAAAAAB;some;content;with;semicolons" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) ( extracted_model_id, @@ -764,9 +704,7 @@ async def test_encrypted_content_affinity_preserves_litellm_metadata_for_respons request_kwargs=request_kwargs, ) - assert ( - request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] is True - ) + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] is True assert request_kwargs["litellm_metadata"]["model_info"] == {"id": "dep-1"} @@ -777,9 +715,7 @@ def test_encrypted_content_wrapping_empty_string(): model_id = "test-deployment" original_content = "" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) assert wrapped.startswith("litellm_enc:") @@ -1132,9 +1068,7 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance(): "api_key": "fake-azure-resource-key-a", } - pydantic_key = EncryptedContentAffinityCheck._encryption_boundary_key( - pydantic_params - ) + pydantic_key = EncryptedContentAffinityCheck._encryption_boundary_key(pydantic_params) plain_key = EncryptedContentAffinityCheck._encryption_boundary_key(plain_params) assert pydantic_key is not None @@ -1161,18 +1095,8 @@ def test_boundary_key_rejects_non_dict_like_inputs(): for bad in (None, [], "not a dict", 42, object()): assert EncryptedContentAffinityCheck._encryption_boundary_key(bad) is None - assert ( - EncryptedContentAffinityCheck._encryption_boundary_key( - {"api_base": "", "api_key": "k"} - ) - is None - ) - assert ( - EncryptedContentAffinityCheck._encryption_boundary_key( - {"api_base": "https://x"} - ) - is None - ) + assert EncryptedContentAffinityCheck._encryption_boundary_key({"api_base": "", "api_key": "k"}) is None + assert EncryptedContentAffinityCheck._encryption_boundary_key({"api_base": "https://x"}) is None # --------------------------------------------------------------------------- @@ -1180,10 +1104,11 @@ def test_boundary_key_rejects_non_dict_like_inputs(): # --------------------------------------------------------------------------- -def _make_originating_mock(api_base: str, api_key: str): +def _make_originating_mock(api_base: str, api_key: str, model_name: str = "gpt-5.4"): from unittest.mock import MagicMock originating = MagicMock() + originating.model_name = model_name originating.litellm_params.model_dump.return_value = { "api_base": api_base, "api_key": api_key, @@ -1192,19 +1117,23 @@ def _make_originating_mock(api_base: str, api_key: str): def _make_router_mock_with_cooldown( - originating, cooldown_entries: Optional[List[tuple]] = None + originating, + cooldown_entries: list[tuple] | None = None, + routed_group_model_ids: list[str] | None = None, ): """ Build a MagicMock router whose ``cooldown_cache.async_get_active_cooldowns`` - returns ``cooldown_entries`` (defaulting to ``[]`` — no active cooldown). + returns ``cooldown_entries`` (defaulting to ``[]`` — no active cooldown), and + whose ``get_candidate_model_ids_for_route`` returns ``routed_group_model_ids`` + (the deployment ids the router resolves for the routed model; defaulting to ``[]`` + — origin absent from the routed group, i.e. a tier change). """ from unittest.mock import AsyncMock, MagicMock mock_router = MagicMock() mock_router.get_deployment.return_value = originating - mock_router.cooldown_cache.async_get_active_cooldowns = AsyncMock( - return_value=list(cooldown_entries or []) - ) + mock_router.cooldown_cache.async_get_active_cooldowns = AsyncMock(return_value=list(cooldown_entries or [])) + mock_router.get_candidate_model_ids_for_route.return_value = frozenset(routed_group_model_ids or []) return mock_router @@ -1235,15 +1164,15 @@ async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_42 }, ) ], + routed_group_model_ids=["deployment-a-cooled", "deployment-b"], ) check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a-cooled", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-cooled", "rs_test") healthy_only_b = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1297,15 +1226,15 @@ async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_fo }, ) ], + routed_group_model_ids=["deployment-a-cooled-429", "deployment-b"], ) check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a-cooled-429", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-cooled-429", "rs_test") healthy_only_b = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1345,15 +1274,16 @@ async def test_affinity_raises_service_unavailable_when_origin_filtered_without_ ) originating = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") - mock_router = _make_router_mock_with_cooldown(originating, cooldown_entries=[]) + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-a-filtered", "deployment-b"] + ) check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a-filtered", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-filtered", "rs_test") healthy_only_b = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1377,15 +1307,18 @@ async def test_affinity_raises_service_unavailable_when_origin_filtered_without_ @pytest.mark.asyncio -async def test_affinity_raises_bad_request_when_origin_removed(): +async def test_affinity_strips_and_dispatches_when_origin_is_unknown_or_removed(): """ - Originating deployment was removed from the router config and no boundary - peer is available. This is permanent (the stale encrypted_content cannot - be honored), so surface a 400 with actionable text. + A removed deployment, or a forged/unknown affinity marker, resolves to no + originating deployment. It is handled like a cross-group origin: the encrypted + reasoning is stripped and the request dispatches with its readable history, + rather than returning a distinguishable error. That uniform handling denies an + authenticated caller a deployment-id existence oracle, an existing cross-group id + and a nonexistent id both strip and proceed, so responses cannot be told apart. + The membership lookup is skipped entirely when the origin is unknown. """ from unittest.mock import MagicMock - from litellm.exceptions import BadRequestError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1394,12 +1327,11 @@ async def test_affinity_raises_bad_request_when_origin_removed(): mock_router.get_deployment.return_value = None check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-removed", "rs_test" - ) - healthy_only_b = [ + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-removed") + routed_pool = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1408,18 +1340,28 @@ async def test_affinity_raises_bad_request_when_origin_removed(): } ] request_kwargs = { - "input": [{"id": encoded_id, "type": "reasoning"}], + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], } - with pytest.raises(BadRequestError) as excinfo: - await check.async_filter_deployments( - model="gpt-5.4", - healthy_deployments=healthy_only_b, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=routed_pool, + messages=None, + request_kwargs=request_kwargs, + ) - assert "deployment-removed" not in str(excinfo.value) + assert result is routed_pool + assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in request_kwargs["input"]) + mock_router.get_candidate_model_ids_for_route.assert_not_called() @pytest.mark.asyncio @@ -1444,9 +1386,7 @@ async def test_affinity_does_not_raise_when_boundary_peer_available(): mock_router.get_deployment.return_value = originating check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a", "rs_test") peer = { "model_info": {"id": "deployment-a-peer"}, "litellm_params": { @@ -1490,9 +1430,7 @@ async def test_model_group_affinity_config_enables_encrypted_content_affinity(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {}, @@ -1536,9 +1474,7 @@ async def test_model_group_affinity_config_does_not_disable_global_encrypted_con }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {}, @@ -1600,15 +1536,9 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen try: callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) assert encrypted_content_callback.enable_global_affinity is False cache_key = DeploymentAffinityCheck.get_affinity_cache_key( @@ -1620,9 +1550,7 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [ { @@ -1643,16 +1571,384 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen ) assert after_deployment_affinity == [deployment_a, deployment_b] - after_encrypted_content_affinity = ( - await encrypted_content_callback.async_filter_deployments( - model=model_group, - healthy_deployments=after_deployment_affinity, - messages=None, - request_kwargs=request_kwargs, - ) + after_encrypted_content_affinity = await encrypted_content_callback.async_filter_deployments( + model=model_group, + healthy_deployments=after_deployment_affinity, + messages=None, + request_kwargs=request_kwargs, ) assert after_encrypted_content_affinity == [deployment_b] assert request_kwargs.get("_encrypted_content_affinity_pinned") is True finally: router.discard() + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_through_the_bridge(): + """ + Claude Code behind /v1/messages replays the encrypted reasoning the bridge packed + into a thinking block's signature (or a redacted block's data). The pin has to be + read from those blocks because the bridge builds the Responses `input` only after + the router has picked a deployment. + """ + check = EncryptedContentAffinityCheck() + deployments = [ + {"model_info": {"id": "openai-org-a"}, "litellm_params": {"model": "openai/gpt-5.1"}}, + {"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5.1"}}, + ] + request_kwargs = {"model": "gpt-5.1"} + + pinned = await check.async_filter_deployments( + model="gpt-5.1", + healthy_deployments=deployments, + messages=_bridge_replayed_anthropic_messages(minted_by="openai-org-b"), + request_kwargs=request_kwargs, + ) + + assert [d["model_info"]["id"] for d in pinned] == ["openai-org-b"] + assert request_kwargs["_encrypted_content_affinity_pinned"] is True + + +def _bridge_replayed_anthropic_messages(minted_by: str) -> list: + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA_turn_one", minted_by) + return [ + {"role": "user", "content": "Solve the zebra puzzle"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + { + "type": "thinking", + "thinking": "The bridge packed this one", + "signature": f"litellm_encrypted_reasoning:{wrapped}", + }, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ], + }, + {"role": "user", "content": "And who drinks water?"}, + ] + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_strips_bridge_reasoning_from_messages_routed_to_another_group(): + """ + The /v1/messages twin of the tier-change case: the routed group holds no deployment + of the org that minted the reasoning, so the bridge-tagged blocks are dropped whole + and the request dispatches to the routed pool. No unsigned thinking block may be left + behind: Anthropic and Bedrock reject a thinking block with a missing signature the + same way they reject a foreign one. + """ + originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["openai-org-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + routed_pool = [{"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5-nano"}}] + messages = _bridge_replayed_anthropic_messages(minted_by="openai-org-a") + assistant_content = messages[1]["content"] + request_kwargs = {"model": "gpt-5.1"} + + result = await check.async_filter_deployments( + model="gpt-simple-tier", + healthy_deployments=routed_pool, + messages=messages, + request_kwargs=request_kwargs, + ) + + assert result is routed_pool + assert "_encrypted_content_affinity_pinned" not in request_kwargs + assert messages[1]["content"] is assistant_content + assert assistant_content == [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ] + assert all(block["signature"] for block in assistant_content if block["type"] == "thinking") + + +class TestStripEncryptedReasoningFromInput: + def test_keeps_summary_and_drops_encrypted_content_and_id(self): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a", "rs_1") + request_input = [ + {"role": "user", "content": "first turn"}, + { + "type": "reasoning", + "id": encoded_id, + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "thought about it"}], + }, + {"type": "reasoning", "id": encoded_id, "encrypted_content": wrapped}, + {"type": "reasoning", "encrypted_content": wrapped, "summary": []}, + {"type": "message", "id": "msg_1", "role": "assistant", "content": "hi"}, + {"role": "user", "content": "second turn"}, + ] + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + assert request_input == [ + {"role": "user", "content": "first turn"}, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "thought about it"}], + }, + {"type": "message", "id": "msg_1", "role": "assistant", "content": "hi"}, + {"role": "user", "content": "second turn"}, + ] + + def test_keeps_string_form_summary_when_stripping(self): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + request_input = [ + {"type": "reasoning", "encrypted_content": wrapped, "summary": "plain string thought"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "content": [{"type": "output_text", "text": "in content"}], + }, + {"type": "reasoning", "encrypted_content": wrapped, "summary": "", "content": []}, + ] + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + assert request_input == [ + {"type": "reasoning", "summary": "plain string thought"}, + {"type": "reasoning", "content": [{"type": "output_text", "text": "in content"}]}, + ] + + def test_leaves_input_untouched_when_no_encrypted_reasoning(self): + request_input = [ + {"role": "user", "content": "first turn"}, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "no blob here"}]}, + {"role": "user", "content": "second turn"}, + ] + before = [dict(item) for item in request_input] + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + assert request_input == before + + +def _cross_group_request_kwargs(): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + return { + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "ZEBRA: why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"type": "message", "role": "assistant", "content": "Rayleigh scattering."}, + {"role": "user", "content": "KIWI: and sunsets?"}, + ], + } + + +@pytest.mark.asyncio +async def test_affinity_strips_encrypted_reasoning_when_routed_to_another_model_group(): + """ + An auto-router tier change (or a model switch with no boundary peer): the + routed pool holds no deployment of the origin's model group. The origin is + healthy, so a 503 would be wrong; the follow-up dispatches to the routed + pool with the origin's encrypted reasoning stripped and its summary kept. + """ + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + routed_pool = [ + { + "model_info": {"id": "deployment-b"}, + "model_name": "gpt-simple-tier", + "litellm_params": { + "api_base": "https://gateway.example/v1", + "api_key": "key-b", + "model": "openai/gpt-5-nano", + }, + } + ] + request_kwargs = _cross_group_request_kwargs() + original_input = request_kwargs["input"] + + result = await check.async_filter_deployments( + model="gpt-simple-tier", + healthy_deployments=routed_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert result is routed_pool + assert "_encrypted_content_affinity_pinned" not in request_kwargs + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] is True + assert request_kwargs["input"] is original_input + assert [item.get("type") or item["role"] for item in original_input] == [ + "user", + "reasoning", + "message", + "user", + ] + assert original_input[1] == { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "scattering"}], + } + assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in original_input) + + +@pytest.mark.asyncio +async def test_affinity_fails_fast_within_the_origins_own_group(): + """ + Negative class for the tier-change discriminator: the routed group IS the + origin's group (a same-group cooldown, not a tier change), so even with a + healthy non-origin sibling that cannot decrypt the content, the request + still fails fast and the encrypted reasoning is left intact rather than + stripped. Preserves the LIT-3051 cooldown contract. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock( + "https://account-a.openai.azure.com/", "key-a", model_name="gpt-reasoning-tier" + ) + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-a", "deployment-a2"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + sibling_pool = [ + { + "model_info": {"id": "deployment-a2"}, + "model_name": "gpt-reasoning-tier", + "litellm_params": { + "api_base": "https://account-a2.openai.azure.com/", + "api_key": "key-a2", + "model": "azure/gpt-5.4", + }, + } + ] + request_kwargs = _cross_group_request_kwargs() + + with pytest.raises(ServiceUnavailableError): + await check.async_filter_deployments( + model="gpt-reasoning-tier", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert request_kwargs["input"][1].get("encrypted_content") + + +@pytest.mark.asyncio +async def test_affinity_does_not_strip_when_group_is_spelled_differently_but_same_by_id(): + """ + The discriminator must key on deployment-id membership, not on the model-group + name string. Here the origin's configured group is spelled ``openai/gpt-5.4-mini`` + while the routed group is the canonical ``gpt-5.4-mini``: same group, different + spelling. A name compare (``originating.model_name != model``) would read this as + a tier change and strip the reasoning it did not have to. Because the origin's id + is a member of the routed group, this is a same-group cooldown instead: the request + fails fast and the encrypted reasoning is left intact. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock(None, "key-a", model_name="openai/gpt-5.4-mini") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-mini-a", "deployment-mini-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-mini-a") + sibling_pool = [ + { + "model_info": {"id": "deployment-mini-b"}, + "model_name": "gpt-5.4-mini", + "litellm_params": { + "api_base": "https://gateway.example/v1", + "api_key": "key-b", + "model": "openai/gpt-5.4-mini", + }, + } + ] + request_kwargs = { + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], + } + + with pytest.raises(ServiceUnavailableError): + await check.async_filter_deployments( + model="gpt-5.4-mini", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert request_kwargs["input"][1].get("encrypted_content") + + +@pytest.mark.asyncio +async def test_affinity_honors_router_candidate_ids_for_team_and_pattern_routes(): + """ + The exact `model_name` index does not include team-public or pattern routes, so a + same-group cooldown reached only through one of those would be misread as a tier change + and stripped. The check asks the router for the candidate ids it resolves for the route + (`get_candidate_model_ids_for_route`), which covers those paths, rather than the bare + index. Here that set marks the origin as a candidate, so the request fails fast with its + reasoning intact, and the routed group and team are passed through to the router. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock(None, "key-a", model_name="model_name_teamA_uuid") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-team-a", "deployment-team-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-team-a") + sibling_pool = [ + { + "model_info": {"id": "deployment-team-b"}, + "model_name": "team-public-model", + "litellm_params": { + "api_base": "https://gateway.example/v1", + "api_key": "key-b", + "model": "openai/gpt-5.4-mini", + }, + } + ] + request_kwargs = { + "litellm_metadata": {"user_api_key_team_id": "teamA"}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], + } + + with pytest.raises(ServiceUnavailableError): + await check.async_filter_deployments( + model="team-public-model", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert request_kwargs["input"][1].get("encrypted_content") + mock_router.get_candidate_model_ids_for_route.assert_called_once_with(model="team-public-model", team_id="teamA") diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 79ae00e155c..333e7b2ff31 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -332,6 +332,67 @@ async def test_per_request_enable_prompt_caching_reaches_the_affinity_key(monkey assert filtered == [deployments[1]] +@pytest.mark.asyncio +async def test_claude_code_one_shot_subagent_does_not_reuse_an_auto_injected_affinity_key(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + messages = cast(List[AllMessageValues], [{"role": "user", "content": "unique " * 3000}]) + request_kwargs = { + "system": [ + { + "type": "text", + "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=true;", + } + ], + "proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.263 (external, cli)"}}, + } + auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections( + messages=messages, + models=(AUTO_CACHING_MODEL,), + ) + assert auto_injected_messages != messages + await PromptCachingCache(cache=cache).async_add_model_id( + model_id="dep-2", messages=auto_injected_messages, tools=None + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs=request_kwargs, + ) + + assert filtered == deployments + + +@pytest.mark.asyncio +async def test_root_cache_control_does_not_reuse_an_auto_injected_affinity_key(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + messages = _auto_caching_messages() + auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections( + messages=messages, + models=(AUTO_CACHING_MODEL,), + ) + assert auto_injected_messages != messages + await PromptCachingCache(cache=cache).async_add_model_id( + model_id="dep-2", messages=auto_injected_messages, tools=None + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs={"cache_control": {"type": "ephemeral"}}, + ) + + assert filtered == deployments + + @pytest.mark.asyncio async def test_tool_marked_cache_control_keeps_routing_off_another_requests_prefix(monkeypatch, local_model_cost_map): """ @@ -416,3 +477,65 @@ async def test_wildcard_route_resolves_underlying_model_minimum(local_model_cost assert deployments[0]["litellm_params"]["model"] == "anthropic/claude-opus-4-6" assert _get_min_token_count_for_deployments(deployments) == 4096 + + +@pytest.mark.asyncio +async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("anthropic/claude-fable-5") + check = PromptCachingDeploymentCheck(cache=DualCache()) + deployments = _deployments("anthropic/claude-fable-5") + messages = cast(List[AllMessageValues], [{"role": "user", "content": text * 100}]) + + result, took, lags = await timed_with_loop_lags( + lambda: check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=messages + ) + ) + + assert result == deployments + assert_loop_stayed_free(took, lags) + + +@pytest.mark.asyncio +async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("anthropic/claude-fable-5") + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": [{"type": "text", "text": text * 100, "cache_control": {"type": "ephemeral"}}]}], + ) + standard_logging_object = { + "call_type": "acompletion", + "model": "anthropic/claude-fable-5", + "messages": messages, + "model_id": "dep-1", + } + + _, took, lags = await timed_with_loop_lags( + lambda: check.async_log_success_event( + kwargs={"standard_logging_object": standard_logging_object}, + response_obj=None, + start_time=None, + end_time=None, + ) + ) + + assert await PromptCachingCache(cache=cache).async_get_model_id(messages=messages, tools=None) == { + "model_id": "dep-1" + } + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py new file mode 100644 index 00000000000..75c115cd3ab --- /dev/null +++ b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py @@ -0,0 +1,308 @@ +"""Behavior pins for the baseline-relative heuristic-v1 tuning gate.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig +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"} +_KEYWORD_DIMENSION: Final = {"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]} +_HISTORICAL_FINGERPRINTS: Final = ( + ({}, "44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a"), + ( + {"custom_dimensions": [_KEYWORD_DIMENSION]}, + "b5c3c3f3be6341a8a16148d68d9067e03f94ed01a0bbfcde7955763042744372", + ), + ( + {"custom_dimensions": [{"name": "sqlDdl", "weight": 0.4, "patterns": [r"\bCREATE\s{1,4}TABLE\b"]}]}, + "814ce0017fc7f60a160b262f658d910e9bdf784e6139a4ba4f1e2657aa203950", + ), + ( + { + "tiers": _TIERS, + "dimension_weights": {"codePresence": 0.3}, + "custom_dimensions": [ + { + "name": "internalFrameworks", + "weight": 0.2, + "keywords": ["orbitmesh", "fluxgate"], + "patterns": [r"\bALTER\s{1,4}TABLE\b"], + } + ], + }, + "38970dc9224e265ab38c89674563d8d0537822591f9239b45251db6f5ca6cc39", + ), +) + + +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}, + "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}], + "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_explicit_empty_tier_model_configs_follow_omission(self) -> None: + assert tuning_fingerprint({"tier_model_configs": {}}) == DEFAULT_TUNING_FINGERPRINT + + @pytest.mark.parametrize(("config", "fingerprint"), _HISTORICAL_FINGERPRINTS) + def test_fingerprints_recorded_before_scoring_mode_existed_are_preserved( + self, config: Mapping[str, object], fingerprint: str + ) -> None: + """Literal hashes captured from the merged implementation at 9bc9104102, before CustomDimension.scoring_mode.""" + assert tuning_fingerprint(config) == fingerprint + + def test_binary_scoring_mode_hashes_like_its_absence(self) -> None: + historical: Final = tuning_fingerprint({"custom_dimensions": [_KEYWORD_DIMENSION]}) + explicit: Final = tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "binary"}]}) + reserialized: Final = ComplexityRouterConfig.model_validate( + {"custom_dimensions": [_KEYWORD_DIMENSION]} + ).model_dump(mode="json", include={"custom_dimensions"}) + assert reserialized["custom_dimensions"][0]["scoring_mode"] == "binary" + assert reserialized["custom_dimensions"][0]["patterns"] == [] + assert historical == explicit == tuning_fingerprint(reserialized) + assert ( + tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]}) + != historical + ) + + 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 + + def test_a_release_changing_a_shipped_default_is_not_an_operator_edit(self, monkeypatch) -> None: + """An omitted setting follows the shipped default and stays off the quota when that default moves; + only what the operator wrote is fingerprinted, so an explicit value equal to the old default still counts.""" + import litellm.router_strategy.complexity_router.config as config_module + + untouched = _router("a", {}) + tiers_only = _router("b", {"tiers": _TIERS}) + pinned = _router("c", {"dimension_weights": dict(config_module.DEFAULT_DIMENSION_WEIGHTS)}) + baselines = snapshot_tuning_baselines([untouched, tiers_only, pinned]) + assert tuning_fingerprint(pinned["litellm_params"]["complexity_router_config"]) != DEFAULT_TUNING_FINGERPRINT + + monkeypatch.setattr( + config_module, + "DEFAULT_DIMENSION_WEIGHTS", + {**config_module.DEFAULT_DIMENSION_WEIGHTS, "codePresence": 0.99}, + ) + monkeypatch.setattr( + config_module, + "DEFAULT_TIER_BOUNDARIES", + {**config_module.DEFAULT_TIER_BOUNDARIES, "simple_medium": 0.42}, + ) + + assert tuning_fingerprint({}) == DEFAULT_TUNING_FINGERPRINT + assert mutable_tuned_identities([untouched, tiers_only, pinned], baselines) == frozenset() + assert mutable_tuned_identities([untouched], snapshot_tuning_baselines([])) == frozenset() + + +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 + ) + + @pytest.mark.parametrize( + "edit", + [ + pytest.param({"weight": 0.9}, id="weight"), + pytest.param({"scoring_mode": "match_count"}, id="scoring-mode"), + ], + ) + def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self, edit: Mapping[str, object]) -> None: + baselines: Final = snapshot_tuning_baselines(()) + original: Final = _router("a", {}) + config: Final = {"custom_dimensions": [_KEYWORD_DIMENSION]} + edited_config: Final = {"custom_dimensions": [{**_KEYWORD_DIMENSION, **edit}]} + added: Final = _router("a", config) + edited: Final = _router("a", edited_config) + second: Final = _router("b", config) + + assert tuning_fingerprint(config) != tuning_fingerprint(edited_config) + assert mutable_tuned_identities((added,), baselines) == {router_identity(original)} + assert tuning_quota_violation(candidate=added, others=(original,), baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=edited, others=(added,), baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=second, others=(edited,), baselines=baselines, limit=1) is not None + assert tuning_quota_violation(candidate=original, others=(edited,), baselines=baselines, limit=1) is None + assert mutable_tuned_identities((original,), baselines) == frozenset() + assert tuning_quota_violation(candidate=second, others=(original,), baselines=baselines, limit=1) is None + + def test_graded_dimension_recorded_at_snapshot_is_its_own_baseline(self) -> None: + graded: Final = _router("a", {"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]}) + baselines: Final = snapshot_tuning_baselines((graded,)) + assert mutable_tuned_identities((graded,), baselines) == frozenset() + reverted_to_binary: Final = _router("a", {"custom_dimensions": [_KEYWORD_DIMENSION]}) + assert mutable_tuned_identities((reverted_to_binary,), baselines) == {router_identity(graded)} + + 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 diff --git a/tests/test_litellm/router_utils/test_cooldown_cache.py b/tests/test_litellm/router_utils/test_cooldown_cache.py index 68e9aeaa4fc..6f90fa8465f 100644 --- a/tests/test_litellm/router_utils/test_cooldown_cache.py +++ b/tests/test_litellm/router_utils/test_cooldown_cache.py @@ -268,12 +268,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired cooldown entry must not appear in active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" + assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" def test_active_entry_is_returned(self): """ @@ -289,7 +289,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -312,14 +312,14 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - (60.0 - remaining), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + cc.in_memory_cache.set_cache(key, value, ttl=600) - before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + before_expiry = cc.in_memory_cache.ttl_dict.get(key) assert before_expiry is not None cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry is not None corrected_remaining = after_expiry - time.time() assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" @@ -340,12 +340,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired entry must not appear in async active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None @pytest.mark.asyncio async def test_async_active_entry_is_returned(self): @@ -363,7 +363,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -389,18 +389,18 @@ class TestCorrectedActiveCooldown: cc = self._make_cooldown_cache() key = "deployment:expired-dep:cooldown" entry = self._entry(timestamp=time.time() - 120.0, cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=600) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) assert result is None - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None def test_active_entry_within_window_returns_value(self): cc = self._make_cooldown_cache() key = "deployment:active-dep:cooldown" entry = self._entry(timestamp=time.time(), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=60) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) @@ -412,12 +412,12 @@ class TestCorrectedActiveCooldown: key = "deployment:backfilled-dep:cooldown" remaining = 30.0 entry = self._entry(timestamp=time.time() - (60.0 - remaining), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=600) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) assert result is not None - corrected_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + corrected_expiry = cc.in_memory_cache.ttl_dict.get(key) assert corrected_expiry is not None assert corrected_expiry - time.time() <= 60.0 @@ -425,10 +425,160 @@ class TestCorrectedActiveCooldown: cc = self._make_cooldown_cache() key = "deployment:normal-dep:cooldown" entry = self._entry(timestamp=time.time(), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60) - original_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=60) + original_expiry = cc.in_memory_cache.ttl_dict.get(key) cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry == original_expiry + + +class SharedRedisDouble: + """ + In-process stand-in for RedisCache, shared by several DualCache instances so that + tests can model two proxy replicas talking to one Redis. + """ + + def __init__(self) -> None: + self.store: dict = {} # mutable-ok: stands in for Redis' own mutable keyspace + + def set_cache(self, key, value, **kwargs): + self.store[key] = value + + async def async_set_cache(self, key, value, **kwargs): + self.store[key] = value + + def batch_get_cache(self, key_list, parent_otel_span=None, **kwargs): + return {key: self.store.get(key) for key in key_list} + + async def async_batch_get_cache(self, key_list, parent_otel_span=None, **kwargs): + return {key: self.store.get(key) for key in key_list} + + +class TestCooldownPropagationBetweenReplicas: + """ + A cooldown written by one replica has to reach its siblings quickly. The router's own + DualCache re-reads a key that is missing from memory only every 10s, so cooldown reads + get their own cache with a much shorter Redis read interval. + """ + + def _make_replica(self, redis: SharedRedisDouble, read_interval: float | None = None) -> CooldownCache: + router_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis) + if read_interval is None: + return CooldownCache(cache=router_cache, default_cooldown_time=60.0) + return CooldownCache( + cache=router_cache, + default_cooldown_time=60.0, + redis_read_interval_seconds=read_interval, + ) + + @pytest.mark.asyncio + async def test_sibling_replica_sees_cooldown_within_configured_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis, read_interval=0.25) + replica_b = self._make_replica(redis, read_interval=0.25) + model_id = "shared-deployment" + + assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(0.3) + + active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "sibling replica must pick up a cooldown written by another replica within the read interval" + ) + + @pytest.mark.asyncio + async def test_sibling_replica_sees_cooldown_within_default_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis) + replica_b = self._make_replica(redis) + model_id = "default-interval-deployment" + + assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(1.2) + + active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "the shipped default read interval must let a sibling replica see a cooldown about a second later" + ) + + def test_sync_read_path_sees_sibling_cooldown_within_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis, read_interval=0.25) + replica_b = self._make_replica(redis, read_interval=0.25) + model_id = "sync-shared-deployment" + + assert replica_b.get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(0.3) + + active = replica_b.get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active] + + @pytest.mark.asyncio + async def test_redis_attached_after_construction_is_still_used(self): + redis = SharedRedisDouble() + router_cache = DualCache(in_memory_cache=InMemoryCache()) + writer = CooldownCache(cache=router_cache, default_cooldown_time=60.0, redis_read_interval_seconds=0.25) + router_cache.attach_redis_cache(redis) + reader = self._make_replica(redis, read_interval=0.25) + model_id = "late-redis-deployment" + + writer.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + active = await reader.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "a router that wires Redis after building its cooldown cache must still publish cooldowns to it" + ) + + +class TestCooldownSurvivesUnrelatedCacheTraffic: + @pytest.mark.asyncio + async def test_unrelated_router_cache_writes_do_not_evict_active_cooldown(self): + router_cache = DualCache(in_memory_cache=InMemoryCache()) + cc = CooldownCache(cache=router_cache, default_cooldown_time=60.0) + model_id = "busy-router-deployment" + + cc.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=30.0, + ) + + for i in range(400): + router_cache.set_cache(key=f"unrelated-router-key-{i}", value={"n": i}) + + active = await cc.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "unrelated router cache traffic must not evict a cooldown that is still running" + ) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index 4768988fc87..7ee0ed3701b 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -1,6 +1,8 @@ from unittest.mock import MagicMock, patch import litellm +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.router_utils.cooldown_handlers import ( _get_deployment_cooldown_policy, _resolve_allowed_fails_from_policy, @@ -269,18 +271,20 @@ class TestShouldCooldownBasedOnDeploymentPolicy: class TestShouldCooldownBasedOnAllowedFailsPolicy: - def _make_router(self, cooldown_time: float = 60.0) -> MagicMock: + def _make_router(self, cooldown_time: float = 60.0, cache: DualCache | None = None) -> MagicMock: router = MagicMock() router.cooldown_time = cooldown_time router.allowed_fails = 0 router.allowed_fails_policy = None router.get_allowed_fails_from_policy.return_value = None - router.failed_calls.get_cache.return_value = None + router.cache = cache if cache is not None else DualCache(in_memory_cache=InMemoryCache()) return router def test_cooldown_time_override_zero_is_not_falsy(self): """cooldown_time_override=0 must be honored; it must not fall through to the router-level value.""" router = self._make_router(cooldown_time=60.0) + router.cache = MagicMock() + router.cache.increment_cache.return_value = 1 exc = litellm.RateLimitError("429", "openai", "gpt-4") should_cooldown_based_on_allowed_fails_policy( @@ -291,12 +295,68 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy: cooldown_time_override=0.0, ) - set_cache_call = router.failed_calls.set_cache.call_args - assert set_cache_call is not None - assert set_cache_call[1]["ttl"] == 0.0, ( + increment_call = router.cache.increment_cache.call_args + assert increment_call is not None + assert increment_call[1]["ttl"] == 0.0, ( "cooldown_time_override=0 should be used as TTL, not the router-level 60.0" ) + def test_fail_counter_is_shared_across_router_instances(self): + """Two workers (two Router objects over one shared cache) must pool their failures toward allowed_fails.""" + shared_cache = DualCache(in_memory_cache=InMemoryCache()) + workers = (self._make_router(cache=shared_cache), self._make_router(cache=shared_cache)) + exc = litellm.AuthenticationError("401", "openai", "gpt-4") + + results = [ + should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=workers[i % 2], + deployment="dep-1", + original_exception=exc, + allowed_fails_override=5, + ) + for i in range(6) + ] + + assert results == [False, False, False, False, False, True] + assert shared_cache.get_cache(key="deployment:dep-1:allowed_fails") == 6 + + def test_fleet_wide_count_from_redis_decides_cooldown(self): + """The Redis (fleet-wide) count decides, even when this process has only seen one failure.""" + redis_cache = MagicMock() + redis_cache.increment_cache.return_value = 6 + router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)) + exc = litellm.AuthenticationError("401", "openai", "gpt-4") + + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="dep-1", + original_exception=exc, + allowed_fails_override=5, + ) + + assert result is True + redis_cache.increment_cache.assert_called_once_with("deployment:dep-1:allowed_fails", 1, ttl=60.0) + + def test_redis_outage_falls_back_to_this_workers_count(self): + """When every Redis increment fails, the worker's own in-memory count must still cool the deployment down.""" + redis_cache = MagicMock() + redis_cache.increment_cache.side_effect = ConnectionError("redis down") + router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)) + exc = litellm.AuthenticationError("401", "openai", "gpt-4") + + results = [ + should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="dep-1", + original_exception=exc, + allowed_fails_override=5, + ) + for _ in range(6) + ] + + assert results == [False, False, False, False, False, True] + assert redis_cache.increment_cache.call_count == 6 + class TestRoutingGroupCooldownAlternatives: def _router(self, routing_groups=None): diff --git a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py index 6effbc5fa7f..9021d842daa 100644 --- a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py +++ b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py @@ -179,7 +179,7 @@ class TestHealthCheckCooldownIntegration: assert result is False # Check counter was incremented - current_fails = router.failed_calls.get_cache(key="deploy-1") + current_fails = router.cache.get_cache(key="deployment:deploy-1:allowed_fails") assert current_fails == 1 def test_health_check_failure_triggers_cooldown_at_threshold(self): @@ -263,7 +263,7 @@ class TestHealthCheckCooldownIntegration: assert "exception" not in healthy_endpoint # Verify failed_calls counter is untouched - current_fails = router.failed_calls.get_cache(key="deploy-1") + current_fails = router.cache.get_cache(key="deployment:deploy-1:allowed_fails") assert current_fails is None def test_disable_cooldowns_prevents_health_check_cooldown(self): diff --git a/tests/test_litellm/router_utils/test_health_state_cache.py b/tests/test_litellm/router_utils/test_health_state_cache.py index ffd031f9b7d..aa976bb1002 100644 --- a/tests/test_litellm/router_utils/test_health_state_cache.py +++ b/tests/test_litellm/router_utils/test_health_state_cache.py @@ -7,6 +7,7 @@ import time import pytest from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.router_utils.health_state_cache import DeploymentHealthCache @@ -145,8 +146,11 @@ class _SharedRedisFake: def __init__(self): self.store = {} self.fail_get = False + self.breaker_open = False def get_cache(self, key, parent_otel_span=None, **kwargs): + if self.breaker_open: + raise RedisCircuitBreakerOpenError("Redis circuit breaker is open - skipping get_cache") if self.fail_get: return None # RedisCache.get_cache swallows connection errors and returns None return self.store.get(key) @@ -192,3 +196,26 @@ def test_failed_redis_read_falls_back_to_local_copy(): {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} ) assert set(redis_fake.store[DeploymentHealthCache.CACHE_KEY]) == {"prod-bad", "internal-bad"} + + +def test_open_circuit_breaker_read_still_merges_into_local_copy(caplog): + """A read refused by the open breaker is a miss, so the merge and local write still happen quietly.""" + redis_fake = _SharedRedisFake() + pod_a = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) + pod_b = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) + pod_a.set_deployment_health_states( + {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} + ) + pod_b.set_deployment_health_states( + {"internal-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "timeout"}} + ) + pod_a.set_deployment_health_states( + {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} + ) + redis_fake.breaker_open = True + with caplog.at_level("ERROR"): + pod_a.set_deployment_health_states( + {"prod-new-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} + ) + assert caplog.records == [] + assert pod_a.get_unhealthy_deployment_ids() == {"prod-bad", "internal-bad", "prod-new-bad"} diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index f181370455d..ccd6766b13a 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -389,14 +389,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", ) diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/test_litellm/router_utils/test_router_utils_common_utils.py index 30f658d7ea2..ac18b4889dd 100644 --- a/tests/test_litellm/router_utils/test_router_utils_common_utils.py +++ b/tests/test_litellm/router_utils/test_router_utils_common_utils.py @@ -12,6 +12,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + provider_for_generic_call, resolve_model_group_alias, truncate_fallback_error_detail, PROVIDER_SCOPED_CREDENTIAL_PARAMS, @@ -756,3 +757,20 @@ class TestWarnOnProviderCredentialMismatch: ) is None ) + + +@pytest.mark.parametrize( + ("litellm_params", "expected"), + [ + ({"model": "azure_ai/gpt-5.4-mini", "custom_llm_provider": "azure"}, "azure"), + ({"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.openai.azure.com"}, "azure_ai"), + ({"model": "cohere/command-r"}, "cohere"), + ({"model": "gpt-5.4-mini"}, "openai"), + ({"model": "no-provider-knows-this-model"}, None), + ({"api_base": "https://my-resource.openai.azure.com"}, None), + ], + ids=["declared_wins", "prefix_beats_host_flip", "prefix_beats_cohere_chat_flip", "unprefixed_inferred", "unknown", "no_model"], +) +def test_provider_for_generic_call(litellm_params, expected, monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") + assert provider_for_generic_call(litellm_params) == expected diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index a7f50a82a99..8d83f4ca8a6 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -73,7 +73,7 @@ def assert_native_request( headers: HTTPMessage, body: object, ) -> None: - if route not in {"ocr", "transcription", "messages", "chat_completions"}: + if route not in {"ocr", "azure_ocr", "azure_di", "transcription", "messages", "chat_completions"}: raise AssertionError(f"unexpected route marker: {route!r}") if outcome not in {"success", "429", "hang"}: raise AssertionError(f"unexpected outcome marker: {outcome!r}") @@ -86,6 +86,19 @@ def assert_native_request( assert body["document"]["document_url"] == "https://example.com/document.pdf" assert body["include_image_base64"] is True return + if route == "azure_ocr": + assert path == "/providers/mistral/azure/ocr" + assert headers.get("authorization") == "Bearer prepared-azure-token" + assert body["model"] == "mistral-ocr-2505" + assert body["document"]["document_url"] == "data:application/pdf;base64,YWJj" + return + if route == "azure_di": + assert path.startswith("/documentintelligence/documentModels/prebuilt-read:analyze?") + assert "api-version=2024-11-30" in path + assert "pages=1%2C3" in path + assert headers.get("ocp-apim-subscription-key") == "di-key" + assert body == {"base64Source": "YWJj"} + return if route == "transcription": assert path == "/model/mistral.voxtral-mini-3b-2507/converse" assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ") @@ -107,8 +120,10 @@ def assert_native_request( def native_response(status: int, route: str | None) -> bytes: if status == 429: return b'{"error":"native-rate-limit"}' - if route == "ocr": + if route in {"ocr", "azure_ocr"}: return b'{"pages":[{"index":0,"markdown":"native-ocr"}]}' + if route == "azure_di": + return b'{"status":"succeeded","analyzeResult":{"pages":[]}}' if route == "transcription": return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}' return ANTHROPIC_RESPONSE @@ -181,6 +196,32 @@ def assert_success(route: str, response: object) -> None: raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}") +def azure_ocr_kwargs(api_base: str) -> dict[str, object]: + return { + "model": "mistral-ocr-2505", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_base": api_base, + "custom_llm_provider": "azure_ai", + "extra_headers": { + "x-test-outcome": "success", + "x-test-route": "azure_ocr", + }, + "optional_params": {"azure_ad_token": "prepared-azure-token"}, + } + + +def azure_di_kwargs(api_base: str) -> dict[str, object]: + return { + "model": "doc-intelligence/prebuilt-read", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_key": "di-key", + "api_base": api_base, + "custom_llm_provider": "azure_ai", + "extra_headers": {"x-test-outcome": "success", "x-test-route": "azure_di"}, + "optional_params": {"req_format": "native", "pages": [0, 2]}, + } + + def success_value(route: str, response: dict[object, object]) -> object: if route == "ocr": return response["pages"][0]["markdown"] @@ -211,6 +252,9 @@ def exercise_sync(native: object, api_base: str) -> None: assert_rate_limit(native, route, error) else: raise AssertionError(f"{route} accepted a 429 response") + assert_success("ocr", native.ocr(**azure_ocr_kwargs(api_base))) + di_response: Final = native.ocr(**azure_di_kwargs(api_base)) + assert di_response["provider_native_response"]["status"] == "succeeded" async def exercise_async(native: object, api_base: str) -> None: @@ -223,16 +267,14 @@ async def exercise_async(native: object, api_base: str) -> None: assert_rate_limit(native, route, error) else: raise AssertionError(f"a{route} accepted a 429 response") + assert_success("ocr", await native.aocr(**azure_ocr_kwargs(api_base))) + di_response: Final = await native.aocr(**azure_di_kwargs(api_base)) + assert di_response["provider_native_response"]["status"] == "succeeded" async def exercise_async_concurrency(native: object, api_base: str) -> None: responses: Final = await asyncio.wait_for( - asyncio.gather( - *( - native.amessages(**route_kwargs("messages", api_base, "success")) - for _ in range(32) - ) - ), + asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))), timeout=15, ) for response in responses: diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 0489f4ff017..b2fd2e6dcc0 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -10,8 +10,8 @@ from __future__ import annotations import pytest import litellm -from litellm.rust_bridge import chat_completions as bridge from litellm.rust_bridge import configuration +from litellm.rust_bridge import chat_completions as bridge from litellm.types.utils import ModelResponse RUST_RESPONSE = { @@ -67,10 +67,11 @@ def _hide_native_bridge(monkeypatch): @pytest.fixture(autouse=True) -def reset_bridge(): +def reset_bridge(monkeypatch): """Every test starts with no injected callables, and leaves none behind.""" bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) configuration.reset_rust_configuration() + monkeypatch.setenv("LITELLM_RUST", "1") yield bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) configuration.reset_rust_configuration() @@ -112,7 +113,7 @@ def _accepts(**overrides) -> bool: "messages": MESSAGES, "optional_params": {"max_tokens": 16}, "custom_llm_provider": "anthropic", - "litellm_params": {"rust": True}, + "litellm_params": {}, "stream": None, } kwargs.update(overrides) @@ -126,23 +127,16 @@ class TestGate: bridge.set_rust_chat_completions(decline=gate) assert _accepts(litellm_params={}) is False assert _accepts(litellm_params=None) is False - assert _accepts(litellm_params={"rust": False}) is False assert gate.calls == [], "the gate must not be consulted before opt-in" def test_accepts_when_the_deployment_opted_in_and_the_core_agrees(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) assert _accepts() is True assert gate.calls[0]["model"] == "claude-sonnet-4-5" assert gate.calls[0]["custom_llm_provider"] == "anthropic" - def test_explicit_false_overrides_process_enable(self): - bridge.set_rust_chat_completions(decline=_RecordingDecline()) - configuration.rust(True) - - assert _accepts(litellm_params={"rust": False}) is False - def test_process_enable_applies_without_request_override(self): bridge.set_rust_chat_completions(decline=_RecordingDecline()) configuration.rust(True) @@ -155,7 +149,7 @@ class TestGate: assert _accepts(litellm_params={}) is True def test_declines_streaming_and_providers_off_the_path(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) assert _accepts(stream=True) is False @@ -170,10 +164,10 @@ class TestGate: handed `optional_params` only, so accepting here would send the request to Anthropic with the abuse-detection attribution silently missing. """ - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) - assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}) is False + assert _accepts(litellm_params={"metadata": {"user_id": "u-123"}}) is False assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of" # Bedrock's Converse transform reads no `user_id`, and an Anthropic request @@ -182,13 +176,13 @@ class TestGate: _accepts( custom_llm_provider="bedrock", model="bedrock/us-east-1/anthropic.claude-v2", - litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}, + litellm_params={"metadata": {"user_id": "u-123"}}, ) is True ) - assert _accepts(litellm_params={"rust": True, "metadata": {"trace_id": "t-1"}}) is True - assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": None}}) is True - assert _accepts(litellm_params={"rust": True, "metadata": None}) is True + assert _accepts(litellm_params={"metadata": {"trace_id": "t-1"}}) is True + assert _accepts(litellm_params={"metadata": {"user_id": None}}) is True + assert _accepts(litellm_params={"metadata": None}) is True def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch): """`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the @@ -196,7 +190,7 @@ class TestGate: evicting a caller-supplied one. The core can do neither, so an operator who armed `bedrock_request_metadata_fields` keeps the Python path. """ - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) bedrock = { @@ -213,17 +207,17 @@ class TestGate: assert _accepts(**bedrock) is True, "the decline follows the operator's opt-in alone" def test_declines_when_the_core_declines(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") bridge.set_rust_chat_completions(decline=_RecordingDecline("streaming")) assert _accepts() is False def test_declines_when_the_bridge_is_unavailable(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") _hide_native_bridge(monkeypatch) assert _accepts() is False def test_declines_when_the_gate_itself_raises(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") def exploding(**_kwargs): raise RuntimeError("boom") diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/test_litellm/rust_bridge/test_configuration.py index 15f69f95335..aff9d5acac1 100644 --- a/tests/test_litellm/rust_bridge/test_configuration.py +++ b/tests/test_litellm/rust_bridge/test_configuration.py @@ -10,7 +10,6 @@ from typing import Final import pytest from litellm.rust_bridge import configuration -from litellm.rust_bridge import ocr as rust_ocr @pytest.fixture(autouse=True) @@ -19,42 +18,31 @@ def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest ) -> Generator[None]: configuration.reset_rust_configuration() monkeypatch.delenv("LITELLM_RUST", raising=False) - monkeypatch.delenv("LITELLM_USE_RUST_OCR", raising=False) - rust_ocr.set_rust_ocr(ocr=None, aocr=None) yield configuration.reset_rust_configuration() - rust_ocr.set_rust_ocr(ocr=None, aocr=None) @pytest.mark.parametrize( - ("request_override", "process", "environment", "legacy_environment", "release_default", "expected"), + ("process", "environment", "release_default", "expected"), ( - (False, True, True, True, True, False), - (True, False, False, False, False, True), - (None, False, True, True, True, False), - (None, True, False, False, False, True), - (None, None, False, True, True, False), - (None, None, True, False, False, True), - (None, None, None, False, True, False), - (None, None, None, True, False, True), - (None, None, None, None, False, False), - (None, None, None, None, True, True), + (False, True, True, False), + (True, False, False, True), + (None, False, True, False), + (None, True, False, True), + (None, None, False, False), + (None, None, True, True), ), ) def test_resolution_precedence( - request_override: bool | None, process: bool | None, environment: bool | None, - legacy_environment: bool | None, release_default: bool, expected: bool, ) -> None: assert ( configuration.resolve_rust_enabled( - request_override=request_override, process_override=process, environment_override=environment, - legacy_environment_override=legacy_environment, release_default=release_default, ) is expected @@ -71,7 +59,6 @@ def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch) configuration.rust(True) assert configuration.rust_enabled() is True - assert configuration.rust_enabled(request_override=False) is False def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPatch) -> None: @@ -83,18 +70,8 @@ def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPat @pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) def test_invalid_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None: monkeypatch.setenv("LITELLM_RUST", value) - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") assert configuration.rust_enabled() is False - assert configuration.rust_ocr_enabled() is False - - -@pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) -def test_invalid_legacy_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None: - monkeypatch.setenv("LITELLM_USE_RUST_OCR", value) - - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_enabled() is False def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytest.MonkeyPatch) -> None: @@ -104,40 +81,20 @@ def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytes assert executor.submit(configuration.rust_enabled).result() is True configuration.rust(False) assert executor.submit(configuration.rust_enabled).result() is False - assert executor.submit(configuration.rust_ocr_enabled).result() is False configuration.reset_rust_configuration() assert executor.submit(configuration.rust_enabled).result() is True - assert executor.submit(configuration.rust_ocr_enabled).result() is True def test_explicit_override_precedes_invalid_environment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_RUST", "sometimes") - assert configuration.rust_enabled(request_override=False) is False configuration.rust(True) assert configuration.rust_enabled() is True -def test_legacy_ocr_environment_is_deprecated_and_global(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_enabled() is True - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_ocr_enabled() is True - - -def test_global_environment_precedes_legacy_ocr_environment(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("LITELLM_RUST", "0") - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - - assert configuration.rust_enabled() is False - - -@pytest.mark.parametrize("environment_name", ("LITELLM_RUST", "LITELLM_USE_RUST_OCR")) @pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False"))) -def test_environment_controls_startup(environment_name: str, value: str, expected: str) -> None: - environment: Final = {**os.environ, environment_name: value} +def test_environment_controls_startup(value: str, expected: str) -> None: + environment: Final = {**os.environ, "LITELLM_RUST": value} result: Final = subprocess.run( ( sys.executable, diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py new file mode 100644 index 00000000000..2df91204390 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_token_counter.py @@ -0,0 +1,242 @@ +"""Tests for the Rust input token counter bridge. + +The native factory is dependency-injected through ``TOKEN_COUNTER.override`` +so the fallback cases run without the compiled extension present. The parity +cases need the extension and are skipped when it is not built. +""" + +from __future__ import annotations + +import json +from typing import Final + +import pytest + +import litellm +from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens +from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge import token_counter as bridge + +MODEL: Final = "claude-sonnet-4-5-20250929" +BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode() + + +class _FakeDeclined(Exception): + pass + + +class _FakeUpstream(Exception): + pass + + +class _FakeNative: + RustBridgeDeclined = _FakeDeclined + RustUpstreamError = _FakeUpstream + + +class _RecordingCounter: + def __init__(self, tokenizer_json: str) -> None: + self.tokenizer_json = tokenizer_json + self.bodies: list[bytes] = [] + + async def acount_request(self, body: bytes) -> object: + self.bodies.append(body) + return {"model": MODEL, "input_tokens": 42} + + +class _DecliningCounter: + def __init__(self, tokenizer_json: str) -> None: + pass + + async def acount_request(self, body: bytes) -> object: + raise _FakeDeclined("request has no messages") + + +class _FailingCounter: + def __init__(self, tokenizer_json: str) -> None: + pass + + async def acount_request(self, body: bytes) -> object: + raise RuntimeError("encode failed") + + +@pytest.fixture(autouse=True) +def _reset_bridge(monkeypatch: pytest.MonkeyPatch): + bridge.TOKEN_COUNTER.reset() + bridge._anthropic_counter.cache_clear() + configuration.reset_rust_configuration() + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) + yield + bridge.TOKEN_COUNTER.reset() + bridge._anthropic_counter.cache_clear() + configuration.reset_rust_configuration() + + +@pytest.mark.asyncio +async def test_disabled_bridge_never_constructs_a_counter() -> None: + constructed: list[str] = [] + + def factory(tokenizer_json: str) -> _RecordingCounter: + constructed.append(tokenizer_json) + return _RecordingCounter(tokenizer_json) + + litellm.rust(False) + bridge.TOKEN_COUNTER.override(factory) + + assert await bridge.count_anthropic_input_tokens(BODY) is None + assert constructed == [] + + +@pytest.mark.asyncio +async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> None: + counters: list[_RecordingCounter] = [] + + def factory(tokenizer_json: str) -> _RecordingCounter: + counter = _RecordingCounter(tokenizer_json) + counters.append(counter) + return counter + + litellm.rust(True) + bridge.TOKEN_COUNTER.override(factory) + + first: Final = await bridge.count_anthropic_input_tokens(BODY) + second: Final = await bridge.count_anthropic_input_tokens(BODY) + + assert first == bridge.InputTokenCount(model=MODEL, input_tokens=42) + assert second == first + assert len(counters) == 1 + assert counters[0].bodies == [BODY, BODY] + assert json.loads(counters[0].tokenizer_json)["model"]["type"] == "BPE" + + +@pytest.mark.asyncio +async def test_missing_native_module_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: + litellm.rust(True) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) + + assert await bridge.count_anthropic_input_tokens(BODY) is None + + +@pytest.mark.asyncio +async def test_declined_request_falls_back() -> None: + litellm.rust(True) + bridge.TOKEN_COUNTER.override(_DecliningCounter) + + assert await bridge.count_anthropic_input_tokens(BODY) is None + + +@pytest.mark.asyncio +async def test_runtime_failure_falls_back() -> None: + litellm.rust(True) + bridge.TOKEN_COUNTER.override(_FailingCounter) + + assert await bridge.count_anthropic_input_tokens(BODY) is None + + +@pytest.mark.parametrize( + ("model", "expected"), + ((MODEL, True), ("claude-3-5-sonnet-20241022", False), ("gpt-4o", False), ("my-router-alias", False)), +) +def test_uses_anthropic_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: bool) -> None: + assert bridge.uses_anthropic_tokenizer(model) is expected + + +@pytest.mark.parametrize("flag", ("disable_hf_tokenizer_download", "disable_token_counter")) +def test_uses_anthropic_tokenizer_respects_python_opt_outs(monkeypatch: pytest.MonkeyPatch, flag: str) -> None: + monkeypatch.setattr(litellm, flag, True) + + assert bridge.uses_anthropic_tokenizer(MODEL) is False + + +PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = ( + {"model": MODEL, "messages": [{"role": "user", "content": "Hello, how are you today?"}]}, + { + "model": MODEL, + "messages": [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "name": "bob", "content": [{"type": "text", "text": "Summarize this."}]}, + {"role": "assistant", "content": "Sure."}, + ], + }, + { + "model": MODEL, + "messages": [{"role": "user", "content": "weather in sf?"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City"}, + "unit": {"type": "string", "enum": ["c", "f"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": {"type": "function", "function": {"name": "get_weather"}}, + }, + { + "model": MODEL, + "messages": [{"role": "user", "content": "x " * 20_000}], + }, + {"model": MODEL, "prompt": "Write a haiku about ships.", "max_tokens": 20}, + {"model": MODEL, "prompt": ["first prompt", "second prompt"]}, + { + "model": MODEL, + "instructions": "be terse", + "input": [ + {"role": "user", "content": [{"type": "input_text", "text": "Summarise caf\u00e9 menus \u2014 \"ok\"?\n"}]}, + {"role": "assistant", "content": "Sure."}, + ], + }, + {"model": MODEL, "input": "a single embedding string"}, + {"model": MODEL, "input": [[101, 2023, 5], [7]], "encoding_format": "float"}, + {"model": MODEL, "query": "best harbour", "documents": ["doc one", {"text": "doc two", "title": "T", "n": 3}]}, + {"model": MODEL, "messages": None, "prompt": "messages key wins even when null"}, + {"prompt": "model comes from the route"}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_body", PARITY_REQUESTS) +async def test_native_count_matches_python_budget_counter( + monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object] +) -> None: + native: Final = pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + litellm.rust(True) + + rust_count: Final = await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode()) + python_count: Final = _count_input_tokens(request_body=request_body, model=MODEL) + + assert rust_count is not None + assert rust_count.model == request_body.get("model") + assert rust_count.input_tokens == python_count + + +DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = ( + { + "model": MODEL, + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}], + }, + {"model": MODEL, "prompt": 1.5}, + {"model": MODEL, "documents": [{"score": 0.5}]}, + {"model": MODEL, "file": "audio.mp3"}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_body", DECLINED_REQUESTS) +async def test_native_declines_shapes_python_prices_differently( + monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object] +) -> None: + native: Final = pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + litellm.rust(True) + + assert await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode()) is None diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py index 7e655b70756..03422d5433c 100644 --- a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py +++ b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py @@ -5,11 +5,68 @@ Tests the write/read/delete cycle for JSON and simple string secrets. """ import json -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx +import litellm from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 +from litellm.types.secret_managers.main import KeyManagementSettings + +_STATIC_CREDENTIALS = {"aws_access_key_id": "test-key", "aws_secret_access_key": "test-secret"} +_CMK_ARN = "arn:aws:kms:us-east-1:123456789012:key/11111111-2222-3333-4444-555555555555" + + +async def _create_secret_body_for_settings( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, settings: KeyManagementSettings +) -> dict[str, object]: + """Boot the manager from settings the way the proxy does and return the CreateSecret body it posts to AWS.""" + monkeypatch.setattr(litellm, "secret_manager_client", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", None) + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + AWSSecretsManagerV2.load_aws_secret_manager(use_aws_secret_manager=True, key_management_settings=settings) + manager = litellm.secret_manager_client + assert isinstance(manager, AWSSecretsManagerV2) + + route = respx_mock.post("https://secretsmanager.us-east-1.amazonaws.com/").respond( + json={"ARN": "arn", "Name": "litellm/test-key"} + ) + await manager.async_write_secret( + secret_name="litellm/test-key", + secret_value="sk-test-value", + optional_params=dict(_STATIC_CREDENTIALS), + ) + assert route.call_count == 1 + request = route.calls.last.request + assert request.headers["X-Amz-Target"] == "secretsmanager.CreateSecret" + return json.loads(request.content) + + +@pytest.mark.asyncio +async def test_create_secret_uses_customer_managed_kms_key_from_settings( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + body = await _create_secret_body_for_settings( + monkeypatch, + respx_mock, + KeyManagementSettings(store_virtual_keys=True, aws_region_name="us-east-1", kms_key_id=_CMK_ARN), + ) + assert body["KmsKeyId"] == _CMK_ARN + assert body["Name"] == "litellm/test-key" + assert body["SecretString"] == "sk-test-value" + + +@pytest.mark.asyncio +async def test_create_secret_omits_kms_key_id_when_not_configured( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + body = await _create_secret_body_for_settings( + monkeypatch, respx_mock, KeyManagementSettings(store_virtual_keys=True, aws_region_name="us-east-1") + ) + assert "KmsKeyId" not in body + assert body["Name"] == "litellm/test-key" @pytest.mark.asyncio diff --git a/tests/test_litellm/skills/test_skills_main.py b/tests/test_litellm/skills/test_skills_main.py new file mode 100644 index 00000000000..e1c66c8d9ea --- /dev/null +++ b/tests/test_litellm/skills/test_skills_main.py @@ -0,0 +1,57 @@ +from unittest.mock import MagicMock + +import litellm.skills.main as skills_main +from litellm.types.utils import LlmProviders + + +def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs( + monkeypatch, +) -> None: + """The REST /v1/skills form endpoint passes description/instructions as top-level + kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch + branch of create_skill() dropped both, so every LiteLLM-hosted skill was created + with description=None and instructions=None regardless of what the caller sent.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Converts files from one language into another" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == ( + "Take an uploaded document and produce it in the target language" + ) + + +def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: + """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under + extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Warehouse SQL Analyst", + extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"}, + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Runs SQL against the inventory database" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results" + + +def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None: + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value) + + assert handler.create_skill_handler.call_args.kwargs["description"] is None + assert handler.create_skill_handler.call_args.kwargs["instructions"] is None diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index bbeb6c38f78..112464bda22 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -44,7 +44,7 @@ class AsyncBridge: def test_enabled_sync_bridge_receives_audio() -> None: bridge = SyncBridge() - rust_bridge.configure_rust_transcription(True, transcription=bridge) + rust_bridge.configure_rust_transcription(transcription=bridge) result = rust_bridge.transcription( model="mistral.voxtral-mini-3b-2507", audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, @@ -61,7 +61,7 @@ def test_enabled_sync_bridge_receives_audio() -> None: @pytest.mark.asyncio async def test_enabled_async_bridge() -> None: - rust_bridge.configure_rust_transcription(True, atranscription=AsyncBridge()) + rust_bridge.configure_rust_transcription(atranscription=AsyncBridge()) result = await rust_bridge.atranscription( model="mistral.voxtral-mini-3b-2507", audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, diff --git a/tests/test_litellm/test_bedrock_extended_beta_models.py b/tests/test_litellm/test_bedrock_extended_beta_models.py deleted file mode 100644 index ebbbd6cab5c..00000000000 --- a/tests/test_litellm/test_bedrock_extended_beta_models.py +++ /dev/null @@ -1,170 +0,0 @@ -""" -Test suite for AWS Bedrock extended beta model support -Tests model configuration, pricing, and regional availability for: -- DeepSeek V3.2 -- Minimax M2.1 -- Moonshot AI Kimi K2.5 -- Qwen3 Coder Next -""" - -import os - -# Set env var to use local model cost map instead of fetching from remote -os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true" - -import pytest - -from litellm import get_model_info - -# Model configurations: (model_name, regions, max_input, max_output) -MODEL_CONFIGS = [ - ( - "deepseek.v3.2", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-north-1", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 163840, - 163840, - ), - ( - "minimax.minimax-m2.1", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-central-1", - "eu-north-1", - "eu-south-1", - "eu-west-1", - "eu-west-2", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 196000, - 8192, - ), - ( - "moonshotai.kimi-k2.5", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-north-1", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 262144, - 262144, - ), - ( - "qwen.qwen3-coder-next", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-central-1", - "eu-south-1", - "eu-west-1", - "eu-west-2", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 262144, - 8192, - ), -] - - -class TestBedrockNewModels: - """Unified test suite for all new Bedrock models""" - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_model_info_primary_region( - self, model_name, regions, max_input, max_output - ): - """Test model configuration in primary region (us-east-1)""" - model = f"bedrock/us-east-1/{model_name}" - model_info = get_model_info(model) - - assert model_info is not None, f"Model {model_name} not found" - assert model_info["max_input_tokens"] == max_input - assert model_info["max_output_tokens"] == max_output - assert model_info["litellm_provider"] == "bedrock" - assert model_info["mode"] == "chat" - assert model_info["supports_function_calling"] is True - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_pricing_configured(self, model_name, regions, max_input, max_output): - """Verify pricing is set for all models""" - model = f"bedrock/us-east-1/{model_name}" - model_info = get_model_info(model) - - assert ( - model_info["input_cost_per_token"] > 0 - ), f"Missing input cost for {model_name}" - assert ( - model_info["output_cost_per_token"] > 0 - ), f"Missing output cost for {model_name}" - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_region_count(self, model_name, regions, max_input, max_output): - """Verify each bedrock/{region}/{model_name} resolves via get_model_info""" - for region in regions: - model = f"bedrock/{region}/{model_name}" - model_info = get_model_info(model) - assert model_info is not None, f"Model {model_name} not found in {region}" - assert model_info["max_input_tokens"] == max_input - assert model_info["max_output_tokens"] == max_output - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_sample_regional_variants(self, model_name, regions, max_input, max_output): - """Test sample regional variants (us-east-1, eu-west-1, ap-northeast-1)""" - for region in ["us-east-1", "ap-northeast-1"]: - if region in regions: - model = f"bedrock/{region}/{model_name}" - model_info = get_model_info(model) - assert ( - model_info is not None - ), f"Model {model_name} not found in {region}" - assert model_info["max_input_tokens"] == max_input - assert model_info["litellm_provider"] == "bedrock" - - -class TestModelSpecificFeatures: - """Model-specific capability tests""" - - def test_deepseek_v3_2_context_window(self): - """DeepSeek V3.2 has 163K context window""" - model_info = get_model_info("bedrock/us-east-1/deepseek.v3.2") - assert model_info["max_input_tokens"] == 163840 - - def test_minimax_m2_1_context_window(self): - """Minimax M2.1 has 196K input, 8K output""" - model_info = get_model_info("bedrock/us-east-1/minimax.minimax-m2.1") - assert model_info["max_input_tokens"] == 196000 - assert model_info["max_output_tokens"] == 8192 - - def test_moonshotai_kimi_k2_5_context_window(self): - """Moonshot AI Kimi K2.5 has 256K context window""" - model_info = get_model_info("bedrock/us-east-1/moonshotai.kimi-k2.5") - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - - def test_qwen3_coder_next_context_window(self): - """Qwen3 Coder Next has 256K input, 8K output""" - model_info = get_model_info("bedrock/us-east-1/qwen.qwen3-coder-next") - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 8192 diff --git a/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py new file mode 100644 index 00000000000..0bb99339435 --- /dev/null +++ b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py @@ -0,0 +1,117 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm.constants import bedrock_embedding_models +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +BASE_MODEL = "twelvelabs.marengo-embed-3-0-v1:0" +PROFILE_MODELS = ("us.twelvelabs.marengo-embed-3-0-v1:0", "eu.twelvelabs.marengo-embed-3-0-v1:0") +ALL_MODELS = (BASE_MODEL, *PROFILE_MODELS) +MARENGO_2_7_MODELS = ( + "twelvelabs.marengo-embed-2-7-v1:0", + "us.twelvelabs.marengo-embed-2-7-v1:0", + "eu.twelvelabs.marengo-embed-2-7-v1:0", +) +PER_REQUEST_MODELS = (*ALL_MODELS, *MARENGO_2_7_MODELS) + +TEXT_REQUEST_COST = 7e-05 +IMAGE_REQUEST_COST = 0.0001 +VIDEO_COST_PER_SECOND = 0.0007 +AUDIO_COST_PER_SECOND = 0.00014 + + +def _load(path): + with open(path) as f: + return json.load(f) + + +@pytest.mark.parametrize("model", ALL_MODELS) +def test_marengo_embed_3_specs(model): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + + assert info["litellm_provider"] == "bedrock" + assert info["mode"] == "embedding" + assert info["input_cost_per_query"] == TEXT_REQUEST_COST + assert info["output_cost_per_token"] == 0.0 + assert info["max_input_tokens"] == 500 + assert info["max_tokens"] == 500 + assert info["output_vector_size"] == 512 + assert info["supports_embedding_image_input"] is True + assert info["supports_image_input"] is True + assert "deprecation_date" not in info + + routed_model, provider, _, _ = get_llm_provider(model=f"bedrock/{model}") + assert routed_model == model + assert provider == "bedrock" + + +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +def test_marengo_prices_are_per_request_not_per_token(model): + info = _load(MAIN_PATH)[model] + assert "input_cost_per_token" not in info + assert info["input_cost_per_query"] == TEXT_REQUEST_COST + assert info["input_cost_per_image"] == IMAGE_REQUEST_COST + assert info["input_cost_per_video_per_second"] == VIDEO_COST_PER_SECOND + assert info["input_cost_per_audio_per_second"] == AUDIO_COST_PER_SECOND + + +@pytest.mark.parametrize("model", ALL_MODELS) +def test_marengo_embed_3_is_visible_to_callers(model, local_model_cost_map): + info = litellm.get_model_info(model=model, custom_llm_provider="bedrock") + assert info["mode"] == "embedding" + assert info["output_vector_size"] == 512 + assert info["max_input_tokens"] == 500 + + +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +@pytest.mark.parametrize( + "details,expected_cost", + [ + (PromptTokensDetailsWrapper(query_count=1), TEXT_REQUEST_COST), + (PromptTokensDetailsWrapper(image_count=1), IMAGE_REQUEST_COST), + (PromptTokensDetailsWrapper(query_count=1, image_count=1), TEXT_REQUEST_COST + IMAGE_REQUEST_COST), + (PromptTokensDetailsWrapper(query_count=1, image_count=2), TEXT_REQUEST_COST + 2 * IMAGE_REQUEST_COST), + (PromptTokensDetailsWrapper(video_length_seconds=10), 10 * VIDEO_COST_PER_SECOND), + (PromptTokensDetailsWrapper(audio_length_seconds=10), 10 * AUDIO_COST_PER_SECOND), + ], +) +def test_marengo_requests_are_billed_per_request(model, details, expected_cost, local_model_cost_map): + usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0, prompt_tokens_details=details) + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, usage_object=usage, custom_llm_provider="bedrock" + ) + assert prompt_cost == pytest.approx(expected_cost) + assert completion_cost == 0.0 + + +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +def test_marengo_token_counts_bill_nothing(model, local_model_cost_map): + usage = Usage(prompt_tokens=128, completion_tokens=0, total_tokens=128) + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, usage_object=usage, custom_llm_provider="bedrock" + ) + assert prompt_cost == 0.0 + assert completion_cost == 0.0 + + +def test_marengo_embed_3_is_a_known_bedrock_embedding_model(): + assert BASE_MODEL in bedrock_embedding_models + + +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +def test_backup_matches_main(model): + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + + assert model in main_cost, f"{model} missing from model_prices_and_context_window.json" + assert model in backup_cost, f"{model} missing from model_prices_and_context_window_backup.json" + assert backup_cost[model] == main_cost[model], f"{model} differs between main and backup model cost maps" diff --git a/tests/test_litellm/test_bedrock_nemotron_super.py b/tests/test_litellm/test_bedrock_nemotron_super.py deleted file mode 100644 index 969db890e84..00000000000 --- a/tests/test_litellm/test_bedrock_nemotron_super.py +++ /dev/null @@ -1,51 +0,0 @@ -""" -Test suite for NVIDIA Nemotron Super 3 120B on AWS Bedrock -Verifies model configuration, pricing, and regional availability. -""" - -import os - -os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true" - -import pytest - -from litellm import get_model_info - - -MODEL_NAME = "nvidia.nemotron-super-3-120b" - - -class TestNemotronSuper3120B: - """Test model definition for nvidia.nemotron-super-3-120b""" - - def test_model_info_primary_region(self): - """Test model resolves in us-east-1""" - model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}") - - assert model_info is not None, f"Model {MODEL_NAME} not found" - assert model_info["max_input_tokens"] == 256000 - assert model_info["max_output_tokens"] == 32768 - assert model_info["litellm_provider"] == "bedrock_converse" - assert model_info["mode"] == "chat" - assert model_info["supports_function_calling"] is True - - def test_pricing_configured(self): - """Verify pricing matches AWS Bedrock rates""" - model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}") - - assert model_info["input_cost_per_token"] == 1.5e-07 - assert model_info["output_cost_per_token"] == 6.5e-07 - - def test_context_window(self): - """Nemotron Super 3 120B has 256K input, 32K output on Bedrock""" - model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}") - - assert model_info["max_input_tokens"] == 256000 - assert model_info["max_output_tokens"] == 32768 - - def test_resolves_without_region(self): - """Test model resolves with just bedrock/ prefix""" - model_info = get_model_info(f"bedrock/{MODEL_NAME}") - - assert model_info is not None, f"Model {MODEL_NAME} not found without region" - assert model_info["max_input_tokens"] == 256000 diff --git a/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py b/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py deleted file mode 100644 index 1312aa110d3..00000000000 --- a/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Validate that AWS GovCloud (Bedrock us-gov-*) Haiku 4.5 entries carry -the 1-hour cache write tier. - -AWS Bedrock GovCloud pricing applies a +20% premium over global -Anthropic rates. Global Haiku 4.5 1h cache write is $2.00/MTok; us-gov -is therefore $2.40/MTok — exactly 1.6x the 5-minute rate of $1.50/MTok. - -Source: https://aws.amazon.com/bedrock/pricing/ -""" - -import json -import os - -import pytest - - -@pytest.fixture(scope="module") -def model_data(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - return json.load(f) - - -HAIKU_USGOV_KEYS = [ - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0", -] - - -@pytest.mark.parametrize("model_key", HAIKU_USGOV_KEYS) -def test_usgov_haiku_4_5_1hr_cache_write(model_data, model_key): - assert model_key in model_data, f"Missing model entry: {model_key}" - info = model_data[model_key] - assert ( - info["cache_creation_input_token_cost"] == 1.5e-06 - ), f"{model_key}: 5m cache write should be $1.50/MTok" - assert ( - info["cache_creation_input_token_cost_above_1hr"] == 2.4e-06 - ), f"{model_key}: 1h cache write should be $2.40/MTok" - ratio = ( - info["cache_creation_input_token_cost_above_1hr"] - / info["cache_creation_input_token_cost"] - ) - assert abs(ratio - 1.6) < 1e-9, f"{model_key}: 1h/5m ratio is {ratio}, expected 1.6" diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py index 3576834dd27..3dfd7350a06 100644 --- a/tests/test_litellm/test_bedrock_usgov_pricing.py +++ b/tests/test_litellm/test_bedrock_usgov_pricing.py @@ -31,34 +31,6 @@ def model_data(): return json.load(f) -SONNET_4_5_USGOV_KEYS = [ - "bedrock/us-gov-east-1/anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0", - "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0", - "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0", -] - - -@pytest.mark.parametrize("model_key", SONNET_4_5_USGOV_KEYS) -def test_usgov_sonnet_4_5_pricing(model_data, model_key): - """Each us-gov sonnet-4-5 entry must carry the +20%-over-global rates - that AWS publishes on the GovCloud pricing page. - """ - assert model_key in model_data, f"Missing model entry: {model_key}" - info = model_data[model_key] - - assert info["input_cost_per_token"] == 3.6e-06, ( - f"{model_key}: input_cost_per_token should be $3.60/MTok (got {info['input_cost_per_token']})" - ) - assert info["output_cost_per_token"] == 1.8e-05, f"{model_key}: output_cost_per_token should be $18.00/MTok" - assert info["cache_creation_input_token_cost"] == 4.5e-06, f"{model_key}: 5m cache write should be $4.50/MTok" - assert info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06, ( - f"{model_key}: 1h cache write should be $7.20/MTok" - ) - assert info["cache_read_input_token_cost"] == 3.6e-07, f"{model_key}: cache read should be $0.36/MTok" - - def test_usgov_carries_20_percent_premium_over_global(model_data): """The us-gov rates must equal 1.2x the global anthropic.* rates, matching AWS's documented GovCloud uplift. @@ -92,18 +64,6 @@ EXPECTED_USGOV_ABOVE_200K = { } -@pytest.mark.parametrize("field,expected", EXPECTED_USGOV_ABOVE_200K.items()) -def test_usgov_cross_region_above_200k_carries_gov_premium(model_data, field, expected): - """The `_above_200k_tokens` tier on the us-gov cross-region inference - profile must also carry the +20% GovCloud uplift. The original PR - corrected the base rates but left the 200k-tier fields at the +10% - commercial-US rates, undercharging long-context requests. - """ - info = model_data[USGOV_CROSS_REGION_KEY] - assert field in info, f"{USGOV_CROSS_REGION_KEY}: missing field {field}" - assert info[field] == expected, f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})" - - def test_usgov_cross_region_above_200k_ratio_to_global(model_data): """Cross-check via the property-based invariant: every `_above_200k_tokens` field on the us-gov cross-region profile must equal 1.2x the global @@ -117,167 +77,6 @@ def test_usgov_cross_region_above_200k_ratio_to_global(model_data): assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2" -CLAUDE_GOV_EXPECTED = { - "anthropic.claude-sonnet-5": { - "input_cost_per_token": 2.4e-06, - "output_cost_per_token": 1.2e-05, - "cache_creation_input_token_cost": 3e-06, - "cache_creation_input_token_cost_above_1hr": 4.8e-06, - "cache_read_input_token_cost": 2.4e-07, - }, - "anthropic.claude-opus-4-8": { - "input_cost_per_token": 6e-06, - "output_cost_per_token": 3e-05, - "cache_creation_input_token_cost": 7.5e-06, - "cache_creation_input_token_cost_above_1hr": 1.2e-05, - "cache_read_input_token_cost": 6e-07, - }, - "anthropic.claude-opus-5": { - "input_cost_per_token": 6e-06, - "output_cost_per_token": 3e-05, - "cache_creation_input_token_cost": 7.5e-06, - "cache_creation_input_token_cost_above_1hr": 1.2e-05, - "cache_read_input_token_cost": 6e-07, - }, - "anthropic.claude-fable-5-1": { - "input_cost_per_token": 1.2e-05, - "output_cost_per_token": 6e-05, - "cache_creation_input_token_cost": 1.5e-05, - "cache_creation_input_token_cost_above_1hr": 2.4e-05, - "cache_read_input_token_cost": 3e-07, - }, -} - - -USGOV_CLAUDE_KEY_TEMPLATES = { - "bedrock/us-gov-east-1/{base_key}": "bedrock", - "bedrock/us-gov-west-1/{base_key}": "bedrock", - "us-gov.{base_key}": "bedrock_converse", -} - - -@pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED) -@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items()) -def test_usgov_claude_pricing(model_data, key_template, expected_provider, base_key): - """Sonnet 5, Opus 4.8, Opus 5, and Fable 5.1 gov entries, both in-region keys - and the us-gov. geo inference profile the model cards list for GovCloud, must - carry the 1.2x GovCloud premium over the global anthropic.* rates. No public - AWS source (offer files, pricing page) lists Claude GovCloud rows; the premium - is the one AWS quotes for Opus 4.8 in GovCloud ($6/$30 per million). - """ - gov_key = key_template.format(base_key=base_key) - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == expected_provider - assert "search_context_cost_per_query" not in info - for field, expected in CLAUDE_GOV_EXPECTED[base_key].items(): - assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" - ratio = info[field] / model_data[base_key][field] - assert abs(ratio - 1.2) < 1e-9, f"{gov_key}: {field} gov/global ratio is {ratio}, expected 1.2" - - -CONVERSE_GOV_EXPECTED = { - "nvidia.nemotron-nano-3-30b": (7.2e-08, 2.88e-07), - "nvidia.nemotron-nano-9b-v2": (7.2e-08, 2.76e-07), - "nvidia.nemotron-nano-12b-v2": (2.4e-07, 7.2e-07), - "nvidia.nemotron-super-3-120b": (1.8e-07, 7.8e-07), - "openai.gpt-oss-20b-1:0": (8.4e-08, 3.6e-07), - "openai.gpt-oss-120b-1:0": (1.8e-07, 7.2e-07), -} - - -@pytest.mark.parametrize("base_key", CONVERSE_GOV_EXPECTED) -@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items()) -def test_usgov_converse_model_pricing(model_data, key_template, expected_provider, base_key): - """Nemotron and gpt-oss gov entries, in-region and the us-gov. geo inference - profile both GovCloud regions list as ACTIVE, must match the AWS Bedrock - offer file, which prices both regions identically at 1.2x commercial. - """ - gov_key = key_template.format(base_key=base_key) - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - expected_input, expected_output = CONVERSE_GOV_EXPECTED[base_key] - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["litellm_provider"] == expected_provider - base = model_data[base_key] - assert abs(info["input_cost_per_token"] / base["input_cost_per_token"] - 1.2) < 1e-9 - assert abs(info["output_cost_per_token"] / base["output_cost_per_token"] - 1.2) < 1e-9 - - -def test_usgov_west_llama3_8b_output_price_fixed(model_data): - """The us-gov-west-1 llama3-8b entry carried the 70B output rate ($2.65/MTok); - the AWS Bedrock offer file prices output at $0.60/MTok. AWS lists the model - in us-gov-west-1 only, so there is no east entry to check. - """ - info = model_data["bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0"] - assert info["input_cost_per_token"] == 3e-07 - assert info["output_cost_per_token"] == 6e-07 - - -MANTLE_GOV_TIERED_EXPECTED = { - "openai.gpt-5.6-luna": { - "input_cost_per_token": 2.64e-07, - "input_cost_per_token_above_272k_tokens": 5.28e-07, - "cache_creation_input_token_cost": 3.3e-07, - "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, - "cache_read_input_token_cost": 2.64e-08, - "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, - "output_cost_per_token": 1.584e-06, - "output_cost_per_token_above_272k_tokens": 2.376e-06, - }, - "openai.gpt-5.6-terra": { - "input_cost_per_token": 2.64e-06, - "input_cost_per_token_above_272k_tokens": 5.28e-06, - "cache_creation_input_token_cost": 3.3e-06, - "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, - "cache_read_input_token_cost": 2.64e-07, - "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, - "output_cost_per_token": 1.584e-05, - "output_cost_per_token_above_272k_tokens": 2.376e-05, - }, -} - - -@pytest.mark.parametrize("model", MANTLE_GOV_TIERED_EXPECTED) -def test_usgov_west_mantle_terra_luna_pricing(model_data, model): - """Terra and Luna carry 1.2x commercial across every tier in the - us-gov-west-1 offer file; the us-gov-east-1 offer file has no SKUs for them. - """ - gov_key = f"bedrock_mantle/us-gov-west-1/{model}" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - for field, expected in MANTLE_GOV_TIERED_EXPECTED[model].items(): - assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" - assert info["litellm_provider"] == "bedrock_mantle" - assert f"bedrock_mantle/us-gov-east-1/{model}" not in model_data - - -@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"]) -def test_usgov_mantle_gpt_5_4_pricing_has_no_long_context_tier(model_data, region): - """gpt-5.4 gov rates come from the offer file, which publishes only the - standard tier in GovCloud: no long-context SKUs exist there, unlike commercial. - """ - gov_key = f"bedrock_mantle/{region}/openai.gpt-5.4" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["input_cost_per_token"] == 3.3e-06 - assert info["cache_read_input_token_cost"] == 3.3e-07 - assert info["output_cost_per_token"] == 1.98e-05 - assert not any(field.endswith("_above_272k_tokens") for field in info) - - -def test_usgov_mantle_grok_4_3_west_only(model_data): - """grok-4.3 is priced in the us-gov-west-1 offer file only; the east offer - file carries grok-4.6 instead. - """ - info = model_data["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] - assert info["input_cost_per_token"] == 1.5e-06 - assert info["output_cost_per_token"] == 3e-06 - assert info["cache_read_input_token_cost"] == 2.4e-07 - assert "bedrock_mantle/us-gov-east-1/xai.grok-4.3" not in model_data - - def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data): """us-gov-east-1 serves claude-3-haiku through the us-gov. inference profile only, so the profile row must bill exactly like the in-region gov row. @@ -290,96 +89,6 @@ def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data): } -GROK_4_6_GOV_KEYS = { - "us-gov.xai.grok-4.6": ("us.xai.grok-4.6", "bedrock_converse"), - "bedrock_mantle/us-gov-west-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"), - "bedrock_mantle/us-gov-east-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"), -} - - -@pytest.mark.parametrize("gov_key", GROK_4_6_GOV_KEYS) -def test_usgov_grok_4_6_pricing(model_data, gov_key): - """Both GovCloud regions serve grok-4.6 through the us-gov. profile only, and - both offer files price its standard SKU at 1.2x the commercial US rate. - """ - base_key, expected_provider = GROK_4_6_GOV_KEYS[gov_key] - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == expected_provider - assert info["input_cost_per_token"] == 2.64e-06 - assert info["output_cost_per_token"] == 7.92e-06 - assert info["cache_read_input_token_cost"] == 6.6e-07 - for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): - assert abs(info[field] / model_data[base_key][field] - 1.2) < 1e-9 - - -NOVA_GOV_WEST_EXPECTED = { - "amazon.nova-lite-v1:0": (7.2e-08, 2.88e-07), - "amazon.nova-micro-v1:0": (4.2e-08, 1.68e-07), -} - - -@pytest.mark.parametrize("base_key", NOVA_GOV_WEST_EXPECTED) -def test_usgov_west_nova_lite_micro_pricing(model_data, base_key): - """Nova Lite and Micro are on-demand in us-gov-west-1 only; the offer file - prices them at 1.2x commercial, like the Nova Pro row that was already there. - """ - gov_key = f"bedrock/us-gov-west-1/{base_key}" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - expected_input, expected_output = NOVA_GOV_WEST_EXPECTED[base_key] - assert info["litellm_provider"] == "bedrock" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert abs(info["input_cost_per_token"] / model_data[base_key]["input_cost_per_token"] - 1.2) < 1e-9 - assert abs(info["output_cost_per_token"] / model_data[base_key]["output_cost_per_token"] - 1.2) < 1e-9 - assert f"bedrock/us-gov-east-1/{base_key}" not in model_data - - -def test_usgov_west_nova_2_multimodal_embeddings_pricing(model_data): - """Every meter of the multimodal embedding model (tokens, images, audio and - video seconds) carries the 1.2x uplift the us-gov-west-1 offer file lists. - """ - gov_key = "bedrock/us-gov-west-1/amazon.nova-2-multimodal-embeddings-v1:0" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == "bedrock" - assert info["mode"] == "embedding" - assert info["input_cost_per_token"] == 1.62e-07 - assert info["input_cost_per_image"] == 7.2e-05 - assert info["input_cost_per_audio_per_second"] == 0.000168 - assert info["input_cost_per_video_per_second"] == 0.00084 - assert "bedrock/us-gov-east-1/amazon.nova-2-multimodal-embeddings-v1:0" not in model_data - - -MANTLE_GOV_FLAT_EXPECTED = { - "google.gemma-4-e2b": (4.8e-08, 9.6e-08, ("us-gov-west-1",)), - "google.gemma-4-26b-a4b": (1.56e-07, 4.8e-07, ("us-gov-west-1",)), - "google.gemma-4-31b": (1.68e-07, 4.8e-07, ("us-gov-west-1",)), - "openai.gpt-oss-20b": (8.4e-08, 3.6e-07, ("us-gov-west-1", "us-gov-east-1")), - "openai.gpt-oss-120b": (1.8e-07, 7.2e-07, ("us-gov-west-1", "us-gov-east-1")), -} - - -@pytest.mark.parametrize("model", MANTLE_GOV_FLAT_EXPECTED) -def test_usgov_mantle_gemma_and_gpt_oss_pricing(model_data, model): - """Gemma 4 is priced in the us-gov-west-1 offer file only and gpt-oss in both; - each Mantle gov row carries the offer file's standard SKU, and no row exists - for a region whose offer file has no SKU. - """ - expected_input, expected_output, regions = MANTLE_GOV_FLAT_EXPECTED[model] - for region in ("us-gov-west-1", "us-gov-east-1"): - gov_key = f"bedrock_mantle/{region}/{model}" - if region not in regions: - assert gov_key not in model_data - continue - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == "bedrock_mantle" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - - GOV_ROW_SOURCES = { "us-gov.anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1", "bedrock/us-gov-west-1/anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1", @@ -417,33 +126,3 @@ def test_usgov_rows_keep_commercial_limits_and_capabilities(model_data, gov_key) assert _non_pricing_fields(gov) == _non_pricing_fields(model_data[GOV_ROW_SOURCES[gov_key]]) assert "search_context_cost_per_query" not in gov assert "source" not in gov - - -AZURE_GOV_EXPECTED = { - "azure/us-gov/gpt-5.1": { - "input_cost_per_token": 1.71875e-06, - "cache_read_input_token_cost": 1.71875e-07, - "output_cost_per_token": 1.375e-05, - }, - "azure/us-gov/o3-mini": { - "input_cost_per_token": 1.513e-06, - "cache_read_input_token_cost": 7.57e-07, - "output_cost_per_token": 6.05e-06, - }, - "azure/us-gov/text-embedding-3-large": {"input_cost_per_token": 1.63e-07}, - "azure/us-gov/text-embedding-3-small": {"input_cost_per_token": 2.5e-08}, -} - - -@pytest.mark.parametrize("gov_key", AZURE_GOV_EXPECTED) -def test_azure_usgov_pricing(model_data, gov_key): - """Azure Government meters from the Azure retail prices API - (usgovvirginia/usgovarizona, serviceName 'Foundry Models'). No Government - retirement schedule is published, so these entries carry no deprecation_date. - """ - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - for field, expected in AZURE_GOV_EXPECTED[gov_key].items(): - assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" - assert info["litellm_provider"] == "azure" - assert "deprecation_date" not in info diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/test_litellm/test_claude_fable_5_config.py index 3ecf94602d9..0473161faac 100644 --- a/tests/test_litellm/test_claude_fable_5_config.py +++ b/tests/test_litellm/test_claude_fable_5_config.py @@ -14,7 +14,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -27,89 +26,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_fable_5_model_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_models = [ - ("claude-fable-5", "anthropic"), - ("anthropic.claude-fable-5", "bedrock_converse"), - ("vertex_ai/claude-fable-5", "vertex_ai-anthropic_models"), - # Unlike Opus 4.8 (200k on Foundry), Fable 5 has the full 1M context - # window on Microsoft Foundry. - ("azure_ai/claude-fable-5", "azure_ai"), - ] - - for model_name, provider in expected_models: - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # $10 / $50 per MTok (2x Opus 4.8), with the standard 1.25x 5m - # cache-write, 2x 1h cache-write, and 0.1x cache-read multipliers. - assert info["input_cost_per_token"] == 1e-05 - assert info["output_cost_per_token"] == 5e-05 - assert info["cache_creation_input_token_cost"] == 1.25e-05 - assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 - assert info["cache_read_input_token_cost"] == 1e-06 - - # Flat-rate across the full 1M context window. - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True - - -def test_fable_5_bedrock_regional_model_pricing(): - model_data = _load_root_cost_map() - - # Fable 5 launched with us/eu geo inference profiles plus a global profile - # (no au/apac/jp). Global uses base pricing; geo profiles carry the - # standard 10% regional premium. - expected_models = { - "global.anthropic.claude-fable-5": { - "input_cost_per_token": 1e-05, - "output_cost_per_token": 5e-05, - "cache_creation_input_token_cost": 1.25e-05, - "cache_read_input_token_cost": 1e-06, - }, - "us.anthropic.claude-fable-5": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 1.1e-06, - }, - "eu.anthropic.claude-fable-5": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 1.1e-06, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in expected.items(): - assert info[key] == value - - def test_fable_5_geo_multiplier_without_fast_mode(): """First-party ``inference_geo='us'`` carries the 1.1x premium, but unlike the Opus line there is no fast-mode variant for Fable 5; a ``fast`` key @@ -144,13 +60,6 @@ def test_fable_5_registered_for_bedrock_converse(): assert "anthropic.claude-fable-5" in BEDROCK_CONVERSE_MODELS -def test_fable_5_provider_resolves_via_model_info(local_model_cost_map): - info = litellm.get_model_info(model="claude-fable-5") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], @@ -222,46 +131,6 @@ FABLE_5_1_VARIANTS = ( ) -def test_fable_5_1_model_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_models = [ - ("claude-fable-5-1", "anthropic"), - ("anthropic.claude-fable-5-1", "bedrock_converse"), - ("vertex_ai/claude-fable-5-1", "vertex_ai-anthropic_models"), - ("azure_ai/claude-fable-5-1", "azure_ai"), - ] - - for model_name, provider in expected_models: - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["input_cost_per_token"] == 1e-05 - assert info["output_cost_per_token"] == 5e-05 - assert info["cache_creation_input_token_cost"] == 1.25e-05 - assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 - - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_forced_tool_use"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True - assert info["prompt_cache_min_tokens"] == 512 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], @@ -280,48 +149,6 @@ def test_fable_5_1_cache_reads_cost_a_quarter_of_fable_5(cost_map): ), model_name -def test_fable_5_1_bedrock_regional_model_pricing(): - model_data = _load_root_cost_map() - - expected_models = { - "global.anthropic.claude-fable-5-1": { - "input_cost_per_token": 1e-05, - "output_cost_per_token": 5e-05, - "cache_creation_input_token_cost": 1.25e-05, - "cache_read_input_token_cost": 2.5e-07, - }, - "us.anthropic.claude-fable-5-1": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 2.75e-07, - }, - "eu.anthropic.claude-fable-5-1": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 2.75e-07, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in expected.items(): - assert info[key] == value - - -def test_fable_5_1_geo_multiplier_without_fast_mode(): - """Fable 5.1 has no fast mode, so a ``fast`` key here would misprice - ``speed='fast'`` requests.""" - model_data = _load_root_cost_map() - assert model_data["claude-fable-5-1"]["provider_specific_entry"] == {"us": 1.1} - - def test_fable_5_1_present_in_bundled_backup(): backup = GetModelCostMap.load_local_model_cost_map() root = _load_root_cost_map() @@ -334,13 +161,6 @@ def test_fable_5_1_registered_for_bedrock_converse(): assert "anthropic.claude-fable-5-1" in BEDROCK_CONVERSE_MODELS -def test_fable_5_1_provider_resolves_via_model_info(local_model_cost_map): - info = litellm.get_model_info(model="claude-fable-5-1") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "model", [ diff --git a/tests/test_litellm/test_claude_haiku_4_5_config.py b/tests/test_litellm/test_claude_haiku_4_5_config.py index 8755e5d156f..9172b6479a5 100644 --- a/tests/test_litellm/test_claude_haiku_4_5_config.py +++ b/tests/test_litellm/test_claude_haiku_4_5_config.py @@ -7,57 +7,6 @@ import json import os -def test_bedrock_haiku_4_5_configuration(): - """Test that all Bedrock Claude Haiku 4.5 models use bedrock_converse provider""" - # Load model configuration - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - # All Bedrock Haiku 4.5 variants that should use bedrock_converse - bedrock_haiku_models = [ - "anthropic.claude-haiku-4-5-20251001-v1:0", - "anthropic.claude-haiku-4-5@20251001", - "us.anthropic.claude-haiku-4-5-20251001-v1:0", - "eu.anthropic.claude-haiku-4-5-20251001-v1:0", - "apac.anthropic.claude-haiku-4-5-20251001-v1:0", - "jp.anthropic.claude-haiku-4-5-20251001-v1:0", - "global.anthropic.claude-haiku-4-5-20251001-v1:0", - "au.anthropic.claude-haiku-4-5-20251001-v1:0", - ] - - for model in bedrock_haiku_models: - assert model in model_data, f"Model {model} not found in config" - model_info = model_data[model] - - # Verify uses bedrock_converse (not legacy bedrock provider) - assert ( - model_info["litellm_provider"] == "bedrock_converse" - ), f"{model} should use bedrock_converse provider, got {model_info['litellm_provider']}" - - # Verify supports vision (key missing capability) - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - - # Verify core capabilities - assert model_info.get("supports_computer_use") is True - assert model_info.get("supports_function_calling") is True - assert model_info.get("supports_tool_choice") is True - assert model_info.get("supports_prompt_caching") is True - assert model_info.get("supports_response_schema") is True - assert model_info.get("supports_pdf_input") is True - assert model_info.get("supports_assistant_prefill") is True - assert model_info.get("supports_reasoning") is True - - # Verify token limits - assert model_info["max_input_tokens"] == 200000 - assert model_info["max_output_tokens"] == 64000 - assert model_info["mode"] == "chat" - - def test_bedrock_haiku_4_5_matches_sonnet_capabilities(): """ Test that Haiku 4.5 has same capabilities as Sonnet 4.5 @@ -97,36 +46,3 @@ def test_bedrock_haiku_4_5_matches_sonnet_capabilities(): assert haiku_info.get(capability) == sonnet_info.get( capability ), f"Capability {capability} mismatch: Haiku={haiku_info.get(capability)}, Sonnet={sonnet_info.get(capability)}" - - -def test_anthropic_api_haiku_4_5_configuration(): - """Test that Anthropic API Claude Haiku 4.5 has correct configuration""" - # Load model configuration - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - # Anthropic API models (not Bedrock) - anthropic_models = [ - "claude-haiku-4-5-20251001", - "claude-haiku-4-5", - ] - - for model in anthropic_models: - assert model in model_data, f"Model {model} not found in config" - model_info = model_data[model] - - # Should use anthropic provider (not bedrock) - assert ( - model_info["litellm_provider"] == "anthropic" - ), f"{model} should use anthropic provider" - - # Should support vision - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - - # Should have larger output token limit (64K for Anthropic API) - assert model_info["max_output_tokens"] == 64000 diff --git a/tests/test_litellm/test_claude_opus_4_6_config.py b/tests/test_litellm/test_claude_opus_4_6_config.py index 89d2cd916e0..9a8632924f2 100644 --- a/tests/test_litellm/test_claude_opus_4_6_config.py +++ b/tests/test_litellm/test_claude_opus_4_6_config.py @@ -71,125 +71,6 @@ def test_claude_4_6_australia_region_uses_au_prefix_not_apac(): ), "apac.anthropic.claude-sonnet-4-6 should not be in bedrock_converse_models" -def test_opus_4_6_model_pricing_and_capabilities(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - expected_models = { - "claude-opus-4-6": { - "provider": "anthropic", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "claude-opus-4-6-20260205": { - "provider": "anthropic", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "anthropic.claude-opus-4-6-v1": { - "provider": "bedrock_converse", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "vertex_ai/claude-opus-4-6": { - "provider": "vertex_ai-anthropic_models", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "azure_ai/claude-opus-4-6": { - "provider": "azure_ai", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - } - - for model_name, config in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == config["provider"] - assert info["mode"] == "chat" - assert info["max_input_tokens"] == config["max_input_tokens"] - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 2.5e-05 - assert info["cache_creation_input_token_cost"] == 6.25e-06 - assert info["cache_read_input_token_cost"] == 5e-07 - - if config["has_long_context_pricing"]: - assert info["input_cost_per_token_above_200k_tokens"] == 1e-05 - assert info["output_cost_per_token_above_200k_tokens"] == 3.75e-05 - assert info["cache_creation_input_token_cost_above_200k_tokens"] == 1.25e-05 - assert info["cache_read_input_token_cost_above_200k_tokens"] == 1e-06 - else: - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - assert "cache_creation_input_token_cost_above_200k_tokens" not in info - assert "cache_read_input_token_cost_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - -def test_opus_4_6_bedrock_regional_model_pricing(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - expected_models = { - "global.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5e-06, - "output_cost_per_token": 2.5e-05, - "cache_creation_input_token_cost": 6.25e-06, - "cache_read_input_token_cost": 5e-07, - }, - "us.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "eu.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "au.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - assert info["supports_assistant_prefill"] is False - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - assert "cache_creation_input_token_cost_above_200k_tokens" not in info - assert "cache_read_input_token_cost_above_200k_tokens" not in info - for key, value in expected.items(): - assert info[key] == value - - def test_opus_4_6_alias_and_dated_metadata_match(): json_path = os.path.join( os.path.dirname(__file__), "../../model_prices_and_context_window.json" diff --git a/tests/test_litellm/test_claude_opus_4_8_config.py b/tests/test_litellm/test_claude_opus_4_8_config.py index 760512ad31b..e75fdba54ed 100644 --- a/tests/test_litellm/test_claude_opus_4_8_config.py +++ b/tests/test_litellm/test_claude_opus_4_8_config.py @@ -16,7 +16,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -29,102 +28,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_opus_4_8_model_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_models = { - "claude-opus-4-8": { - "provider": "anthropic", - "max_input_tokens": 1000000, - }, - "anthropic.claude-opus-4-8": { - "provider": "bedrock_converse", - "max_input_tokens": 1000000, - }, - "vertex_ai/claude-opus-4-8": { - "provider": "vertex_ai-anthropic_models", - "max_input_tokens": 1000000, - }, - "azure_ai/claude-opus-4-8": { - "provider": "azure_ai", - "max_input_tokens": 1000000, - }, - } - - for model_name, config in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == config["provider"] - assert info["mode"] == "chat" - assert info["max_input_tokens"] == config["max_input_tokens"] - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # Base pricing matches Opus 4.7: $5 / $25 per MTok, with the standard - # 1.25x cache-write and 0.1x cache-read multipliers. - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 2.5e-05 - assert info["cache_creation_input_token_cost"] == 6.25e-06 - assert info["cache_read_input_token_cost"] == 5e-07 - - # Opus 4.x flagships are flat-rate across the full context window. - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - assert model_data["claude-opus-4-8"]["supports_native_structured_output"] is True - - -def test_opus_4_8_bedrock_regional_model_pricing(): - model_data = _load_root_cost_map() - - # Global endpoints use base pricing; regional endpoints carry a 10% premium. - expected_models = { - "global.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5e-06, - "output_cost_per_token": 2.5e-05, - "cache_creation_input_token_cost": 6.25e-06, - "cache_read_input_token_cost": 5e-07, - }, - "us.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "eu.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "au.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in expected.items(): - assert info[key] == value - - def test_opus_4_8_fast_mode_multiplier(): """Opus 4.8 dropped fast-mode pricing to 2x base ($10/$50 per MTok); Opus 4.7 was 6x ($30/$150).""" @@ -134,44 +37,10 @@ def test_opus_4_8_fast_mode_multiplier(): assert entry["fast"] == 2.0 -def test_opus_4_8_present_in_bundled_backup(): - """The bundled backup is the runtime fallback (and what tests load with - ``LITELLM_LOCAL_MODEL_COST_MAP=True``) — it must carry the same entries as - the root cost map, otherwise the model resolves on one path but not the - other.""" - backup = GetModelCostMap.load_local_model_cost_map() - for model_name in ( - "claude-opus-4-8", - "anthropic.claude-opus-4-8", - "global.anthropic.claude-opus-4-8", - "us.anthropic.claude-opus-4-8", - "eu.anthropic.claude-opus-4-8", - "au.anthropic.claude-opus-4-8", - "vertex_ai/claude-opus-4-8", - "vertex_ai/claude-opus-4-8@default", - "azure_ai/claude-opus-4-8", - ): - assert model_name in backup, f"Missing from backup cost map: {model_name}" - assert backup["claude-opus-4-8"]["supports_native_structured_output"] is True - - def test_opus_4_8_registered_for_bedrock_converse(): assert "anthropic.claude-opus-4-8" in BEDROCK_CONVERSE_MODELS -def test_opus_4_8_provider_resolves_via_model_info(local_model_cost_map): - """Regression: ``claude-opus-4-8`` must resolve to provider ``anthropic``. - - Before the cost-map entry existed, the model was unknown to LiteLLM, so it - could not be tied to the ``anthropic`` provider and an ``anthropic/*`` - wildcard deployment would not match it. - """ - info = litellm.get_model_info(model="claude-opus-4-8") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/test_litellm/test_claude_opus_5_config.py index 34744aad17b..285d556ef2b 100644 --- a/tests/test_litellm/test_claude_opus_5_config.py +++ b/tests/test_litellm/test_claude_opus_5_config.py @@ -17,7 +17,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -52,91 +51,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_opus_5_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_providers = { - "claude-opus-5": "anthropic", - "anthropic.claude-opus-5": "bedrock_converse", - "vertex_ai/claude-opus-5": "vertex_ai-anthropic_models", - "azure_ai/claude-opus-5": "azure_ai", - } - - for model_name, provider in expected_providers.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # Opus 5 ships at Opus 4.8's rates: $5 / $25 per MTok, with the standard - # 1.25x cache-write, 2x 1-hour cache-write, and 0.1x cache-read multipliers. - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 2.5e-05 - assert info["cache_creation_input_token_cost"] == 6.25e-06 - assert info["cache_creation_input_token_cost_above_1hr"] == 1e-05 - assert info["cache_read_input_token_cost"] == 5e-07 - - # Flat rate across the full 1M window, no long-context premium. - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no - # assistant prefill. - assert info["supports_adaptive_thinking"] is True - assert info["supports_reasoning"] is True - assert info["supports_sampling_params"] is False - assert info["supports_assistant_prefill"] is False - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - -def test_opus_5_bedrock_regional_pricing(): - """Global/base endpoints use base pricing; the us./eu./au./jp. regional - cross-region inference profiles carry a 10% premium.""" - model_data = _load_root_cost_map() - - base_pricing = { - "input_cost_per_token": 5e-06, - "output_cost_per_token": 2.5e-05, - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_1hr": 1e-05, - "cache_read_input_token_cost": 5e-07, - } - regional_pricing = { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - } - - expected = { - "anthropic.claude-opus-5": base_pricing, - "global.anthropic.claude-opus-5": base_pricing, - "us.anthropic.claude-opus-5": regional_pricing, - "eu.anthropic.claude-opus-5": regional_pricing, - "au.anthropic.claude-opus-5": regional_pricing, - "jp.anthropic.claude-opus-5": regional_pricing, - } - - for model_name, pricing in expected.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - for key, value in pricing.items(): - assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}" - - @pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS) def test_opus_5_bedrock_entries_declare_no_effort_ceiling(model_name): """Bedrock accepts every effort level for Opus 5, so no clamp belongs here. @@ -216,18 +130,6 @@ def test_opus_5_registered_for_bedrock_converse(): assert "anthropic.claude-opus-5" in BEDROCK_CONVERSE_MODELS -def test_opus_5_provider_resolves_via_model_info(local_model_cost_map): - """Regression: ``claude-opus-5`` must resolve to provider ``anthropic``. - - Without the cost-map entry the model is unknown to LiteLLM, so it cannot be - tied to the ``anthropic`` provider and an ``anthropic/*`` wildcard deployment - would not match it.""" - info = litellm.get_model_info(model="claude-opus-5") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_claude_sonnet_5_config.py b/tests/test_litellm/test_claude_sonnet_5_config.py index 8504326cd21..8c6d2cd1851 100644 --- a/tests/test_litellm/test_claude_sonnet_5_config.py +++ b/tests/test_litellm/test_claude_sonnet_5_config.py @@ -15,7 +15,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -41,96 +40,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_sonnet_5_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_providers = { - "claude-sonnet-5": "anthropic", - "anthropic.claude-sonnet-5": "bedrock_converse", - "vertex_ai/claude-sonnet-5": "vertex_ai-anthropic_models", - "azure_ai/claude-sonnet-5": "azure_ai", - } - - for model_name, provider in expected_providers.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # Introductory Sonnet 5 pricing through 2026-08-31: $2 / $10 per MTok, - # with the 1.25x cache-write and 0.1x cache-read multipliers. On - # 2026-09-01 flip these five fields back to the sticker rate, here and - # in both cost-map JSON files (all ten claude-sonnet-5 entries): - # input_cost_per_token: 3e-06 - # output_cost_per_token: 1.5e-05 - # cache_creation_input_token_cost: 3.75e-06 - # cache_creation_input_token_cost_above_1hr: 6e-06 - # cache_read_input_token_cost: 3e-07 - # Regional Bedrock profiles (us./eu./au./jp.) stay at 1.1x those values: - # 3.3e-06 / 1.65e-05 / 4.125e-06 / 6.6e-06 / 3.3e-07 (see - # test_sonnet_5_bedrock_regional_pricing below). - assert info["input_cost_per_token"] == 2e-06 - assert info["output_cost_per_token"] == 1e-05 - assert info["cache_creation_input_token_cost"] == 2.5e-06 - assert info["cache_creation_input_token_cost_above_1hr"] == 4e-06 - assert info["cache_read_input_token_cost"] == 2e-07 - - # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no - # assistant prefill. - assert info["supports_adaptive_thinking"] is True - assert info["supports_reasoning"] is True - assert info["supports_sampling_params"] is False - assert info["supports_assistant_prefill"] is False - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - -def test_sonnet_5_bedrock_regional_pricing(): - """Global/base endpoints use base pricing; the us./eu./au./jp. regional - cross-region inference profiles carry a 10% premium.""" - model_data = _load_root_cost_map() - - base_pricing = { - "input_cost_per_token": 2e-06, - "output_cost_per_token": 1e-05, - "cache_creation_input_token_cost": 2.5e-06, - "cache_creation_input_token_cost_above_1hr": 4e-06, - "cache_read_input_token_cost": 2e-07, - } - regional_pricing = { - "input_cost_per_token": 2.2e-06, - "output_cost_per_token": 1.1e-05, - "cache_creation_input_token_cost": 2.75e-06, - "cache_creation_input_token_cost_above_1hr": 4.4e-06, - "cache_read_input_token_cost": 2.2e-07, - } - - expected = { - "anthropic.claude-sonnet-5": base_pricing, - "global.anthropic.claude-sonnet-5": base_pricing, - "us.anthropic.claude-sonnet-5": regional_pricing, - "eu.anthropic.claude-sonnet-5": regional_pricing, - "au.anthropic.claude-sonnet-5": regional_pricing, - "jp.anthropic.claude-sonnet-5": regional_pricing, - } - - for model_name, pricing in expected.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in pricing.items(): - assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}" - - def test_sonnet_5_present_in_bundled_backup(): """The bundled backup is the runtime fallback (and what tests load with ``LITELLM_LOCAL_MODEL_COST_MAP=True``); it must carry the same entries as the @@ -144,18 +53,6 @@ def test_sonnet_5_registered_for_bedrock_converse(): assert "anthropic.claude-sonnet-5" in BEDROCK_CONVERSE_MODELS -def test_sonnet_5_provider_resolves_via_model_info(local_model_cost_map): - """Regression: ``claude-sonnet-5`` must resolve to provider ``anthropic``. - - Before the cost-map entry existed, the model was unknown to LiteLLM, so it - could not be tied to the ``anthropic`` provider and an ``anthropic/*`` - wildcard deployment would not match it.""" - info = litellm.get_model_info(model="claude-sonnet-5") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py index e33bcfb8378..4e770be7c3e 100644 --- a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py +++ b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py @@ -27,17 +27,6 @@ BACKUP_MAP = os.path.join( ) -@pytest.fixture(autouse=True) -def _use_local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - try: - yield - finally: - litellm.model_cost = original_model_cost - - def _load(path: str) -> dict: with open(path, encoding="utf-8") as f: return json.load(f) @@ -47,50 +36,6 @@ def _cloudflare_keys(data: dict) -> set: return {k for k in data if k.startswith("cloudflare/")} -def test_glm_5_2_entry_is_present_and_well_formed(): - entry = litellm.model_cost["cloudflare/@cf/zai-org/glm-5.2"] - assert entry["litellm_provider"] == "cloudflare" - assert entry["mode"] == "chat" - assert entry["supports_function_calling"] is True - assert entry["input_cost_per_token"] > 0 - assert entry["output_cost_per_token"] > 0 - - -def test_vision_model_is_flagged_supports_vision(): - entry = litellm.model_cost["cloudflare/@cf/meta/llama-3.2-11b-vision-instruct"] - assert entry["litellm_provider"] == "cloudflare" - assert entry.get("supports_vision") is True - - -def test_additional_current_models_are_present(): - for key in ( - "cloudflare/@cf/openai/gpt-oss-120b", - "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast", - ): - entry = litellm.model_cost[key] - assert entry["litellm_provider"] == "cloudflare" - assert entry["mode"] == "chat" - assert entry["supports_function_calling"] is True - assert entry["input_cost_per_token"] > 0 - assert entry["output_cost_per_token"] > 0 - - -@pytest.mark.parametrize( - "key, published_price_per_audio_minute", - [ - ("cloudflare/@cf/openai/whisper", 0.00045), - ("cloudflare/@cf/openai/whisper-large-v3-turbo", 0.00051), - ], -) -def test_whisper_transcription_pricing_is_stored_per_second(key, published_price_per_audio_minute): - entry = litellm.model_cost[key] - assert entry["litellm_provider"] == "cloudflare" - assert entry["mode"] == "audio_transcription" - assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"] - assert entry["output_cost_per_second"] == 0.0 - assert entry["input_cost_per_second"] == pytest.approx(published_price_per_audio_minute / 60) - - def test_root_and_backup_have_identical_cloudflare_keys(): if not os.path.exists(ROOT_MAP): pytest.skip("root cost map only ships in source checkouts") diff --git a/tests/test_litellm/test_component_entrypoint.py b/tests/test_litellm/test_component_entrypoint.py index b0969e2c694..0c2a533b8bc 100644 --- a/tests/test_litellm/test_component_entrypoint.py +++ b/tests/test_litellm/test_component_entrypoint.py @@ -38,12 +38,16 @@ _STUB_TEMPLATE = """#!/bin/sh _ENTRYPOINT_RE = re.compile(r"^ENTRYPOINT\s+(\[.*\])\s*$", re.MULTILINE) _CMD_RE = re.compile(r"^CMD\s+(\[.*\])\s*$", re.MULTILINE) _COPY_RE = re.compile(r"^COPY\s+(?!--from)(\S+)\s+(\S+)\s*$", re.MULTILINE) -_APP_TARGET_RE = re.compile(r"(?:gateway|backend)\.main:app") +_APP_TARGET_RE = re.compile(r"(?:gateway|backend)\.main:app|gateway\.launch") _TF_STRING_LOCAL_RE = re.compile(r'^\s*(\w+)\s*=\s*"((?:[^"\\]|\\.)*)"\s*$', re.MULTILINE) _TF_INTERPOLATION_RE = re.compile(r"\$\{(local|var)\.(\w+)\}") TERRAFORM_LAUNCH_SITES = {TERRAFORM_ECS: 2, TERRAFORM_CLOUDRUN: 2} TERRAFORM_VAR_STUBS = {"gateway_num_workers": "2"} +COMPONENT_LAUNCHERS = { + "gateway": ("python", "-m", "gateway.launch"), + "backend": ("uvicorn", "backend.main:app"), +} _MAX_INTERPOLATION_PASSES = 5 @@ -63,7 +67,7 @@ def _run_entrypoint( """Run `script` with stubbed executables on PATH and return the recorded lines.""" bin_dir = tmp_path / "bin" bin_dir.mkdir(parents=True) - _write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "litellm")) + _write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "python", "litellm")) record = tmp_path / "record.txt" env = { @@ -223,6 +227,59 @@ def test_gating_matches_the_monolithic_entrypoint_and_get_secret_bool( assert monolith[1] == ("args=litellm --port 4000" if traced else "args=--port 4000") +def test_wipes_the_prometheus_multiproc_dir_before_uvicorn_forks(tmp_path: Path) -> None: + """A restarted container inherits the emptyDir of its predecessor, whose worker pids it may reuse, so the + stale .db files must be gone before any worker opens the one carrying its own pid.""" + multiproc_dir = tmp_path / "multiproc" + multiproc_dir.mkdir() + (multiproc_dir / "gauge_livesum_7.db").write_bytes(b"stale") + (multiproc_dir / "counter_7.db").write_bytes(b"stale") + (multiproc_dir / "keep.txt").write_text("not a sample") + + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + _write_stubs(bin_dir, ("uvicorn",)) + record = tmp_path / "record.txt" + env = { + **os.environ, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "RECORD": str(record), + "PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir), + } + env.pop("USE_DDTRACE", None) + result = subprocess.run( + ["sh", str(COMPONENT_ENTRYPOINT), "uvicorn", "gateway.main:app"], + env=env, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}" + assert sorted(p.name for p in multiproc_dir.iterdir()) == ["keep.txt"] + assert record.read_text().splitlines()[0] == "exec=uvicorn" + + +def test_creates_a_missing_prometheus_multiproc_dir(tmp_path: Path) -> None: + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + _write_stubs(bin_dir, ("uvicorn",)) + missing = tmp_path / "multiproc" + env = { + **os.environ, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "RECORD": str(tmp_path / "record.txt"), + "PROMETHEUS_MULTIPROC_DIR": str(missing), + } + env.pop("USE_DDTRACE", None) + result = subprocess.run( + ["sh", str(COMPONENT_ENTRYPOINT), "uvicorn", "gateway.main:app"], env=env, capture_output=True, text=True + ) + + assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}" + assert missing.is_dir() + + def _copied_script(dockerfile: Path, image_path: str) -> Path: """Resolve the repo file a Dockerfile `COPY`s to `image_path`, so tests run what the image ships.""" matches = _COPY_RE.findall(dockerfile.read_text()) @@ -277,22 +334,63 @@ def test_entrypoint_script_has_no_carriage_returns() -> None: @pytest.mark.parametrize( - "dockerfile, app_target", + "dockerfile, launcher", [ - (GATEWAY_DOCKERFILE, "gateway.main:app"), - (BACKEND_DOCKERFILE, "backend.main:app"), + (GATEWAY_DOCKERFILE, "python -m gateway.launch"), + (BACKEND_DOCKERFILE, "uvicorn backend.main:app"), ], ) -def test_component_images_launch_uvicorn_through_the_entrypoint(dockerfile: Path, app_target: str) -> None: +def test_component_images_launch_uvicorn_through_the_entrypoint(dockerfile: Path, launcher: str) -> None: entrypoint = " ".join(_entrypoint_argv(dockerfile)) assert IMAGE_ENTRYPOINT_PATH in entrypoint, f"{dockerfile} bypasses the ddtrace-aware entrypoint" - assert app_target in entrypoint - assert entrypoint.index(IMAGE_ENTRYPOINT_PATH) < entrypoint.index("uvicorn"), ( + assert launcher in entrypoint + assert entrypoint.index(IMAGE_ENTRYPOINT_PATH) < entrypoint.index(launcher), ( f"{dockerfile} must invoke uvicorn through the entrypoint, not the other way around" ) +@pytest.mark.parametrize( + "use_ddtrace, num_workers, expected_exec, expected_args", + [ + (None, "4", "exec=python", "args=-m gateway.launch --workers 4 --host 0.0.0.0 --port 4000"), + (None, None, "exec=python", "args=-m gateway.launch --workers 1 --host 0.0.0.0 --port 4000"), + ("true", "4", "exec=ddtrace-run", "args=python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000"), + ], +) +def test_gateway_image_execs_the_supervisor_with_its_worker_count( + use_ddtrace: str | None, num_workers: str | None, expected_exec: str, expected_args: str, tmp_path: Path +) -> None: + """Run the gateway image's ENTRYPOINT + CMD and record what the container execs. + + The Dockerfile's `/app/...` script path is resolved to the checked-in script and `python` + is stubbed on PATH, so the assertion is on the argv `gateway.launch` receives, not on the + Dockerfile text. + """ + entrypoint = tuple( + part.replace(IMAGE_ENTRYPOINT_PATH, str(COMPONENT_ENTRYPOINT)) for part in _entrypoint_argv(GATEWAY_DOCKERFILE) + ) + bin_dir = tmp_path / "bin" + bin_dir.mkdir(parents=True) + _write_stubs(bin_dir, ("ddtrace-run", "python", "uvicorn")) + record = tmp_path / "record.txt" + overrides = {"USE_DDTRACE": use_ddtrace, "NUM_WORKERS": num_workers} + env = { + **{k: v for k, v in os.environ.items() if k not in ("DD_TRACE_OPENAI_ENABLED", *overrides)}, + **{k: v for k, v in overrides.items() if v is not None}, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "RECORD": str(record), + "PYTHONPATH": PYTHONPATH_SENTINEL, + } + + result = subprocess.run( + [*entrypoint, *_cmd_argv(GATEWAY_DOCKERFILE)], env=env, capture_output=True, text=True, check=False + ) + + assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}" + assert tuple(record.read_text().splitlines())[:2] == (expected_exec, expected_args) + + @pytest.mark.parametrize("dockerfile", [GATEWAY_DOCKERFILE, BACKEND_DOCKERFILE]) def test_component_images_make_the_entrypoint_executable(dockerfile: Path) -> None: body = dockerfile.read_text() @@ -314,17 +412,18 @@ def test_terraform_launch_command_matches_the_script_contract( implementations under the same environment and asserts they agree on which binary is exec'd and on whether the openai integration is disabled. """ - app_target = f"{component}.main:app" + launcher = COMPONENT_LAUNCHERS[component] + app_target = " ".join(launcher[1:]) command = _resolve_tf_local(terraform_file, f"{component}_launch_cmd") bin_dir = tmp_path / "bin" bin_dir.mkdir(parents=True) - _write_stubs(bin_dir, ("ddtrace-run", "uvicorn")) + _write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "python")) from_terraform = _run_shell_command(command, bin_dir, tmp_path / "terraform.txt", use_ddtrace) from_script = _run_entrypoint( COMPONENT_ENTRYPOINT, - ("uvicorn", app_target), + launcher, use_ddtrace=use_ddtrace, tmp_path=tmp_path / "script", ) @@ -334,12 +433,14 @@ def test_terraform_launch_command_matches_the_script_contract( ) assert from_terraform[2] == from_script[2], f"{terraform_file} disagrees with the script on the openai integration" assert app_target in from_terraform[1] + assert "gateway.main:app" not in from_terraform[1], f"{terraform_file} bypasses the gateway.launch supervisor" if use_ddtrace in TRUTHY_USE_DDTRACE: assert from_terraform[0] == "exec=ddtrace-run" + assert from_terraform[1].startswith(f"args={launcher[0]} ") assert from_terraform[2] == "DD_TRACE_OPENAI_ENABLED=False" else: - assert from_terraform[0] == "exec=uvicorn" + assert from_terraform[0] == f"exec={launcher[0]}" assert from_terraform[2] == "DD_TRACE_OPENAI_ENABLED=" diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 2046695f151..f610821e06a 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1,6 +1,7 @@ import json from pathlib import Path +from typing import Final import pytest @@ -17,6 +18,7 @@ from litellm.cost_calculator import ( response_cost_calculator, ) from litellm.types.llms.openai import OpenAIRealtimeStreamList +from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CacheCreationTokenDetails, ModelInfo, @@ -127,6 +129,22 @@ def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_co assert cost > 0, "Cost should be calculated using response model" +def test_jina_rerank_bills_total_tokens_at_input_rate_only(_local_model_cost_map): + response: Final = RerankResponse( + id="rerank-1", + results=[{"index": 0, "relevance_score": 0.9}], + meta={"billed_units": {"total_tokens": 1000}}, + ) + + cost: Final = completion_cost( + completion_response=response, + model="jina_ai/jina-reranker-v2-base-multilingual", + call_type="rerank", + ) + + assert cost == pytest.approx(1000 * 5e-08) + + def test_cost_calculator_with_response_cost_in_additional_headers(): class MockResponse(BaseModel): _hidden_params = { @@ -146,6 +164,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 = { @@ -3690,6 +3753,112 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(100 * 3e-08) +def test_completion_cost_logs_the_rates_it_billed_at(monkeypatch): + """A caller reporting the cost lines beside their per-token rates reads both off this one call. + completion_cost infers the provider, and xai's inclusive tier thresholds put a request sitting + exactly on 200k at the tier rate, which a lookup made without that inferred provider would miss. + """ + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setitem( + litellm.model_cost, + "xai/tiered-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "litellm_provider": "xai", + "mode": "chat", + }, + ) + logging_obj = Logging( + model="xai/tiered-model", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="billed-rates", + function_id="f", + ) + usage = Usage( + prompt_tokens=200_000, + completion_tokens=1_000, + total_tokens=201_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100_000), + ) + + litellm.completion_cost( + completion_response=ModelResponse(model="xai/tiered-model", usage=usage), + model="xai/tiered-model", + custom_llm_provider=None, + litellm_logging_obj=logging_obj, + ) + + rates = logging_obj.billed_token_rates + assert rates is not None + assert rates.input_cost_per_token == pytest.approx(6e-6) + assert rates.cache_read_input_token_cost == pytest.approx(6e-7) + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx( + 100_000 * rates.cache_read_input_token_cost + ) + assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(1_000 * rates.output_cost_per_token) + + +def test_completion_cost_logs_cache_and_reasoning_breakdown_for_custom_pricing(): + """ + A custom-priced deployment bills cache tokens at its custom cache rates, but the + breakdown stored for the spend logs carried no cache or reasoning lines for it. + """ + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CompletionTokensDetailsWrapper, CostPerToken + + logging_obj = Logging( + model="openai/onprem-model", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="custom-pricing-breakdown", + function_id="f", + ) + response = ModelResponse( + model="openai/onprem-model", + usage=Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800, cache_creation_tokens=100), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=200), + ), + ) + + total = completion_cost( + completion_response=response, + model="openai/onprem-model", + custom_llm_provider="openai", + custom_cost_per_token=CostPerToken( + input_cost_per_token=1e-6, + output_cost_per_token=2e-6, + cache_read_input_token_cost=1e-7, + cache_creation_input_token_cost=1.25e-6, + ), + litellm_logging_obj=logging_obj, + ) + + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(800 * 1e-7) + assert logging_obj.cost_breakdown["cache_creation_cost"] == pytest.approx(100 * 1.25e-6) + assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(200 * 2e-6) + assert total == pytest.approx(100 * 1e-6 + 800 * 1e-7 + 100 * 1.25e-6 + 500 * 2e-6) + + def test_cost_per_token_per_second_pricing(monkeypatch): """ Models priced by duration (input/output_cost_per_second) with no per-token rates @@ -4480,6 +4649,18 @@ def test_explicit_pricing_precedes_private_provider_response_model( assert selected == expected +def test_cost_per_token_mistral_voxtral_tts_bills_per_input_character(_local_model_cost_map): + prompt_usd, completion_usd = cost_per_token( + model="voxtral-mini-tts-2603", + custom_llm_provider="mistral", + call_type="speech", + prompt_characters=1000, + ) + + assert prompt_usd == pytest.approx(1000 * 1.6e-05) + assert completion_usd == 0.0 + + def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_model_cost_map): """gpt-6-astra batch pricing is 50% off the standard $10 input and $50 output rates per 1M tokens.""" from litellm.cost_calculator import batch_cost_calculator @@ -4492,3 +4673,98 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m assert prompt_cost == pytest.approx(1000 * 5e-6) assert completion_cost == pytest.approx(500 * 2.5e-5) + + +def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once( + _local_model_cost_map: None, +) -> None: + """Realtime response.done nests reasoning_tokens inside text_tokens, so they are billed once.""" + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}}, + { + "type": "response.done", + "response": { + "usage": { + "total_tokens": 260, + "input_tokens": 237, + "output_tokens": 23, + "input_token_details": { + "text_tokens": 43, + "audio_tokens": 0, + "image_tokens": 194, + "cached_tokens": 0, + "cached_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "image_tokens": 0}, + }, + "output_token_details": {"text_tokens": 23, "audio_tokens": 0, "reasoning_tokens": 18}, + } + }, + }, + ] + combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results, + ) + + total_cost = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=combined_usage_object, + custom_llm_provider="azure", + litellm_model_name="azure/gpt-realtime-2.1-mini", + ) + + info = litellm.get_model_info(model="azure/gpt-realtime-2.1-mini", custom_llm_provider="azure") + expected = ( + 43 * info["input_cost_per_token"] + + 194 * info["input_cost_per_image_token"] + + 23 * info["output_cost_per_token"] + ) + assert total_cost == pytest.approx(expected) + assert total_cost == pytest.approx(0.0002362) + + +def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> None: + """The combined usage that lands in spend logs keeps reasoning out of text_tokens for every turn.""" + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}}, + { + "type": "response.done", + "response": { + "usage": { + "total_tokens": 307, + "input_tokens": 237, + "output_tokens": 70, + "input_token_details": { + "text_tokens": 43, + "audio_tokens": 0, + "image_tokens": 194, + "cached_tokens": 0, + }, + "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52}, + } + }, + }, + { + "type": "response.done", + "response": { + "usage": { + "total_tokens": 363, + "input_tokens": 300, + "output_tokens": 63, + "input_token_details": { + "text_tokens": 106, + "audio_tokens": 0, + "image_tokens": 194, + "cached_tokens": 0, + }, + "output_token_details": {"text_tokens": 63, "audio_tokens": 0, "reasoning_tokens": 43}, + } + }, + }, + ] + + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) + + assert combined.completion_tokens == 133 + assert combined.completion_tokens_details is not None + assert combined.completion_tokens_details.reasoning_tokens == 95 + assert combined.completion_tokens_details.text_tokens == 38 + assert combined.completion_tokens_details.audio_tokens == 0 diff --git a/tests/test_litellm/test_cost_map_guard.py b/tests/test_litellm/test_cost_map_guard.py new file mode 100644 index 00000000000..1b4330ed62c --- /dev/null +++ b/tests/test_litellm/test_cost_map_guard.py @@ -0,0 +1,195 @@ +import importlib.util +import json +import subprocess +import sys +from pathlib import Path +from types import ModuleType +from typing import Final + +import pytest + +ROOT: Final = Path(__file__).resolve().parents[2] +CI_CD: Final = ROOT / "ci_cd" + + +def _load(name: str) -> ModuleType: + spec = importlib.util.spec_from_file_location(name, CI_CD / f"{name}.py") + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +schema_module: Final = _load("generate_model_prices_schema") +guard: Final = _load("cost_map_guard") + +MAP_FILES: Final = (guard.COST_MAP_PATH,) +BOT_REF: Final = "litellm_cost_map_sync_2026-09-04T12-00Z" + + +def _entry(price: float = 1e-06, **extra: object) -> dict[str, object]: + return { + "input_cost_per_token": price, + "output_cost_per_token": price * 2, + "litellm_provider": "openrouter", + "mode": "chat", + "max_tokens": 4096, + **extra, + } + + +BASE_MAP: Final = { + "sample_spec": {"input_cost_per_token": "USD per prompt token"}, + "fallback_generalizations": {"rules": [{"name": "r", "pattern": "^x"}]}, + "openrouter/a": _entry(supports_vision=True), + "openrouter/b": _entry(2e-06), +} + + +def _serialize(cost_map: dict[str, object]) -> str: + return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n" + + +def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: str | None = None) -> object: + text = _serialize(cost_map) + rendered = schema_module.render(schema_module.build_schema(cost_map)) + return guard.Snapshot( + cost_map=text, backup=text if backup is None else backup, schema=rendered if schema is None else schema + ) + + +BASE: Final = _snapshot(BASE_MAP) + + +def _failures(head: object, changed_files: tuple[str, ...] = MAP_FILES, bot: bool = True) -> tuple[str, ...]: + return guard.guard_failures(BASE, head, changed_files, bot) + + +def test_in_sync_files_pass_for_humans_and_bots() -> None: + assert _failures(BASE, bot=False) == () + assert _failures(BASE, bot=True) == () + + +def test_bot_may_add_and_reprice_models() -> None: + head = _snapshot({**BASE_MAP, "openrouter/a": _entry(9e-06, supports_vision=True), "openrouter/c": _entry()}) + assert _failures(head) == () + + +def test_broken_json_is_reported() -> None: + head = guard.Snapshot(cost_map="{not json", backup="{not json", schema="{}") + assert _failures(head, bot=False) == ( + f"{guard.COST_MAP_PATH} is not valid JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1)", + ) + + +def test_non_object_root_is_reported() -> None: + head = guard.Snapshot(cost_map="[]", backup="[]", schema="{}") + assert _failures(head, bot=False) == (f"{guard.COST_MAP_PATH} must be a JSON object at the root",) + + +def test_backup_drift_is_reported() -> None: + head = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)})) + assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.BACKUP_PATH)] + + +def test_schema_out_of_sync_is_reported() -> None: + head = _snapshot({**BASE_MAP, "openrouter/c": _entry(supports_audio_input=True)}, schema=BASE.schema) + assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.SCHEMA_PATH)] + + +def test_schema_validation_errors_are_reported() -> None: + head = _snapshot({**BASE_MAP, "openrouter/c": _entry(-1e-06)}) + prefix = f"{guard.COST_MAP_PATH} does not validate against its schema: openrouter/c." + assert [failure.removeprefix(prefix).split(":")[0] for failure in _failures(head, bot=False)] == [ + "input_cost_per_token", + "output_cost_per_token", + ] + + +def test_unclassified_entry_key_is_reported() -> None: + text = _serialize({**BASE_MAP, "openrouter/c": _entry(weird_thing=1)}) + head = guard.Snapshot(cost_map=text, backup=text, schema=BASE.schema) + (failure,) = _failures(head, bot=False) + assert "Unclassified keys" in failure and "weird_thing" in failure + + +def test_bot_may_only_touch_the_cost_map_files() -> None: + changed = (*guard.GUARDED_PATHS, "litellm/utils.py", ".github/workflows/cost-map-guard.yml") + assert _failures(BASE, changed_files=changed, bot=False) == () + assert _failures(BASE, changed_files=changed) == ( + "bot PRs may only change the cost map files, not litellm/utils.py", + "bot PRs may only change the cost map files, not .github/workflows/cost-map-guard.yml", + ) + + +def test_bot_may_not_remove_models() -> None: + head = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}) + assert _failures(head, bot=False) == () + assert _failures(head) == ("bot PRs may not remove models: openrouter/b",) + + +def test_bot_may_not_remove_fields() -> None: + head = _snapshot({**BASE_MAP, "openrouter/a": _entry()}) + assert _failures(head, bot=False) == () + assert _failures(head) == ("bot PRs may not remove fields: openrouter/a.supports_vision",) + + +def test_bot_may_not_change_special_root_keys() -> None: + head = _snapshot({**BASE_MAP, "fallback_generalizations": {"rules": []}}) + assert _failures(head, bot=False) == () + assert _failures(head) == ("bot PRs may not change fallback_generalizations",) + + +def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str: + text = _serialize(cost_map) + (repo / guard.COST_MAP_PATH).write_text(text) + (repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True) + (repo / guard.BACKUP_PATH).write_text(text) + (repo / guard.SCHEMA_PATH).write_text(schema_module.render(schema_module.build_schema(cost_map))) + subprocess.run(("git", "add", "-A"), cwd=repo, check=True) + subprocess.run( + ("git", "-c", "user.name=t", "-c", "user.email=t@example.com", "commit", "-q", "-m", message), + cwd=repo, + check=True, + ) + return subprocess.run( + ("git", "rev-parse", "HEAD"), cwd=repo, check=True, capture_output=True, text=True + ).stdout.strip() + + +def _run_guard(repo: Path, base: str, head: str, head_ref: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + (sys.executable, str(CI_CD / "cost_map_guard.py"), "--base", base, "--head", head, "--head-ref", head_ref), + cwd=repo, + capture_output=True, + text=True, + check=False, + ) + + +@pytest.mark.parametrize( + ("head_ref", "expected_code", "expected_line"), + [ + (BOT_REF, 1, "- bot PRs may not remove models: openrouter/b"), + ("litellm_fix_pricing", 0, "cost map guard passed (human PR, file checks only)"), + ], +) +def test_main_reads_both_revisions_from_git( + tmp_path: Path, head_ref: str, expected_code: int, expected_line: str +) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + base = _commit(tmp_path, BASE_MAP, "base") + head = _commit(tmp_path, {key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}, "head") + result = _run_guard(tmp_path, base, head, head_ref) + assert result.returncode == expected_code, result.stdout + result.stderr + assert expected_line in result.stdout.splitlines() + + +def test_main_rejects_a_bot_pr_that_edits_code(tmp_path: Path) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + base = _commit(tmp_path, BASE_MAP, "base") + (tmp_path / "litellm" / "utils.py").write_text("print('hi')\n") + head = _commit(tmp_path, {**BASE_MAP, "openrouter/c": _entry()}, "head") + assert _run_guard(tmp_path, base, head, BOT_REF).returncode == 1 + assert _run_guard(tmp_path, base, head, "litellm_fix_pricing").returncode == 0 diff --git a/tests/test_litellm/test_daybreak_model_metadata.py b/tests/test_litellm/test_daybreak_model_metadata.py index dbb7ecdffac..c3bac14dbbd 100644 --- a/tests/test_litellm/test_daybreak_model_metadata.py +++ b/tests/test_litellm/test_daybreak_model_metadata.py @@ -32,22 +32,6 @@ def _load(path): return json.load(f) -@pytest.mark.parametrize("model", DAYBREAK_MODELS) -def test_daybreak_capability_contract(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "openai" - assert info["mode"] == "chat" - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] - - assert info["supports_computer_use"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - assert info["supports_vision"] is True - - def test_blue_alias_matches_its_snapshot_computer_use(): cost_map = _load(MAIN_PATH) diff --git a/tests/test_litellm/test_deepseek_model_metadata.py b/tests/test_litellm/test_deepseek_model_metadata.py index b9eb33f0972..9cbd14ebd1e 100644 --- a/tests/test_litellm/test_deepseek_model_metadata.py +++ b/tests/test_litellm/test_deepseek_model_metadata.py @@ -39,26 +39,6 @@ class TestDeepSeekModelCostEntries: """Verify that provider-prefixed DeepSeek entries contain the same capability flags as their bare-name counterparts in the JSON files.""" - def test_deepseek_chat_supports_response_schema_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-chat", {}) - assert entry.get("supports_response_schema") is True - - def test_deepseek_reasoner_supports_response_schema_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-reasoner", {}) - assert entry.get("supports_response_schema") is True - - def test_deepseek_chat_supports_system_messages_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-chat", {}) - assert entry.get("supports_system_messages") is True - - def test_deepseek_reasoner_supports_system_messages_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-reasoner", {}) - assert entry.get("supports_system_messages") is True - def test_deepseek_chat_max_input_tokens_matches_bare_in_backup(self): data = _load_backup_json() bare = data.get("deepseek-chat", {}) @@ -71,26 +51,6 @@ class TestDeepSeekModelCostEntries: prefixed = data.get("deepseek/deepseek-reasoner", {}) assert prefixed.get("max_output_tokens") == bare.get("max_output_tokens") - def test_main_json_deepseek_chat_supports_response_schema(self): - main_path = os.path.join( - os.path.dirname(os.path.dirname(litellm.__file__)), - "model_prices_and_context_window.json", - ) - with open(main_path, encoding="utf-8") as f: - data = json.load(f) - entry = data.get("deepseek/deepseek-chat", {}) - assert entry.get("supports_response_schema") is True - - def test_main_json_deepseek_reasoner_supports_response_schema(self): - main_path = os.path.join( - os.path.dirname(os.path.dirname(litellm.__file__)), - "model_prices_and_context_window.json", - ) - with open(main_path, encoding="utf-8") as f: - data = json.load(f) - entry = data.get("deepseek/deepseek-reasoner", {}) - assert entry.get("supports_response_schema") is True - # --------------------------------------------------------------------------- # API-level tests – verify supports_response_schema returns True diff --git a/tests/test_litellm/test_default_branch.py b/tests/test_litellm/test_default_branch.py new file mode 100644 index 00000000000..a1b2a8c5c91 --- /dev/null +++ b/tests/test_litellm/test_default_branch.py @@ -0,0 +1,240 @@ +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Final + +import pytest + +ROOT: Final = Path(__file__).resolve().parents[2] + + +def _git(repo: Path, *args: str) -> str: + return subprocess.run(["git", *args], cwd=repo, check=True, capture_output=True, text=True).stdout.strip() + + +def _commit(repo: Path, message: str) -> None: + _git(repo, "add", ".") + _git(repo, "-c", "user.name=Test", "-c", "user.email=test@example.com", "commit", "-qm", message) + + +@pytest.fixture +def remote_and_clone(tmp_path: Path) -> tuple[Path, Path]: + seed: Final = tmp_path / "seed" + seed.mkdir() + _git(seed, "init", "-q", "-b", "litellm_internal_staging") + (seed / "scripts").mkdir() + for name in ( + "default_branch.py", + "budget_ratchet_check.py", + "ruff_strict_gate.py", + "type_discipline_gate.py", + "test_quality_gate.py", + "type_check_gate.py", + "gate_slot_lock.py", + ): + shutil.copyfile(ROOT / "scripts" / name, seed / "scripts" / name) + shutil.copyfile(ROOT / "Makefile", seed / "Makefile") + (seed / "litellm").mkdir() + (seed / "litellm" / "example.py").write_text("value = 0\n") + (seed / "ruff-strict-budget.json").write_text('{"C901": {"limit": 1}}\n') + _commit(seed, "staging base") + _git(seed, "checkout", "-qb", "main") + (seed / "litellm" / "example.py").write_text("value = 1\n") + (seed / "ruff-strict-budget.json").write_text('{"C901": {"limit": 0}}\n') + _commit(seed, "main base") + remote: Final = tmp_path / "remote.git" + _git(tmp_path, "clone", "-q", "--bare", str(seed), str(remote)) + _git(remote, "symbolic-ref", "HEAD", "refs/heads/litellm_internal_staging") + repo: Final = tmp_path / "clone" + _git(tmp_path, "clone", "-q", "--single-branch", str(remote), str(repo)) + return remote, repo + + +def _resolve(repo: Path, *args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(ROOT / "scripts" / "default_branch.py"), *args], + cwd=repo, + capture_output=True, + text=True, + check=False, + ) + + +def _make(repo: Path, target: str, *args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["make", target, "LINT_DEP_INSTALL=", "LINT_DEP_BASE=", *args], + cwd=repo, + capture_output=True, + text=True, + check=False, + env={key: value for key, value in os.environ.items() if key != "BASE_REF"}, + ) + + +def test_existing_single_branch_clone_follows_remote_switch(remote_and_clone: tuple[Path, Path]) -> None: + remote, repo = remote_and_clone + before: Final = _resolve(repo) + assert before.returncode == 0, before.stderr + assert before.stdout.strip() == "origin/litellm_internal_staging" + _git(remote, "symbolic-ref", "HEAD", "refs/heads/main") + after: Final = _resolve(repo) + assert after.returncode == 0, after.stderr + assert after.stdout.strip() == "origin/main" + assert _git(repo, "rev-parse", "origin/main") == _git(remote, "rev-parse", "main") + assert _git(repo, "symbolic-ref", "refs/remotes/origin/HEAD").endswith("/litellm_internal_staging") + + +@pytest.mark.parametrize("missing_head", [False, True]) +def test_unverifiable_default_never_uses_cached_head( + remote_and_clone: tuple[Path, Path], + missing_head: bool, +) -> None: + remote, repo = remote_and_clone + if missing_head: + _git(remote, "symbolic-ref", "HEAD", "refs/heads/missing") + else: + _git(repo, "remote", "set-url", "origin", str(remote / "missing")) + result: Final = _resolve(repo) + assert result.returncode != 0 + assert not result.stdout + assert "explicit base ref" in result.stderr + checked: Final = _make(repo, "lint-format-check-changed") + assert checked.returncode != 0 + assert "No changed" not in checked.stdout + + +@pytest.mark.parametrize("base_ref", ["HEAD", "origin/litellm_internal_staging"]) +def test_explicit_base_works_without_remote_access( + remote_and_clone: tuple[Path, Path], + base_ref: str, +) -> None: + remote, repo = remote_and_clone + _git(repo, "remote", "set-url", "origin", str(remote / "missing")) + result: Final = _resolve(repo, "--base", base_ref) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == base_ref + checked: Final = _make(repo, "lint-format-check-changed", f"BASE_REF={base_ref}") + assert checked.returncode == 0, checked.stderr + assert "No changed litellm Python files" in checked.stdout + + +def test_budget_ratchet_compares_against_new_default(remote_and_clone: tuple[Path, Path]) -> None: + remote, repo = remote_and_clone + _git(remote, "symbolic-ref", "HEAD", "refs/heads/main") + resolved: Final = _resolve(repo) + assert resolved.returncode == 0, resolved.stderr + _git(repo, "checkout", "-qb", "litellm_feature", "origin/main") + (repo / "ruff-strict-budget.json").write_text('{"C901": {"limit": 1}}\n') + command: Final = [sys.executable, "scripts/budget_ratchet_check.py"] + checked: Final = subprocess.run(command, cwd=repo, capture_output=True, text=True, check=False) + assert checked.returncode == 1 + assert "limit raised 0 -> 1" in checked.stdout + assert "base origin/main" in checked.stdout + overridden: Final = subprocess.run( + [*command, "--base", "origin/litellm_internal_staging"], + cwd=repo, + capture_output=True, + text=True, + check=False, + ) + assert overridden.returncode == 0, overridden.stdout + overridden.stderr + + +def _freshness(repo: Path, *args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [ + sys.executable, + "-c", + "import sys; from pathlib import Path; " + "from ci_cd.run_migration import _check_branch_freshness; " + "_check_branch_freshness(Path(sys.argv[1]), sys.argv[2] if len(sys.argv) > 2 else None)", + str(repo), + *args, + ], + cwd=ROOT, + capture_output=True, + text=True, + check=False, + ) + + +def test_migration_freshness_refuses_stale_branch_after_switch(remote_and_clone: tuple[Path, Path]) -> None: + remote, repo = remote_and_clone + before: Final = _freshness(repo) + assert before.returncode == 0, before.stderr + assert "Branch freshness OK" in before.stdout + _git(remote, "symbolic-ref", "HEAD", "refs/heads/main") + after: Final = _freshness(repo) + assert after.returncode == 3 + assert "1 commit(s) behind origin/main" in after.stderr + overridden: Final = _freshness(repo, "litellm_internal_staging") + assert overridden.returncode == 0, overridden.stderr + _git(repo, "merge", "--ff-only", "origin/main") + updated: Final = _freshness(repo) + assert updated.returncode == 0, updated.stderr + assert "up to date with origin/main" in updated.stdout + + +def test_migration_freshness_refuses_unavailable_remote(remote_and_clone: tuple[Path, Path]) -> None: + remote, repo = remote_and_clone + _git(repo, "remote", "set-url", "origin", str(remote / "missing")) + result: Final = _freshness(repo) + assert result.returncode == 3 + assert "Could not discover origin's default branch" in result.stderr + explicit: Final = _freshness(repo, "litellm_internal_staging") + assert explicit.returncode == 3 + assert "git fetch origin litellm_internal_staging" in explicit.stderr + + +@pytest.mark.parametrize( + "gate", + [ + "budget_ratchet_check", + "ruff_strict_gate", + "type_discipline_gate", + "test_quality_gate", + "type_check_gate", + ], +) +def test_each_gate_refuses_an_unverifiable_default(remote_and_clone: tuple[Path, Path], gate: str) -> None: + remote, repo = remote_and_clone + _git(repo, "remote", "set-url", "origin", str(remote / "missing")) + result: Final = subprocess.run( + [sys.executable, f"scripts/{gate}.py"], + cwd=repo, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode != 0 + assert "Cannot verify the base branch against origin" in result.stderr + + +@pytest.mark.parametrize( + "target", ["lint-format-check-changed", "lint-test-quality", "lint-test-quality-budget-update"] +) +def test_direct_make_target_fetches_default_once(remote_and_clone: tuple[Path, Path], target: str) -> None: + _, repo = remote_and_clone + trace: Final = repo.parent / "git-trace.jsonl" + shutil.copyfile(ROOT / "scripts" / "check_test_quality.py", repo / "scripts" / "check_test_quality.py") + shutil.copyfile(ROOT / "test-quality-budget.json", repo / "test-quality-budget.json") + (repo / "tests").mkdir() + result: Final = subprocess.run( + ["make", "-o", "install-dev", target, "LINT_DEP_INSTALL=", "UV_RUN=env"], + cwd=repo, + capture_output=True, + text=True, + check=False, + env={**{key: value for key, value in os.environ.items() if key != "BASE_REF"}, "GIT_TRACE2_EVENT": str(trace)}, + ) + assert result.returncode == 0, result.stdout + result.stderr + commands: Final = tuple( + event["argv"][1:] + for line in trace.read_text().splitlines() + if (event := json.loads(line)).get("event") == "start" + ) + assert sum(command[0] == "ls-remote" for command in commands) == 1 + assert sum(command[0] == "fetch" for command in commands) == 1 diff --git a/tests/test_litellm/test_drop_params_env_var.py b/tests/test_litellm/test_drop_params_env_var.py new file mode 100644 index 00000000000..1e0b7801ef1 --- /dev/null +++ b/tests/test_litellm/test_drop_params_env_var.py @@ -0,0 +1,33 @@ +import os +import subprocess +import sys + +import pytest + + +def _import_litellm_with(configured: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, "-c", "import litellm; print(litellm.drop_params)"], + env={**os.environ, "LITELLM_DROP_PARAMS": configured}, + capture_output=True, + text=True, + check=True, + ) + + +@pytest.mark.parametrize("configured, expected", [("false", "False"), ("true", "True"), ("", "False")]) +def test_litellm_drop_params_env_var_is_parsed_as_a_flag(configured, expected): + result = _import_litellm_with(configured) + + assert result.stdout.strip() == expected + assert "is not a flag value" not in result.stderr + + +def test_litellm_drop_params_env_var_non_flag_value_stays_on_with_a_warning(): + result = _import_litellm_with("temperature") + + assert result.stdout.strip() == "True" + assert ( + "LITELLM_DROP_PARAMS='temperature' is not a flag value, treating it as on. Set it to true or false" + in result.stderr + ) diff --git a/tests/test_litellm/test_exception_header_preservation.py b/tests/test_litellm/test_exception_header_preservation.py index 6ea478c633b..dd142d9d40b 100644 --- a/tests/test_litellm/test_exception_header_preservation.py +++ b/tests/test_litellm/test_exception_header_preservation.py @@ -18,6 +18,7 @@ from litellm.exceptions import ( ImageFetchError, MidStreamFallbackError, RateLimitError, + ServiceUnavailableError, ) @@ -312,3 +313,85 @@ class TestProxyHeaderExtraction: # Verify headers are extracted and prefixed correctly assert headers.get("llm_provider-x-request-id") == "req-abc123" assert headers.get("llm_provider-x-ms-region") == "eastus" + + +class TestBedrockErrorHeaders: + """A BedrockError built with headers but no response still exposes them (LIT-5428).""" + + def test_synthesized_response_carries_headers(self): + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError( + status_code=500, + message="Amazon Bedrock is unable to process your request.", + headers={"x-amzn-RequestId": "req-base-500"}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-base-500" + assert str(error.request.url) == str(BedrockError(status_code=500, message="boom").request.url) + assert str(error.response.request.url) == str(error.request.url) + + def test_synthesized_response_without_headers_stays_empty(self): + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError(status_code=500, message="boom") + + assert dict(error.response.headers) == {} + + def test_explicit_response_is_kept(self): + from litellm.llms.bedrock.common_utils import BedrockError + + provider_response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "from-response"}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + error = BedrockError( + status_code=500, + message="boom", + headers={"x-amzn-RequestId": "from-headers"}, + response=provider_response, + ) + + assert error.response is provider_response + + def test_proxy_extraction_surfaces_bedrock_request_id(self): + """End-to-end shape the proxy error handler returns to the caller.""" + from litellm.litellm_core_utils.exception_mapping_utils import exception_type + from litellm.litellm_core_utils.llm_response_utils.get_headers import ( + get_response_headers, + ) + from litellm.llms.bedrock.common_utils import BedrockError + + provider_response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-proxy-500"}, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + with pytest.raises(ServiceUnavailableError) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=BedrockError( + status_code=500, + message=provider_response.text, + headers=provider_response.headers, + response=provider_response, + ), + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + # Mirrors ProxyBaseLLMRequestProcessing._handle_llm_api_exception + error = exc_info.value + headers = getattr(error, "headers", None) or {} + if not headers: + _response = getattr(error, "response", None) + if _response is not None: + _response_headers = getattr(_response, "headers", None) + if _response_headers: + headers = get_response_headers(dict(_response_headers)) + + assert headers.get("llm_provider-x-amzn-requestid") == "req-proxy-500" diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index a7a9e0fc37d..701938f5677 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -14,26 +14,9 @@ import os import pytest -import litellm from litellm.utils import get_model_info -@pytest.fixture(scope="module", autouse=True) -def _local_model_cost_map(): - """ - Point litellm at the bundled cost map for the duration of this module - only. ``mp.undo()`` restores both the environment variable and - ``litellm.model_cost`` so nothing leaks into later tests. - """ - mp = pytest.MonkeyPatch() - mp.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - mp.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - get_model_info.cache_clear() - yield - mp.undo() - get_model_info.cache_clear() - - NEW_ENTRIES = { "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "input_cost_per_token": 1.32e-06, @@ -54,19 +37,6 @@ def model_data(): return json.load(f) -def test_fireworks_serverless_entries_exist(model_data): - """The new prefixed entry carries the pricing and metadata from #37274.""" - for key, expected in NEW_ENTRIES.items(): - assert key in model_data, f"{key} is missing from model_prices_and_context_window.json" - entry = model_data[key] - for field, value in expected.items(): - assert entry[field] == pytest.approx(value), f"{key}.{field}" - assert entry["litellm_provider"] == "fireworks_ai" - assert entry["mode"] == "chat" - assert entry["supports_function_calling"] is True - assert entry["supports_vision"] is False - - def test_bare_fireworks_ids_resolve_through_prefixed_entries(): """Bare IDs from #37274 resolve via the provider-prefix lookup path.""" for bare_id, prefixed_key in [ diff --git a/tests/test_litellm/test_gate_slot_lock.py b/tests/test_litellm/test_gate_slot_lock.py index 17fa8547ce7..c80e876700b 100644 --- a/tests/test_litellm/test_gate_slot_lock.py +++ b/tests/test_litellm/test_gate_slot_lock.py @@ -1,6 +1,8 @@ import fcntl import importlib.util +import json import os +import shlex import signal import subprocess import sys @@ -301,33 +303,47 @@ def test_held_slot_context_manager_releases_on_exit(tmp_path: Path, monkeypatch: fcntl.flock(probe, fcntl.LOCK_UN) -def _make_rule(target: str) -> tuple[list[str], list[str]]: - database = subprocess.run( - ["make", "--dry-run", "--print-data-base", "info"], - cwd=ROOT, - capture_output=True, - text=True, - check=True, - ).stdout - lines = database.splitlines() - for index, line in enumerate(lines): - if line != f"{target}:" and not line.startswith(f"{target}: "): - continue - recipe: list[str] = [] - for follower in lines[index + 1 :]: - if follower.startswith("#"): - continue - if not follower.startswith("\t"): - break - recipe.append(follower.strip()) - return line.split(":", 1)[1].split(), recipe - raise AssertionError(f"target {target} not found in make database") - - -def test_direct_make_lint_takes_a_slot_before_any_setup() -> None: - lint_prerequisites, lint_recipe = _make_rule("lint") - assert lint_prerequisites == [] - assert any("$(GATE_SLOT_LOCK)" in line for line in lint_recipe) - inner_prerequisites, _ = _make_rule("lint-inner") - assert "lint-install" in inner_prerequisites - assert "lint-fetch-base" in inner_prerequisites +def test_direct_make_lint_takes_a_slot_before_any_setup(tmp_path: Path) -> None: + lock_dir = tmp_path / "locks" + lock_dir.mkdir() + events_file = tmp_path / "setup.jsonl" + stderr_file = tmp_path / "make.stderr" + probe = tmp_path / "probe.py" + probe.write_text( + "import fcntl, json, os, pathlib, sys\n" + "with (pathlib.Path(os.environ['LITELLM_GATE_SLOT_DIR']) / 'slot-0.lock').open('wb') as slot:\n" + " try:\n" + " fcntl.flock(slot, fcntl.LOCK_EX | fcntl.LOCK_NB)\n" + " locked = False\n" + " except BlockingIOError:\n" + " locked = True\n" + "with open(os.environ['EVENTS_FILE'], 'a') as events:\n" + " events.write(json.dumps({'phase': sys.argv[1], 'locked': locked}) + '\\n')\n" + "if sys.argv[1] == 'base':\n" + " print('HEAD')\n" + ) + (tmp_path / "Makefile").write_text((ROOT / "Makefile").read_text()) + command = [ + "make", "-o", "lint-checks", "lint", "MAKE=make -o lint-checks", + f"GATE_SLOT_LOCK={shlex.join([sys.executable, str(HELPER)])}", + f"UV={shlex.join([sys.executable, str(probe), 'setup'])}", + f"UV_RUN={shlex.join([sys.executable, str(probe), 'setup'])}", + f"RESOLVE_BASE={shlex.join([sys.executable, str(probe), 'base'])}", + ] + with (lock_dir / "slot-0.lock").open("wb") as held, stderr_file.open("wb") as stderr: + fcntl.flock(held, fcntl.LOCK_EX) + process = subprocess.Popen( + command, cwd=tmp_path, stdout=subprocess.DEVNULL, stderr=stderr, + env={**_env(lock_dir, "1"), "EVENTS_FILE": str(events_file)}, + ) + try: + assert _wait_until(lambda: "queueing" in stderr_file.read_text(), 10) + assert not events_file.exists() + fcntl.flock(held, fcntl.LOCK_UN) + assert process.wait(timeout=30) == 0, stderr_file.read_text() + finally: + fcntl.flock(held, fcntl.LOCK_UN) + _reap(process) + events = tuple(json.loads(line) for line in events_file.read_text().splitlines()) + assert {event["phase"] for event in events} == {"setup", "base"} + assert all(event["locked"] for event in events) diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/test_litellm/test_gpt_5_5_model_metadata.py index a60fa9466e6..e07efbcc913 100644 --- a/tests/test_litellm/test_gpt_5_5_model_metadata.py +++ b/tests/test_litellm/test_gpt_5_5_model_metadata.py @@ -1,54 +1,6 @@ import json from pathlib import Path -import pytest - -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - -@pytest.mark.parametrize("model", ["azure_ai/gpt-5.5", "azure_ai/gpt-5.5-2026-04-23"]) -def test_azure_ai_gpt_5_5_model_info(model): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "azure_ai" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 3e-05 - assert info["cache_read_input_token_cost"] == 5e-07 - - assert info["input_cost_per_token_above_272k_tokens"] == 1e-05 - assert info["output_cost_per_token_above_272k_tokens"] == 4.5e-05 - assert info["cache_read_input_token_cost_above_272k_tokens"] == 1e-06 - - assert info["input_cost_per_token_priority"] == 1e-05 - assert info["output_cost_per_token_priority"] == 6e-05 - - assert info["max_input_tokens"] == 1050000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_web_search"] is True - # gpt-5.5 dropped minimal reasoning effort support (true on gpt-5.4) - assert info["supports_minimal_reasoning_effort"] is False - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.split("/", 1)[1] - assert provider == "azure_ai" - def test_azure_ai_gpt_5_5_backup_matches_main(): """Ensure the bundled model cost map stays in sync with the canonical file.""" diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/test_litellm/test_gpt_realtime_mode.py index 314fd63c4cc..8c41e474486 100644 --- a/tests/test_litellm/test_gpt_realtime_mode.py +++ b/tests/test_litellm/test_gpt_realtime_mode.py @@ -1,10 +1,8 @@ import json from pathlib import Path -import pytest from typing_extensions import get_args, get_type_hints -import litellm from litellm.types.utils import ModelInfoBase REALTIME_ONLY_GPT_MODELS = ( @@ -43,44 +41,11 @@ REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS = ( ALL_REALTIME_ONLY_GPT_MODELS = REALTIME_ONLY_GPT_MODELS + REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS -def _load_cost_map() -> dict: - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - return json.load(f) - - def test_realtime_is_a_valid_mode_literal(): hints = get_type_hints(ModelInfoBase, include_extras=False) assert "realtime" in get_args(hints["mode"]) -@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS) -def test_realtime_only_gpt_models_are_mode_realtime(model): - """These models only serve /v1/realtime and are rejected by /v1/chat/completions - ("This is not a chat model ..."), so they must not be tagged mode=chat.""" - info = _load_cost_map()[model] - assert info["supported_endpoints"] == ["/v1/realtime"] - assert info["mode"] == "realtime" - - -@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS) -def test_realtime_only_gpt_4o_models_are_mode_realtime(model): - """gpt-4o(-mini)-realtime-preview are realtime-only and must not be mode=chat.""" - assert _load_cost_map()[model]["mode"] == "realtime" - - -def test_get_model_info_reports_realtime_mode(monkeypatch): - """get_model_info must resolve the retag against the bundled cost map, not the - hosted map fetched from main, which lags this repo until the next promotion.""" - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - litellm.get_model_info.cache_clear() - try: - assert litellm.get_model_info("gpt-realtime-mini")["mode"] == "realtime" - finally: - litellm.get_model_info.cache_clear() - - def test_backup_matches_main_for_realtime_models(): repo_root = Path(__file__).parents[2] with open(repo_root / "model_prices_and_context_window.json") as f: diff --git a/tests/test_litellm/test_lint_workflow_diff_gates.py b/tests/test_litellm/test_lint_workflow_diff_gates.py new file mode 100644 index 00000000000..62e67cdaaab --- /dev/null +++ b/tests/test_litellm/test_lint_workflow_diff_gates.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import re +import shlex +import subprocess +from pathlib import Path +from typing import Final + +import pytest + +WORKFLOW: Final = Path(__file__).resolve().parents[2] / ".github" / "workflows" / "test-linting.yml" +DIFF_GATE: Final = re.compile(r'git diff --name-only --diff-filter=\w+ "\$GATE_BASE_SHA" HEAD -- (.+?) \|') +GATES: Final = tuple(tuple(shlex.split(gate.group(1))) for gate in DIFF_GATE.finditer(WORKFLOW.read_text())) + + +def _git(cwd: Path, *args: str) -> str: + return subprocess.run(["git", *args], cwd=cwd, check=True, capture_output=True, text=True).stdout + + +def _scoped_root(pathspec: str) -> str: + return re.sub(r"^:\([^)]*\)", "", pathspec).split("*", 1)[0] + + +def _changed_files_selected_by(tmp_path: Path, pathspecs: tuple[str, ...], files: tuple[str, ...]) -> frozenset[str]: + _git(tmp_path, "init", "-q", "-b", "main") + _git(tmp_path, "config", "user.email", "t@t") + _git(tmp_path, "config", "user.name", "t") + _git(tmp_path, "commit", "-q", "--allow-empty", "-m", "base") + for name in files: + target = tmp_path / name + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text("x = 1\n") + _git(tmp_path, "add", "-A") + _git(tmp_path, "commit", "-qm", "change") + return frozenset( + _git(tmp_path, "diff", "--name-only", "--diff-filter=ACMRD", "HEAD~1", "HEAD", "--", *pathspecs).split() + ) + + +def test_workflow_still_carries_the_ruff_format_and_e2e_basedpyright_diff_gates() -> None: + assert frozenset(_scoped_root(gate[0]) for gate in GATES) == frozenset({"litellm/", "tests/e2e/"}) + + +@pytest.mark.parametrize("pathspecs", GATES, ids=" ".join) +def test_diff_gate_selects_top_level_and_nested_python_files_only(tmp_path: Path, pathspecs: tuple[str, ...]) -> None: + root = _scoped_root(pathspecs[0]) + top_level = f"{root}top_level_module.py" + nested = f"{root}pkg/sub/nested_module.py" + selected = _changed_files_selected_by( + tmp_path, + pathspecs, + (top_level, nested, f"{root}notes.md", "elsewhere/top_level_module.py", "elsewhere/pkg/nested_module.py"), + ) + assert selected == frozenset({top_level, nested}) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 715ca8672b2..f71225c6fc5 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -120,7 +120,7 @@ def test_completion_missing_role(openai_api_response): print(f"openai_api_response: {openai_api_response}") with patch.object( - client.chat.completions.with_raw_response, "create", mock_raw_response + client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) ) as mock_create: litellm.completion( model="gpt-4o-mini", @@ -1367,6 +1367,78 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( } +@pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) +def test_responses_bridge_preserves_reasoning_effort_with_drop_params( + reasoning_effort, + restore_model_registry, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + response_body: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "test-responses-bridge", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, + } + response_route: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").respond(json=response_body) + model: Final = "perplexity/test-responses-bridge" + litellm.register_model( + { + model: { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_reasoning": False, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + api_base="https://api.perplexity.ai", + ) + + request_body: Final = json.loads(response_route.calls[0].request.content) + assert request_body["reasoning"] == {"effort": "high"} + + @pytest.mark.parametrize( "model, model_info, expected_model_param, expected_base_model_param", [ @@ -2322,6 +2394,280 @@ def test_image_edit_merges_headers_and_extra_headers(): assert "extra_headers" not in handler_kwargs["image_edit_optional_request_params"] +@pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata")) +@pytest.mark.parametrize("input_tokens", (51234, 0)) +def test_mock_completion_usage_reports_admission_input_tokens(metadata_key: str, input_tokens: int): + response = litellm.completion( + model="anthropic/claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + api_key="mock", + **{metadata_key: {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}}}, + ) + + assert response.usage.prompt_tokens == input_tokens + assert response.usage.total_tokens == input_tokens + response.usage.completion_tokens + + +def test_mock_completion_usage_falls_back_to_default_without_admission_count(): + response = litellm.completion( + model="anthropic/claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + api_key="mock", + metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, + ) + + assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + + +_ADMISSION_INPUT_TOKENS: Final = 51234 + + +def _admission_metadata(input_tokens: int) -> dict[str, object]: # mutable-ok: logging writes into metadata + return {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}} + + +_ADMISSION_METADATA: Final = _admission_metadata(_ADMISSION_INPUT_TOKENS) +_MOCK_STREAM_MESSAGES: Final = [{"role": "user", "content": "hello " * 200}] +_STREAM_CHUNK_BUILDER_TOKEN_COUNTER: Final = "litellm.litellm_core_utils.streaming_chunk_builder_utils.token_counter" + + +def _prompt_token_counter_calls(token_counter: MagicMock) -> list[object]: + return [call for call in token_counter.call_args_list if call.kwargs.get("messages") is not None] + + +def _client_usage_chunks(chunks: list[ModelResponseStream]) -> list[Usage]: + return [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] + + +@pytest.mark.parametrize("n", (None, 2)) +def test_mock_completion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback(n: int | None): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + metadata=_ADMISSION_METADATA, + ) + ) + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].completion_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("n", (None, 2)) +async def test_mock_acompletion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback( + n: int | None, +): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + litellm_metadata=_ADMISSION_METADATA, + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +def test_mock_completion_stream_without_include_usage_hides_usage_chunk_but_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + metadata=_ADMISSION_METADATA, + ) + ) + + assert _client_usage_chunks(chunks) == [] + assert all(len(chunk.choices) == 1 for chunk in chunks) + assert chunks[-1]._hidden_params["usage"].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_completion_stream_with_empty_stream_options_completes_and_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={}, + metadata=_ADMISSION_METADATA, + ) + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert _client_usage_chunks(chunks) == [] + assert _prompt_token_counter_calls(token_counter) == [] + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_with_empty_stream_options_completes_and_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={}, + litellm_metadata=_ADMISSION_METADATA, + ) + chunks: Final = [chunk async for chunk in response] + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert _client_usage_chunks(chunks) == [] + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_completion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, + ) + ) + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert usage_chunks[0].total_tokens == expected_prompt_tokens + usage_chunks[0].completion_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +def _usage_triple(usage: Usage) -> tuple[int, int, int]: + return (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) + + +@pytest.mark.parametrize("input_tokens", (_ADMISSION_INPUT_TOKENS, 0)) +def test_mock_completion_stream_and_non_stream_report_the_same_admission_usage(input_tokens: int): + metadata: Final = _admission_metadata(input_tokens) + non_stream: Final = litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + metadata=metadata, + ) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata=metadata, + ) + ) + + assert _usage_triple(non_stream.usage) == _usage_triple(_client_usage_chunks(chunks)[0]) + assert non_stream.usage.prompt_tokens == input_tokens + assert _prompt_token_counter_calls(token_counter) == [] + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_reports_zero_admission_input_tokens_without_tokenizer_fallback(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=[{"role": "user", "content": ""}], + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + litellm_metadata=_admission_metadata(0), + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert _usage_triple(usage_chunks[0]) == (0, usage_chunks[0].completion_tokens, usage_chunks[0].completion_tokens) + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_text_completion_stream_and_non_stream_report_the_same_zero_admission_usage(): + metadata: Final = _admission_metadata(0) + non_stream: Final = litellm.text_completion( + model="openai/gpt-5.4-mini", prompt="", mock_response="ok", api_key="mock", metadata=metadata + ) + chunks: Final = list( + litellm.text_completion( + model="openai/gpt-5.4-mini", + prompt="", + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata=metadata, + ) + ) + + stream_usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None) + assert len(stream_usages) == 1 + assert _usage_triple(non_stream.usage) == _usage_triple(stream_usages[0]) + assert non_stream.usage.prompt_tokens == 0 + + def test_mock_completion_stream_with_model_response(): """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" from litellm import completion @@ -3351,6 +3697,52 @@ def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeyp assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) +def test_speech_mistral_dispatches_and_decodes_audio(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-fake-mp3-bytes" + mock_route: Final = respx_mock.post("https://api.mistral.ai/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + response_format="wav", + speed=2, + instructions="sound cheerful", + ) + + assert mock_route.called + request_body: Final = json.loads(mock_route.calls.last.request.content) + assert request_body == { + "model": "voxtral-mini-tts-2603", + "input": "hello from litellm", + "voice_id": "en_paul_neutral", + "response_format": "wav", + } + assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-mistral-test" + assert response.content == audio_bytes + + +def test_speech_mistral_routes_to_configured_api_base(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-gateway-bytes" + gateway_route: Final = respx_mock.post("https://mistral.gateway.internal/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + api_base="https://mistral.gateway.internal", + ) + + assert gateway_route.called + assert response.content == audio_bytes + + FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com" diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py index 6f1ba702d8d..d73311baae9 100644 --- a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py +++ b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py @@ -3,8 +3,6 @@ from pathlib import Path import pytest -import litellm -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider REPO_ROOT = Path(__file__).parents[2] MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -27,56 +25,6 @@ def _load(path): return json.load(f) - -@pytest.mark.parametrize("model", MEDIUM_3_5_MODELS) -def test_medium_3_5_specs(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "mistral" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 1.5e-06 - assert info["output_cost_per_token"] == 7.5e-06 - - assert info["max_input_tokens"] == 262144 - assert info["max_output_tokens"] == 262144 - assert info["max_tokens"] == 262144 - - assert info["supports_reasoning"] is True - assert info["supports_vision"] is True - assert info["supports_function_calling"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_assistant_prefill"] is True - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.split("/", 1)[1] - assert provider == "mistral" - - -def test_mistral_medium_latest_resolves_to_medium_3_5(local_model_cost_map): - """LIT-3883: the -latest alias was retargeted to Medium 3.5; get_model_info must - return the 3.5 pricing/context/reasoning, not the stale Medium 3.1 values.""" - info = litellm.get_model_info(model="mistral/mistral-medium-latest") - - assert info["input_cost_per_token"] == 1.5e-06 - assert info["output_cost_per_token"] == 7.5e-06 - assert info["max_input_tokens"] == 262144 - assert info["supports_reasoning"] is True - - -def test_mistral_medium_2508_keeps_medium_3_1_specs(): - """The date-pinned 2508 alias is Medium 3.1 and must not inherit 3.5 pricing.""" - info = _load(MAIN_PATH).get("mistral/mistral-medium-2508") - assert info is not None, "mistral/mistral-medium-2508 missing from cost map" - - assert info["input_cost_per_token"] == 4e-07 - assert info["output_cost_per_token"] == 2e-06 - assert info["max_input_tokens"] == 131072 - assert info.get("supports_reasoning") is not True - - @pytest.mark.parametrize("model", SYNCED_MODELS) def test_backup_matches_main(model): """Ensure the bundled (backup) cost map stays in sync with the canonical file.""" diff --git a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py index 0442321ba0b..182c444bac9 100644 --- a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py +++ b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py @@ -18,29 +18,6 @@ def _load(path): return json.load(f) -@pytest.mark.parametrize("model", SMALL_4_0_MODELS) -def test_small_4_0_specs(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "mistral" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 1.5e-07 - assert info["output_cost_per_token"] == 6e-07 - - assert info["max_input_tokens"] == 262144 - assert info["max_output_tokens"] == 262144 - assert info["max_tokens"] == 262144 - - assert info["supports_reasoning"] is True - assert info["supports_vision"] is True - assert info["supports_function_calling"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_assistant_prefill"] is True - - @pytest.mark.parametrize("model", SMALL_4_0_MODELS) def test_backup_matches_main(model): main_cost = _load(MAIN_PATH) diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index 6114d1d8aba..0b9dbd23097 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -4,10 +4,15 @@ import importlib.util import json import re from pathlib import Path +from types import MappingProxyType +from typing import Final import jsonschema import pytest +from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name +from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + REPO_ROOT = Path(__file__).parents[2] GENERATOR_PATH = REPO_ROOT / "ci_cd" / "generate_model_prices_schema.py" PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -173,3 +178,91 @@ def test_dated_variants_carry_base_alias_service_tier_pricing(prices: dict): "sync the tier keys so service-tier requests against pinned snapshots are not " "billed at standard rates:\n" + "\n".join(drifted) ) + + +OPENAI_REASONING_FAMILY_MARKERS = ("codex", "deep-research", "chat-latest") + + +def is_openai_o_series(name: str) -> bool: + return len(name) > 1 and name[0] == "o" and name[1].isdigit() + + +def is_openai_reasoning_family(name: str) -> bool: + base = name.split("/")[-1].removeprefix("ft:") + if "search-api" in base: + return False + return ( + is_openai_o_series(base) + or is_gpt_reasoning_series_name(base) + or any(marker in base for marker in OPENAI_REASONING_FAMILY_MARKERS) + ) + + +def test_openai_reasoning_family_entries_carry_supports_reasoning(prices: dict): + unflagged = [ + name + for name, entry in prices.items() + if isinstance(entry, dict) + and entry.get("litellm_provider") == "openai" + and is_openai_reasoning_family(name) + and entry.get("supports_reasoning") is not True + ] + assert unflagged == [], ( + "OpenAI o-series, gpt-5+, codex, deep-research, and chat-latest models are reasoning " + "models, and the Responses API drops the `reasoning` param for any mapped OpenAI model " + "whose entry lacks supports_reasoning; flag these entries:\n" + "\n".join(unflagged) + ) + + +def test_chat_latest_declares_the_one_effort_openai_accepts(prices: dict): + """OpenAI rejects every reasoning.effort on chat-latest except medium, and a reasoning entry + with no declared levels resolves to None, which lets /model_group/info and the dashboard offer + levels the upstream will 400 on.""" + assert resolve_supported_reasoning_efforts(prices["chat-latest"], deployment_is_mapped=True) == ("medium",) + + +BEDROCK_OPENAI_GPT_MARKERS: Final = ("openai.gpt-5.4", "openai.gpt-5.5", "openai.gpt-5.6", "openai.gpt-6-astra") +BEDROCK_PROVIDERS: Final = frozenset(("bedrock", "bedrock_converse", "bedrock_mantle")) +BEDROCK_ROW_PREFIXES: Final = ("bedrock_mantle/", "us.", "global.") +GPT_5_4_BEDROCK_LADDER: Final = ("none", "low", "medium", "high", "xhigh") +GPT_5_6_BEDROCK_LADDER: Final = ("none", "low", "medium", "high", "xhigh", "max") +GPT_6_ASTRA_BEDROCK_LADDER: Final = ("low", "medium", "high", "xhigh", "max") +BEDROCK_OPENAI_GPT_LADDERS: Final = MappingProxyType( + { + "bedrock_mantle/openai.gpt-5.4": GPT_5_4_BEDROCK_LADDER, + "bedrock_mantle/openai.gpt-5.5": GPT_5_4_BEDROCK_LADDER, + **{ + f"{prefix}openai.gpt-5.6-{variant}": GPT_5_6_BEDROCK_LADDER + for prefix in BEDROCK_ROW_PREFIXES + for variant in ("luna", "sol", "terra") + }, + **{f"{prefix}openai.gpt-6-astra": GPT_6_ASTRA_BEDROCK_LADDER for prefix in BEDROCK_ROW_PREFIXES}, + } +) + + +@pytest.mark.parametrize( + ("name", "ladder"), tuple(BEDROCK_OPENAI_GPT_LADDERS.items()), ids=tuple(BEDROCK_OPENAI_GPT_LADDERS) +) +def test_bedrock_openai_gpt_rows_advertise_the_ladder_bedrock_accepts(prices: dict, name: str, ladder: tuple[str, ...]): + """Each ladder is the set of levels Bedrock answered 200 to for that row through the proxy on + 2026-09-11 (PR #40740): the Mantle rows go out over its Responses endpoint and the Converse rows + over inference profiles. Bedrock differs from the direct OpenAI rows in two places, gpt-5.6 and + gpt-6-astra take max there, and gpt-6-astra refuses none; minimal is refused on every row. + xhigh and max are opt-in for the resolver, so a row missing either flag silently drops that + level from every group it belongs to.""" + assert resolve_supported_reasoning_efforts(prices[name], deployment_is_mapped=True) == ladder + + +def test_every_bedrock_openai_gpt_row_advertises_xhigh(prices: dict): + """The GovCloud and gpt-5.6-cyber rows cannot be called from our account, so they carry the + family's xhigh flag rather than a measured ladder.""" + missing: Final = [ + name + for name, entry in prices.items() + if isinstance(entry, dict) + and entry.get("litellm_provider") in BEDROCK_PROVIDERS + and any(marker in name for marker in BEDROCK_OPENAI_GPT_MARKERS) + and "xhigh" not in (resolve_supported_reasoning_efforts(entry, deployment_is_mapped=True) or ()) + ] + assert missing == [] diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py index 0587883aa44..02527a98711 100644 --- a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py +++ b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py @@ -23,46 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di return json.load(f) - -@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) -def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): - info = _load_cost_map().get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "meta" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cached_cost - - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 131072 - assert info["max_tokens"] == 131072 - - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_pdf_input"] is True - assert info["supports_web_search"] is True - assert info["supports_minimal_reasoning_effort"] is True - assert info["supports_xhigh_reasoning_effort"] is True - - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] - assert info["supported_modalities"] == ["text", "image", "video"] - assert info["supported_output_modalities"] == ["text"] - - assert info["search_context_cost_per_query"] == { - "search_context_size_high": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_low": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY, - } - - @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) def test_muse_spark_1_2_cost_per_token( local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float diff --git a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py index 1ecd9490f78..92b099fc780 100644 --- a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py +++ b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py @@ -23,46 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di return json.load(f) - -@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) -def test_muse_spark_1_3_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): - info = _load_cost_map().get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "meta" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cached_cost - - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 131072 - assert info["max_tokens"] == 131072 - - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_pdf_input"] is True - assert info["supports_web_search"] is True - assert info["supports_minimal_reasoning_effort"] is True - assert info["supports_xhigh_reasoning_effort"] is True - - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] - assert info["supported_modalities"] == ["text", "image", "video"] - assert info["supported_output_modalities"] == ["text"] - - assert info["search_context_cost_per_query"] == { - "search_context_size_high": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_low": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY, - } - - @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) def test_muse_spark_1_3_cost_per_token( local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index aa2260e89ea..e12da0833dc 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -11,6 +11,12 @@ import pytest ROOT = Path(__file__).resolve().parents[2] SCRIPT = ROOT / "scripts" / "pre_commit_lint.sh" +WHOLE_TREE_RUFF = "run --no-sync ruff check --config ruff-tests.toml tests" +TEST_TREE_RAN = "ran: test-tree lint (ruff-tests.toml + test-quality budget)" +TEST_TREE_SKIPPED = ( + "skipped: test-tree lint (ruff-tests.toml + test-quality budget) " + "(no tests/ Python files or test-tree lint inputs in scope)" +) BARRIER_HELPER = """barrier_sync() { touch "$STUB_BARRIER_DIR/$1.started" @@ -30,6 +36,7 @@ BARRIER_HELPER = """barrier_sync() { MAKE_STUB = """#!/bin/sh . "$STUB_BIN/barrier.sh" +[ -n "${STUB_ARGS_DIR:-}" ] && echo "$*" >> "$STUB_ARGS_DIR/make.args" case "$*" in lint) [ "${STUB_FAIL:-}" = "make-lint" ] && exit 1 @@ -40,6 +47,9 @@ case "$*" in sleep 60 fi ;; + lint-test-quality) + [ "${STUB_FAIL:-}" = "test-quality" ] && exit 1 + ;; esac exit 0 """ @@ -69,6 +79,10 @@ case "$*" in *orjson*) [ -n "${STUB_BARRIER_DIR:-}" ] && barrier_sync genapi "python dashboard" ;; + "run --no-sync ruff check --config ruff-tests.toml"*) + [ -n "${STUB_ARGS_DIR:-}" ] && echo "$*" >> "$STUB_ARGS_DIR/ruff_tests.args" + [ "${STUB_FAIL:-}" = "tests-ruff" ] && exit 1 + ;; esac exit 0 """ @@ -145,18 +159,26 @@ def _commit_all(repo: Path, message: str) -> None: ) -def _set_base_ref(repo: Path) -> None: - subprocess.run( - ["git", "update-ref", "refs/remotes/origin/litellm_internal_staging", "HEAD"], - cwd=repo, - check=True, - ) +def _set_base_ref(repo: Path, branch: str = "litellm_internal_staging") -> None: + remote = repo.parent / "remote.git" + subprocess.run(["git", "clone", "-q", "--bare", str(repo), str(remote)], check=True) + subprocess.run(["git", "update-ref", f"refs/heads/{branch}", "HEAD"], cwd=remote, check=True) + subprocess.run(["git", "symbolic-ref", "HEAD", f"refs/heads/{branch}"], cwd=remote, check=True) + subprocess.run(["git", "remote", "add", "origin", str(remote)], cwd=repo, check=True) -def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path) -> None: +def _stage_file(repo: Path, relative: str, body: str) -> None: + path = repo / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(body) + subprocess.run(["git", "add", relative], cwd=repo, check=True) + + +@pytest.mark.parametrize("branch", ["litellm_internal_staging", "main"]) +def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path, branch: str) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") - _set_base_ref(repo) + _set_base_ref(repo, branch) (repo / "litellm" / "foo.py").write_text("x = 2\n") proc = _run(repo, bin_dir, {}) assert proc.returncode == 0, proc.stdout + proc.stderr @@ -240,8 +262,8 @@ def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Pat _commit_all(repo, "base") proc = _run(repo, bin_dir, {}) assert proc.returncode == 1 - assert "cannot resolve the merge base" in proc.stdout - assert "git fetch origin litellm_internal_staging" in proc.stdout + assert "Cannot verify the base branch against origin" in proc.stdout + assert "explicit base ref" in proc.stdout assert "check: FAIL" in proc.stdout @@ -405,6 +427,7 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No assert "ran: dashboard lint (prettier + eslint + lint budgets)" in proc.stdout assert "ran: dashboard API-type sync (npm run gen:api)" in proc.stdout assert "skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope)" in proc.stdout + assert TEST_TREE_SKIPPED in proc.stdout assert "check: PASS" in proc.stdout assert "check: FAIL" not in proc.stdout @@ -412,20 +435,146 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty_log(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") - tests_dir = repo / "tests" / "test_litellm" - tests_dir.mkdir(parents=True) - (tests_dir / "test_x.py").write_text("def test_x() -> None: ...\n") - subprocess.run(["git", "add", "tests"], cwd=repo, check=True) + _stage_file(repo, "scripts/tool.py", "def main() -> None: ...\n") proc = _run(repo, bin_dir, {}) assert proc.returncode == 0, proc.stdout + proc.stderr assert "no gating lint check matches the files in scope, so nothing ran" in proc.stdout - assert "tests/test_litellm/test_x.py" in proc.stdout + assert "scripts/tool.py" in proc.stdout assert "a no-op, not a lint verdict" in proc.stdout assert "check: PASS" in proc.stdout assert "linting Python" not in proc.stdout log = (repo / ".git" / "pre_commit_lint.log").read_text() assert "check: summary" in log assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log + assert TEST_TREE_SKIPPED in log + + +def _recorded(args_dir: Path, name: str) -> list[str]: + path = args_dir / name + return path.read_text().splitlines() if path.exists() else [] + + +def test_tests_only_change_runs_the_whole_test_tree_ruff_and_the_quality_gate(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _stage_file(repo, "tests/fixtures/data.json", "{}\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + assert TEST_TREE_RAN in proc.stdout + assert "no gating lint check matches" not in proc.stdout + assert "linting Python" not in proc.stdout + assert "check: PASS" in proc.stdout + + +@pytest.mark.parametrize( + "changed", + [ + "ruff-tests.toml", + "test-quality-budget.json", + "scripts/check_test_quality.py", + "scripts/test_quality_gate.py", + "tests/e2e/test_x.py", + ], +) +def test_test_tree_lint_inputs_trigger_the_test_tree_checks(tmp_path: Path, changed: str) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, changed, "x = 1\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert "lint-test-quality" in _recorded(args_dir, "make.args") + assert TEST_TREE_RAN in proc.stdout + + +def test_nothing_staged_tests_only_working_tree_change_runs_the_test_tree_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _commit_all(repo, "base") + _set_base_ref(repo) + args_dir = tmp_path / "args" + args_dir.mkdir() + (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n assert True\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged; scoping to the working tree's diff" in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + + +def test_a_failing_test_tree_ruff_fails_the_run_and_still_runs_the_quality_gate(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "tests-ruff"}) + assert proc.returncode == 1 + assert "Test-tree ruff failed" in proc.stdout + proc.stderr + assert "check: FAIL" in proc.stdout + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + + +def test_a_failing_quality_gate_fails_a_tests_only_run(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + proc = _run(repo, bin_dir, {"STUB_FAIL": "test-quality"}) + assert proc.returncode == 1 + assert "Test-quality budget failed" in proc.stdout + proc.stderr + assert "check: FAIL" in proc.stdout + + +def test_tests_changed_alongside_litellm_files_defer_to_make_lint(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "litellm/foo.py", "x = 2\n") + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "linting Python" in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [] + assert _recorded(args_dir, "make.args") == ["lint"] + assert TEST_TREE_RAN in proc.stdout + + +def test_deleted_test_file_still_runs_the_test_tree_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + subprocess.run(["git", "rm", "-q", "tests/test_a.py"], cwd=repo, check=True) + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + assert TEST_TREE_RAN in proc.stdout + + +def test_partial_staging_warns_when_test_files_are_left_unstaged(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "notes.md", "hi\n") + (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n assert True\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "SKIPPED test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout + assert "tests/test_a.py" in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [] + assert _recorded(args_dir, "make.args") == [] def test_run_queues_through_the_machine_wide_gate_slot_lock(tmp_path: Path) -> None: @@ -474,3 +623,28 @@ def test_failing_run_ends_with_a_fail_verdict(tmp_path: Path) -> None: assert proc.returncode == 1 assert "check: FAIL" in proc.stdout assert "check: PASS" not in proc.stdout + + + +def test_explicit_base_scopes_offline_without_a_remote(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + (repo / "litellm" / "foo.py").write_text("x = 2\n") + proc = _run(repo, bin_dir, {"BASE_REF": "HEAD"}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "merge base with HEAD" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_symlinked_hook_can_resolve_default_branch(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo, "main") + hook = repo / ".git" / "hooks" / "pre-commit" + hook.symlink_to(SCRIPT) + proc = subprocess.run( + [str(hook)], cwd=repo, capture_output=True, text=True, + env=_env(repo, bin_dir, {}), timeout=120, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "no branch changes vs origin/main" in proc.stdout diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index a96e8541e06..0e8c86d26df 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -10,13 +10,16 @@ from redis.credentials import CredentialProvider import litellm from litellm._redis import ( + _AWS_IAM_KWARG_NAMES, _async_auth_kwargs, + _coerce_redis_kwargs_types, _get_redis_client_logic, _get_redis_cluster_kwargs, _get_redis_env_kwarg_mapping, _get_redis_kwargs, _get_redis_url_kwargs, _pretty_print_redis_config, + _uses_tls, get_redis_async_client, get_redis_client, get_redis_connection_pool, @@ -24,12 +27,14 @@ from litellm._redis import ( ) from litellm._redis_credential_provider import ( AzureADCredentialProvider, + ElastiCacheIAMCredentialProvider, GCPIAMCredentialProvider, _token_cache, ) from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import REDIS_CLUSTER_HEALTH_CHECK_INTERVAL +from litellm.proxy._types import CoordinationRedisParams class _StubCredentialProvider(CredentialProvider): @@ -78,6 +83,8 @@ def clean_redis_environment(monkeypatch): "REDIS_URL", "REDIS_CLUSTER_NODES", "REDIS_SENTINEL_NODES", + "AWS_REGION", + "AWS_DEFAULT_REGION", *_get_redis_env_kwarg_mapping(), ): monkeypatch.delenv(var, raising=False) @@ -110,6 +117,29 @@ def test_credential_provider_is_not_environment_derived(): assert "credential_provider" not in mapping.values() +_AWS_IAM_SETTINGS = { + "aws_iam_auth", + "aws_iam_user_name", + "aws_iam_cache_name", + "aws_iam_region", + "aws_iam_serverless", +} + + +def test_aws_iam_settings_are_environment_derived(): + allowed = _get_redis_kwargs() + mapping = _get_redis_env_kwarg_mapping() + + assert _AWS_IAM_SETTINGS <= allowed + assert set(_AWS_IAM_KWARG_NAMES) == _AWS_IAM_SETTINGS + assert {f for f in CoordinationRedisParams.model_fields if f.startswith("aws_iam_")} == _AWS_IAM_SETTINGS + assert mapping["REDIS_AWS_IAM_AUTH"] == "aws_iam_auth" + assert mapping["REDIS_AWS_IAM_USER_NAME"] == "aws_iam_user_name" + assert mapping["REDIS_AWS_IAM_CACHE_NAME"] == "aws_iam_cache_name" + assert mapping["REDIS_AWS_IAM_REGION"] == "aws_iam_region" + assert mapping["REDIS_AWS_IAM_SERVERLESS"] == "aws_iam_serverless" + + def test_sync_direct_preserves_credential_provider_identity(clean_redis_environment): provider = _StubCredentialProvider() @@ -300,6 +330,333 @@ def test_gcp_kwargs_never_survive_client_logic(clean_redis_environment, override assert "gcp_ssl_ca_certs" not in redis_kwargs +def test_aws_iam_environment_settings_install_provider(clean_redis_environment, monkeypatch): + monkeypatch.setenv("REDIS_AWS_IAM_AUTH", "true") + monkeypatch.setenv("REDIS_AWS_IAM_USER_NAME", "iam-user") + monkeypatch.setenv("REDIS_AWS_IAM_CACHE_NAME", "cache.example.com") + monkeypatch.setenv("REDIS_AWS_IAM_REGION", "us-east-1") + monkeypatch.setenv("REDIS_AWS_IAM_SERVERLESS", "1") + monkeypatch.setenv("REDIS_SSL", "1") + + redis_kwargs = _get_redis_client_logic(host="cache.example.com", port=6379) + + provider = redis_kwargs["credential_provider"] + assert isinstance(provider, ElastiCacheIAMCredentialProvider) + assert provider._is_serverless is True + assert not _AWS_IAM_SETTINGS & redis_kwargs.keys() + + +@pytest.mark.parametrize( + "transport", + [ + pytest.param({"host": "cache.example.com", "port": 6379}, id="host_without_ssl"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": False}, id="host_ssl_false"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "false"}, id="host_ssl_false_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "0"}, id="host_ssl_zero_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "no"}, id="host_ssl_no_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "off"}, id="host_ssl_off_string"), + pytest.param({"url": "redis://cache.example.com:6379", "ssl": True}, id="plaintext_url"), + pytest.param( + {"startup_nodes": [{"host": "cache.example.com", "port": 6379}]}, + id="cluster_without_ssl", + ), + pytest.param( + { + "url": "rediss://cache.example.com:6379", + "startup_nodes": [{"host": "cache.example.com", "port": 6379}], + }, + id="cluster_without_ssl_ignores_url_scheme", + ), + pytest.param( + { + "sentinel_nodes": [("sentinel.example.com", 26379)], + "service_name": "cache", + }, + id="sentinel_without_ssl", + ), + ], +) +def test_aws_iam_auth_rejects_non_tls_connections(clean_redis_environment, transport): + with pytest.raises(ValueError, match="requires TLS"): + _get_redis_client_logic( + **transport, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + +@pytest.mark.parametrize( + "transport", + [ + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": True}, id="host"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "true"}, id="host_ssl_true_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "True"}, id="host_ssl_true_capitalized"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "1"}, id="host_ssl_one_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "yes"}, id="host_ssl_yes_string"), + pytest.param( + {"startup_nodes": [{"host": "cache.example.com", "port": 6379}], "ssl": "1"}, + id="cluster_ssl_one_string", + ), + pytest.param({"url": "rediss://cache.example.com:6379"}, id="url"), + pytest.param( + { + "startup_nodes": [{"host": "cache.example.com", "port": 6379}], + "ssl": True, + }, + id="cluster", + ), + pytest.param( + { + "sentinel_nodes": [("sentinel.example.com", 26379)], + "service_name": "cache", + "ssl": True, + }, + id="sentinel", + ), + ], +) +def test_aws_iam_auth_accepts_tls_connections(clean_redis_environment, transport): + redis_kwargs = _get_redis_client_logic( + **transport, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) + + +@pytest.mark.parametrize("ssl", ["true", "True", "TRUE", "1", "yes", "YES", "false", "0", "no", "off", "", "maybe"]) +def test_tls_detection_agrees_with_the_ssl_kwarg_coercion(ssl): + assert _uses_tls({"ssl": ssl}) is _coerce_redis_kwargs_types({"ssl": ssl})["ssl"] + + +def test_aws_iam_settings_are_removed_for_url_and_static_credentials(clean_redis_environment): + redis_kwargs = _get_redis_client_logic( + url="rediss://url-user:url-pass@cache.example.com:6380", + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + username="static-user", + password="static-password", + ) + + assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) + assert redis_kwargs["url"] == "rediss://cache.example.com:6380" + assert "username" not in redis_kwargs + assert "password" not in redis_kwargs + assert not _AWS_IAM_SETTINGS & redis_kwargs.keys() + + +@pytest.mark.parametrize("missing", ["aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"]) +def test_aws_iam_missing_setting_fails_closed(clean_redis_environment, missing): + settings = { + "aws_iam_auth": True, + "aws_iam_user_name": "iam-user", + "aws_iam_cache_name": "cache.example.com", + "aws_iam_region": "us-east-1", + "ssl": True, + } + settings[missing] = None + + with pytest.raises(ValueError, match=missing): + _get_redis_client_logic(host="cache.example.com", port=6379, **settings) + + +@pytest.mark.parametrize("region_var", ["AWS_REGION", "AWS_DEFAULT_REGION"]) +def test_aws_iam_region_falls_back_to_environment(clean_redis_environment, monkeypatch, region_var): + monkeypatch.setenv(region_var, "sa-east-1") + + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + ssl=True, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + ) + + provider = redis_kwargs["credential_provider"] + assert isinstance(provider, ElastiCacheIAMCredentialProvider) + assert provider._region == "sa-east-1" + + +def test_aws_iam_region_prefers_aws_region_over_default_region(clean_redis_environment, monkeypatch): + monkeypatch.setenv("AWS_REGION", "sa-east-1") + monkeypatch.setenv("AWS_DEFAULT_REGION", "eu-west-1") + + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + ssl=True, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + ) + + assert redis_kwargs["credential_provider"]._region == "sa-east-1" + + +def test_aws_iam_region_prefers_explicit_over_environment(clean_redis_environment, monkeypatch): + monkeypatch.setenv("AWS_REGION", "sa-east-1") + + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + ssl=True, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="explicit-region", + ) + + assert redis_kwargs["credential_provider"]._region == "explicit-region" + + +def test_aws_iam_settings_map_to_distinct_provider_fields(clean_redis_environment): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + ssl=True, + aws_iam_auth=True, + aws_iam_user_name="iam-user-value", + aws_iam_cache_name="iam-cache-value", + aws_iam_region="iam-region-value", + ) + + provider = redis_kwargs["credential_provider"] + assert isinstance(provider, ElastiCacheIAMCredentialProvider) + assert provider._user_name == "iam-user-value" + assert provider._cache_name == "iam-cache-value" + assert provider._region == "iam-region-value" + + +@pytest.mark.parametrize("aws_iam_auth", [None, False, "", "false", "0", "no", "off"]) +def test_aws_iam_auth_disabled_does_not_install_provider(clean_redis_environment, aws_iam_auth): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + aws_iam_auth=aws_iam_auth, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert "credential_provider" not in redis_kwargs + assert not _AWS_IAM_SETTINGS & redis_kwargs.keys() + + +@pytest.mark.parametrize("aws_iam_auth", [True, "true", "True", "TRUE", "1", "yes"]) +def test_aws_iam_auth_enabled_by_any_truthy_flag(clean_redis_environment, aws_iam_auth): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + ssl=True, + aws_iam_auth=aws_iam_auth, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache-name", + aws_iam_region="us-east-1", + ) + + assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) + + +@pytest.mark.parametrize( + "aws_iam_serverless, expected", + [ + pytest.param(None, False, id="unset"), + pytest.param(False, False, id="bool_false"), + pytest.param("false", False, id="string_false"), + pytest.param("0", False, id="string_zero"), + pytest.param(True, True, id="bool_true"), + pytest.param("true", True, id="string_true"), + pytest.param("1", True, id="string_one"), + ], +) +def test_aws_iam_serverless_flag_reaches_the_provider(clean_redis_environment, aws_iam_serverless, expected): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + ssl=True, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache-name", + aws_iam_region="us-east-1", + aws_iam_serverless=aws_iam_serverless, + ) + + provider = redis_kwargs["credential_provider"] + assert isinstance(provider, ElastiCacheIAMCredentialProvider) + assert provider._is_serverless is expected + assert "aws_iam_serverless" not in redis_kwargs + + +def test_explicit_provider_wins_over_aws_iam(clean_redis_environment): + provider = _StubCredentialProvider() + + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + credential_provider=provider, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert redis_kwargs["credential_provider"] is provider + assert "aws_iam_auth" not in redis_kwargs + + +def test_gcp_wins_over_aws_iam(clean_redis_environment): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + gcp_service_account="sa@example.com", + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert "credential_provider" not in redis_kwargs + assert redis_kwargs["redis_connect_func"]._gcp_service_account == "sa@example.com" + + +def test_azure_wins_over_aws_iam(clean_redis_environment): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + azure_redis_ad_token="true", + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert "credential_provider" not in redis_kwargs + assert redis_kwargs["redis_connect_func"]._azure_redis_ad_token is True + + +def test_async_cluster_installs_aws_iam_provider(clean_redis_environment): + startup_nodes = [{"host": "cluster-node", "port": 6379}] + + client = get_redis_async_client( + startup_nodes=startup_nodes, + ssl=True, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert isinstance(client.connection_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) + + def test_provider_keeps_the_rest_of_the_url_intact(clean_redis_environment): provider = _StubCredentialProvider() diff --git a/tests/test_litellm/test_redis_credential_provider.py b/tests/test_litellm/test_redis_credential_provider.py new file mode 100644 index 00000000000..96b1933b0e3 --- /dev/null +++ b/tests/test_litellm/test_redis_credential_provider.py @@ -0,0 +1,236 @@ +import asyncio +import builtins +import sys +from types import SimpleNamespace +from urllib.parse import parse_qs, urlsplit + +import pytest + +from litellm._redis_credential_provider import ElastiCacheIAMCredentialProvider + + +class _FakeCredentials: + def __init__(self, access_key: str) -> None: + self.access_key = access_key + self.secret_key = "synthetic-secret" + self.token = "synthetic-session-token" + + def get_frozen_credentials(self): + return self + + +class _RotatingFakeCredentials: + def __init__(self) -> None: + self.calls = 0 + + def __bool__(self) -> bool: + return False + + def get_frozen_credentials(self): + self.calls += 1 + return SimpleNamespace( + access_key=f"AKIA-SYNTHETIC-{self.calls}", + secret_key="synthetic-secret", + token="synthetic-session-token", + ) + + +class _FakeResolver: + def __init__(self, credentials: _FakeCredentials | _RotatingFakeCredentials | None) -> None: + self.credentials = credentials + self.calls = 0 + + def __call__(self): + self.calls += 1 + return self.credentials + + +def test_elasticache_provider_signs_expected_query(): + resolver = _FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=resolver, + ) + + user_name, token = provider.get_credentials() + parsed = urlsplit("https://" + token) + query = parse_qs(parsed.query) + + assert user_name == "iam-user" + assert parsed.netloc == "cache.example.com" + assert query["Action"] == ["connect"] + assert query["User"] == ["iam-user"] + assert query["X-Amz-Expires"] == ["900"] + assert "elasticache" in query["X-Amz-Credential"][0] + assert query["X-Amz-Credential"][0].split("/")[2] == "us-east-1" + assert not token.startswith("https://") + + +def test_elasticache_provider_resolves_credentials_once_but_refreshes_signature(): + rotating_credentials = _RotatingFakeCredentials() + resolver = _FakeResolver(rotating_credentials) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=resolver, + ) + + first = provider.get_credentials() + second = provider.get_credentials() + async_result = asyncio.run(provider.get_credentials_async()) + + assert first[0] == second[0] == async_result[0] == "iam-user" + assert first[1] != second[1] + assert async_result[1] != second[1] + assert resolver.calls == 1 + assert rotating_credentials.calls == 3 + + +def test_elasticache_provider_uses_botocore_session_credentials(monkeypatch): + credentials = _FakeCredentials("AKIA-SYNTHETIC") + monkeypatch.setattr("botocore.session.get_session", lambda: SimpleNamespace(get_credentials=lambda: credentials)) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + ) + + user_name, token = provider.get_credentials() + + assert user_name == "iam-user" + assert "AKIA-SYNTHETIC" in token + + +def test_elasticache_provider_reports_missing_botocore(monkeypatch): + original_import = builtins.__import__ + + def import_without_botocore(name, *args, **kwargs): + if name == "botocore.session": + raise ImportError("synthetic missing dependency") + return original_import(name, *args, **kwargs) + + monkeypatch.delitem(sys.modules, "botocore.session", raising=False) + monkeypatch.setattr(builtins, "__import__", import_without_botocore) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + ) + + with pytest.raises(ImportError, match="pip install boto3"): + provider.get_credentials() + + +def test_elasticache_provider_reports_missing_credentials(): + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=_FakeResolver(None), + ) + + with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"): + provider.get_credentials() + + +def test_elasticache_provider_reports_missing_signing_dependency(monkeypatch): + original_import = builtins.__import__ + + def import_without_botocore_auth(name, *args, **kwargs): + if name == "botocore.auth": + raise ImportError("synthetic missing dependency") + return original_import(name, *args, **kwargs) + + monkeypatch.delitem(sys.modules, "botocore.auth", raising=False) + monkeypatch.setattr(builtins, "__import__", import_without_botocore_auth) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")), + ) + + with pytest.raises(ImportError, match="pip install boto3"): + provider.get_credentials() + + +def test_elasticache_provider_recovers_after_a_failed_resolution(): + resolver = _FakeResolver(None) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=resolver, + ) + + with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"): + provider.get_credentials() + + resolver.credentials = _FakeCredentials("AKIA-SYNTHETIC") + user_name, token = provider.get_credentials() + + assert user_name == "iam-user" + assert token + assert resolver.calls == 2 + + +@pytest.mark.parametrize( + "provider_kwargs, expected_operation_params", + [ + pytest.param({}, frozenset({"Action", "User"}), id="default_is_self_designed"), + pytest.param({"is_serverless": False}, frozenset({"Action", "User"}), id="self_designed"), + pytest.param({"is_serverless": True}, frozenset({"Action", "User", "ResourceType"}), id="serverless"), + ], +) +def test_elasticache_provider_signs_resource_type_only_for_serverless(provider_kwargs, expected_operation_params): + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache-name", + region="us-east-1", + credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")), + **provider_kwargs, + ) + + _, token = provider.get_credentials() + query_string = urlsplit("https://" + token).query + param_names = tuple(pair.split("=", 1)[0] for pair in query_string.split("&")) + first_auth_param = next(i for i, name in enumerate(param_names) if name.startswith("X-Amz-")) + query = parse_qs(query_string) + + assert frozenset(param_names[:first_auth_param]) == expected_operation_params + assert all(name.startswith("X-Amz-") for name in param_names[first_auth_param:]) + assert query.get("ResourceType") == (["ServerlessCache"] if "ResourceType" in expected_operation_params else None) + assert query["X-Amz-Signature"] + + +def test_elasticache_provider_lowercases_the_cache_name(): + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="Mixed-Case-Cache", + region="us-east-1", + credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")), + ) + + _, token = provider.get_credentials() + + assert urlsplit("https://" + token).netloc == "mixed-case-cache" + + +def test_elasticache_provider_encodes_reserved_characters_in_the_user_name(): + user_name = "iam user/with+reserved&chars" + provider = ElastiCacheIAMCredentialProvider( + user_name=user_name, + cache_name="cache-name", + region="us-east-1", + credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")), + ) + + returned_user_name, token = provider.get_credentials() + query = parse_qs(urlsplit("https://" + token).query) + + assert returned_user_name == user_name + assert query["User"] == [user_name] + assert query["Action"] == ["connect"] diff --git a/tests/test_litellm/test_replicate_model_key_format.py b/tests/test_litellm/test_replicate_model_key_format.py index 8c2b72f8ed2..77ae5e1b069 100644 --- a/tests/test_litellm/test_replicate_model_key_format.py +++ b/tests/test_litellm/test_replicate_model_key_format.py @@ -20,14 +20,6 @@ def test_replicate_models_have_valid_key_prefix(model_cost: dict[str, Any]) -> N ) -def test_replicate_openai_gpt_oss_20b_key_exists(model_cost: dict[str, Any]) -> None: - assert "replicate/openai/gpt-oss-20b" in model_cost - info = model_cost["replicate/openai/gpt-oss-20b"] - assert info["litellm_provider"] == "replicate" - assert info["mode"] == "chat" - assert info["supports_function_calling"] is True - - def test_replicate_backup_matches_main() -> None: repo_root = Path(__file__).parents[2] main_path = repo_root / "model_prices_and_context_window.json" diff --git a/tests/test_litellm/test_responses_id_security.py b/tests/test_litellm/test_responses_id_security.py index d35b9563888..a6081670172 100644 --- a/tests/test_litellm/test_responses_id_security.py +++ b/tests/test_litellm/test_responses_id_security.py @@ -855,3 +855,242 @@ class TestAsyncPostCallSuccessHook: ) assert result == mock_response + + + +_FABRICATED_PROVIDER_RESPONSE_ID = "resp_fabricatedprovideridaaaaaaaaaaaaaaaa" +_FABRICATED_UNMANAGED_ID = "resp_fabricatedunmanagedidbbbbbbbbbbbbbbbb" +_UNIT_TEST_SALT_KEY = "lit6837-unit-test-salt-key" +_ADDRESSED_ID_FIELD_BY_CALL_TYPE = { + "aresponses": "previous_response_id", + "aget_responses": "response_id", + "adelete_responses": "response_id", + "acancel_responses": "response_id", + "alist_input_items": "response_id", +} + + +@pytest.fixture +def salt_key_env(monkeypatch): + """Give the encrypt/decrypt helpers a real salt key so ids round-trip for real.""" + monkeypatch.setenv("LITELLM_SALT_KEY", _UNIT_TEST_SALT_KEY) + return _UNIT_TEST_SALT_KEY + + +def _hook(general_settings=None, signing_key=_UNIT_TEST_SALT_KEY): + settings = general_settings if general_settings is not None else {} + return ResponsesIDSecurity( + general_settings_reader=lambda: settings, + signing_key_reader=lambda: signing_key, + ) + + +def _auth(user_id="owner-user", team_id="owner-team", user_role=None): + from litellm.proxy._types import UserAPIKeyAuth + + return UserAPIKeyAuth(user_id=user_id, team_id=team_id, user_role=user_role) + + +def _issue_managed_id(hook, owner, provider_response_id=_FABRICATED_PROVIDER_RESPONSE_ID): + """Mint an id exactly the way the proxy hands one to a client on create.""" + issued = hook._encrypt_response_id( + ResponsesAPIResponse( + id=provider_response_id, created_at=1234567890, output=[], status="completed" + ), + owner, + ) + return issued.id + + +class TestUnrecognizedResponseIdIsRejected: + """An id this proxy never issued carries no owner, so it must not reach the provider.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE)) + async def test_unmanaged_id_is_rejected_and_not_forwarded(self, mock_cache, salt_key_env, call_type): + field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type] + data = {field: _FABRICATED_UNMANAGED_ID} + + with pytest.raises(HTTPException) as exc_info: + await _hook().async_pre_call_hook( + user_api_key_dict=_auth(), + cache=mock_cache, + data=data, + call_type=call_type, + ) + + assert exc_info.value.status_code == 403 + assert "allow_unmanaged_response_ids" in exc_info.value.detail + assert data[field] == _FABRICATED_UNMANAGED_ID + + @pytest.mark.asyncio + async def test_owner_can_still_address_the_id_the_proxy_issued_it(self, mock_cache, salt_key_env): + hook = _hook() + owner = _auth() + data = {"response_id": _issue_managed_id(hook, owner)} + + result = await hook.async_pre_call_hook( + user_api_key_dict=owner, + cache=mock_cache, + data=data, + call_type="aget_responses", + ) + + assert result["response_id"] == _FABRICATED_PROVIDER_RESPONSE_ID + + @pytest.mark.asyncio + async def test_stranger_cannot_address_an_id_issued_to_someone_else(self, mock_cache, salt_key_env): + hook = _hook() + issued_id = _issue_managed_id(hook, _auth()) + data = {"response_id": issued_id} + + with pytest.raises(HTTPException) as exc_info: + await hook.async_pre_call_hook( + user_api_key_dict=_auth(user_id="stranger-user", team_id="stranger-team"), + cache=mock_cache, + data=data, + call_type="aget_responses", + ) + + assert exc_info.value.status_code == 403 + assert data["response_id"] == issued_id + + @pytest.mark.asyncio + async def test_unmanaged_previous_response_id_cannot_seed_a_new_response(self, mock_cache, salt_key_env): + data = {"model": "gpt-fake", "previous_response_id": _FABRICATED_UNMANAGED_ID} + + with pytest.raises(HTTPException) as exc_info: + await _hook().async_pre_call_hook( + user_api_key_dict=_auth(), + cache=mock_cache, + data=data, + call_type="aresponses", + ) + + assert exc_info.value.status_code == 403 + assert data["previous_response_id"] == _FABRICATED_UNMANAGED_ID + + @pytest.mark.asyncio + async def test_re_entering_the_hook_on_the_same_request_does_not_reject(self, mock_cache, salt_key_env): + """The rate-limit fallback retry runs pre-call twice over one already-rewritten dict.""" + hook = _hook() + owner = _auth() + data = {"model": "gpt-fake", "previous_response_id": _issue_managed_id(hook, owner)} + + first = await hook.async_pre_call_hook( + user_api_key_dict=owner, cache=mock_cache, data=data, call_type="aresponses" + ) + second = await hook.async_pre_call_hook( + user_api_key_dict=owner, cache=mock_cache, data=first, call_type="aresponses" + ) + + assert second["previous_response_id"] == _FABRICATED_PROVIDER_RESPONSE_ID + + +class TestUnmanagedResponseIdEscapeHatches: + """Deployments that pass provider ids through on purpose must keep working.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "general_settings", + [{"allow_unmanaged_response_ids": True}, {"disable_responses_id_security": True}], + ) + async def test_opted_in_settings_forward_the_id_untouched(self, mock_cache, salt_key_env, general_settings): + data = {"response_id": _FABRICATED_UNMANAGED_ID} + + result = await _hook(general_settings=general_settings).async_pre_call_hook( + user_api_key_dict=_auth(), + cache=mock_cache, + data=data, + call_type="aget_responses", + ) + + assert result["response_id"] == _FABRICATED_UNMANAGED_ID + + @pytest.mark.asyncio + async def test_proxy_without_a_signing_key_forwards_the_id_untouched(self, mock_cache, monkeypatch): + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + data = {"response_id": _FABRICATED_UNMANAGED_ID} + + result = await _hook(signing_key=None).async_pre_call_hook( + user_api_key_dict=_auth(), + cache=mock_cache, + data=data, + call_type="aget_responses", + ) + + assert result["response_id"] == _FABRICATED_UNMANAGED_ID + + @pytest.mark.asyncio + async def test_proxy_admin_may_address_an_unmanaged_id(self, mock_cache, salt_key_env): + from litellm.proxy._types import LitellmUserRoles + + data = {"response_id": _FABRICATED_UNMANAGED_ID} + + result = await _hook().async_pre_call_hook( + user_api_key_dict=_auth(user_role=LitellmUserRoles.PROXY_ADMIN), + cache=mock_cache, + data=data, + call_type="aget_responses", + ) + + assert result["response_id"] == _FABRICATED_UNMANAGED_ID + + +class TestClientSuppliedRetainedIdCannotBypassAuthorization: + """The retained-id key travels in the request body, so it is re-authorized, never trusted.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE)) + async def test_forged_retained_id_is_still_authorized(self, mock_cache, salt_key_env, call_type): + field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type] + data = { + field: _FABRICATED_UNMANAGED_ID, + "_litellm_addressed_response_id": _FABRICATED_UNMANAGED_ID, + } + + with pytest.raises(HTTPException) as exc_info: + await _hook().async_pre_call_hook( + user_api_key_dict=_auth(), + cache=mock_cache, + data=data, + call_type=call_type, + ) + + assert exc_info.value.status_code == 403 + assert data[field] == _FABRICATED_UNMANAGED_ID + + @pytest.mark.asyncio + @pytest.mark.parametrize("forged", [{"nested": "value"}, ["list"], 42, "", None]) + async def test_non_string_retained_id_falls_back_to_the_addressed_field(self, mock_cache, salt_key_env, forged): + data = {"response_id": _FABRICATED_UNMANAGED_ID, "_litellm_addressed_response_id": forged} + + with pytest.raises(HTTPException) as exc_info: + await _hook().async_pre_call_hook( + user_api_key_dict=_auth(), + cache=mock_cache, + data=data, + call_type="aget_responses", + ) + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_stranger_forging_their_own_id_never_reaches_someone_elses_response( + self, mock_cache, salt_key_env + ): + hook = _hook() + stranger = _auth(user_id="stranger-user", team_id="stranger-team") + stranger_id = _issue_managed_id(hook, stranger, provider_response_id="resp_strangerownprovideridcccccccc") + victim_provider_id = "resp_victimprovideriddddddddddddddddddddd" + data = {"response_id": victim_provider_id, "_litellm_addressed_response_id": stranger_id} + + result = await hook.async_pre_call_hook( + user_api_key_dict=stranger, + cache=mock_cache, + data=data, + call_type="aget_responses", + ) + + assert result["response_id"] == "resp_strangerownprovideridcccccccc" + assert result["response_id"] != victim_provider_id diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7c044310e14..2a97e92396a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -6,6 +6,7 @@ import logging import os import threading from datetime import datetime +from collections.abc import Awaitable, Callable, Mapping from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -18,6 +19,8 @@ import respx import litellm +from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm import Router from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_guardrail import CustomGuardrail @@ -27,6 +30,7 @@ from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, ) +from litellm.types.llms.openai import ChatCompletionRequest from litellm.router import ( MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, FallbackAwareAnthropicMessagesStream, @@ -39,7 +43,8 @@ from litellm.router import ( _anthropic_stream_should_drop_pre_content_ping, _is_retriable_anthropic_status, ) -from litellm.types.router import DeploymentTypedDict +from litellm.router_strategy import simple_shuffle +from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, RetryPolicy def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -2386,6 +2391,580 @@ async def test_acompletion_streaming_iterator_preserves_hidden_params(): assert result._hidden_params.get("_response_ms") == 500.0 +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_preserves_response_headers(): + """LIT-6767: the returned wrapper must carry the provider's raw response headers. + + Proxy callbacks read ``_response_headers`` off the object the router hands + back. The wrapper used to be built without it, so every streaming chat + completion reported zero raw provider headers while the non-streaming path + reported the full set. + """ + from unittest.mock import MagicMock + + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + async def _empty(): + return + yield # make it an async generator + + provider_headers = { + "x-request-id": "req-provider-123", + "x-ratelimit-remaining-requests": "42", + # a provider must never be able to spoof an internal header + "x-litellm-model-id": "spoofed", + } + source = CustomStreamWrapper( + completion_stream=_empty(), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers=provider_headers, + ) + + result = await router._acompletion_streaming_iterator( + model_response=source, + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + + assert result._response_headers == provider_headers + additional_headers = result._hidden_params["additional_headers"] + assert additional_headers["llm_provider-x-request-id"] == "req-provider-123" + assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42" + # internal-header protection survives: the provider value is namespaced, never promoted + assert additional_headers["llm_provider-x-litellm-model-id"] == "spoofed" + assert "x-litellm-model-id" not in additional_headers + + +def test_completion_streaming_iterator_preserves_response_headers(): + """LIT-6767, sync counterpart of the async header-preservation test.""" + from unittest.mock import MagicMock + + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + provider_headers = {"x-request-id": "req-provider-sync", "openai-organization": "org-real"} + source = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers=provider_headers, + ) + + result = router._completion_streaming_iterator( + model_response=source, + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + + assert result._response_headers == provider_headers + assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-provider-sync" + + +def test_adopt_fallback_response_headers_replaces_rather_than_merges(): + """LIT-6767: direct unit for FallbackAwareStreamWrapper.adopt_fallback_response_headers. + + Values from the failed attempt must not survive, so the wrapper replaces both + ``_response_headers`` and ``_hidden_params`` instead of merging them. + """ + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers={"x-request-id": "req-FAILED"}, + ) + wrapper._hidden_params = { + "model_id": "failed-deployment", + "only_on_failed_attempt": "stale", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + } + + fallback = MagicMock() + fallback._response_headers = {"x-request-id": "req-FALLBACK"} + fallback._hidden_params = { + "model_id": "fallback-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + + wrapper.adopt_fallback_response_headers( + fallback, Router._prepare_fallback_hidden_params(fallback) + ) + + assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} + assert wrapper._hidden_params["model_id"] == "fallback-deployment" + assert "only_on_failed_attempt" not in wrapper._hidden_params + assert wrapper._hidden_params is not fallback._hidden_params + # the snapshot CustomStreamWrapper caches at init has to follow, or a chunk built + # from it would still be stamped with the deployment that failed + assert wrapper._base_hidden_params["model_id"] == "fallback-deployment" + # the nested header dict is copied too, so a later mutation on the fallback + # response cannot reach headers the proxy has already published + assert wrapper._hidden_params["additional_headers"] is not fallback._hidden_params["additional_headers"] + fallback._hidden_params["additional_headers"]["llm_provider-x-request-id"] = "req-MUTATED" + assert wrapper._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} + + +def test_adopt_fallback_response_headers_survives_a_collected_wrapper(): + """LIT-6767: adoption still returns the fallback's params once the wrapper is gone.""" + import weakref + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + fallback: Final = MagicMock() + fallback._response_headers = {"x-request-id": "req-FALLBACK"} + fallback._hidden_params = { + "model_id": "fallback-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + ) + live_ref: Final = weakref.ref(wrapper) + prepared: Final = Router._adopt_fallback_response_headers(live_ref, fallback) + assert prepared == (fallback._hidden_params, fallback._hidden_params["additional_headers"]) + assert wrapper.fallback_headers_adopted is True + assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} + + dead_ref: Final = weakref.ref(wrapper) + del wrapper + assert dead_ref() is None + assert Router._adopt_fallback_response_headers(dead_ref, fallback) == prepared + + +def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_replace(): + """LIT-6767: a fallback that carries no raw provider headers publishes none. + + Keeping the failed attempt's raw headers would hand the client and the callbacks a + provider ``x-request-id`` for a request that deployment never served, which is the + leak this fix exists to close. + """ + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers={"x-request-id": "req-FAILED"}, + ) + wrapper._hidden_params = {"model_id": "failed-deployment", "additional_headers": {}} + + fallback = MagicMock() + fallback._response_headers = None + fallback._hidden_params = {"model_id": "fallback-deployment"} + + wrapper.adopt_fallback_response_headers( + fallback, Router._prepare_fallback_hidden_params(fallback) + ) + + assert wrapper._response_headers is None + assert wrapper._hidden_params["model_id"] == "fallback-deployment" + assert wrapper.fallback_headers_adopted is True + + +def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none(): + """A fallback response carrying no hidden params keeps the identity headers. + + Publishing no ``x-litellm-*`` header at all for a request the fallback served is + worse than keeping what is there, so only the raw provider headers are dropped. + """ + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers={"x-request-id": "req-FAILED"}, + ) + hidden_params_before = wrapper._hidden_params + + fallback = object() + wrapper.adopt_fallback_response_headers( + fallback, Router._prepare_fallback_hidden_params(fallback) + ) + + assert wrapper._response_headers is None + assert wrapper._hidden_params is hidden_params_before + assert wrapper.fallback_headers_adopted is True + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_adopts_fallback_response_headers(): + """LIT-6767: after a successful pre-first-chunk fallback, the wrapper must + describe the deployment that served the stream, with no value left over + from the attempt that failed.""" + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "api_base": "https://failed.example", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + "only_on_failed_attempt": "stale", + } + + def __aiter__(self): + return self + + async def __anext__(self): + raise failed_error + + class FallbackStream: + def __init__(self): + self._response_headers = {"x-request-id": "req-FALLBACK"} + self._hidden_params = { + "model_id": "fallback-deployment", + "api_base": "https://fallback.example", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._chunks) + except StopIteration: + raise StopAsyncIteration from None + + fallback_stream = FallbackStream() + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=fallback_stream, + ): + result = await router._acompletion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + # the failed attempt is what the wrapper is built from + assert result._response_headers == {"x-request-id": "req-FAILED"} + # the very first chunk the fallback produces must already be published under + # the fallback's identity: the proxy commits response headers once that chunk + # is buffered, so adopting any later is adopting too late + await result.__anext__() + assert result._response_headers == {"x-request-id": "req-FALLBACK"} + assert result._hidden_params["model_id"] == "fallback-deployment" + async for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-FALLBACK"} + assert result._hidden_params["model_id"] == "fallback-deployment" + assert result._hidden_params["api_base"] == "https://fallback.example" + assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} + # stale values are removed, not merged over + assert "only_on_failed_attempt" not in result._hidden_params + # and the wrapper holds its own copy, so later fallback mutations cannot leak in + assert result._hidden_params is not fallback_stream._hidden_params + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): + """LIT-6767: a fallback that itself fails over before its first chunk. + + The selected fallback still describes its own failed attempt at selection time, so + the wrapper has to re-read it once a chunk exists or it publishes a deployment that + produced no output. + """ + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error: Final = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + } + + def __aiter__(self): + return self + + async def __anext__(self): + raise failed_error + + class NestedFallbackStream: + """A fallback that repoints itself at a third deployment as it yields.""" + + def __init__(self): + self._response_headers = {"x-request-id": "req-MIDDLE"} + self._hidden_params = { + "model_id": "middle-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, + } + self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + chunk = next(self._chunks) + except StopIteration: + raise StopAsyncIteration from None + self._response_headers = {"x-request-id": "req-SERVED"} + self._hidden_params = { + "model_id": "served-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, + } + return chunk + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=NestedFallbackStream(), + ): + result = await router._acompletion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + first_chunk: Final = await result.__anext__() + # the proxy commits response headers once this chunk is buffered + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-SERVED"} + # and the chunk itself carries the same deployment + assert first_chunk._hidden_params["model_id"] == "served-deployment" + async for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + + +def test_completion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): + """LIT-6767, sync counterpart of the nested-fallback adoption test.""" + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error: Final = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + } + + def __iter__(self): + return self + + def __next__(self): + raise failed_error + + class NestedFallbackStream: + """A fallback that repoints itself at a third deployment as it yields.""" + + def __init__(self): + self._response_headers = {"x-request-id": "req-MIDDLE"} + self._hidden_params = { + "model_id": "middle-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, + } + self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) + + def __iter__(self): + return self + + def __next__(self): + chunk = next(self._chunks) + self._response_headers = {"x-request-id": "req-SERVED"} + self._hidden_params = { + "model_id": "served-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, + } + return chunk + + with patch.object(router, "function_with_fallbacks", return_value=NestedFallbackStream()): + result = router._completion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + first_chunk: Final = next(result) + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + assert first_chunk._hidden_params["model_id"] == "served-deployment" + for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + + +def test_completion_streaming_iterator_adopts_fallback_response_headers(): + """LIT-6767, sync counterpart of the fallback-adoption test.""" + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + "only_on_failed_attempt": "stale", + } + + def __iter__(self): + return self + + def __next__(self): + raise failed_error + + class FallbackStream: + def __init__(self): + self._response_headers = {"x-request-id": "req-FALLBACK"} + self._hidden_params = { + "model_id": "fallback-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + + def __iter__(self): + return iter([]) + + with patch.object(router, "function_with_fallbacks", return_value=FallbackStream()): + result = router._completion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + assert result._response_headers == {"x-request-id": "req-FAILED"} + for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-FALLBACK"} + assert result._hidden_params["model_id"] == "fallback-deployment" + assert "only_on_failed_attempt" not in result._hidden_params + + def test_completion_streaming_iterator_fallback_on_429(): """Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback. @@ -4472,6 +5051,41 @@ def test_get_deployment_model_info_base_model_merge_priority(): print("✓ Base model merge priority test passed!") +@pytest.mark.parametrize( + "model, litellm_params, endpoint, expected", + [ + ( + "gpt", + {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.services.ai.azure.com", "api_key": "key"}, + "gpt/openai/deployments/gpt-5.4-mini/chat/completions", + "gpt-5.4-mini/openai/deployments/gpt-5.4-mini/chat/completions", + ), + ( + "aws/anthropic/bedrock-claude", + {"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"}, + "/model/aws/anthropic/bedrock-claude/invoke", + "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", + ), + ( + "my-gemini", + {"model": "gemini/gemini-3.1-pro-preview", "api_key": "key"}, + "v1beta/models/my-gemini:streamGenerateContent", + "v1beta/models/gemini-3.1-pro-preview:streamGenerateContent", + ), + ], +) +def test_add_deployment_model_to_endpoint_rewrites_the_model_group_only_as_whole_path_segments( + model, litellm_params, endpoint, expected +): + router = litellm.Router(model_list=[{"model_name": model, "litellm_params": litellm_params}]) + + result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( + kwargs={"endpoint": endpoint}, model=model, model_name=litellm_params["model"] + ) + + assert result["endpoint"] == expected + + def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): """ Test that _add_deployment_model_to_endpoint_for_llm_passthrough_route correctly strips bedrock provider prefix @@ -5557,6 +6171,50 @@ def test_update_kwargs_with_deployment_no_tags(): assert "tags" not in kwargs["metadata"] +@pytest.mark.asyncio +async def test_retry_does_not_narrow_tag_filtered_group_to_failed_deployments_tags(): + router = Router( + model_list=[ + { + "model_name": "tagged-group", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "fake-key", + "tags": ["free"], + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000001, + "mock_response": "litellm.ContextWindowExceededError", + }, + "model_info": {"id": "tagged-failing"}, + }, + { + "model_name": "tagged-group", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "fake-key", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.001, + "mock_response": "ok", + }, + "model_info": {"id": "untagged-healthy"}, + }, + ], + routing_strategy="cost-based-routing", + enable_tag_filtering=True, + num_retries=2, + retry_after=0, + retry_policy=RetryPolicy(BadRequestErrorRetries=2), + ) + metadata: Final[dict[str, object]] = {} + + response = await router.acompletion( + model="tagged-group", messages=[{"role": "user", "content": "hi"}], metadata=metadata + ) + + assert response._hidden_params["model_id"] == "untagged-healthy" + assert metadata["tags"] == ["free"] + + def test_update_kwargs_with_deployment_merges_tools(): """ Test that when both deployment litellm_params and request have tools, @@ -6723,6 +7381,63 @@ async def test_async_get_fully_unhealthy_model_names_marks_name_when_all_unhealt assert await router.async_get_fully_unhealthy_model_names() == {"gpt-4o"} +@pytest.mark.asyncio +@pytest.mark.parametrize("health_check_probe", [False, True]) +@pytest.mark.parametrize( + "state, health_routing, fails_policy, scoped, strict_ids", + [ + ("absent", True, False, False, ("dep-0", "dep-1")), + ("partial", True, False, False, ("dep-1",)), + ("all", True, False, False, ()), + ("stale", True, False, False, ("dep-0", "dep-1")), + ("all", False, False, False, ("dep-0", "dep-1")), + ("all", True, True, False, ("dep-0", "dep-1")), + ("all", True, True, True, ()), + ], +) +async def test_health_probe_preserves_normal_caller_policy( + health_check_probe: bool, + state: str, + health_routing: bool, + fails_policy: bool, + scoped: bool, + strict_ids: tuple[str, ...], +) -> None: + import time + from litellm.types.router import AllowedFailsPolicy, RouterRateLimitError + + router: Final = Router( + model_list=[ + { + "model_name": "health-group", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "test-only"}, + "model_info": {"id": model_id}, + } + for model_id in ("dep-0", "dep-1") + ], + enable_health_check_routing=health_routing, + allowed_fails_policy=AllowedFailsPolicy(ServiceUnavailableErrorAllowedFails=2) if fails_policy else None, + background_health_check_model_groups=["health-group"] if scoped else None, + ) + if state != "absent": + _seed_unhealthy_states( + router, + ("dep-0",) if state == "partial" else ("dep-0", "dep-1"), + time.time() - router.health_state_cache.staleness_threshold - 10 if state == "stale" else None, + ) + expected: Final = strict_ids if strict_ids or health_check_probe else ("dep-0", "dep-1") + if not expected: + with pytest.raises(RouterRateLimitError, match="No deployments available"): + await router.async_get_healthy_deployments(model="health-group", request_kwargs={}, health_check_probe=True) + else: + deployments: Final = await router.async_get_healthy_deployments( + model="health-group", request_kwargs={}, health_check_probe=health_check_probe + ) + assert {d["model_info"]["id"] for d in deployments} == set(expected) + assert await router.cooldown_cache.async_get_active_cooldowns(["dep-0", "dep-1"], parent_otel_span=None) == [] + + + @pytest.mark.asyncio async def test_async_get_fully_unhealthy_model_names_keeps_name_when_partial(): router = _router_with_two_deployments([False, False]) @@ -7552,6 +8267,204 @@ def test_get_configured_token_limits_coerces_numeric_strings(): assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000) +def test_get_model_listing_info_prefers_base_model_over_litellm_params_model(): + """The cost-map key comes from base_model when set, so a deployment pointing at an + opaque backend name still resolves the real catalog entry.""" + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude-opus-5", + "litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"}, + "model_info": {"base_model": "eu.anthropic.claude-opus-5"}, + } + ] + ) + + info = router.get_model_listing_info("bedrock-claude-opus-5") + assert info is not None + assert info.cost_map_keys == ("eu.anthropic.claude-opus-5",) + + +def test_get_model_listing_info_falls_back_to_litellm_params_model(): + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude-opus-5", + "litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"}, + } + ] + ) + + info = router.get_model_listing_info("bedrock-claude-opus-5") + assert info is not None + assert info.cost_map_keys == ("bedrock/eu.anthropic.claude-opus-5",) + + +def test_get_model_listing_info_ignores_blank_base_model(): + """A base_model set to an empty string is absent, not a cost-map key.""" + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude-opus-5", + "litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"}, + "model_info": {"base_model": ""}, + } + ] + ) + + info = router.get_model_listing_info("bedrock-claude-opus-5") + assert info is not None + assert info.cost_map_keys == ("bedrock/eu.anthropic.claude-opus-5",) + + +def test_get_model_listing_info_returns_none_for_unknown_name(): + router = litellm.Router( + model_list=[ + { + "model_name": "no-limits-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + } + ] + ) + + assert router.get_model_listing_info("not-a-real-model") is None + + +def test_get_model_listing_info_carries_configured_limits(): + router = litellm.Router( + model_list=[ + { + "model_name": "my-custom-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"max_input_tokens": 32000, "max_output_tokens": 8000}, + } + ] + ) + + info = router.get_model_listing_info("my-custom-model") + assert info is not None + assert (info.max_input_tokens, info.max_output_tokens) == (32000, 8000) + + +def test_widest_configured_limit_ignores_absent_and_malformed_values(): + model_infos = ( + {"max_input_tokens": 32000}, + {}, + {"max_input_tokens": "not-a-number"}, + {"max_input_tokens": "128000"}, + {"max_output_tokens": 4096}, + ) + + assert litellm.Router._widest_configured_limit(model_infos, "max_input_tokens") == 128000 + assert litellm.Router._widest_configured_limit(model_infos, "max_output_tokens") == 4096 + assert litellm.Router._widest_configured_limit((), "max_input_tokens") is None + + +def test_get_model_listing_info_dedupes_interchangeable_deployments(): + """The ordinary group is N deployments of one model, so it yields exactly one key.""" + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-a"}, + }, + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-b"}, + }, + ] + ) + + info = router.get_model_listing_info("gpt-4o") + assert info is not None + assert info.cost_map_keys == ("openai/gpt-4o",) + + +def test_get_model_listing_info_collects_every_model_in_a_mixed_group(): + router = litellm.Router( + model_list=[ + { + "model_name": "house-claude", + "litellm_params": {"model": "anthropic/claude-3-haiku-20240307"}, + }, + { + "model_name": "house-claude", + "litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"}, + }, + ] + ) + + info = router.get_model_listing_info("house-claude") + assert info is not None + assert info.cost_map_keys == ( + "anthropic/claude-3-haiku-20240307", + "bedrock/eu.anthropic.claude-opus-5", + ) + + +def test_get_model_listing_info_reports_widest_configured_limits_in_a_mixed_group(): + """Matches how get_model_group_info aggregates for the Admin UI, so the two agree.""" + router = litellm.Router( + model_list=[ + { + "model_name": "house-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"max_input_tokens": 32000, "max_output_tokens": 4096}, + }, + { + "model_name": "house-model", + "litellm_params": {"model": "openai/another-unmapped-model"}, + "model_info": {"max_input_tokens": 128000, "max_output_tokens": 16384}, + }, + ] + ) + + info = router.get_model_listing_info("house-model") + assert info is not None + assert (info.max_input_tokens, info.max_output_tokens) == (128000, 16384) + + +def test_get_model_listing_info_reads_base_model_from_litellm_params(): + """base_model resolution mirrors get_router_model_info, which also accepts it there.""" + router = litellm.Router( + model_list=[ + { + "model_name": "azure-deployment", + "litellm_params": { + "model": "azure/my-azure-deployment-name", + "base_model": "azure/gpt-4o", + "api_key": "sk-a", + "api_base": "https://example.openai.azure.com", + }, + } + ] + ) + + info = router.get_model_listing_info("azure-deployment") + assert info is not None + assert info.cost_map_keys == ("azure/gpt-4o",) + + +def test_get_model_listing_info_skips_wildcard_pattern_matching(): + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + "model_info": {"max_input_tokens": 12345}, + } + ] + ) + + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) + + def test_get_configured_mode_reads_deployment_model_info(): router = litellm.Router( model_list=[ @@ -9382,13 +10295,6 @@ class TestTaggedAutoRouterOnSharedModelName: def test_deployment_without_litellm_params_mapping_is_not_a_marker(self): assert litellm.Router._is_strategy_marker_deployment({"model_name": "gpt4o"}) is False - def test_model_name_has_plain_deployments_reflects_the_pool(self): - mixed = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=True) - marker_only = self._router(marker_tags=["route"], include_plain_sibling=False, enable_tag_filtering=True) - - assert mixed._model_name_has_plain_deployments("gpt4o") is True - assert marker_only._model_name_has_plain_deployments("gpt4o") is False - class TestAutoRouterSharedModelNameConnectionParams: """A plain deployment sharing its model_name with an `auto_router/` marker must not have @@ -9997,6 +10903,300 @@ class TestModelGroupAliasReachesPreRoutingStrategies: ) +class TestTeamPublicNameReachesPreRoutingStrategies: + """A team-scoped strategy router is stored under an internal `model_name_{team}_{uuid}` with the + caller-facing name in `model_info.team_public_model_name`, and the four registries key on that + internal name. A team key asks for the public name, so the hook has to resolve it to the team's + marker through the same team-first resolution the deployment path uses, and a resolution that + yields only markers is not callable on any path (LIT-7363).""" + + MARKER_TIMEOUT = 42.0 + REGISTRY_NAMES = ("auto_routers", "complexity_routers", "adaptive_routers", "quality_routers") + TEAM = "team-a" + OTHER_TEAM = "team-b" + PUBLIC_NAME = "smart-route" + INTERNAL_NAME = "model_name_team-a_0b3c" + SIBLING_INTERNAL_NAME = "model_name_team-a_9e1d" + + class _RewriteStrategy: + def __init__(self, rewrite_to: str = "gemini-flash"): + self.rewrite_to = rewrite_to + + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + return PreRoutingHookResponse(model=self.rewrite_to, messages=messages) + + @classmethod + def _team_marker(cls, internal_name: str, tags: list[str] | None = None) -> dict: + tiers = dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "gemini-flash") + return { + "model_name": internal_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": tiers}, + "complexity_router_default_model": "gemini-flash", + "timeout": cls.MARKER_TIMEOUT, + **({"tags": tags} if tags else {}), + }, + "model_info": {"team_id": cls.TEAM, "team_public_model_name": cls.PUBLIC_NAME}, + } + + @classmethod + def _router( + cls, + registrations: dict[str, "TestTeamPublicNameReachesPreRoutingStrategies._RewriteStrategy"], + registry_name: str = "complexity_routers", + extra_deployments: tuple[dict, ...] = (), + markers: tuple[dict, ...] | None = None, + enable_tag_filtering: bool = False, + ) -> "litellm.Router": + from litellm.types.router import TaggedPreRoutingStrategy + + markers = markers if markers is not None else (cls._team_marker(cls.INTERNAL_NAME),) + tier = { + "model_name": "gemini-flash", + "litellm_params": {"model": "gemini/gemini-3.6-flash", "mock_response": "routed by the tier"}, + } + router = litellm.Router( + model_list=[*markers, tier, *extra_deployments], + enable_tag_filtering=enable_tag_filtering, + ) + tags_by_name = {m["model_name"]: tuple(m["litellm_params"].get("tags") or ()) for m in markers} + for name in cls.REGISTRY_NAMES: + setattr(router, name, {}) + setattr( + router, + registry_name, + { + name: [TaggedPreRoutingStrategy(tags=tags_by_name[name], strategy=strategy)] + for name, strategy in registrations.items() + }, + ) + return router + + @staticmethod + def _messages() -> list[dict[str, str]]: + return [{"role": "user", "content": "What is the capital of France?"}] + + @classmethod + def _team_request(cls, team_id: str | None = "team-a", tags: list[str] | None = None) -> dict: + metadata = {**({"user_api_key_team_id": team_id} if team_id else {}), **({"tags": tags} if tags else {})} + return {"metadata": metadata} + + @pytest.mark.parametrize("registry_name", REGISTRY_NAMES) + @pytest.mark.asyncio + async def test_team_key_dispatches_to_the_strategy_registered_under_the_internal_name(self, registry_name): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}, registry_name=registry_name) + + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + + assert response is not None + assert response.model == "gemini-flash" + + @pytest.mark.asyncio + async def test_team_key_deployment_selection_lands_on_the_tier_and_forwards_the_marker_params(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + request_kwargs = self._team_request() + + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + + assert deployment["litellm_params"]["model"] == "gemini/gemini-3.6-flash" + assert request_kwargs["timeout"] == self.MARKER_TIMEOUT + + @pytest.mark.asyncio + async def test_another_team_never_reaches_the_strategy_or_the_marker(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + + assert ( + await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(self.OTHER_TEAM), messages=self._messages() + ) + is None + ) + with pytest.raises(litellm.BadRequestError): + await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(self.OTHER_TEAM), messages=self._messages() + ) + + @pytest.mark.asyncio + async def test_sibling_team_markers_select_by_request_tag_then_default(self): + router = self._router( + { + self.INTERNAL_NAME: self._RewriteStrategy("cn-model"), + self.SIBLING_INTERNAL_NAME: self._RewriteStrategy("us-model"), + }, + markers=( + self._team_marker(self.INTERNAL_NAME, tags=["cn"]), + self._team_marker(self.SIBLING_INTERNAL_NAME, tags=["us", "default"]), + ), + ) + + async def routed(tags: list[str] | None) -> str | None: + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(tags=tags), messages=self._messages() + ) + return response.model if response else None + + assert await routed(["cn"]) == "cn-model" + assert await routed(["us"]) == "us-model" + assert await routed(None) == "us-model" + + @pytest.mark.asyncio + async def test_team_public_name_shadows_a_global_model_for_that_team_only(self): + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + extra_deployments=({"model_name": self.PUBLIC_NAME, "litellm_params": {"model": "openai/gpt-4o"}},), + ) + + async def routed(request_kwargs: dict) -> str | None: + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + return response.model if response else None + + async def selected(request_kwargs: dict) -> str: + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + return deployment["litellm_params"]["model"] + + assert await routed(self._team_request()) == "gemini-flash" + assert await selected(self._team_request()) == "gemini/gemini-3.6-flash" + for request_kwargs in (self._team_request(None), self._team_request(self.OTHER_TEAM)): + assert await routed(request_kwargs) is None + assert await selected(request_kwargs) == "openai/gpt-4o" + + @pytest.mark.asyncio + async def test_tag_filtering_hands_untagged_team_requests_to_the_team_plain_sibling(self): + plain_sibling = { + "model_name": self.SIBLING_INTERNAL_NAME, + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"team_id": self.TEAM, "team_public_model_name": self.PUBLIC_NAME}, + } + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + markers=(self._team_marker(self.INTERNAL_NAME, tags=["route"]),), + extra_deployments=(plain_sibling,), + enable_tag_filtering=True, + ) + + tagged = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(tags=["route"]), messages=self._messages() + ) + assert tagged is not None and tagged.model == "gemini-flash" + for _ in range(20): + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + assert deployment["litellm_params"]["model"] == "openai/gpt-4o" + + @pytest.mark.asyncio + async def test_marker_only_team_resolution_is_rejected_as_uncallable(self): + import re + + from litellm.types.router import RouterErrors + + router = self._router({}) + + with pytest.raises( + litellm.BadRequestError, match=re.escape(RouterErrors.only_strategy_marker_deployments.value) + ): + await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + + @pytest.mark.asyncio + async def test_proxy_admin_without_a_team_reaches_the_team_strategy_by_public_name(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + request_kwargs = {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}} + + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + + assert response is not None + assert response.model == "gemini-flash" + @pytest.mark.asyncio + async def test_strategy_resolution_agrees_with_the_deployment_path_for_every_principal(self): + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + extra_deployments=({"model_name": "shared-name", "litellm_params": {"model": "openai/gpt-4o"}},), + ) + principals = { + "team": self._team_request(), + "other-team": self._team_request(self.OTHER_TEAM), + "no-team": self._team_request(None), + "admin": {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}}, + } + for principal, request_kwargs in principals.items(): + for model in (self.PUBLIC_NAME, "shared-name", "gemini-flash", "missing"): + resolved = [d["model_name"] for d in router.deployments_for_request(model, request_kwargs)] + callable_names = [ + name + for name, deployment in zip(resolved, router.deployments_for_request(model, request_kwargs)) + if not router._is_strategy_marker_deployment(deployment) + ] + if resolved and not callable_names: + with pytest.raises(litellm.BadRequestError, match="strategy router marker"): + router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) + elif not resolved: + with pytest.raises(litellm.BadRequestError): + router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) + else: + _, deployments = router._common_checks_available_deployment( + model=model, request_kwargs=request_kwargs + ) + assert [d["model_name"] for d in deployments] == callable_names, (principal, model) + + def test_drop_strategy_markers_keeps_plain_deployments_and_rejects_marker_only_sets(self): + router = self._router({}) + marker = router.model_list[0] + plain = {"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o"}} + + assert router._drop_strategy_markers("x", [marker, plain]) == [plain] + assert router._drop_strategy_markers("x", [plain]) == [plain] + assert router._drop_strategy_markers("x", []) == [] + with pytest.raises(litellm.BadRequestError, match="strategy router marker"): + router._drop_strategy_markers("x", [marker]) + + def test_team_deployments_across_teams_unions_one_team_and_rejects_two(self): + other_team_marker = { + **self._team_marker(self.SIBLING_INTERNAL_NAME), + "model_info": {"team_id": self.OTHER_TEAM, "team_public_model_name": self.PUBLIC_NAME}, + } + one_team = self._router({}) + two_teams = self._router({}, markers=(self._team_marker(self.INTERNAL_NAME), other_team_marker)) + + assert [d["model_name"] for d in one_team._team_deployments_across_teams(self.PUBLIC_NAME)] == [ + self.INTERNAL_NAME + ] + assert one_team._team_deployments_across_teams("missing") == [] + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + two_teams._team_deployments_across_teams(self.PUBLIC_NAME) + + + def test_compression_policy_follows_the_same_resolution_for_every_principal(self): + from litellm.proxy.guardrails.auto_router_compression import AutoRouterCompressionPolicy, policy_for_model + + marker = self._team_marker(self.INTERNAL_NAME) + marker["litellm_params"]["auto_router_routing_compression"] = "headroom-team" + router = self._router({}, markers=(marker,)) + admin = {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}} + expected = AutoRouterCompressionPolicy(routing="headroom-team", model=None) + + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(), ()) == expected + assert policy_for_model(router, self.PUBLIC_NAME, admin, ()) == expected + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(self.OTHER_TEAM), ()) is None + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(None), ()) is None + + class TestAutoRouterCompressionDecoupling: """An auto router's `auto_router_routing_compression` / `auto_router_model_compression` decouple what the routing decision sees from what the model call sees. The one @@ -12805,6 +14005,82 @@ class TestTierParamsTheTargetAccepts: assert accepted == {"reasoning_effort": "max"} +@pytest.mark.asyncio +async def test_router_aspeech_without_voice_dispatches_ref_audio_cloning(respx_mock, monkeypatch): + import base64 + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + audio_bytes = b"RIFFfake-wav-bytes" + respx_mock.post("https://api.mistral.ai/v1/audio/speech").respond( + json={"audio_data": base64.b64encode(audio_bytes).decode()} + ) + router = Router( + model_list=[ + { + "model_name": "voxtral-tts", + "litellm_params": {"model": "mistral/voxtral-mini-tts-2603"}, + } + ] + ) + + response = await router.aspeech(model="voxtral-tts", input="clone me", ref_audio="ZmFrZQ==") + + request_body = json.loads(respx_mock.calls.last.request.content) + assert request_body == {"model": "voxtral-mini-tts-2603", "input": "clone me", "ref_audio": "ZmFrZQ=="} + assert response.content == audio_bytes + + +@pytest.mark.asyncio +async def test_router_aspeech_without_voice_keeps_deployment_default_voice(respx_mock, monkeypatch): + import base64 + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + audio_bytes = b"RIFFfake-wav-bytes" + respx_mock.post("https://api.mistral.ai/v1/audio/speech").respond( + json={"audio_data": base64.b64encode(audio_bytes).decode()} + ) + router = Router( + model_list=[ + { + "model_name": "voxtral-tts", + "litellm_params": {"model": "mistral/voxtral-mini-tts-2603", "voice": "en_paul_neutral"}, + } + ] + ) + + await router.aspeech(model="voxtral-tts", input="use my default") + + request_body = json.loads(respx_mock.calls.last.request.content) + assert request_body["voice_id"] == "en_paul_neutral" + + +@pytest.mark.asyncio +async def test_router_aspeech_request_voice_overrides_deployment_default(respx_mock, monkeypatch): + import base64 + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + audio_bytes = b"RIFFfake-wav-bytes" + respx_mock.post("https://api.mistral.ai/v1/audio/speech").respond( + json={"audio_data": base64.b64encode(audio_bytes).decode()} + ) + router = Router( + model_list=[ + { + "model_name": "voxtral-tts", + "litellm_params": {"model": "mistral/voxtral-mini-tts-2603", "voice": "en_paul_neutral"}, + } + ] + ) + + await router.aspeech(model="voxtral-tts", input="override me", voice="gb_oliver_neutral") + + request_body = json.loads(respx_mock.calls.last.request.content) + assert request_body["voice_id"] == "gb_oliver_neutral" + + class TestRequestReasoningEffortOverride: def test_drop_effort_from_nested_carrier_preserves_other_nested_values(self): params: dict[str, object] = {"output_config": {"effort": "high", "format": "json"}} @@ -13116,6 +14392,7 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo ({"DefaultRetries": 0}, 502, litellm.BadGatewayError, 1), ({"DefaultRetries": 0, "ServiceUnavailableErrorRetries": 1}, 503, litellm.ServiceUnavailableError, 2), ({"ServiceUnavailableErrorRetries": 0}, 502, litellm.BadGatewayError, 3), + ({"BadRequestErrorRetries": 2}, 400, litellm.BadRequestError, 3), ], ) async def test_router_retry_policy_controls_upstream_attempt_count( @@ -13152,6 +14429,591 @@ async def test_router_retry_policy_controls_upstream_attempt_count( assert upstream.call_count == expected_upstream_calls +@pytest.mark.asyncio +async def test_generic_call_keeps_the_deployment_name_of_an_azure_ai_model_on_an_azure_openai_host(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + { + "model_name": "aoai-gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://my-resource.openai.azure.com", + "api_key": "deployment-key", + }, + } + ] + ) + + with respx.mock(assert_all_called=True) as respx_mock: + upstream = respx_mock.post(host="my-resource.openai.azure.com", path__regex=r"^/openai/.*responses$").mock( + return_value=httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + }, + ) + ) + await router.aresponses(model="aoai-gpt", input="hi") + + assert json.loads(upstream.calls.last.request.content)["model"] == "gpt-5.4-mini" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retry_policy,upstream_error", + [ + ( + {"BadRequestErrorRetries": 2}, + { + "message": "This model's maximum context length is 16385 tokens", + "type": "invalid_request_error", + "code": "context_length_exceeded", + }, + ), + ( + {"ContentPolicyViolationErrorRetries": 2}, + { + "message": "Your request was rejected as a result of our safety system", + "type": "invalid_request_error", + "code": "content_policy_violation", + }, + ), + ], +) +async def test_router_retry_policy_400_retries_on_sibling_deployment( + monkeypatch: pytest.MonkeyPatch, retry_policy, upstream_error +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": "https://rejecting.local/v1", + "weight": 1, + }, + "model_info": {"id": "rejecting"}, + }, + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": "https://accepting.local/v1", + "weight": 0, + }, + "model_info": {"id": "accepting"}, + }, + ], + num_retries=2, + retry_policy=retry_policy, + disable_cooldowns=True, + ) + + with respx.mock(assert_all_called=True) as respx_mock: + rejecting = respx_mock.post("https://rejecting.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": upstream_error}) + ) + accepting = respx_mock.post("https://accepting.local/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-lit-7036", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "hi back"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + }, + ) + ) + response = await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + + assert rejecting.call_count == 1 + assert accepting.call_count == 1 + assert response.choices[0].message.content == "hi back" + assert response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] == 1 + + +_UPSTREAM_400 = {"message": "upstream refused this request", "type": "invalid_request_error", "code": "bad_request"} + + +def _retry_skip_deployment(deployment_id, host, litellm_params=None, model_info=None): + return { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": f"https://{host}.local/v1", + **(litellm_params or {}), + }, + "model_info": {"id": deployment_id, **(model_info or {})}, + } + + +@pytest.mark.parametrize( + "status_code,failed_deployment_id,already_skipped,expected", + [ + (400, "rejecting", None, ("rejecting",)), + (403, "rejecting", None, ("rejecting",)), + (400, "second", ("first",), ("first", "second")), + (400, "first", ("first",), ("first",)), + (429, "rejecting", None, ()), + (503, "rejecting", None, ()), + (408, "rejecting", None, ()), + (400, None, None, ()), + (None, "rejecting", None, ()), + ("400", "rejecting", None, ()), + (400, "second", 7, ("second",)), + (400, "second", "first", ("second",)), + (400, "second", ["first"], ("second",)), + (400, "second", ("first", 7), ("first", "second")), + ], +) +def test_router_deployment_ids_to_skip_on_retry(status_code, failed_deployment_id, already_skipped, expected): + exception = Exception("upstream refused this request") + exception.status_code = status_code + exception.failed_deployment_id = failed_deployment_id + + assert litellm.Router._deployment_ids_to_skip_on_retry(exception, already_skipped) == expected + + +@pytest.mark.parametrize( + "kwargs,failed_deployment_id,expected", + [ + ({"model_info": {"id": "rejecting"}}, None, ("rejecting",)), + ({"model_info": {"id": "rejecting"}}, "cooldown-target", ("rejecting",)), + ({"model_info": {"id": ""}}, None, ()), + ({"model_info": {"id": 7}}, None, ()), + ({"model_info": "rejecting"}, None, ()), + ({}, None, ()), + ({}, "cooldown-target", ("cooldown-target",)), + ], +) +def test_router_retry_skip_stamp_feeds_deployment_ids_to_skip_on_retry( + kwargs: Mapping[str, object], failed_deployment_id: str | None, expected: tuple[str, ...] +): + exception = Exception("upstream refused this request") + exception.status_code = 400 + exception.failed_deployment_id = failed_deployment_id + + litellm.Router._stamp_retry_skip_deployment_id(exception, kwargs) + + assert litellm.Router._deployment_ids_to_skip_on_retry(exception, None) == expected + assert exception.failed_deployment_id == failed_deployment_id + + +@pytest.mark.parametrize( + "value,expected", + [ + (("first", "second"), ("first", "second")), + ((), ()), + (("first", 7, None, "second"), ("first", "second")), + (None, ()), + (7, ()), + ("first", ()), + (["first"], ()), + ({"first": True}, ()), + (object(), ()), + ], +) +def test_router_as_retry_skipped_deployment_ids_keeps_only_a_tuple_of_strings(value, expected): + from litellm.router import _as_retry_skipped_deployment_ids + + assert _as_retry_skipped_deployment_ids(value) == expected + + +@pytest.mark.parametrize( + "deployment_ids,skipped,expected", + [ + (["rejecting", "sibling"], ("rejecting",), ["sibling"]), + (["rejecting"], ("rejecting",), ["rejecting"]), + (["rejecting", "sibling"], ("rejecting", "sibling"), ["rejecting", "sibling"]), + (["rejecting", "sibling"], (), ["rejecting", "sibling"]), + (["rejecting", "sibling"], None, ["rejecting", "sibling"]), + (["rejecting", "sibling"], ("absent",), ["rejecting", "sibling"]), + (["rejecting", "sibling"], 7, ["rejecting", "sibling"]), + (["rejecting", "sibling"], "rejecting", ["rejecting", "sibling"]), + (["rejecting", "sibling"], ["rejecting"], ["rejecting", "sibling"]), + (["rejecting", "sibling"], {"rejecting": True}, ["rejecting", "sibling"]), + (["rejecting", "sibling"], ("rejecting", 7), ["sibling"]), + ], +) +@pytest.mark.asyncio +async def test_router_healthy_deployments_keep_the_last_candidate_a_retry_skipped(deployment_ids, skipped, expected): + router = litellm.Router( + model_list=[_retry_skip_deployment(deployment_id, deployment_id) for deployment_id in deployment_ids], + disable_cooldowns=True, + ) + request_kwargs = {"_retry_skipped_deployment_ids": skipped} + + healthy_deployments = await router.async_get_healthy_deployments(model="gpt-5.6", request_kwargs=request_kwargs) + + assert sorted(deployment["model_info"]["id"] for deployment in healthy_deployments) == sorted(expected) + assert "_retry_skipped_deployment_ids" not in request_kwargs + + +@pytest.mark.parametrize("client_supplied", [7, "rejecting", ["rejecting"], {"rejecting": True}, object()]) +@pytest.mark.asyncio +async def test_router_retry_policy_400_keeps_upstream_error_when_a_client_forges_the_skip_list( + monkeypatch: pytest.MonkeyPatch, client_supplied +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[_retry_skip_deployment("rejecting", "rejecting"), _retry_skip_deployment("sibling", "sibling")], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + respx_mock.post("https://rejecting.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + respx_mock.post("https://sibling.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + _retry_skipped_deployment_ids=client_supplied, + ) + + assert "upstream refused this request" in str(raised.value) + + +@pytest.mark.asyncio +async def test_router_retry_policy_400_keeps_upstream_error_on_order_fallback_hop(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + _retry_skip_deployment("order1", "order1", litellm_params={"order": 1}), + _retry_skip_deployment("order2", "order2", litellm_params={"order": 2}), + ], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + order1 = respx_mock.post("https://order1.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + order2 = respx_mock.post("https://order2.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + + assert "upstream refused this request" in str(raised.value) + assert "No deployments available" not in str(raised.value) + assert order1.call_count >= 1 + assert order2.call_count >= 1 + + +@pytest.mark.asyncio +async def test_router_retry_policy_400_keeps_upstream_error_when_tags_narrow_the_group( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + _retry_skip_deployment( + "tagged", "tagged", litellm_params={"tags": ["free"]}, model_info={"enable_tag_filtering": True} + ), + _retry_skip_deployment("untagged", "untagged", model_info={"enable_tag_filtering": True}), + ], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + tagged = respx_mock.post("https://tagged.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + untagged = respx_mock.post("https://untagged.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["free"]}, + ) + + assert "upstream refused this request" in str(raised.value) + assert "No deployments available" not in str(raised.value) + assert tagged.call_count == 3 + assert untagged.call_count == 0 + + +@pytest.mark.asyncio +async def test_router_retry_policy_400_never_returns_to_a_deployment_that_already_refused( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(simple_shuffle.random, "choice", lambda deployments: deployments[0]) + router = litellm.Router( + model_list=[ + _retry_skip_deployment("first-refuser", "first-refuser", litellm_params={"weight": 1}), + _retry_skip_deployment("second-refuser", "second-refuser", litellm_params={"weight": 0}), + _retry_skip_deployment("accepting", "accepting", litellm_params={"weight": 0}), + ], + num_retries=3, + retry_policy={"BadRequestErrorRetries": 3}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + first = respx_mock.post("https://first-refuser.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + second = respx_mock.post("https://second-refuser.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + accepting = respx_mock.post("https://accepting.local/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-lit-7036", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "hi back"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + }, + ) + ) + response = await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + + assert first.call_count == 1 + assert second.call_count == 1 + assert accepting.call_count == 1 + assert response.choices[0].message.content == "hi back" + + +_LIT_7114_CHAT_OK = { + "id": "chatcmpl-lit-7114", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi back"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, +} +_LIT_7114_EMBEDDING_OK = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}], + "model": "text-embedding-3-large", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, +} +_LIT_7114_IMAGE_OK = {"created": 1, "data": [{"b64_json": "aGk="}]} +_LIT_7114_BATCH_OK = { + "id": "batch_lit_7114", + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": "file-lit-7114", + "completion_window": "24h", + "status": "validating", + "created_at": 1, +} + + +class _PassthroughAdapter(CustomLogger): + def translate_completion_input_params(self, kwargs: ChatCompletionRequest) -> ChatCompletionRequest: + return ChatCompletionRequest(**kwargs) + + def translate_completion_output_params(self, response: litellm.ModelResponse) -> litellm.ModelResponse: + return response + + +def _lit_7114_router(litellm_model: str) -> litellm.Router: + api_base_suffix: Final = "" if litellm_model.startswith("cohere/") else "/v1" + return litellm.Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": litellm_model, + "api_key": "sk-fake", + "api_base": f"https://{host}.local{api_base_suffix}", + "weight": weight, + }, + "model_info": {"id": host}, + } + for host, weight in (("rejecting", 1), ("accepting", 0)) + ], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + +def _lit_7114_mock_upstreams( + respx_mock: respx.MockRouter, path: str, refusal_status: int, success_body: Mapping[str, object] | bytes +) -> tuple[respx.Route, respx.Route]: + ok_response: Final = ( + httpx.Response(200, content=success_body) + if isinstance(success_body, bytes) + else httpx.Response(200, json=success_body) + ) + rejecting: Final = respx_mock.post(f"https://rejecting.local{path}").mock( + return_value=httpx.Response( + refusal_status, json={"error": _UPSTREAM_400, "message": "upstream refused this request"} + ) + ) + accepting: Final = respx_mock.post(f"https://accepting.local{path}").mock(return_value=ok_response) + return rejecting, accepting + + +_LIT_7114_ASYNC_ENTRYPOINTS: Final[ + Mapping[str, tuple[str, str, int, Mapping[str, object] | bytes, Callable[[litellm.Router], Awaitable[object]]]] +] = { + "aembedding": ( + "openai/text-embedding-3-large", + "/v1/embeddings", + 400, + _LIT_7114_EMBEDDING_OK, + lambda router: router.aembedding(model="gpt-5.6", input="hi"), + ), + "aimage_generation": ( + "openai/gpt-image-1", + "/v1/images/generations", + 400, + _LIT_7114_IMAGE_OK, + lambda router: router.aimage_generation(model="gpt-5.6", prompt="a cat"), + ), + "atext_completion": ( + "text-completion-openai/gpt-3.5-turbo-instruct", + "/v1/completions", + 400, + {"id": "c", "object": "text_completion", "created": 1, "model": "i", "choices": [{"text": "hi", "index": 0}]}, + lambda router: router.atext_completion(model="gpt-5.6", prompt="hi"), + ), + "aspeech": ( + "openai/gpt-4o-mini-tts", + "/v1/audio/speech", + 400, + b"RIFF", + lambda router: router.aspeech(model="gpt-5.6", input="hi", voice="alloy"), + ), + "atranscription": ( + "openai/gpt-4o-transcribe", + "/v1/audio/transcriptions", + 400, + {"text": "hi"}, + lambda router: router.atranscription(model="gpt-5.6", file=("hi.wav", b"RIFF", "audio/wav")), + ), + "arerank": ( + "cohere/rerank-v3.5", + "/v2/rerank", + 400, + {"id": "r", "results": [{"index": 0, "relevance_score": 0.9}], "meta": {}}, + lambda router: router.arerank(model="gpt-5.6", query="hi", documents=["hi"]), + ), + "aadapter_completion": ( + "openai/gpt-5.6", + "/v1/chat/completions", + 400, + _LIT_7114_CHAT_OK, + lambda router: router.aadapter_completion( + adapter_id="lit-7114", model="gpt-5.6", messages=[{"role": "user", "content": "hi"}] + ), + ), + "acreate_batch": ( + "openai/gpt-5.6", + "/v1/batches", + 401, + _LIT_7114_BATCH_OK, + lambda router: router.acreate_batch( + model="gpt-5.6", completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-lit-7114" + ), + ), + "acancel_batch": ( + "openai/gpt-5.6", + "/v1/batches/batch_lit_7114/cancel", + 401, + {**_LIT_7114_BATCH_OK, "status": "cancelling"}, + lambda router: router.acancel_batch(model="gpt-5.6", batch_id="batch_lit_7114"), + ), +} + + +@pytest.mark.parametrize("entrypoint", sorted(_LIT_7114_ASYNC_ENTRYPOINTS)) +@pytest.mark.asyncio +async def test_router_retry_moves_off_the_refusing_deployment_on_every_async_entrypoint( + monkeypatch: pytest.MonkeyPatch, entrypoint: str +): + litellm_model, path, refusal_status, success_body, call = _LIT_7114_ASYNC_ENTRYPOINTS[entrypoint] + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "adapters", [{"id": "lit-7114", "adapter": _PassthroughAdapter()}]) + router: Final = _lit_7114_router(litellm_model) + + with respx.mock as respx_mock: + rejecting, accepting = _lit_7114_mock_upstreams(respx_mock, path, refusal_status, success_body) + response: Final = await call(router) + + assert response is not None + assert rejecting.call_count == 1 + assert accepting.call_count == 1 + + +_LIT_7114_SYNC_ENTRYPOINTS: Final[ + Mapping[str, tuple[str, str, Mapping[str, object], Callable[[litellm.Router], object]]] +] = { + "embedding": ( + "openai/text-embedding-3-large", + "/v1/embeddings", + _LIT_7114_EMBEDDING_OK, + lambda router: router.embedding(model="gpt-5.6", input="hi"), + ), + "image_generation": ( + "openai/gpt-image-1", + "/v1/images/generations", + _LIT_7114_IMAGE_OK, + lambda router: router.image_generation(model="gpt-5.6", prompt="a cat"), + ), +} + + +@pytest.mark.parametrize("entrypoint", sorted(_LIT_7114_SYNC_ENTRYPOINTS)) +def test_router_retry_policy_400_moves_off_the_refusing_deployment_on_every_sync_entrypoint( + monkeypatch: pytest.MonkeyPatch, entrypoint: str +): + litellm_model, path, success_body, call = _LIT_7114_SYNC_ENTRYPOINTS[entrypoint] + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _lit_7114_router(litellm_model) + + with respx.mock as respx_mock: + rejecting, accepting = _lit_7114_mock_upstreams(respx_mock, path, 400, success_body) + response: Final = call(router) + + assert response is not None + assert rejecting.call_count == 1 + assert accepting.call_count == 1 + + def _make_failure_logging_obj(): return LiteLLMLogging( model="gpt-5.6", @@ -13421,3 +15283,167 @@ async def test_router_max_parallel_requests_slot_released_when_stream_closed_ear assert tracker.peak == 1 assert tracker.current == 0 + + +@pytest.mark.asyncio +async def test_router_deployment_drop_params_string_true_is_honored(monkeypatch): + from litellm import Router + + monkeypatch.setattr(litellm, "drop_params", False) + router = Router( + model_list=[ + { + "model_name": "gpt-5-nano", + "litellm_params": { + "model": "openai/gpt-5-nano", + "api_key": "sk-fake", + "temperature": 1, + "reasoning_effort": "minimal", + "drop_params": "true", + "mock_response": "Hello, world!", + }, + } + ], + num_retries=0, + ) + + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-nano") + assert deployment is not None + assert deployment.litellm_params.drop_params is True + + response = await router.acompletion( + model="gpt-5-nano", + messages=[{"role": "user", "content": "hi"}], + temperature=0.1, + ) + assert response.choices[0].message.content == "Hello, world!" + + +@pytest.mark.parametrize("value", ["ture", "enabled"]) +def test_router_warns_when_a_deployment_drop_params_string_is_not_a_flag(value, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + router = Router( + model_list=[ + { + "model_name": "gpt-5-nano", + "litellm_params": {"model": "openai/gpt-5-nano", "api_key": "sk-fake", "drop_params": value}, + } + ] + ) + + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-nano") + assert deployment is not None + assert deployment.litellm_params.drop_params == value + assert f"model=gpt-5-nano drop_params={value!r} is not a flag value, treating it as unset" in caplog.text + + +@pytest.mark.parametrize("value", [True, "true", "off", None]) +def test_router_stays_quiet_when_a_deployment_drop_params_is_a_flag(value, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + Router( + model_list=[ + { + "model_name": "gpt-5-nano", + "litellm_params": {"model": "openai/gpt-5-nano", "api_key": "sk-fake", "drop_params": value}, + } + ] + ) + + assert "is not a flag value" not in caplog.text + + +def test_get_candidate_model_ids_for_route_covers_model_name_and_pattern(): + """ + get_candidate_model_ids_for_route resolves a route the way the router does, so a + pre-call check can tell a genuine cross-group route from same-group unavailability. + A concrete model group returns its member ids; a wildcard/pattern deployment is + included for a concrete model it matches, which the bare model_name index misses. + The unprefixed-name case must resolve through get_deployments_by_pattern (which retries + the provider-qualified form), not a bare pattern_router.route that only sees the literal + name. Regression guard for the LIT-7195 tier-change discriminator's team/pattern gaps. + """ + router = Router( + model_list=[ + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-a", "api_base": "https://x.invalid"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-b", "api_base": "https://x.invalid"}, + "model_info": {"id": "dep-b"}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "sk-c", "api_base": "https://x.invalid"}, + "model_info": {"id": "dep-wild"}, + }, + ] + ) + + assert router.get_candidate_model_ids_for_route(model="grp") == frozenset({"dep-a", "dep-b"}) + assert "dep-wild" in router.get_candidate_model_ids_for_route(model="openai/gpt-4o-some-new-model") + # unprefixed name whose provider resolves to openai: only get_deployments_by_pattern's + # provider-qualified retry matches "openai/*"; a bare route() on the literal name misses it + assert "dep-wild" in router.get_candidate_model_ids_for_route(model="gpt-5") + + +def test_deployment_ids_stringifies_ids_and_skips_entries_without_a_model_info_id(): + deployments = ( + {"model_info": {"id": "a"}}, + {"model_info": {"id": 2}}, + {"model_info": {}}, + {"no_model_info": True}, + ) + assert Router._deployment_ids(deployments) == frozenset({"a", "2"}) + + +def test_cached_model_info_lookups_match_uncached_and_reset_on_model_list_change(): + def deployment(max_output_tokens: int) -> Deployment: + return Deployment( + model_name="grp", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_key="sk-a"), + model_info=ModelInfo(id="dep-a", max_output_tokens=max_output_tokens), + ) + + router = Router(model_list=[deployment(100).model_dump()]) + + assert router.cached_model_group_info("grp") == router.get_model_group_info("grp") + first = router.cached_deployment_model_info("dep-a", "openai/gpt-4o") + assert first == router.get_deployment_model_info(model_id="dep-a", model_name="openai/gpt-4o") + assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o") is first + + router.upsert_deployment(deployment(200)) + + assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["max_output_tokens"] == 200 + assert router.cached_model_group_info("grp").max_output_tokens == 200 + + +class _OpenBreakerRedis: + def __init__(self) -> None: + from litellm.caching.redis_cache import RedisCircuitBreaker + + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + for _ in range(3): + self._circuit_breaker.record_failure() + + @_redis_circuit_breaker_guard + async def async_get_cache(self, key, **kwargs): + raise AssertionError("never reached") + + +@pytest.mark.asyncio +async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warning(caplog): + router = litellm.Router( + model_list=[{"model_name": "haiku", "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"}}] + ) + router._claude_code_session_router_cache = DualCache(redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + caplog.clear() + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): + binding = await router._get_claude_code_session_router_binding("quiet-session") + + assert binding is None + assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] + assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py index a5a68271111..3cef1c7bb63 100644 --- a/tests/test_litellm/test_router/test_io_token_rate_limits.py +++ b/tests/test_litellm/test_router/test_io_token_rate_limits.py @@ -1039,3 +1039,31 @@ class TestContextSlotRetention: assert deployment is not None router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) assert get_io_token_rate_limit_request_kwargs() is kwargs + + +@pytest.mark.asyncio +async def test_the_deployment_itpm_reservation_counts_the_request_off_the_event_loop(): + from litellm.utils import get_utc_datetime + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + warm_tokenizer("anthropic/claude-fable-5") + deployment = { + "litellm_params": {"model": "anthropic/claude-fable-5", "itpm": 10_000_000}, + "model_info": {"id": "io-loop-id"}, + "model_name": "claude", + } + set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": text * 100}], "metadata": {}}) + + _, took, lags = await timed_with_loop_lags(lambda: check.async_pre_call_check(deployment)) + + minute = get_utc_datetime().strftime("%H-%M") + reserved = await dual_cache.async_get_cache(key=f"global_router:io-loop-id:anthropic/claude-fable-5:itpm:{minute}") + assert reserved > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 30b265905f3..bd38eecd1c6 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -2361,3 +2361,31 @@ def test_a_config_deployment_dropped_for_a_permanent_reason_is_not_retried_on_re assert router.get_model_names() == ["control-model"] assert router.deployment_names == names_after_boot + + +def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info(monkeypatch): + """ + Budget reservation reads pricing through the router's lru-cached group and + deployment lookups. A reload swaps the catalog without touching model_list, so + unless the replay clears those caches the next reservation prices against the + old catalog until some unrelated model-list change happens to evict it. + """ + router = Router( + model_list=[ + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, + "model_info": {"id": "dep-a"}, + } + ] + ) + old_price = router.cached_model_group_info("grp").input_cost_per_token + assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == old_price + + new_price = old_price * 10 + fresh_catalog = copy.deepcopy(litellm.model_cost) + fresh_catalog["gpt-4o"]["input_cost_per_token"] = new_price + _simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog) + + assert router.cached_model_group_info("grp").input_cost_per_token == new_price + assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price diff --git a/tests/test_litellm/test_router_order_fallback.py b/tests/test_litellm/test_router_order_fallback.py index fde870e5abe..93895bbde08 100644 --- a/tests/test_litellm/test_router_order_fallback.py +++ b/tests/test_litellm/test_router_order_fallback.py @@ -18,10 +18,10 @@ from litellm import Router from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.prompt_caching_cache import PromptCachingCache from litellm.types.router import RouterRateLimitError -from litellm.utils import _get_deployment_order, _get_order_filtered_deployments +from litellm.utils import _get_deployment_order, get_order_filtered_deployments # --------------------------------------------------------------------------- -# Unit tests for _get_order_filtered_deployments +# Unit tests for get_order_filtered_deployments # --------------------------------------------------------------------------- @@ -42,7 +42,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(2, "b"), self._make_deployment(1, "c"), ] - result = _get_order_filtered_deployments(deps) + result = get_order_filtered_deployments(deps) assert len(result) == 2 assert all(d["model_info"]["id"] in ("a", "c") for d in result) @@ -52,7 +52,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(2, "b"), self._make_deployment(3, "c"), ] - result = _get_order_filtered_deployments(deps, target_order=2) + result = get_order_filtered_deployments(deps, target_order=2) assert len(result) == 1 assert result[0]["model_info"]["id"] == "b" @@ -61,7 +61,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(1, "a"), self._make_deployment(2, "b"), ] - result = _get_order_filtered_deployments(deps, target_order=99) + result = get_order_filtered_deployments(deps, target_order=99) assert result == [] def test_target_order_no_match_does_not_reselect_lower_order(self): @@ -70,7 +70,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(2, "b"), ] remaining_after_pre_call = [deps[0]] - result = _get_order_filtered_deployments(remaining_after_pre_call, target_order=2) + result = get_order_filtered_deployments(remaining_after_pre_call, target_order=2) assert result == [] def test_no_order_set_returns_all(self): @@ -78,11 +78,11 @@ class TestGetOrderFilteredDeployments: self._make_deployment(None, "a"), self._make_deployment(None, "b"), ] - result = _get_order_filtered_deployments(deps) + result = get_order_filtered_deployments(deps) assert len(result) == 2 def test_empty_list(self): - result = _get_order_filtered_deployments([]) + result = get_order_filtered_deployments([]) assert result == [] def test_single_order_returns_all_with_that_order(self): @@ -90,7 +90,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(1, "a"), self._make_deployment(1, "b"), ] - result = _get_order_filtered_deployments(deps) + result = get_order_filtered_deployments(deps) assert len(result) == 2 diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/test_litellm/test_router_weighted_failover.py index 162312a8c67..9f05654f23f 100644 --- a/tests/test_litellm/test_router_weighted_failover.py +++ b/tests/test_litellm/test_router_weighted_failover.py @@ -15,11 +15,11 @@ import pytest import litellm from litellm import Router -from litellm.utils import _get_excluded_filtered_deployments +from litellm.utils import get_excluded_filtered_deployments # --------------------------------------------------------------------------- -# Unit tests for _get_excluded_filtered_deployments +# Unit tests for get_excluded_filtered_deployments # --------------------------------------------------------------------------- @@ -37,17 +37,17 @@ def _make_dep(dep_id: str, weight: Optional[int] = None) -> dict: class TestGetExcludedFilteredDeployments: def test_no_excluded_returns_all(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=None) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=None) assert len(result) == 2 def test_empty_excluded_returns_all(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=[]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=[]) assert len(result) == 2 def test_drops_excluded(self): deps = [_make_dep("a"), _make_dep("b"), _make_dep("c")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) ids = sorted(d["model_info"]["id"] for d in result) assert ids == ["a", "c"] @@ -57,12 +57,12 @@ class TestGetExcludedFilteredDeployments: # error. Returning the original list here would re-include the # just-failed deployment and let weighted failover re-pick it. deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["a", "b"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["a", "b"]) assert result == [] def test_excluded_set_with_unknown_ids(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["zzz"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["zzz"]) assert len(result) == 2 def test_handles_missing_model_info(self): @@ -70,7 +70,7 @@ class TestGetExcludedFilteredDeployments: {"model_name": "x", "litellm_params": {"model": "gpt-4o"}}, # no model_info _make_dep("b"), ] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) assert len(result) == 1 diff --git a/tests/test_litellm/test_select_ui_test_scope.py b/tests/test_litellm/test_select_ui_test_scope.py index c395e07bc79..bc11fb495aa 100644 --- a/tests/test_litellm/test_select_ui_test_scope.py +++ b/tests/test_litellm/test_select_ui_test_scope.py @@ -29,7 +29,7 @@ SCOPE_SCRIPT = REPO_ROOT / ".github" / "scripts" / "select_ui_test_scope.sh" WORKFLOW = REPO_ROOT / ".github" / "workflows" / "test-litellm-ui-unit.yml" STEP_NAME = "Run UI unit tests (Vitest)" -FULL_SUITE_ARGV = ["run", "test", "--", "--run", "--pool", "forks", "--poolOptions.forks.maxForks=14"] +FULL_SUITE_ARGV = ["run", "test", "--", "--run", "--pool", "forks", "--maxWorkers=14"] NON_SRC_FILES = [ "package.json", @@ -135,7 +135,7 @@ def _related_argv(changed: list[str]) -> list[str]: "--passWithNoTests", "--pool", "forks", - "--poolOptions.forks.maxForks=14", + "--maxWorkers=14", ] diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py index 8cce6bc735a..6652211a828 100644 --- a/tests/test_litellm/test_test_quality_gate.py +++ b/tests/test_litellm/test_test_quality_gate.py @@ -1,16 +1,21 @@ """Tests for scripts/test_quality_gate.py. -The gate's whole value is that it blames a change only for what it adds, that a limit -can never rise, and that a limit cannot stay above a count the branch pushed below it. -All three live in pure functions, so they are tested directly: `evaluate` for the blame -rule, `ratcheted_budget` for the one-way ratchet, `unratcheted` for the ceiling a branch -left behind, and `parse_changed_lines` for the diff scan that turns a breach into -file:line. +The gate's whole value is that it blames a change only for what it adds and that a +limit can never rise. Both live in pure functions, so they are tested directly: +`evaluate` for the blame rule, `ratcheted_budget` for the one-way ratchet, and +`parse_changed_lines` for the diff scan that turns a breach into file:line. """ import importlib.util +import os +import signal +import subprocess import sys +import time +from collections.abc import Callable +from contextlib import suppress from pathlib import Path +from typing import NamedTuple _REPO_ROOT = Path(__file__).resolve().parents[2] _MODULE_PATH = _REPO_ROOT / "scripts" / "test_quality_gate.py" @@ -23,6 +28,16 @@ _spec.loader.exec_module(gate) _BUDGET = {"TQ001": {"limit": 10}, "TQ003": {"limit": 5}} +_SCAN_BASE = ( + "import importlib.util, pathlib, sys\n" + "spec = importlib.util.spec_from_file_location('test_quality_gate', sys.argv[1])\n" + "gate = importlib.util.module_from_spec(spec)\n" + "sys.modules[spec.name] = gate\n" + "spec.loader.exec_module(gate)\n" + "gate.base_counts('HEAD', repo_root=pathlib.Path(sys.argv[2]), checker=pathlib.Path(sys.argv[3]))\n" +) +_SCAN_BASE_WITH_SIGHUP_IGNORED = "import signal\nsignal.signal(signal.SIGHUP, signal.SIG_IGN)\n" + _SCAN_BASE + def test_a_rule_within_its_limit_is_not_a_breach(): assert gate.evaluate({"TQ001": 10}, {"TQ001": 10}, _BUDGET) == () @@ -73,29 +88,6 @@ def test_ratchet_lowers_a_rule_introduced_on_this_branch_like_any_other(): assert updated["TQ001"]["limit"] == 4 -def test_a_branch_that_cleared_violations_must_lower_the_ceiling(): - stale = gate.unratcheted({"TQ001": 6}, {"TQ001": 10}, _BUDGET) - assert [(b.rule, b.total, b.cap, b.added) for b in stale] == [("TQ001", 6, 10, -4)] - - -def test_headroom_already_in_the_base_is_not_blamed_on_this_branch(): - assert gate.unratcheted({"TQ001": 6}, {"TQ001": 6}, _BUDGET) == () - - -def test_a_branch_that_cleared_down_to_the_ceiling_exactly_is_clean(): - assert gate.unratcheted({"TQ001": 10}, {"TQ001": 12}, _BUDGET) == () - - -def test_a_branch_that_added_violations_is_not_a_ratchet_finding(): - assert gate.unratcheted({"TQ001": 14}, {"TQ001": 10}, _BUDGET) == () - - -def test_the_ratchet_finding_survives_the_update_that_answers_it(): - cleared = {"TQ001": 6} - updated = gate.ratcheted_budget(_BUDGET, cleared, {"TQ001": 10}) - assert gate.unratcheted(cleared, {"TQ001": 10}, updated) == () - - def test_parse_changed_lines_groups_hunks_under_their_own_file(): diff = ( "diff --git a/tests/a.py b/tests/a.py\n" @@ -146,3 +138,99 @@ def test_the_shipped_budget_covers_every_rule_the_checker_can_emit(): budget = json.loads((_REPO_ROOT / "test-quality-budget.json").read_text()) assert set(budget) == {"TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008"} assert all(spec["limit"] >= 0 for spec in budget.values()) + + +def _git(cwd: Path, *args: str) -> str: + proc = subprocess.run(["git", *args], cwd=cwd, capture_output=True, text=True) + assert proc.returncode == 0, proc.stderr + return proc.stdout.strip() + + +def _committed_repo(tmp_path: Path) -> Path: + repo = tmp_path / "repo" + (repo / "tests").mkdir(parents=True) + (repo / "tests" / "test_seed.py").write_text("def test_seed():\n assert True\n") + _git(repo, "init", "-q", "-b", "main") + _git(repo, "config", "user.email", "gate@example.com") + _git(repo, "config", "user.name", "gate") + _git(repo, "config", "commit.gpgsign", "false") + _git(repo, "add", "-A") + _git(repo, "commit", "-q", "-m", "seed") + return repo + + +def _wait_until(predicate: Callable[[], bool], timeout_seconds: float) -> bool: + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return predicate() + + +def _reap(process: subprocess.Popen[bytes]) -> None: + with suppress(subprocess.TimeoutExpired): + process.wait(timeout=10) + if process.poll() is None: + process.kill() + process.wait(timeout=10) + + +def _registered_worktrees(repo: Path) -> int: + listing = _git(repo, "worktree", "list", "--porcelain") + return sum(line.startswith("worktree ") for line in listing.splitlines()) + + +class _StalledScan(NamedTuple): + process: subprocess.Popen[bytes] + repo: Path + release: Path + temp_dir: Path + + +def _base_scan_stalled_in_its_checker(tmp_path: Path, driver: str) -> _StalledScan: + repo = _committed_repo(tmp_path) + scanning = tmp_path / "scanning" + release = tmp_path / "release" + slow_checker = tmp_path / "slow_checker.py" + slow_checker.write_text( + "import pathlib, time\n" + f"pathlib.Path({str(scanning)!r}).touch()\n" + f"while not pathlib.Path({str(release)!r}).exists():\n" + " time.sleep(0.05)\n" + ) + temp_dir = tmp_path / "tmp" + temp_dir.mkdir() + scan = subprocess.Popen( + [sys.executable, "-c", driver, str(_MODULE_PATH), str(repo), str(slow_checker)], + env={**os.environ, "TMPDIR": str(temp_dir)}, + ) + if not _wait_until(scanning.exists, 30): + _reap(scan) + raise AssertionError("the base scan never reached the checker") + return _StalledScan(scan, repo, release, temp_dir) + + +def test_a_terminated_base_scan_still_removes_its_worktree(tmp_path: Path) -> None: + stalled = _base_scan_stalled_in_its_checker(tmp_path, _SCAN_BASE) + try: + stalled.process.send_signal(signal.SIGTERM) + assert stalled.process.wait(timeout=30) == 128 + signal.SIGTERM + finally: + _reap(stalled.process) + assert _registered_worktrees(stalled.repo) == 1 + assert list(stalled.temp_dir.iterdir()) == [] + + +def test_a_base_scan_keeps_ignoring_the_hangup_its_parent_ignored(tmp_path: Path) -> None: + stalled = _base_scan_stalled_in_its_checker(tmp_path, _SCAN_BASE_WITH_SIGHUP_IGNORED) + try: + stalled.process.send_signal(signal.SIGHUP) + time.sleep(1) + assert stalled.process.poll() is None, "a hangup the parent ignored killed the scan" + stalled.release.touch() + assert stalled.process.wait(timeout=30) == 0 + finally: + _reap(stalled.process) + assert _registered_worktrees(stalled.repo) == 1 + assert list(stalled.temp_dir.iterdir()) == [] diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py index c9e2863d240..b9764eca2f8 100644 --- a/tests/test_litellm/test_together_ai_model_metadata.py +++ b/tests/test_litellm/test_together_ai_model_metadata.py @@ -5,7 +5,6 @@ from typing import Final import pytest from pydantic import TypeAdapter -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider REPO_ROOT: Final = Path(__file__).parents[2] @@ -77,59 +76,6 @@ def cost_map() -> CostMap: return COST_MAP_ADAPTER.validate_python(json.load(f)) -@pytest.mark.parametrize("model", SERVERLESS_CHAT_MODELS) -def test_together_serverless_chat_model_is_mapped(cost_map: CostMap, model: str): - info = cost_map.get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - assert info["litellm_provider"] == "together_ai" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] >= 0 - assert info["output_cost_per_token"] >= info["input_cost_per_token"] - assert "deprecation_date" not in info - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.removeprefix("together_ai/") - assert provider == "together_ai" - - -def test_together_kimi_k3_pricing_and_capabilities(cost_map: CostMap): - info = cost_map["together_ai/moonshotai/Kimi-K3"] - assert info["input_cost_per_token"] == 3e-06 - assert info["output_cost_per_token"] == 1.5e-05 - assert info["max_input_tokens"] == 1048576 - assert info["supports_function_calling"] is True - assert info["supports_tool_choice"] is True - assert info["supports_response_schema"] is True - assert info["supports_vision"] is True - assert info["supports_reasoning"] is True - - -def test_together_glm_52_pricing(cost_map: CostMap): - info = cost_map["together_ai/zai-org/GLM-5.2"] - assert info["input_cost_per_token"] == 1.4e-06 - assert info["output_cost_per_token"] == 4.4e-06 - assert info["max_input_tokens"] == 1048575 - assert info["max_output_tokens"] == 128000 - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - - -def test_together_glm_53_flash_pricing_and_capabilities(cost_map: CostMap): - info = cost_map["together_ai/zai-org/GLM-5.3-Flash"] - assert info["input_cost_per_token"] == 1.5e-07 - assert info["output_cost_per_token"] == 5e-07 - assert info["cache_read_input_token_cost"] == 3e-08 - assert info["max_input_tokens"] == 1048575 - assert info["max_output_tokens"] == 128000 - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_tool_choice"] is True - assert info["supports_response_schema"] is True - assert info["supports_vision"] is True - assert info["supports_reasoning"] is True - - def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost_map: CostMap): inflated = sorted( model @@ -142,21 +88,6 @@ def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost assert inflated == [] -def test_together_multilingual_e5_embedding_entry(cost_map: CostMap): - info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] - assert info["mode"] == "embedding" - assert info["input_cost_per_token"] == 2e-08 - assert info["max_input_tokens"] == 514 - assert info["output_vector_size"] == 1024 - - -def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: CostMap): - info = cost_map["together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"] - assert info["input_cost_per_token"] == 1.04e-06 - assert info["output_cost_per_token"] == 1.04e-06 - assert info["max_input_tokens"] == 131072 - - @pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS)) def test_together_deprecated_model_carries_deprecation_date(cost_map: CostMap, model: str): info = cost_map.get(model) @@ -210,32 +141,7 @@ CACHED_INPUT_MODELS: Final = ( ) -@pytest.mark.parametrize("model", CACHED_INPUT_MODELS) -def test_together_cached_input_model_carries_cache_read_pricing(cost_map: CostMap, model: str): - info = cost_map.get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - assert info.get("supports_prompt_caching") is True - cache_read = info.get("cache_read_input_token_cost") - assert isinstance(cache_read, float) - assert 0 < cache_read < info["input_cost_per_token"] - assert "cache_creation_input_token_cost" not in info - - def test_together_prompt_caching_flag_implies_cache_read_rate(cost_map: CostMap): for model, info in cost_map.items(): if model.startswith("together_ai/") and info.get("supports_prompt_caching"): assert "cache_read_input_token_cost" in info, f"{model} flags caching without a cache read rate" - - -def test_together_deepseek_v4_flash_cache_read_rate(cost_map: CostMap): - info = cost_map["together_ai/deepseek-ai/DeepSeek-V4-Flash-0731"] - assert info["input_cost_per_token"] == 1.4e-07 - assert info["cache_read_input_token_cost"] == 3e-08 - assert info["output_cost_per_token"] == 2.8e-07 - - -def test_together_qwen_37_max_repriced_to_current_together_rate(cost_map: CostMap): - info = cost_map["together_ai/Qwen/Qwen3.7-Max"] - assert info["input_cost_per_token"] == 2.5e-06 - assert info["output_cost_per_token"] == 7.5e-06 - assert info["cache_read_input_token_cost"] == 5e-07 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 14907e17b1b..c7e46829aba 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,11 +1,16 @@ import asyncio +import contextlib import json import logging import os +import queue +import threading from datetime import datetime, timedelta, timezone +from collections.abc import Iterator from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from jsonschema import validate @@ -13,6 +18,7 @@ from jsonschema import validate import litellm from litellm._internal_context import is_internal_call +from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT from litellm._logging import ( CorrelationContextFilter, JsonFormatter, @@ -54,6 +60,12 @@ from litellm.utils import ( # Adds the parent directory to the system path +def test_cloudflare_model_info_includes_rpm(local_model_cost_map: None) -> None: + assert litellm.get_model_info("cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8")["rpm"] == 300 + assert litellm.get_model_info("cloudflare/@cf/moonshotai/kimi-k2.6")["rpm"] == 20 + assert litellm.get_model_info("cloudflare/@cf/openai/whisper-large-v3-turbo")["rpm"] == 720 + + def test_get_utc_datetime_returns_current_aware_utc_time() -> None: before: Final = datetime.now(timezone.utc) result: Final = litellm.utils.get_utc_datetime() @@ -186,6 +198,8 @@ def test_get_model_info_resolves_provider_prefixed_model_ids(local_model_cost_ma ("perplexity/perplexity/kimi-k3", True), ("perplexity/perplexity/deepseek-v4-flash-0731", True), ("perplexity/perplexity/kimi-k2.7-code", False), + ("perplexity/perplexity/nemotron-3.5-lightning-30b-a3b", True), + ("perplexity/perplexity/nemotron-3-ultra-550b-a55b", True), ): assert litellm.supports_reasoning(model=model) is reasoning, model @@ -197,6 +211,23 @@ def test_get_model_info_resolves_provider_prefixed_model_ids(local_model_cost_ma assert via_provider["output_cost_per_token"] == 4.4e-06 assert via_provider["mode"] == "responses" + lightning = litellm.get_model_info( + model="perplexity/nemotron-3.5-lightning-30b-a3b", custom_llm_provider="perplexity" + ) + assert lightning["key"] == "perplexity/perplexity/nemotron-3.5-lightning-30b-a3b" + assert lightning["input_cost_per_token"] == 1.15e-08 + assert lightning["output_cost_per_token"] == 1.7e-07 + assert lightning["cache_read_input_token_cost"] == 1.15e-09 + assert lightning["mode"] == "responses" + + ultra = litellm.get_model_info(model="perplexity/perplexity/nemotron-3-ultra-550b-a55b") + assert ultra["key"] == "perplexity/perplexity/nemotron-3-ultra-550b-a55b" + + +def test_get_model_info_strips_openai_finetune_ids_without_a_custom_suffix(local_model_cost_map): + info = litellm.get_model_info(model="ft:gpt-4o-2024-08-06:my-org::abc123", custom_llm_provider="openai") + assert info["key"] == "ft:gpt-4o-2024-08-06" + def test_provider_prefixed_lookup_never_outranks_an_existing_row(local_model_cost_map): """The provider-prefixed candidate is tried last, after every candidate that @@ -803,6 +834,7 @@ def validate_model_cost_values(model_data, exceptions=None): "input_cost_per_second", "output_cost_per_second", "output_cost_per_second_480p", + "output_cost_per_second_720p", "output_cost_per_second_1080p", "output_cost_per_second_4k", "input_cost_per_query", @@ -1026,6 +1058,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_pixel": {"type": "number"}, "output_cost_per_second": {"type": "number"}, "output_cost_per_second_480p": {"type": "number"}, + "output_cost_per_second_720p": {"type": "number"}, "output_cost_per_second_1080p": {"type": "number"}, "output_cost_per_second_4k": {"type": "number"}, "output_cost_per_token": {"type": "number"}, @@ -1091,6 +1124,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 +1157,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", @@ -1386,23 +1431,35 @@ def test_supports_tool_choice_simple_tests(): is True ) - assert ( - litellm.utils.supports_tool_choice(model="us.amazon.nova-micro-v1:0") is False - ) - assert ( - litellm.utils.supports_tool_choice(model="bedrock/us.amazon.nova-micro-v1:0") - is False - ) - assert ( - litellm.utils.supports_tool_choice( - model="us.amazon.nova-micro-v1:0", custom_llm_provider="bedrock_converse" - ) - is False - ) - assert litellm.utils.supports_tool_choice(model="perplexity/sonar") is False +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "model", + [ + "amazon.nova-lite-v1:0", + "amazon.nova-micro-v1:0", + "amazon.nova-pro-v1:0", + "apac.amazon.nova-lite-v1:0", + "apac.amazon.nova-micro-v1:0", + "apac.amazon.nova-pro-v1:0", + "bedrock/us-gov-east-1/amazon.nova-pro-v1:0", + "bedrock/us-gov-west-1/amazon.nova-lite-v1:0", + "bedrock/us-gov-west-1/amazon.nova-micro-v1:0", + "bedrock/us-gov-west-1/amazon.nova-pro-v1:0", + "eu.amazon.nova-lite-v1:0", + "eu.amazon.nova-micro-v1:0", + "eu.amazon.nova-pro-v1:0", + "us.amazon.nova-lite-v1:0", + "us.amazon.nova-micro-v1:0", + "us.amazon.nova-pro-v1:0", + ], +) +def test_amazon_nova_v1_understanding_models_support_tool_choice(model: str) -> None: + assert litellm.utils.supports_tool_choice(model=model) is True + + def test_check_provider_match(): """ Test the _check_provider_match function for various provider scenarios @@ -2365,6 +2422,28 @@ def test_register_model_with_scientific_notation(): _invalidate_model_cost_lowercase_map() +@respx.mock +def test_register_model_url_fetch_uses_single_attempt(monkeypatch): + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + before = dict(litellm.model_cost) + threads_before = {thread.name for thread in threading.enumerate()} + route = respx.get("https://example.invalid/custom_pricing.json").mock( + return_value=httpx.Response(503) + ) + + litellm.register_model(model_cost="https://example.invalid/custom_pricing.json") + + threads_after = {thread.name for thread in threading.enumerate()} + assert route.call_count == 1 + assert not (threads_after - threads_before) & {"litellm-model-cost-map-retry"} + assert not any( + thread.name == "litellm-model-cost-map-retry" and thread.is_alive() + for thread in threading.enumerate() + ) + assert litellm.model_cost.keys() >= before.keys() + + def test_register_model_openrouter_without_slash(): """ Test that register_model handles openrouter models without '/' in the name. @@ -2879,6 +2958,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/) @@ -3956,7 +4089,7 @@ def test_deepseek_v4_models_in_cost_map(): configured in model_prices_and_context_window.json. Prices sourced from https://api-docs.deepseek.com/quick_start/pricing: - - deepseek-v4-flash: $0.44/M input, $1.32/M output + - deepseek-v4-flash: $0.30/M input, $1.20/M output - deepseek-v4-pro: $1.32/M input, $3.96/M output Closes https://github.com/BerriAI/litellm/issues/26709 @@ -3969,9 +4102,9 @@ def test_deepseek_v4_models_in_cost_map(): model_cost = json.load(f) # --- bare model names --- - for key, expected_input, expected_output, expected_cache in [ - ("deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), - ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), + for key, expected_input, expected_output, expected_cache, expected_vision in [ + ("deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09, True), + ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08, False), ]: info = model_cost.get(key) assert info is not None, f"{key} missing from model_prices_and_context_window.json" @@ -3983,11 +4116,12 @@ def test_deepseek_v4_models_in_cost_map(): assert info["max_input_tokens"] == 1_000_000 assert info["supports_function_calling"] is True assert info["supports_tool_choice"] is True + assert info.get("supports_vision", False) is expected_vision # --- provider-prefixed names --- - for key, expected_input, expected_output, expected_cache in [ - ("deepseek/deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), - ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), + for key, expected_input, expected_output, expected_cache, expected_vision in [ + ("deepseek/deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09, True), + ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08, False), ]: info = model_cost.get(key) assert info is not None, f"{key} missing from model_prices_and_context_window.json" @@ -3998,6 +4132,7 @@ def test_deepseek_v4_models_in_cost_map(): assert info["cache_read_input_token_cost"] == expected_cache assert info["supports_function_calling"] is True assert info["supports_tool_choice"] is True + assert info.get("supports_vision", False) is expected_vision def test_deepseek_v4_models_in_backup_cost_map(): @@ -4013,9 +4148,9 @@ def test_deepseek_v4_models_in_backup_cost_map(): model_cost = json.load(f) # --- bare model names --- - for key, expected_input, expected_output, expected_cache in [ - ("deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), - ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), + for key, expected_input, expected_output, expected_cache, expected_vision in [ + ("deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09, True), + ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08, False), ]: info = model_cost.get(key) assert info is not None, f"{key} missing from backup JSON" @@ -4025,11 +4160,12 @@ def test_deepseek_v4_models_in_backup_cost_map(): assert info["output_cost_per_token"] == expected_output assert info["cache_read_input_token_cost"] == expected_cache assert info["max_input_tokens"] == 1_000_000 + assert info.get("supports_vision", False) is expected_vision # --- provider-prefixed names --- - for key, expected_input, expected_output, expected_cache in [ - ("deepseek/deepseek-v4-flash", 4.4e-07, 1.32e-06, 1.4e-08), - ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08), + for key, expected_input, expected_output, expected_cache, expected_vision in [ + ("deepseek/deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09, True), + ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08, False), ]: info = model_cost.get(key) assert info is not None, f"{key} missing from backup JSON" @@ -4038,6 +4174,43 @@ def test_deepseek_v4_models_in_backup_cost_map(): assert info["input_cost_per_token"] == expected_input assert info["output_cost_per_token"] == expected_output assert info["cache_read_input_token_cost"] == expected_cache + assert info.get("supports_vision", False) is expected_vision + + +def test_deprecation_dates_for_retired_xai_and_groq_models(): + 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) + + assert model_cost["xai/grok-imagine-image-quality"]["deprecation_date"] == "2026-11-02" + assert model_cost["xai/grok-imagine-image-quality-latest"]["deprecation_date"] == "2026-11-02" + assert model_cost["xai/grok-imagine-image-quality-20260403"]["deprecation_date"] == "2026-11-02" + assert model_cost["groq/gemma-7b-it"]["deprecation_date"] == "2024-12-18" + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_deepseek_flash_completion_cost(): + from litellm.types.utils import ModelResponse + + response = ModelResponse( + model="deepseek-flash", + usage=Usage( + prompt_tokens=1_000_000, + completion_tokens=1_000_000, + total_tokens=2_000_000, + ), + ) + + cost = litellm.completion_cost( + completion_response=response, + model="deepseek-flash", + custom_llm_provider="deepseek", + ) + + assert cost == pytest.approx(1.50, abs=1e-9) _FIREWORKS_MODELS = [ @@ -5173,52 +5346,6 @@ def test_client_side_timeout_marker_never_reaches_the_provider(): ) -def test_rust_flag_not_forwarded_as_provider_param(): - forwarded = get_non_default_completion_params({"rust": True, "temperature": 0.5}) - assert "rust" not in forwarded - - -def test_completion_does_not_leak_rust_flag_into_provider_request_body(): - mock_response = MagicMock() - mock_response.model_dump.return_value = { - "id": "chatcmpl-1", - "object": "chat.completion", - "created": 1234567890, - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "hi"}, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 1, - "completion_tokens": 1, - "total_tokens": 2, - }, - } - - mock_raw_response = MagicMock() - mock_raw_response.headers = {} - mock_raw_response.parse.return_value = mock_response - - mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = mock_raw_response - - litellm.completion( - model="openai/gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - rust=True, - api_key="sk-test", - client=mock_client, - ) - - create_kwargs = mock_client.chat.completions.with_raw_response.create.call_args.kwargs - assert "rust" not in create_kwargs - assert "rust" not in (create_kwargs.get("extra_body") or {}) - - class _RecordingDeploymentFailureLogger(CustomLogger): def __init__(self) -> None: super().__init__() @@ -6067,3 +6194,207 @@ class TestFinalOptionalParamsLineRedaction: assert "'max_tokens': 17" in printed assert "'temperature': 0.25" in printed + + +class TestDropParamsStringCoercion: + @pytest.mark.parametrize("drop_params", ["true", "True", True]) + def test_truthy_drop_params_drops_unsupported_temperature(self, drop_params, monkeypatch): + from litellm.utils import get_optional_params + + monkeypatch.setattr(litellm, "drop_params", False) + result = get_optional_params( + model="gpt-5-nano", + custom_llm_provider="openai", + temperature=0.1, + drop_params=drop_params, + ) + assert "temperature" not in result + + @pytest.mark.parametrize("drop_params", ["false", False, None]) + def test_falsy_drop_params_still_raises(self, drop_params, monkeypatch): + from litellm.utils import get_optional_params + + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError): + get_optional_params( + model="gpt-5-nano", + custom_llm_provider="openai", + temperature=0.1, + drop_params=drop_params, + ) + + +def _credential_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [record.getMessage() for record in caplog.records if "litellm_credential_name=" in record.getMessage()] + + +def test_load_credentials_from_list_warns_when_the_named_credential_is_not_loaded( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.utils import load_credentials_from_list + + monkeypatch.setattr(litellm, "credential_list", []) + request_kwargs = {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"} + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + load_credentials_from_list(request_kwargs) + + assert request_kwargs == {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"} + assert _credential_warnings(caplog) == [ + "litellm_credential_name=openai-cred matched none of the 0 loaded credentials; the request runs without it" + ] + + +def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_without_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.types.utils import CredentialItem + from litellm.utils import load_credentials_from_list + + loaded = CredentialItem( + credential_name="openai-cred", + credential_values={"api_key": "sk-from-db", "api_base": "https://credential.example"}, + credential_info={}, + ) + monkeypatch.setattr(litellm, "credential_list", [loaded]) + request_kwargs = {"litellm_credential_name": "openai-cred", "api_base": "https://request.example"} + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + load_credentials_from_list(request_kwargs) + + assert request_kwargs == { + "litellm_credential_name": "openai-cred", + "api_base": "https://request.example", + "api_key": "sk-from-db", + } + assert _credential_warnings(caplog) == [] + + +_MOCK_STREAM_ID: Final = "chatcmpl-mock-stream" +_ChunkSnapshot = tuple[str, tuple[str | None, ...], Usage | None] + + +def _snapshot(chunk: ModelResponseStream) -> _ChunkSnapshot: + return chunk.id, tuple(choice.delta.content for choice in chunk.choices), getattr(chunk, "usage", None) + + +def _mock_stream_snapshots(mock_response: object, prompt_tokens: int | None) -> list[_ChunkSnapshot]: + from litellm.utils import mock_completion_streaming_obj + + return [ + _snapshot(chunk) + for chunk in mock_completion_streaming_obj( + ModelResponseStream(id=_MOCK_STREAM_ID, model="gpt-5.4-mini"), + mock_response=mock_response, + model="gpt-5.4-mini", + prompt_tokens=prompt_tokens, + ) + ] + + +async def _async_mock_stream_snapshots(mock_response: object, prompt_tokens: int | None) -> list[_ChunkSnapshot]: + from litellm.utils import async_mock_completion_streaming_obj + + return [ + _snapshot(chunk) + async for chunk in async_mock_completion_streaming_obj( + ModelResponseStream(id=_MOCK_STREAM_ID, model="gpt-5.4-mini"), + mock_response=mock_response, + model="gpt-5.4-mini", + prompt_tokens=prompt_tokens, + ) + ] + + +_CONTENT_SNAPSHOTS: Final = [(_MOCK_STREAM_ID, (content,), None) for content in ("hel", "lo ", "wor", "ld")] + + +def _assert_trailing_usage_chunk(snapshots: list[_ChunkSnapshot], prompt_tokens: int) -> None: + assert snapshots[:-1] == _CONTENT_SNAPSHOTS + chunk_id, choices, usage = snapshots[-1] + assert chunk_id == _MOCK_STREAM_ID + assert choices == () + assert usage is not None + assert usage.prompt_tokens == prompt_tokens + assert usage.completion_tokens == DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT + assert usage.total_tokens == prompt_tokens + usage.completion_tokens + + +@pytest.mark.parametrize("prompt_tokens", (51234, 0)) +def test_mock_completion_streaming_obj_emits_usage_chunk_with_admission_prompt_tokens(prompt_tokens: int) -> None: + _assert_trailing_usage_chunk(_mock_stream_snapshots("hello world", prompt_tokens), prompt_tokens) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prompt_tokens", (51234, 0)) +async def test_async_mock_completion_streaming_obj_emits_usage_chunk_with_admission_prompt_tokens( + prompt_tokens: int, +) -> None: + _assert_trailing_usage_chunk(await _async_mock_stream_snapshots("hello world", prompt_tokens), prompt_tokens) + + +def test_mock_completion_streaming_obj_emits_no_usage_chunk_without_admission_prompt_tokens() -> None: + assert _mock_stream_snapshots("hello world", None) == _CONTENT_SNAPSHOTS + + +@pytest.mark.asyncio +async def test_async_mock_completion_streaming_obj_emits_no_usage_chunk_without_admission_prompt_tokens() -> None: + assert await _async_mock_stream_snapshots("hello world", None) == _CONTENT_SNAPSHOTS + + +def test_mock_completion_streaming_obj_passes_prebuilt_stream_chunk_through_without_usage_chunk() -> None: + prebuilt: Final = ModelResponseStream( + model="gpt-5.4-mini", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="prebuilt"))] + ) + + assert _mock_stream_snapshots(prebuilt, 51234) == [(prebuilt.id, ("prebuilt",), None)] + + +@pytest.mark.asyncio +async def test_async_mock_completion_streaming_obj_raises_mock_exception_before_usage_chunk() -> None: + mock_exception: Final = litellm.MockException( + status_code=500, message="boom", llm_provider="openai", model="gpt-5.4-mini" + ) + with pytest.raises(litellm.MockException): + await _async_mock_stream_snapshots(mock_exception, 51234) + + + +@contextlib.contextmanager +def _recording_hidden_params_at_submit(submit_target: str) -> "Iterator[queue.SimpleQueue[dict[str, object]]]": + seen: Final = queue.SimpleQueue() + + def record_submit(_fn, *args, **_kwargs): + response: Final = next(arg for arg in args if isinstance(arg, litellm.ModelResponse)) + seen.put(dict(response._hidden_params)) + return MagicMock() + + with patch(submit_target, side_effect=record_submit): + yield seen + + +@pytest.mark.asyncio +async def test_acompletion_finishes_response_metadata_before_handing_the_response_to_the_logging_thread(monkeypatch): + monkeypatch.setattr(litellm, "success_callback", [lambda kwargs, response, start_time, end_time: None]) + with _recording_hidden_params_at_submit("litellm.litellm_core_utils.litellm_logging.executor.submit") as seen: + await litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + num_retries=0, + ) + snapshot: Final = seen.get_nowait() + assert snapshot["litellm_call_id"] + assert snapshot["response_cost"] is not None + assert snapshot["api_base"] + + +def test_completion_finishes_response_metadata_before_handing_the_response_to_the_logging_thread(): + with _recording_hidden_params_at_submit("litellm.utils.executor.submit") as seen: + litellm.completion( + model="gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + ) + snapshot: Final = seen.get_nowait() + assert snapshot["litellm_call_id"] + assert snapshot["response_cost"] is not None + assert snapshot["api_base"] diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index 2a60ff9c4b5..f3cd4618078 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -532,6 +532,32 @@ class TestVideoGeneration: assert abs(cost_for("runwayml/seedance2_5", "480p", 8.0) - 1.6) < 0.001 assert abs(cost_for("runwayml/gen4.5", None, 8.0) - 0.96) < 0.001 + def test_completion_cost_xai_imagine_video_720p_tier_from_cost_map(self, monkeypatch): + """720p xAI Imagine Video requests bill the published 720p rate, not the 480p base rate.""" + from litellm.cost_calculator import completion_cost + + local_map_path = os.path.join( + os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json" + ) + with open(local_map_path, "r") as f: + monkeypatch.setattr(litellm, "model_cost", json.load(f)) + + def cost_for(model: str, resolution: str, duration: float) -> float: + mock_response = MagicMock() + mock_response.usage = {"duration_seconds": duration, "video_resolution": resolution} + type(mock_response)._hidden_params = {} + return completion_cost( + completion_response=mock_response, + model=model, + call_type="create_video", + custom_llm_provider="xai", + ) + + assert abs(cost_for("xai/grok-imagine-video", "720p", 10.0) - 0.7) < 0.001 + assert abs(cost_for("xai/grok-imagine-video-1.5", "720p", 10.0) - 1.4) < 0.001 + assert abs(cost_for("xai/grok-imagine-video-1.5", "480p", 10.0) - 0.8) < 0.001 + assert abs(cost_for("xai/grok-imagine-video-1.5", "1080p", 10.0) - 2.5) < 0.001 + def test_completion_cost_veo_31_tiers_pin_published_rates(self, monkeypatch): """The gemini and vertex_ai veo 3.1 entries bill Google's published per-second tier rates.""" from litellm.cost_calculator import completion_cost diff --git a/tests/test_litellm/types/test_guardrails_case_normalization.py b/tests/test_litellm/types/test_guardrails_case_normalization.py index 3e7a573ea8e..26c1d395320 100644 --- a/tests/test_litellm/types/test_guardrails_case_normalization.py +++ b/tests/test_litellm/types/test_guardrails_case_normalization.py @@ -2,6 +2,8 @@ Test case normalization in LitellmParams for all guardrail types """ +from typing import Literal + import pytest from pydantic import ValidationError @@ -93,6 +95,34 @@ class TestLitellmParamsCaseNormalization: assert params.on_disallowed_action.islower() +class TestOnViolationAcceptedValues: + """on_violation is shared by /v1/realtime guardrails and the mcp_security guardrail""" + + @pytest.mark.parametrize("action", ["block", "alert"]) + def test_mcp_security_policy_template_on_violation_is_accepted(self, action: Literal["block", "alert"]): + params = LitellmParams( + guardrail="mcp_security", + mode="pre_call", + default_on=True, + on_violation=action, + ) + assert params.on_violation == action + + @pytest.mark.parametrize("action", ["warn", "end_session"]) + def test_realtime_on_violation_still_accepted(self, action: Literal["warn", "end_session"]): + params = LitellmParams(guardrail="presidio", mode="pre_call", on_violation=action) + assert params.on_violation == action + + @pytest.mark.parametrize("action", ["block", "alert"]) + def test_mcp_only_on_violation_is_rejected_for_other_guardrails(self, action: Literal["block", "alert"]): + with pytest.raises(ValidationError, match="only supported by guardrail='mcp_security'"): + LitellmParams(guardrail="presidio", mode="pre_call", on_violation=action) + + def test_unknown_on_violation_is_rejected(self): + with pytest.raises(ValidationError): + LitellmParams(guardrail="mcp_security", mode="pre_call", on_violation="ignore") + + class TestSensitiveDataRoutingValidation: """on_sensitive_data='route' requires a target model to be set""" diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index accd3b32a0d..7fcb76b638f 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -1,8 +1,12 @@ +import logging + import pytest +from pydantic import ValidationError from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, + GenericLiteLLMParams, LiteLLM_Params, ModelInfo, ) @@ -89,3 +93,56 @@ def test_pricing_strings_are_coerced_to_float(): def test_invalid_pricing_is_rejected(): with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", input_cost_per_token="free") + + +@pytest.mark.parametrize( + "value, expected", + [ + (True, True), + ("true", True), + (" False ", False), + ("yes", True), + (None, None), + ("os.environ/DROP_PARAMS", "os.environ/DROP_PARAMS"), + ("v2:gcm:ciphertext-from-a-pre-fix-row", "v2:gcm:ciphertext-from-a-pre-fix-row"), + ], +) +def test_drop_params_coerces_flags_and_keeps_unresolved_strings(value, expected): + assert GenericLiteLLMParams(drop_params=value).drop_params == expected + + +@pytest.mark.parametrize("value", [2, 2.5, [], {}]) +def test_drop_params_ignores_non_flag_non_string_values_with_a_warning(value, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert GenericLiteLLMParams(drop_params=value).drop_params is None + assert f"drop_params={value!r} is not a flag value" in caplog.text + + +@pytest.mark.parametrize("value", [True, "true", None, "os.environ/DROP_PARAMS", "v2:gcm:ciphertext-from-a-pre-fix-row"]) +def test_drop_params_flags_and_strings_log_nothing(value, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + GenericLiteLLMParams(drop_params=value) + assert caplog.text == "" + + +def test_aws_session_tags_round_trip_as_sts_shaped_pairs(): + """The deployment field keeps the exact Key/Value shape STS AssumeRole expects.""" + params = LiteLLM_Params( + model="bedrock/anthropic.claude-opus-5", + aws_session_tags=[{"Key": "team", "Value": "genai"}, {"Key": "env", "Value": "prod"}], + ) + + assert params.model_dump(exclude_none=True)["aws_session_tags"] == [ + {"Key": "team", "Value": "genai"}, + {"Key": "env", "Value": "prod"}, + ] + + +@pytest.mark.parametrize( + "aws_session_tags", + ["team=genai", {"team": "genai"}, [{"key": "team", "value": "genai"}], [{"Key": "team"}]], + ids=["string", "flat-dict", "lowercase-keys", "missing-value"], +) +def test_aws_session_tags_reject_shapes_sts_would_refuse(aws_session_tags): + with pytest.raises(ValidationError, match="aws_session_tags"): + LiteLLM_Params(model="bedrock/anthropic.claude-opus-5", aws_session_tags=aws_session_tags) diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index 554604ab200..5f44ba1773e 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -3,7 +3,7 @@ from typing import Final import pytest -from litellm.types.utils import HiddenParams, all_litellm_params +from litellm.types.utils import HiddenParams, all_litellm_params, text_tokens_without_nested_reasoning def test_rust_is_a_known_litellm_param(): @@ -768,3 +768,32 @@ def test_image_response_keeps_background(): response = ImageResponse(created=1, data=[{"b64_json": "aGk="}], background="transparent", output_format="png") assert response.background == "transparent" assert response.model_dump()["background"] == "transparent" + + +@pytest.mark.parametrize( + ("completion_tokens", "text_tokens", "reasoning_tokens", "other_modality_tokens", "expected_text_tokens"), + ( + pytest.param(50, 30, 20, 0, 30, id="details_sum_to_completion_is_a_no_op"), + pytest.param(34, 30, 24, 0, 10, id="strip_is_capped_at_the_over_sum"), + pytest.param(100, 100, 10, 70, 90, id="only_the_reasoning_share_is_stripped_when_text_over_reports_further"), + pytest.param(10, 5, 20, 0, 0, id="text_never_goes_negative_when_reasoning_exceeds_it"), + ), +) +def test_text_tokens_without_nested_reasoning_clamps( + completion_tokens: int, + text_tokens: int, + reasoning_tokens: int, + other_modality_tokens: int, + expected_text_tokens: int, +) -> None: + """The strip never exceeds the reasoning share, the reported text, or the over-sum past completion_tokens.""" + + assert ( + text_tokens_without_nested_reasoning( + completion_tokens=completion_tokens, + text_tokens=text_tokens, + reasoning_tokens=reasoning_tokens, + other_modality_tokens=other_modality_tokens, + ) + == expected_text_tokens + ) diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md new file mode 100644 index 00000000000..4c117fb846b --- /dev/null +++ b/tests/test_litellm_rust/README.md @@ -0,0 +1,13 @@ +# Rust OCR bridge tests + +This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR behavior tests live under `ocr/`; reusable OCR request, callback, and recording-server fixtures live under `support/` + +A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions + +`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules call the Rust bridge directly. `ocr/test_dispatch.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch. `test_ocr.py` is a strict smoke test of the compiled Rust OCR transport + +Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports + +Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture isolates callback and configuration state but does not select a backend. Native contract tests call `litellm.rust_bridge.ocr` directly, while the strict dispatch test explicitly enables and disables Rust and records which OCR entrypoint runs + +The OCR contract modules are non-strict expected failures until the retained callback implementation from #40070 lands. The public dispatch test remains strict. Passing contract cases appear as XPASS so staging coverage stays visible diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py new file mode 100644 index 00000000000..b0c75d9d2f5 --- /dev/null +++ b/tests/test_litellm_rust/conftest.py @@ -0,0 +1,132 @@ +import asyncio +import os +from collections.abc import AsyncIterator, Generator, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager +from types import ModuleType +from typing import Final, cast + +import pytest +import pytest_asyncio + +import litellm +from litellm import utils +from litellm.litellm_core_utils import litellm_logging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivateUsage] # preserve raw configuration state in test isolation + _CONFIGURATION, + _parse_env_bool, +) +from tests.test_litellm_rust.support.callback_recorder import drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service + +CALLBACK_ATTRIBUTES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", +) +EXPECTED_FAILURE_REASONS: Final = { + "ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070", + "ocr/test_guardrails.py": "requires the OCR guardrail lifecycle implementation from #40070", + "ocr/test_requests.py": "requires the OCR request and Azure authentication implementation from #40070", +} + + +def _list_attribute(container: ModuleType, attribute: str) -> list[object]: + value: Final = getattr(container, attribute) + if not isinstance(value, list): + raise AssertionError(f"{container.__name__}.{attribute} is not a list") + return cast(list[object], value) + + +@contextmanager +def _isolated_list(container: ModuleType, attribute: str) -> Iterator[None]: + source: Final = _list_attribute(container, attribute) + original: Final = list(source) + source.clear() # mutable-ok: test isolation mutates global registries by design + try: + yield + finally: + source.clear() + source.extend(original) + setattr(container, attribute, source) + + +@contextmanager +def _rebound(container: object, attribute: str, value: object) -> Iterator[None]: + original: Final[object] = getattr(container, attribute) + setattr(container, attribute, value) + try: + yield + finally: + setattr(container, attribute, original) + + +@pytest_asyncio.fixture(autouse=True, loop_scope="function") +async def isolate_ocr_test_state() -> AsyncIterator[None]: + with ExitStack() as stack: + for attribute in CALLBACK_ATTRIBUTES: + stack.enter_context(_isolated_list(litellm, attribute)) + stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # no public callback-cache accessor + stack.enter_context(_rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry + stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache + stack.enter_context(_rebound(_CONFIGURATION, "override", None)) + executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging") + stack.enter_context(_rebound(utils, "executor", executor)) + try: + yield + finally: + try: + await drain_logging() + finally: + await asyncio.to_thread(executor.shutdown, wait=True) + await GLOBAL_LOGGING_WORKER.stop() + + +@pytest.fixture +def recording_server() -> Generator[RecordingServer]: + with recording_service() as server: + yield server + + +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + for item in items: + if "test_litellm_rust" not in item.path.parts: + continue + relative_path: Final = "/".join(item.path.parts[item.path.parts.index("test_litellm_rust") + 1 :]) + reason: Final = EXPECTED_FAILURE_REASONS.get(relative_path) + if reason is not None: + item.add_marker(pytest.mark.xfail(reason=reason, strict=False)) + + if not _parse_env_bool(os.environ.get("LITELLM_RUST")): + skip: Final = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") + for item in items: + if "test_litellm_rust" in item.path.parts: + item.add_marker(skip) + return + + try: + from litellm.rust_bridge import _native # noqa: F401 # validates the installed extension + except ImportError as error: + raise pytest.UsageError("LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension") from error + + +@pytest.fixture +def isolated_azure_auth(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "AZURE_AI_API_KEY", + "AZURE_AI_API_BASE", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_USERNAME", + "AZURE_PASSWORD", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False) diff --git a/tests/test_litellm_rust/ocr/__init__.py b/tests/test_litellm_rust/ocr/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm_rust/ocr/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py new file mode 100644 index 00000000000..b08446412c0 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -0,0 +1,513 @@ +import asyncio +import copy +import queue +import threading +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.requests import ( + OCR_DOCUMENT, + OCR_RESPONSE, + call_native_aocr, + call_native_ocr, + request_body, + request_headers, +) +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +def call_native_ocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): + return call_native_ocr(server, callbacks=callbacks, **kwargs) + + +async def call_native_aocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): + return await call_native_aocr(server, callbacks=callbacks, **kwargs) + + +def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_server: RecordingServer) -> None: + observations: Final = [] + + class Observe(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + observations.append((model, copy.deepcopy(kwargs["additional_args"]))) + + call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0]) + + assert len(observations) == 1 + model, additional_args = observations[0] + assert model == "mistral-ocr-latest" + assert additional_args["api_base"] == f"{ocr_server.base_url}/v1/ocr" + assert additional_args["complete_input_dict"] == { + "model": "mistral-ocr-latest", + "document": OCR_DOCUMENT, + "pages": [0], + } + + +@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["callback-returns", "callback-raises"]) +def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider( + ocr_server: RecordingServer, raise_after_edit: bool +) -> None: + observed: Final = [] + + class Edit(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + request_body(kwargs)["include_image_base64"] = True + if raise_after_edit: + raise RuntimeError("pre-call callback failed") + + class Observe(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + observed.append(copy.deepcopy(request_body(kwargs))) + + call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False) + + assert observed[0]["include_image_base64"] is True + assert ocr_server.requests[0].body["include_image_base64"] is True + + +def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_server: RecordingServer) -> None: + observed: Final = [] + + class Edit(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + request_headers(kwargs)["x-audit-tag"] = "reviewed" + + class Observe(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + observed.append(dict(request_headers(kwargs))) + + call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()]) + + assert observed[0]["x-audit-tag"] == "reviewed" + assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + original: Final = dict(OCR_DOCUMENT) + replacement_url: Final = "data:application/pdf;base64,ZGVm" + retained: Final = [] + aliases: Final = [] + + class Retain(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + aliases.append(request_body(kwargs)["document"] is original) + retained.append(request_body(kwargs)["document"]) + + class Edit(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + original["document_url"] = replacement_url + + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": original, + "api_key": "test-key", + "api_base": ocr_server.base_url, + "callbacks": [Retain(), Edit()], + } + response: Final = ( + await call_native_aocr(ocr_server, **arguments) + if asynchronous + else call_native_ocr(ocr_server, **arguments) + ) + + assert aliases == [True] + assert retained[0]["document_url"] == replacement_url + assert original["document_url"] == replacement_url + assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url + assert response.pages[0].markdown == "native OCR response" + + +def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_document( + ocr_server: RecordingServer, +) -> None: + original: Final = dict(OCR_DOCUMENT) + replacement: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,ZGVm"} + retained: Final = [] + + class RetainAndReplace(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + body = request_body(kwargs) + retained.append(body["document"]) + body["document"] = replacement + + call_native_ocr( + ocr_server, + document=original, + callbacks=[RetainAndReplace()], + ) + + assert retained[0] is original + assert original["document_url"] == OCR_DOCUMENT["document_url"] + assert ocr_server.requests[0].body["document"] == replacement + + +def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_provider( + ocr_server: RecordingServer, +) -> None: + observed: Final = [] + + class Rebind(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + kwargs["additional_args"]["complete_input_dict"] = {"replacement": True} + + class Observe(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + observed.append(request_body(kwargs)) + + call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()]) + + assert observed == [{"replacement": True}] + assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT} + + +def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_server: RecordingServer) -> None: + queued: Final = [] + + class QueuePayload(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + queued.append(request_body(kwargs)) + + class Edit(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + request_body(kwargs)["queued-edit"] = True + + call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()]) + + assert queued[0]["queued-edit"] is True + + +def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(ocr_server: RecordingServer) -> None: + token: Final = object() + terminal_tokens: queue.SimpleQueue[object] = queue.SimpleQueue() + finished: Final = threading.Event() + + class Stash(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + kwargs["test-token"] = token + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + terminal_tokens.put(kwargs["test-token"]) + finished.set() + + call_native_ocr_with_callbacks(ocr_server, [Stash()]) + + assert finished.wait(10) + assert terminal_tokens.get_nowait() is token + + +@pytest.mark.asyncio +async def test_native_aocr_success_callback_receives_call_id_metadata_and_response( + ocr_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + await call_native_aocr_with_callbacks( + ocr_server, + [recorder], + litellm_call_id="ocr-success", + metadata={"source": "callback-test"}, + ) + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert len(events) == 1 + assert events[0].call_type == "aocr" + assert events[0].kwargs["litellm_call_id"] == "ocr-success" + assert events[0].kwargs["litellm_params"]["metadata"]["source"] == "callback-test" + assert events[0].response.pages[0].markdown == "native OCR response" + + +@pytest.mark.asyncio +async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_response( + ocr_server: RecordingServer, +) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + observations: Final = [] + + class Observe(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observations.append(("sync", kwargs["call_type"], kwargs["exception"], response_obj)) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observations.append(("async", kwargs["call_type"], kwargs["exception"], response_obj)) + + with pytest.raises(litellm.InternalServerError): + await call_native_aocr_with_callbacks(ocr_server, [Observe()]) + + assert [observation[0] for observation in observations] == ["sync", "async"] + assert all(observation[1] == "aocr" for observation in observations) + assert all(isinstance(observation[2], litellm.InternalServerError) for observation in observations) + assert all(observation[3] is None for observation in observations) + + +@pytest.mark.asyncio +async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_server: RecordingServer) -> None: + caller_loop: Final = asyncio.get_running_loop() + caller_thread: Final = threading.current_thread() + recorder: Final = RecordingLogger() + + await call_native_aocr_with_callbacks(ocr_server, [recorder]) + + events: Final = await recorder.wait_for_async("log_pre_api_call") + assert len(events) == 1 + assert events[0].loop is caller_loop + assert events[0].thread is caller_thread + + +@pytest.mark.asyncio +async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_callback( + ocr_server: RecordingServer, +) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + token: Final = object() + observed: Final = [] + + class TrackInFlightRequest(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + kwargs["request-token"] = token + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("sync", kwargs["request-token"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("async", kwargs["request-token"])) + + with pytest.raises(litellm.InternalServerError): + await call_native_aocr_with_callbacks(ocr_server, [TrackInFlightRequest()]) + + assert [event for event, _ in observed] == ["sync", "async"] + assert all(observed_token is token for _, observed_token in observed) + + +@pytest.mark.asyncio +async def test_native_aocr_callback_error_does_not_mask_provider_error_or_skip_later_failure_callbacks( + ocr_server: RecordingServer, +) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + recorder: Final = RecordingLogger() + + class FailingCallback(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("failure callback failed") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("failure callback failed") + + with pytest.raises(litellm.InternalServerError) as caught: + await call_native_aocr_with_callbacks(ocr_server, [FailingCallback(), recorder]) + + sync_events: Final = tuple(event for event in recorder.events if event.name == "log_failure_event") + async_events: Final = tuple(event for event in recorder.events if event.name == "async_log_failure_event") + assert len(sync_events) == 1 + assert len(async_events) == 1 + assert sync_events[0].kwargs["exception"] is caught.value + assert async_events[0].kwargs["exception"] is caught.value + assert "async_log_success_event" not in recorder.names + + +def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registered_multiple_times( + ocr_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + call_native_ocr_with_callbacks( + ocr_server, + [recorder, recorder], + success_callback=[recorder], + failure_callback=[recorder], + ) + recorder.wait_for("log_success_event") + + assert recorder.names.count("log_pre_api_call") == 1 + assert recorder.names.count("logging_hook") == 1 + assert recorder.names.count("log_success_event") == 1 + assert "log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, +) -> None: + from contextvars import ContextVar + context: Final = ContextVar("azure-token-context", default="missing") + context.set("caller") + caller_thread: Final = threading.current_thread() + caller_loop: Final = asyncio.get_running_loop() + observations: Final = [] + + class Provider: + def __call__(self) -> str: + assert context.get() == "caller" + assert threading.current_thread() is caller_thread + assert asyncio.get_running_loop() is caller_loop + observations.append("token") + return "caller-token" + + class Edit(CustomLogger): + def log_pre_api_call(self, model, _messages, kwargs): + assert request_headers(kwargs)["Authorization"] == "Bearer caller-token" + observations.append("pre_call") + request_headers(kwargs)["Authorization"] = "Bearer edited" + + provider: Final = Provider() + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": provider, + "callbacks": [Edit()], + } + response: Final = ( + await call_native_aocr(ocr_server, **arguments) + if asynchronous + else call_native_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR response" + assert observations == ["token", "pre_call"] + assert ocr_server.requests[0].headers["authorization"] == "Bearer edited" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, +) -> None: + ocr_server.expected_requests = 2 + calls: Final = [] + + def provider() -> str: + calls.append("token") + nested: Final = call_native_ocr(ocr_server) + assert nested.pages[0].markdown == "native OCR response" + return "outer-token" + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": provider, + } + response: Final = ( + await call_native_aocr(ocr_server, **arguments) + if asynchronous + else call_native_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR response" + assert calls == ["token"] + assert [request.headers["authorization"] for request in ocr_server.requests] == [ + "Bearer test-key", + "Bearer outer-token", + ] + + +@pytest.mark.asyncio +async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + ocr_server.expected_requests = 2 + + async def request(token: str, fail: bool) -> object: + def provider() -> str: + if fail: + raise ValueError(token) + return token + + return await call_native_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token_provider=provider, + ) + + responses: Final = await asyncio.gather( + request("first", False), + request("failed", True), + request("second", False), + return_exceptions=True, + ) + assert isinstance(responses[0], OCRResponse) + assert isinstance(responses[1], litellm.APIConnectionError) + assert "Failed to get Azure AD token: failed" in str(responses[1]) + assert isinstance(responses[2], OCRResponse) + assert sorted(request.headers["authorization"] for request in ocr_server.requests) == [ + "Bearer first", + "Bearer second", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["success", "failure", "cancellation"]) +async def test_native_azure_ocr_releases_token_provider_after_terminal_outcome( + ocr_server: RecordingServer, + isolated_azure_auth: None, + outcome: str, +) -> None: + import gc + import weakref + from tests.test_litellm_rust.support.callback_recorder import drain_logging + class Provider: + def __call__(self) -> str: + if outcome == "failure": + raise ValueError("unavailable") + return "caller-token" + + async def invoke() -> weakref.ReferenceType[Provider]: + provider: Final = Provider() + reference: Final = weakref.ref(provider) + if outcome == "failure": + ocr_server.expected_requests = 0 + with pytest.raises(litellm.APIConnectionError): + await call_native_aocr( + ocr_server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider + ) + elif outcome == "cancellation": + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1)) + task: Final = asyncio.create_task( + call_native_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token_provider=provider, + ) + ) + await ocr_server.wait_for_requests(1) + assert reference() is provider + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + response: Final = await call_native_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token_provider=provider, + ) + assert response.pages[0].markdown == "native OCR response" + return reference + + reference: Final = await invoke() + await drain_logging() + await asyncio.sleep(0) + gc.collect() + assert reference() is None diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py new file mode 100644 index 00000000000..a6c76bc5d0e --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -0,0 +1,44 @@ +from typing import Final +from unittest.mock import Mock + +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.ocr import main as ocr_main +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +@pytest.mark.parametrize("rust_enabled", [True, False], ids=["enabled", "disabled"]) +def test_public_ocr_dispatches_according_to_rust_setting( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + rust_enabled: bool, +) -> None: + rust_call: Final = Mock(wraps=ocr_main.rust_ocr_bridge.ocr) + python_call: Final = Mock(wraps=ocr_main.base_llm_http_handler.ocr) + monkeypatch.setattr(ocr_main.rust_ocr_bridge, "ocr", rust_call) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", python_call) + litellm.rust(rust_enabled) + + response: Final = litellm.ocr( + model=OCR_MODEL, + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "native OCR response" + assert rust_call.call_count == int(rust_enabled) + assert python_call.call_count == int(not rust_enabled) + assert len(ocr_server.requests) == 1 diff --git a/tests/test_litellm_rust/ocr/test_guardrails.py b/tests/test_litellm_rust/ocr/test_guardrails.py new file mode 100644 index 00000000000..f6fc1c7cb8d --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_guardrails.py @@ -0,0 +1,74 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.utils import CallTypes +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_native_aocr + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +class ReplaceOCRMarkdown(CustomGuardrail): + def __init__(self) -> None: + super().__init__( + guardrail_name="replace-ocr-markdown", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + self.call_types: list[CallTypes] = [] + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self.call_types.append(call_type) + reviewed_page: Final = response.pages[0].model_copy(update={"markdown": "Reviewed OCR"}) + return response.model_copy(update={"pages": [reviewed_page]}) + + +@pytest.mark.asyncio +async def test_native_aocr_post_call_content_filter_blocks_matching_markdown( + ocr_server: RecordingServer, +) -> None: + guardrail: Final = ContentFilterGuardrail( + guardrail_name="block-native-ocr-markdown", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="native OCR response", action=ContentFilterAction.BLOCK)], + ) + litellm.callbacks.append(guardrail) + + with pytest.raises(HTTPException, match="Content blocked") as blocked: + await call_native_aocr(ocr_server, guardrails=[guardrail.guardrail_name]) + + assert blocked.value.status_code == 400 + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +async def test_native_aocr_post_call_replacement_reaches_caller_and_success_callback( + ocr_server: RecordingServer, +) -> None: + guardrail: Final = ReplaceOCRMarkdown() + recorder: Final = RecordingLogger() + litellm.callbacks.append(guardrail) + + response: Final = await call_native_aocr( + ocr_server, + callbacks=[recorder], + guardrails=[guardrail.guardrail_name], + ) + success_events: Final = await recorder.wait_for_async("async_log_success_event") + + assert guardrail.call_types == [CallTypes.aocr] + assert response.pages[0].markdown == "Reviewed OCR" + assert len(success_events) == 1 + assert success_events[0].response.pages[0].markdown == "Reviewed OCR" + assert "guardrails" not in ocr_server.requests[0].body diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py new file mode 100644 index 00000000000..d241fe08fc8 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -0,0 +1,434 @@ +from typing import Final + +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.requests import ( + OCR_DOCUMENT, + OCR_RESPONSE, + call_native_aocr, + call_native_ocr, +) +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token( + ocr_server: RecordingServer, isolated_azure_auth: None, asynchronous: bool +) -> None: + calls: Final = [] + + def token_provider() -> str: + calls.append("token") + return "callback-token" + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": token_provider, + } + response: Final = ( + await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) + ) + + assert calls == ["token"] + assert response.pages[0].markdown == "native OCR response" + assert_native_request(ocr_server) + assert ocr_server.requests[0].headers["authorization"] == "Bearer callback-token" + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +def assert_native_request(server: RecordingServer) -> None: + assert len(server.requests) == 1 + assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx") + + +def test_native_ocr_sends_model_and_document_to_mistral_ocr_path(ocr_server: RecordingServer) -> None: + response: Final = call_native_ocr(ocr_server) + + assert response.pages[0].markdown == "native OCR response" + assert_native_request(ocr_server) + assert ocr_server.requests[0].path == "/v1/ocr" + assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT} + + +def test_native_ocr_prepares_file_document_like_python(ocr_server: RecordingServer) -> None: + response: Final = call_native_ocr( + ocr_server, + document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, + ) + + assert response.pages[0].markdown == "native OCR response" + assert_native_request(ocr_server) + assert ocr_server.requests[0].body == { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=", + }, + } + + +def test_native_ocr_sends_pages_and_image_options(ocr_server: RecordingServer) -> None: + call_native_ocr(ocr_server, pages=[0, 2], include_image_base64=True) + + assert ocr_server.requests[0].body["pages"] == [0, 2] + assert ocr_server.requests[0].body["include_image_base64"] is True + + +def test_native_ocr_merges_custom_headers_with_authorization(ocr_server: RecordingServer) -> None: + call_native_ocr(ocr_server, extra_headers={"x-trace-id": "trace-1"}) + + assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key" + assert ocr_server.requests[0].headers["x-trace-id"] == "trace-1" + + +def test_native_mistral_ocr_uses_environment_api_key_when_argument_is_missing( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + + call_native_ocr(ocr_server, api_key=None) + + assert ocr_server.requests[0].headers["authorization"] == "Bearer environment-key" + + +def test_native_mistral_ocr_prefers_explicit_api_key_over_environment( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + + call_native_ocr(ocr_server) + + assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key" + + +def test_native_azure_ocr_uses_environment_endpoint_and_api_key( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("AZURE_AI_API_KEY", "azure-key") + monkeypatch.setenv("AZURE_AI_API_BASE", ocr_server.base_url) + + call_native_ocr(ocr_server, model="azure_ai/pixtral-12b-2409", api_key=None, api_base=None) + + assert_native_request(ocr_server) + assert ocr_server.requests[0].path == "/providers/mistral/azure/ocr" + assert ocr_server.requests[0].headers["authorization"] == "Bearer azure-key" + + +def test_native_vertex_ocr_builds_path_from_project_and_location(ocr_server: RecordingServer) -> None: + call_native_ocr( + ocr_server, + model="vertex_ai/mistral-ocr-2505", + api_key="vertex-token", + vertex_project="project-1", + vertex_location="us-central1", + ) + + assert_native_request(ocr_server) + assert ocr_server.requests[0].path == ( + "/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-2505:rawPredict" + ) + + +def test_native_ocr_normalizes_provider_response_model_and_usage(ocr_server: RecordingServer) -> None: + response: Final = call_native_ocr(ocr_server) + + assert isinstance(response, OCRResponse) + assert response.model == "mistral-ocr-latest" + assert response.usage_info.pages_processed == 1 + + +def test_native_ocr_maps_provider_400_without_exposing_response_body(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400)) + + with pytest.raises(litellm.BadRequestError) as caught: + call_native_ocr(ocr_server) + + assert caught.value.status_code == 400 + assert caught.value.model == "mistral-ocr-latest" + assert caught.value.llm_provider == "mistral" + assert "invalid OCR request" not in str(caught.value) + + +def test_native_ocr_raises_transport_error_when_request_exceeds_timeout(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) + + with pytest.raises(RuntimeError, match="OCR transport failed"): + call_native_ocr(ocr_server, timeout=0.01) + + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "credentials, expected_token, expected_calls", + [ + ({"api_key": "resource-key"}, "resource-key", 0), + ({"azure_ad_token": "static-token"}, "callback-1", 1), + ({"extra_headers": {"Authorization": "Bearer override"}}, "override", 1), + ], + ids=["api-key-skips-provider", "provider-overrides-static-token", "header-overrides-provider"], +) +async def test_native_azure_ocr_applies_python_credential_precedence( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, + credentials: dict[str, object], + expected_token: str, + expected_calls: int, +) -> None: + calls: Final = [] + + def token_provider() -> str: + calls.append("token") + return f"callback-{len(calls)}" + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": token_provider, + **credentials, + } + response: Final = ( + await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR response" + assert len(calls) == expected_calls + assert len(ocr_server.requests) == 1 + assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_token}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_azure_ocr_calls_token_provider_for_each_request( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, +) -> None: + calls: Final = [] + ocr_server.expected_requests = 2 + + def token_provider() -> str: + calls.append("token") + return f"callback-{len(calls)}" + + for _ in range(2): + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": token_provider, + } + response: Final = ( + await call_native_aocr(ocr_server, **arguments) + if asynchronous + else call_native_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR response" + assert len(calls) == 2 + assert [request.headers["authorization"] for request in ocr_server.requests] == [ + "Bearer callback-1", + "Bearer callback-2", + ] + + +class TokenAbort(BaseException): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "failure", + ["non_string", "type_error", "ordinary", "abort"], + ids=["non-string-result", "type-error", "value-error", "base-exception"], +) +async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request( + ocr_server: RecordingServer, + isolated_azure_auth: None, + asynchronous: bool, + failure: str, +) -> None: + ocr_server.expected_requests = 0 + calls: Final = [] + recorder: Final = RecordingLogger() + original: Final = { + "type_error": TypeError("token type"), + "ordinary": ValueError("token unavailable"), + "abort": TokenAbort("abort"), + } + + def token_provider() -> object: + calls.append("token") + if failure == "non_string": + return 123 + raise original[failure] + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": token_provider, + "callbacks": [recorder], + } + expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError + with pytest.raises(expected) as caught: + await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) + assert calls == ["token"] + assert ocr_server.requests == [] + assert "log_pre_api_call" not in recorder.names + if failure == "ordinary": + assert "Failed to get Azure AD token: token unavailable" in str(caught.value) + assert isinstance(caught.value.__context__, RuntimeError) + assert caught.value.__context__.__cause__ is original[failure] + elif failure == "abort": + assert caught.value is original[failure] + elif failure == "type_error": + assert caught.value.__context__ is original[failure] + else: + assert isinstance(caught.value.__context__, TypeError) + + +@pytest.mark.parametrize( + "configuration", + [ + {"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}, + {"model": "azure_ai/doc-intelligence/prebuilt-read"}, + ], + ids=["oidc-assertion", "document-intelligence-model"], +) +def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_callbacks( + ocr_server: RecordingServer, + isolated_azure_auth: None, + configuration: dict[str, object], +) -> None: + ocr_server.expected_requests = 0 + calls: Final = [] + recorder: Final = RecordingLogger() + + def provider() -> str: + calls.append("token") + return "unused" + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": provider, + "callbacks": [recorder], + **configuration, + } + with pytest.raises(NotImplementedError): + call_native_ocr(ocr_server, **arguments) + assert calls == [] + assert recorder.events == () + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +async def test_native_azure_ocr_validates_endpoint_before_calling_token_provider( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + ocr_server.expected_requests = 0 + calls: Final = [] + + def provider() -> str: + calls.append("token") + return "unused" + + with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI API Base"): + await call_native_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + api_base=None, + azure_ad_token_provider=provider, + ) + assert calls == [] + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +async def test_native_azure_ocr_does_not_fall_back_to_static_token_after_empty_provider_result( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + ocr_server.expected_requests = 0 + + def provider() -> str: + return "" + + with pytest.raises(litellm.APIConnectionError, match="Missing Azure AI credentials"): + await call_native_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token="static-token", + azure_ad_token_provider=provider, + ) + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +async def test_native_azure_ocr_ignores_falsey_token_provider_and_uses_static_token( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + calls: Final = [] + + class Provider: + def __bool__(self) -> bool: + return False + + def __call__(self) -> str: + calls.append("token") + return "unused" + + response: Final = await call_native_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + api_key=None, + azure_ad_token="static-token", + azure_ad_token_provider=Provider(), + ) + assert response.pages[0].markdown == "native OCR response" + assert calls == [] + assert ocr_server.requests[0].headers["authorization"] == "Bearer static-token" + + +@pytest.mark.asyncio +async def test_native_azure_ocr_rejects_coroutine_returned_by_sync_token_provider( + ocr_server: RecordingServer, + isolated_azure_auth: None, +) -> None: + ocr_server.expected_requests = 0 + calls: Final = [] + + async def acquire() -> str: + calls.append("awaited") + return "unused" + + coroutine: Final = acquire() + + def provider() -> object: + return coroutine + + try: + with pytest.raises(litellm.APIConnectionError, match="Azure AD token must be a string"): + await call_native_aocr( + ocr_server, model="azure_ai/mistral-ocr-latest", api_key=None, azure_ad_token_provider=provider + ) + finally: + coroutine.close() + assert calls == [] + assert ocr_server.requests == [] diff --git a/tests/test_litellm_rust/support/__init__.py b/tests/test_litellm_rust/support/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm_rust/support/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm_rust/support/callback_recorder.py b/tests/test_litellm_rust/support/callback_recorder.py new file mode 100644 index 00000000000..6de011b1414 --- /dev/null +++ b/tests/test_litellm_rust/support/callback_recorder.py @@ -0,0 +1,101 @@ +import asyncio +import copy +import threading +import time +from dataclasses import dataclass +from typing import Final + +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker + + +async def drain_logging(worker: LoggingWorker = GLOBAL_LOGGING_WORKER) -> None: + await asyncio.sleep(0) + worker.start() + await asyncio.wait_for(worker.flush(), timeout=10) + + +@dataclass(frozen=True, slots=True) +class HookEvent: + name: str + call_type: str | None + thread: threading.Thread + loop: asyncio.AbstractEventLoop | None + kwargs: object + response: object + + +class RecordingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self._events: list[HookEvent] = [] + self._condition = threading.Condition() + + @property + def events(self) -> tuple[HookEvent, ...]: + with self._condition: + return tuple(self._events) + + @property + def names(self) -> tuple[str, ...]: + return tuple(event.name for event in self.events) + + def _record(self, name: str, kwargs: object = None, response: object = None) -> None: + details: Final = kwargs if isinstance(kwargs, dict) else {} + try: + snapshot: Final = copy.deepcopy(details) + except Exception: + snapshot = dict(details) + if "exception" in details: + snapshot["exception"] = details["exception"] + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + loop = None + event: Final = HookEvent( + name=name, + call_type=details.get("call_type"), + thread=threading.current_thread(), + loop=loop, + kwargs=snapshot, + response=response, + ) + with self._condition: + self._events.append(event) + self._condition.notify_all() + + def wait_for(self, name: str, count: int = 1, timeout: float = 10) -> tuple[HookEvent, ...]: + deadline: Final = time.monotonic() + timeout + with self._condition: + while sum(event.name == name for event in self._events) < count: + remaining: Final = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Timed out waiting for {count} {name} events; saw {self.names}") + self._condition.wait(remaining) + return tuple(event for event in self._events if event.name == name) + + async def wait_for_async(self, name: str, count: int = 1, timeout: float = 10) -> tuple[HookEvent, ...]: + await asyncio.wait_for(asyncio.to_thread(self.wait_for, name, count, timeout), timeout=timeout + 1) + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=timeout) + return tuple(event for event in self.events if event.name == name) + + def log_pre_api_call(self, model, _messages, kwargs): + self._record("log_pre_api_call", kwargs) + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record("log_success_event", kwargs, response_obj) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record("async_log_success_event", kwargs, response_obj) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._record("log_failure_event", kwargs, response_obj) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._record("async_log_failure_event", kwargs, response_obj) + + def logging_hook(self, kwargs, result, call_type): + self._record("logging_hook", kwargs, result) + return kwargs, result diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py new file mode 100644 index 00000000000..5a9b9497c6e --- /dev/null +++ b/tests/test_litellm_rust/support/recording_server.py @@ -0,0 +1,108 @@ +import asyncio +import copy +import json +import threading +import time +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final + + +@dataclass +class RecordedRequest: + method: str + path: str + headers: dict[str, str] + raw_body: bytes + body: object | None + + +@dataclass +class ResponseSpec: + body: object + status: int = 200 + headers: dict[str, str] = field(default_factory=dict) + delay: float = 0 + + +@dataclass +class RecordingServer: + server: ThreadingHTTPServer + requests: list[RecordedRequest] + responses: list[ResponseSpec] + default_response: ResponseSpec + expected_requests: int | None = 1 + + @property + def base_url(self) -> str: + host, port = self.server.server_address + return f"http://{host}:{port}" + + def enqueue(self, response: ResponseSpec) -> None: + self.responses.append(response) + + async def wait_for_requests(self, count: int) -> None: + async with asyncio.timeout(2): + while len(self.requests) < count: + await asyncio.sleep(0.01) + + +@contextmanager +def recording_service() -> Iterator[RecordingServer]: + requests: list[RecordedRequest] = [] + responses: list[ResponseSpec] = [] + + class Handler(BaseHTTPRequestHandler): + def _handle(self) -> None: + content_length: Final = int(self.headers.get("Content-Length", "0")) + raw_body: Final = self.rfile.read(content_length) if content_length else b"" + body: Final = json.loads(raw_body) if raw_body else None + requests.append( + RecordedRequest( + method=self.command, + path=self.path, + headers={name.lower(): value for name, value in self.headers.items()}, + raw_body=raw_body, + body=body, + ) + ) + response: Final = responses.pop(0) if responses else copy.deepcopy(recording_server.default_response) + if response.delay: + time.sleep(response.delay) + payload: Final = json.dumps(response.body).encode() + self.send_response(response.status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + for name, value in response.headers.items(): + self.send_header(name, value) + self.end_headers() + try: + self.wfile.write(payload) + except (BrokenPipeError, ConnectionResetError): + pass + + do_POST = _handle + + def log_message(self, format: str, *args: object) -> None: + pass + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + try: + recording_server = RecordingServer( + server=server, + requests=requests, + responses=responses, + default_response=ResponseSpec(body={}), + ) + yield recording_server + finally: + server.shutdown() + server.server_close() + thread.join() + if recording_server.expected_requests is not None: + assert len(recording_server.requests) == recording_server.expected_requests + assert recording_server.responses == [] diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py new file mode 100644 index 00000000000..d681d752ffd --- /dev/null +++ b/tests/test_litellm_rust/support/requests.py @@ -0,0 +1,59 @@ +from typing import Final + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge import ocr as native_ocr +from tests.test_litellm_rust.support.recording_server import RecordingServer + +OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} +OCR_MODEL: Final = "mistral/mistral-ocr-latest" +OCR_RESPONSE: Final = { + "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, +} + + +def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: + return { + "model": OCR_MODEL, + "document": dict(OCR_DOCUMENT), + "api_key": "test-key", + "api_base": server.base_url, + **kwargs, + } + + +def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: + response: Final = litellm.ocr(**ocr_arguments(server, **kwargs)) + if not isinstance(response, OCRResponse): + raise TypeError(f"Expected OCRResponse, got {type(response).__name__}") + return response + + +async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: + return await litellm.aocr(**ocr_arguments(server, **kwargs)) + + +def call_native_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: + return native_ocr.ocr(ocr_arguments(server, **kwargs)) + + +async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: + return await native_ocr.aocr(ocr_arguments(server, **kwargs)) + + +def request_body(kwargs: dict[str, object]) -> dict[str, object]: + additional_args = kwargs["additional_args"] + assert isinstance(additional_args, dict) + body = additional_args["complete_input_dict"] + assert isinstance(body, dict) + return body + + +def request_headers(kwargs: dict[str, object]) -> dict[str, object]: + additional_args = kwargs["additional_args"] + assert isinstance(additional_args, dict) + headers = additional_args["headers"] + assert isinstance(headers, dict) + return headers diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py new file mode 100644 index 00000000000..ad1c8c652bb --- /dev/null +++ b/tests/test_litellm_rust/test_ocr.py @@ -0,0 +1,234 @@ +import json +import threading +from collections.abc import Generator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge import ocr as rust_ocr_bridge + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[dict[str, object]]]]: + requests: Final[list[dict[str, object]]] = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + requests.append( + { + "headers": {name.lower(): value for name, value in self.headers.items()}, + "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))), + } + ) + if self.headers.get("x-test-stall") == "true": + self.connection.settimeout(2) + try: + self.rfile.read(1) + except TimeoutError: + pass + return + if self.headers.get("User-Agent", "").startswith("python-httpx"): + self.send_response(418) + self.end_headers() + return + status = int(self.headers.get("x-test-status", "200")) + if status != 200: + body = b'{"error":"provider unavailable"}' + self.send_response(status) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + return + response: Final = json.dumps( + { + "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, + } + ).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(response))) + self.end_headers() + self.wfile.write(response) + + def log_message(self, format: str, *args: object) -> None: + pass + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + try: + yield server, requests + finally: + server.shutdown() + server.server_close() + thread.join() + + +def test_native_ocr_with_compiled_rust_extension( + ocr_server: tuple[ThreadingHTTPServer, list[dict[str, object]]], +) -> None: + server, requests = ocr_server + address: Final = server.server_address + host: Final = str(address[0]) + port: Final = int(address[1]) + + response: Final = rust_ocr_bridge.ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + api_key="test-key", + api_base=f"http://{host}:{port}", + custom_llm_provider="mistral", + extra_headers=None, + optional_params={}, + timeout=None, + ) + + assert response is not None + assert response["pages"][0]["markdown"] == "native OCR response" + assert len(requests) == 1 + assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") + assert requests[0]["body"] == { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + } + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/doc-intelligence/prebuilt-read"]) +@pytest.mark.asyncio +async def test_native_public_ocr_matches_python(model, asynchronous): + import json + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + from threading import Thread + from typing import Final + from urllib.parse import parse_qsl, urlsplit + + from litellm.rust_bridge import _native + + assert callable(_native.ocr) + calls: Final = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + body: Final = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + target: Final = urlsplit(self.path) + calls.append( + ( + target.path, + parse_qsl(target.query), + self.headers.get("Authorization"), + self.headers.get("Ocp-Apim-Subscription-Key"), + body, + ) + ) + payload: Final = ( + {"status": "succeeded", "analyzeResult": {"pages": []}} + if "doc-intelligence" in model + else {"pages": [{"index": 0, "markdown": "hello"}]} + ) + encoded: Final = json.dumps(payload).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, *_args): + pass + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + responses: Final = [] + try: + for enabled in (False, True): + litellm.rust(enabled) + arguments: Final = { + "model": model, + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{server.server_port}", + "pages": [0, 2], + "timeout": 3.0, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + responses.append(response.model_dump()) + assert len(calls) == 2 + assert calls[0] == calls[1] + for key in ("model", "pages", "object"): + assert responses[0][key] == responses[1][key] + finally: + server.shutdown() + server.server_close() + thread.join(timeout=3) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_native_ocr_failures_do_not_retry_on_python(ocr_server, asynchronous): + server, requests = ocr_server + arguments = { + "model": "mistral-ocr-latest", + "custom_llm_provider": "mistral", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{server.server_port}", + "extra_headers": {"x-test-status": "503"}, + "num_retries": 0, + } + litellm.rust(True) + with pytest.raises(litellm.ServiceUnavailableError) as caught: + await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert caught.value.status_code == 503 + assert len(requests) == 1 + assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") + + +@pytest.mark.parametrize("custom_provider", ["mistral", "not-a-provider"]) +def test_native_ocr_rejects_invalid_input_before_network(ocr_server, custom_provider): + from litellm.rust_bridge import _native + + server, requests = ocr_server + with pytest.raises(ValueError, match=r"invalid (OCR request field|provider)|invalid request"): + _native.ocr( + model="mistral-ocr-latest", + custom_llm_provider=custom_provider, + document={"type": "document_url"}, + api_key="test-key", + api_base=f"http://127.0.0.1:{server.server_port}", + ) + assert requests == [] + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_native_ocr_enforces_request_deadline_without_fallback(ocr_server, asynchronous): + import asyncio + import time + + server, requests = ocr_server + litellm.rust(True) + arguments = { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{server.server_port}", + "extra_headers": {"x-test-stall": "true"}, + "timeout": 0.1, + "num_retries": 0, + } + started = time.monotonic() + with pytest.raises(litellm.APIConnectionError): + await asyncio.wait_for( + litellm.aocr(**arguments) if asynchronous else asyncio.to_thread(litellm.ocr, **arguments), + timeout=3, + ) + assert 0.09 <= time.monotonic() - started < 3 + assert len(requests) == 1 + assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e7186dfe186..0c0952289e2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22180 + "limit": 22174 }, "LIT002": { - "limit": 26729 + "limit": 26715 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16426 + "limit": 16398 }, "LIT011": { - "limit": 5506 + "limit": 5504 }, "LIT012": { "limit": 4486 diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index e98cea9261e..44294b5fa97 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -3,8 +3,8 @@ "no-console": { "max": 12, "target": 0 }, "complexity": { "max": 140, "target": 80 }, "max-depth": { "max": 70, "target": 30 }, - "local/no-large-inline-object-arg": { "max": 554, "target": 300 }, - "local/no-long-condition-chain": { "max": 265, "target": 120 }, + "local/no-large-inline-object-arg": { "max": 551, "target": 300 }, + "local/no-long-condition-chain": { "max": 196, "target": 120 }, "testing-library/no-container": { "max": 133, "target": 50 }, "testing-library/no-node-access": { "max": 707, "target": 500 }, "testing-library/prefer-screen-queries": { "max": 18, "target": 18 } diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index d7475173b90..e2a08a40bcb 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1619,14 +1619,6 @@ "count": 1 } }, - "src/components/common_components/fetch_teams.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "max-params": { - "count": 1 - } - }, "src/components/common_components/simple_table.tsx": { "local/filename-pascal-case": { "count": 1 @@ -1823,7 +1815,7 @@ "count": 5 }, "no-restricted-syntax": { - "count": 152 + "count": 150 }, "prefer-const": { "count": 32 @@ -1871,9 +1863,6 @@ "src/components/per_user_usage.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 } }, "src/components/permissions/MCPServerPermissions.tsx": { @@ -2303,17 +2292,6 @@ "count": 1 } }, - "src/components/user_dashboard.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "prefer-const": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/vector_store_management/types.tsx": { "local/filename-pascal-case": { "count": 1 @@ -2347,11 +2325,6 @@ "count": 4 } }, - "src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx": { - "no-nested-ternary": { - "count": 3 - } - }, "src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx": { "no-nested-ternary": { "count": 2 diff --git a/ui/litellm-dashboard/next.config.mjs b/ui/litellm-dashboard/next.config.mjs index 876df2b49cf..128ce0a84a7 100644 --- a/ui/litellm-dashboard/next.config.mjs +++ b/ui/litellm-dashboard/next.config.mjs @@ -7,6 +7,9 @@ const __dirname = path.dirname(__filename); const nextConfig = { output: "export", + experimental: { + useTypeScriptCli: false, + }, compiler: { removeConsole: process.env.NODE_ENV === "production" ? { exclude: ["error", "warn"] } : false, }, diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 4e5b0c1dda6..1f459bc50ea 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -24,7 +24,7 @@ "jwt-decode": "4.0.0", "lucide-react": "0.513.0", "moment": "2.30.1", - "next": "16.2.11", + "next": "16.3.3", "next-themes": "^0.4.6", "nuqs": "^2.9.4", "openai": "4.104.0", @@ -58,10 +58,10 @@ "@types/react-copy-to-clipboard": "5.0.7", "@types/react-dom": "19.2.4", "@types/react-syntax-highlighter": "15.5.13", - "@vitest/coverage-v8": "3.2.6", - "@vitest/ui": "3.2.6", + "@vitest/coverage-v8": "4.1.11", + "@vitest/ui": "4.1.11", "eslint": "9.39.2", - "eslint-config-next": "16.2.11", + "eslint-config-next": "16.3.3", "eslint-config-prettier": "10.1.8", "eslint-plugin-jest-dom": "5.10.1", "eslint-plugin-testing-library": "7.16.2", @@ -75,7 +75,8 @@ "tw-animate-css": "1.4.0", "typescript": "5.9.3", "typescript-eslint": "8.60.1", - "vitest": "3.2.6" + "vite": "7.3.5", + "vitest": "4.1.11" }, "engines": { "node": ">=24.14.1", @@ -109,20 +110,6 @@ "url": "https://github.com/sponsors/sindresorhus" } }, - "node_modules/@ampproject/remapping": { - "version": "2.3.0", - "resolved": "https://registry.npmjs.org/@ampproject/remapping/-/remapping-2.3.0.tgz", - "integrity": "sha512-30iZtAPgz+LTIYoeivqYo853f02jBYSd5uGnGpkFV0M3xOt9aN73erkgYAmZU43x4VfqcnLxW9Kpg3R5LC4YYw==", - "dev": true, - "license": "Apache-2.0", - "dependencies": { - "@jridgewell/gen-mapping": "^0.3.5", - "@jridgewell/trace-mapping": "^0.3.24" - }, - "engines": { - "node": ">=6.0.0" - } - }, "node_modules/@anthropic-ai/sdk": { "version": "0.92.0", "resolved": "https://registry.npmjs.org/@anthropic-ai/sdk/-/sdk-0.92.0.tgz", @@ -693,9 +680,9 @@ } }, "node_modules/@emnapi/runtime": { - "version": "1.11.2", - "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.11.2.tgz", - "integrity": "sha512-kyOl3X0DuTiT1h2ft8r2fYO8JYtU9a9Xis/zBSiGArNaagCOWx90N1k2wxp18czFDH+OgcWGb5ZP/XMt3dcyPA==", + "version": "1.11.3", + "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.11.3.tgz", + "integrity": "sha512-Xz4Tpyki7XyrpbUK1jR1AhdAdaXyhhY4lZ3neLodmhpuWfy2PAQN5B46sAiU4liOXGLkHypn/qU+jvfWSCYYLA==", "license": "MIT", "optional": true, "dependencies": { @@ -1465,9 +1452,9 @@ } }, "node_modules/@img/sharp-darwin-arm64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-darwin-arm64/-/sharp-darwin-arm64-0.35.3.tgz", - "integrity": "sha512-RMnFX7YQsMoh7lWfcM4NEHHymBX/rLuKNPVM84XE9ONPcaSCDgE7CHIHpSgPcO2xcRthgBy1HfNO319mwhIAkg==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-darwin-arm64/-/sharp-darwin-arm64-0.35.4.tgz", + "integrity": "sha512-Uhfl4V4lhP2nbUVF9+hyH1+luj86f1gUFeo8ALYxFoULoU+G87D43BfeMP8XHsk9boxAnCY/bf2EHwhA7MuGsA==", "cpu": [ "arm64" ], @@ -1483,13 +1470,13 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-darwin-arm64": "1.3.2" + "@img/sharp-libvips-darwin-arm64": "1.3.3" } }, "node_modules/@img/sharp-darwin-x64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-darwin-x64/-/sharp-darwin-x64-0.35.3.tgz", - "integrity": "sha512-Xo+5uFBtLN0BKqieTxiFzFPQAUlBbbH5iBKyRX/z1JrbnYsHTfKJnUfL8+p2TPXr1pXqao4eeL4Rl144uDpK9w==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-darwin-x64/-/sharp-darwin-x64-0.35.4.tgz", + "integrity": "sha512-hWniXY3bG5qKpkKrAwPe4y+VTPmf086YQAnkxWh7uA1YrlRouWGa0M0Mxj3ZjnXFkv7/TD1bTy9lGUK26vRvWw==", "cpu": [ "x64" ], @@ -1505,20 +1492,20 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-darwin-x64": "1.3.2" + "@img/sharp-libvips-darwin-x64": "1.3.3" } }, "node_modules/@img/sharp-freebsd-wasm32": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-freebsd-wasm32/-/sharp-freebsd-wasm32-0.35.3.tgz", - "integrity": "sha512-lUxcqWIj2wMQ9BrwNjngcr1gWUr5xgaGThBRqPPalIC2n67Cqj1uPh8NnA/ZhAg8hUbKl+kVHKwgUIwe6ZYPrg==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-freebsd-wasm32/-/sharp-freebsd-wasm32-0.35.4.tgz", + "integrity": "sha512-lIsKw/BU+kjB4eZjxrYrZmwOJYi3Ajrv66iAlBmUPyKc3HpnloevB1g3wxGD9P/5BbQ1brBGl65VRRrCvQDEqA==", "license": "Apache-2.0", "optional": true, "os": [ "freebsd" ], "dependencies": { - "@img/sharp-wasm32": "0.35.3" + "@img/sharp-wasm32": "0.35.4" }, "engines": { "node": ">=20.9.0" @@ -1528,9 +1515,9 @@ } }, "node_modules/@img/sharp-libvips-darwin-arm64": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-darwin-arm64/-/sharp-libvips-darwin-arm64-1.3.2.tgz", - "integrity": "sha512-9J6ypZFpQBj4YnePGoq/S38w6nz+vqg5WZLrLGY4YuSemdMq47GMLBPO42MzwdGwpg/agZ7xzZcFHa48xlywfg==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-darwin-arm64/-/sharp-libvips-darwin-arm64-1.3.3.tgz", + "integrity": "sha512-suTBPTDGrI9WodccaDdwZItTSaBYASlBk1NSfElSHrUfzu3szG6lvIF58+WiFvnfzuK8ZBFS5zE00PxqxnRiPg==", "cpu": [ "arm64" ], @@ -1544,9 +1531,9 @@ } }, "node_modules/@img/sharp-libvips-darwin-x64": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-darwin-x64/-/sharp-libvips-darwin-x64-1.3.2.tgz", - "integrity": "sha512-m2pW1n6cns9VaubNwsZ+c3CRYjxNQWgJ5gPlnL1nbBcpkBvFm6SCFN5o0psFHI8w9n11NKhFkeEDns98tiqbEw==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-darwin-x64/-/sharp-libvips-darwin-x64-1.3.3.tgz", + "integrity": "sha512-FVJZ5mITMobmXIz/hPDTw0EintTW5H3WfrxwLqEqjiIihlu+hVRyGrFQ60xl0Lxn7Bt3zdpevPaQi0HEzqz9fw==", "cpu": [ "x64" ], @@ -1560,12 +1547,15 @@ } }, "node_modules/@img/sharp-libvips-linux-arm": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-arm/-/sharp-libvips-linux-arm-1.3.2.tgz", - "integrity": "sha512-1eMLzy92I4J6rmi4mAT8yC3HxOtniyGELlzGbNMLLeqe052ahFQ0h6LFq+lh5DsDIdYViIDst08abvSbcEdLXQ==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-arm/-/sharp-libvips-linux-arm-1.3.3.tgz", + "integrity": "sha512-3rbU4vqXXc3hY/OiXdl52xZvT0F1yEngWfvqudtPJg/KkyiaQw2DRsFrNzpmLvfavbwOq3qXn36GP8obHRULQA==", "cpu": [ "arm" ], + "libc": [ + "glibc" + ], "license": "LGPL-3.0-or-later", "optional": true, "os": [ @@ -1576,12 +1566,15 @@ } }, "node_modules/@img/sharp-libvips-linux-arm64": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-arm64/-/sharp-libvips-linux-arm64-1.3.2.tgz", - "integrity": "sha512-dqVSFynCox4C/J8kT16V7SIFAns0IjgLwkvYT7p8LQVmJ5OS5b6tI9IGflxTeuBS//zXeFIUbwt5dwxyZ17cnA==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-arm64/-/sharp-libvips-linux-arm64-1.3.3.tgz", + "integrity": "sha512-0DaL0A6Xu6sQSQFwe4iVCrKWU2cCTItnRsYsCdxAMm9NF6twAA9BKnoqy4hqz4+azQ0JHuA26qiUKsf1XJ/v5A==", "cpu": [ "arm64" ], + "libc": [ + "glibc" + ], "license": "LGPL-3.0-or-later", "optional": true, "os": [ @@ -1592,12 +1585,15 @@ } }, "node_modules/@img/sharp-libvips-linux-ppc64": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-ppc64/-/sharp-libvips-linux-ppc64-1.3.2.tgz", - "integrity": "sha512-3z0NHDxD6n5I9gc05U1eW1AyRm+Gznzq3naMrthPNqE6oYykcogW0l/jfpJdjYnuNl8R7yI9pNbE1XiUeyq0Aw==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-ppc64/-/sharp-libvips-linux-ppc64-1.3.3.tgz", + "integrity": "sha512-cdn1OvUBwsXhbC0zSzJnNzf5MZ/mTrobawDvNXBTxe8VtqKAm0sRuEY2Evzovb/w9JMk4TvRxqt1mekSuJz64w==", "cpu": [ "ppc64" ], + "libc": [ + "glibc" + ], "license": "LGPL-3.0-or-later", "optional": true, "os": [ @@ -1608,12 +1604,15 @@ } }, "node_modules/@img/sharp-libvips-linux-riscv64": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-riscv64/-/sharp-libvips-linux-riscv64-1.3.2.tgz", - "integrity": "sha512-bsb4rI+NldGOsXuej2r8OdSS8+zXDVaCWxyWrcv6kneTOlgAHtZABRzBBCwdsPiD90J4myNJuHpg6kA20ImW/w==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-riscv64/-/sharp-libvips-linux-riscv64-1.3.3.tgz", + "integrity": "sha512-HjPVx7yKz+0lqdhDlTw1tt90wamBoxhiXpvl1XZpJLiHH4RCJ5yDTqH+VlYPv2fwFs89JFw4c1IexYOcQUi4IQ==", "cpu": [ "riscv64" ], + "libc": [ + "glibc" + ], "license": "LGPL-3.0-or-later", "optional": true, "os": [ @@ -1624,12 +1623,15 @@ } }, "node_modules/@img/sharp-libvips-linux-s390x": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-s390x/-/sharp-libvips-linux-s390x-1.3.2.tgz", - "integrity": "sha512-/ABshyj8gCpyIrNXnHn4LorDJ0HHm1VhXPBlxZ8zAtfVPAaSafXPGn+sUSIRiwaSBy0mmFjSjiXI5mkcwdChKQ==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-s390x/-/sharp-libvips-linux-s390x-1.3.3.tgz", + "integrity": "sha512-neWLh+3yCNThxnfy3c4BbVBeGgt9aftno+XbT56iK28RgeDs3UOFWviLWlUu0bArYVYJaFDK+RRohbicUNCm8Q==", "cpu": [ "s390x" ], + "libc": [ + "glibc" + ], "license": "LGPL-3.0-or-later", "optional": true, "os": [ @@ -1640,12 +1642,15 @@ } }, "node_modules/@img/sharp-libvips-linux-x64": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-x64/-/sharp-libvips-linux-x64-1.3.2.tgz", - "integrity": "sha512-ITPEtgffGJ0S6G9dRyw/366tJQqFRcHWPHhC+Stpg3Z8AEMrDrTr2lhdz4f/Y/HMbRh//7Z5mBzEpVdi62Oc3w==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-x64/-/sharp-libvips-linux-x64-1.3.3.tgz", + "integrity": "sha512-4vKmvAst9nrowcqquKFAyZJUDolUaIp8uRiN0mWFguJ1IplC9/pitXtlnnlU4aa/eJw3J7i67V+pwUL+wZGdsA==", "cpu": [ "x64" ], + "libc": [ + "glibc" + ], "license": "LGPL-3.0-or-later", "optional": true, "os": [ @@ -1656,12 +1661,15 @@ } }, "node_modules/@img/sharp-libvips-linuxmusl-arm64": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linuxmusl-arm64/-/sharp-libvips-linuxmusl-arm64-1.3.2.tgz", - "integrity": "sha512-zE9EdiUzUmg5mDT5a1rk5fYJ6GWPloTwWBYDS14naqHsL+EaMpDj1AWnpLgh3u0YCORv2Tt50wrcrpYqkP97Kw==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linuxmusl-arm64/-/sharp-libvips-linuxmusl-arm64-1.3.3.tgz", + "integrity": "sha512-Y9kQaLMuNoB0bPYOOdcZMaseNrFpPodIWWMrx+CZyydf2xn68j9WYc6sWWRrDwNkzCQjKYfc68L7jKjGlHMibw==", "cpu": [ "arm64" ], + "libc": [ + "musl" + ], "license": "LGPL-3.0-or-later", "optional": true, "os": [ @@ -1672,12 +1680,15 @@ } }, "node_modules/@img/sharp-libvips-linuxmusl-x64": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linuxmusl-x64/-/sharp-libvips-linuxmusl-x64-1.3.2.tgz", - "integrity": "sha512-m0lrLiUt+lBYnCFr8qV/65yMR4E/c7/wf78I5eKTdkEakFAlZ9QlzEM3QIhhAwVeUhLAHLcCq7a7Vszq/oFNZQ==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linuxmusl-x64/-/sharp-libvips-linuxmusl-x64-1.3.3.tgz", + "integrity": "sha512-fj8Mv0HHfD1Rr+4I68+3agJynxDWtBFgicTbSOb9Bke6pIwzGcJ+RX/yHjmiEGFMCavY/dxvem7MyNaJF+wDiw==", "cpu": [ "x64" ], + "libc": [ + "musl" + ], "license": "LGPL-3.0-or-later", "optional": true, "os": [ @@ -1688,12 +1699,15 @@ } }, "node_modules/@img/sharp-linux-arm": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-arm/-/sharp-linux-arm-0.35.3.tgz", - "integrity": "sha512-affVWCTLooy8TSxbDx2qkzuDeaWLNVBA+P//FNBirHsXpP2fuBhk5AuboYUnrDnzoXes8GFjpTx0SBFOCRg+FA==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-arm/-/sharp-linux-arm-0.35.4.tgz", + "integrity": "sha512-7OAS8gI0EReKGVN2HssHlM6umJgxF5VI3xN0p9FA91p/YO+ou5hiNghLdZ5BEHztwaaK5+bLKRf8x/o2L2nk9A==", "cpu": [ "arm" ], + "libc": [ + "glibc" + ], "license": "Apache-2.0", "optional": true, "os": [ @@ -1706,16 +1720,19 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-linux-arm": "1.3.2" + "@img/sharp-libvips-linux-arm": "1.3.3" } }, "node_modules/@img/sharp-linux-arm64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-arm64/-/sharp-linux-arm64-0.35.3.tgz", - "integrity": "sha512-QgKDspHPnrU+GQ55XPhGwyhC8acLVOOSyAvo1oVfFmrIXLkDNmGWzAfDZ4xK8oSA1qBQrALcHX0G5UZni/SuFQ==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-arm64/-/sharp-linux-arm64-0.35.4.tgz", + "integrity": "sha512-De4jpEnAU8Hd5oT0j1G3uL4ZvTuipVMn7YC6vPaJhy6/7EwEae0SVAoBrUMYQbkLGDm85taVWwuPc1a44LTzCQ==", "cpu": [ "arm64" ], + "libc": [ + "glibc" + ], "license": "Apache-2.0", "optional": true, "os": [ @@ -1728,16 +1745,19 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-linux-arm64": "1.3.2" + "@img/sharp-libvips-linux-arm64": "1.3.3" } }, "node_modules/@img/sharp-linux-ppc64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-ppc64/-/sharp-linux-ppc64-0.35.3.tgz", - "integrity": "sha512-sMd8rDxmpLOwv/7N44klFjOD5DUO7FLdjiXDI0hoxYaf7Ar262dQIEkosE98bps+5HPLtp/EvNqeqQtOycP/IA==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-ppc64/-/sharp-linux-ppc64-0.35.4.tgz", + "integrity": "sha512-2oYZJeIl4kCcMGk4ouZVjnkCtFrpQFlNEtJ6GbxzhHQchwH0NH/qEb9ykmOl29dqwMq+JhFdZn+1ak2FKhI9fQ==", "cpu": [ "ppc64" ], + "libc": [ + "glibc" + ], "license": "Apache-2.0", "optional": true, "os": [ @@ -1750,16 +1770,19 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-linux-ppc64": "1.3.2" + "@img/sharp-libvips-linux-ppc64": "1.3.3" } }, "node_modules/@img/sharp-linux-riscv64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-riscv64/-/sharp-linux-riscv64-0.35.3.tgz", - "integrity": "sha512-0Eob78yjlYPfL5vMNWAW55l3R9Y6BQS/gOfe0ZcP9mEz9ohhKSt4im1hayiknXgf8AWrFqMvJcKIdmLmEe7yeQ==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-riscv64/-/sharp-linux-riscv64-0.35.4.tgz", + "integrity": "sha512-cPbNChoRURAWdebDIHSenxRpgEdy7JkPydSnUxRm9VvKD7m0/xVaR/8Fzlu81pk5nHEvHH87UZUA7cTtwnbJSA==", "cpu": [ "riscv64" ], + "libc": [ + "glibc" + ], "license": "Apache-2.0", "optional": true, "os": [ @@ -1772,16 +1795,19 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-linux-riscv64": "1.3.2" + "@img/sharp-libvips-linux-riscv64": "1.3.3" } }, "node_modules/@img/sharp-linux-s390x": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-s390x/-/sharp-linux-s390x-0.35.3.tgz", - "integrity": "sha512-KgAxQ0DxpNOq1rG2t5cgTgShJFGSuU7XO45cqC+1NVOuZnP6tlgZRuSYOfNupGkHID0o3cJOsw4DVeJpMovcGw==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-s390x/-/sharp-linux-s390x-0.35.4.tgz", + "integrity": "sha512-RY0JFY8Fd6RonCBtHz+DvadaPkXDSI1AUn6yWL9TipqkZ1vY8w8evqdgyDFnkm4/K1ve1TvZiaePP5oSd4+WVQ==", "cpu": [ "s390x" ], + "libc": [ + "glibc" + ], "license": "Apache-2.0", "optional": true, "os": [ @@ -1794,16 +1820,19 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-linux-s390x": "1.3.2" + "@img/sharp-libvips-linux-s390x": "1.3.3" } }, "node_modules/@img/sharp-linux-x64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-linux-x64/-/sharp-linux-x64-0.35.3.tgz", - "integrity": "sha512-8pqvxubL2PGdhlPy6GLqzDYMUjyRmKAwKHYKixpdJYBUK7PJ0C029XdsnpFIdgRZG68fZiGdHVWcKPvtiPB4cA==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-x64/-/sharp-linux-x64-0.35.4.tgz", + "integrity": "sha512-9qvvEAuk8k89TfWUoX2htWjbAMX8p+NxCppjpcg5k6xMsjhBQPTsoIh36h9Qde4WRuGpJeYnOjdosDn/cnv+OA==", "cpu": [ "x64" ], + "libc": [ + "glibc" + ], "license": "Apache-2.0", "optional": true, "os": [ @@ -1816,16 +1845,19 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-linux-x64": "1.3.2" + "@img/sharp-libvips-linux-x64": "1.3.3" } }, "node_modules/@img/sharp-linuxmusl-arm64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-linuxmusl-arm64/-/sharp-linuxmusl-arm64-0.35.3.tgz", - "integrity": "sha512-Vz0iQjzzcSX3HCbfwFfCSG/9SCIqyO0mH2sXyiHaAYfBk0cRsCWXRyQYX0ovCK/PAQBbTzQ0dsPQHh5MAFL59w==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-linuxmusl-arm64/-/sharp-linuxmusl-arm64-0.35.4.tgz", + "integrity": "sha512-KB5jxpfWQTr0nc3xdHtWChdbifHrBGsd2SM62Eyxrl8afikm+f5qGBU75SJIZBT/S1MC8XyacdlXBMSWq6OURA==", "cpu": [ "arm64" ], + "libc": [ + "musl" + ], "license": "Apache-2.0", "optional": true, "os": [ @@ -1838,16 +1870,19 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-linuxmusl-arm64": "1.3.2" + "@img/sharp-libvips-linuxmusl-arm64": "1.3.3" } }, "node_modules/@img/sharp-linuxmusl-x64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-linuxmusl-x64/-/sharp-linuxmusl-x64-0.35.3.tgz", - "integrity": "sha512-6O1NPKcDVj9QEdg7Hx549EX8U0rp6yXQERqru6yRN7fGBn32UvIRJUlWnk+8xDCiG76hXVBbX82NZ/ZKr0euIg==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-linuxmusl-x64/-/sharp-linuxmusl-x64-0.35.4.tgz", + "integrity": "sha512-f+eZJZIQNEEd26RPSW+76chwOf1XtA2Y/O+5ocVyLliHkeih3e+jhLVBdNTd2rS3IbNXK8+ug93Vf5ZXtF5Lxg==", "cpu": [ "x64" ], + "libc": [ + "musl" + ], "license": "Apache-2.0", "optional": true, "os": [ @@ -1860,17 +1895,17 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-libvips-linuxmusl-x64": "1.3.2" + "@img/sharp-libvips-linuxmusl-x64": "1.3.3" } }, "node_modules/@img/sharp-wasm32": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-wasm32/-/sharp-wasm32-0.35.3.tgz", - "integrity": "sha512-cZ0XkcYGpHZkqW6iCkqTcmUC0CD9DhD5d/qeZlZkfRBn6GnHniZXLUo5+9xw8Iv76YE6LQFN9YNBlKREcCG76w==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-wasm32/-/sharp-wasm32-0.35.4.tgz", + "integrity": "sha512-zQnl4Kwp7Q6NHsENtU2T/00Zi+w3AQNwz3+UaTyVBy2FpXrzXzGjndpK61onhZjRtRpQXxCTeqw19bVyXOh7jA==", "license": "Apache-2.0 AND LGPL-3.0-or-later AND MIT", "optional": true, "dependencies": { - "@emnapi/runtime": "^1.11.1" + "@emnapi/runtime": "^1.11.3" }, "engines": { "node": ">=20.9.0" @@ -1880,16 +1915,16 @@ } }, "node_modules/@img/sharp-webcontainers-wasm32": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-webcontainers-wasm32/-/sharp-webcontainers-wasm32-0.35.3.tgz", - "integrity": "sha512-2rnq7bX3NzeR2T4YWgz8qiG4h3TSdMe+vN1iQXpJleSJ3SM5zQ8Fy2SyyXAWlbxpEZ2Y+Z4u1BePgJEYbSy80Q==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-webcontainers-wasm32/-/sharp-webcontainers-wasm32-0.35.4.tgz", + "integrity": "sha512-ESfNkywmCfPNyaZjxooddJQiQ+l/nTpGEOGthxiLnIHXC/CmcBixnfwUleX9mCz9ovrUUvKMap/pm8RYbzfwaA==", "cpu": [ "wasm32" ], "license": "Apache-2.0", "optional": true, "dependencies": { - "@img/sharp-wasm32": "0.35.3" + "@img/sharp-wasm32": "0.35.4" }, "engines": { "node": ">=20.9.0" @@ -1899,9 +1934,9 @@ } }, "node_modules/@img/sharp-win32-arm64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-win32-arm64/-/sharp-win32-arm64-0.35.3.tgz", - "integrity": "sha512-4bPwFdMbeC4JQ8L8LOyWp6nsHcboP5fxkp6iPOXz2Vg49R42TuMs2whkJ5OAP4/Ul035qOzy0AecOF9VOscn4w==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-win32-arm64/-/sharp-win32-arm64-0.35.4.tgz", + "integrity": "sha512-iNdlBX9gLVvqe2I3uIJSIKTq6wckP/DYxZtcqxm09x5Gi24DnFBmPAWZmr60ZyYMG0xlzo6goG3670ar+RXvRw==", "cpu": [ "arm64" ], @@ -1918,9 +1953,9 @@ } }, "node_modules/@img/sharp-win32-ia32": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-win32-ia32/-/sharp-win32-ia32-0.35.3.tgz", - "integrity": "sha512-r53mXsBN6lFUDiST764SvgwUdHAqM4rPAiDzAmf4fLoB6X/rkfyTrLCg6+g17wJJiCmB3JYgHuUldCWUIRFSXw==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-win32-ia32/-/sharp-win32-ia32-0.35.4.tgz", + "integrity": "sha512-kqRsbaa5CS6KHlpxnN7WhE6vAAugXyZButpRdvDWetlv6Qv4N9WTcrWzF7tXfB9T7MsoadqdI8hmwLq6UlLvtw==", "cpu": [ "ia32" ], @@ -1937,9 +1972,9 @@ } }, "node_modules/@img/sharp-win32-x64": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/@img/sharp-win32-x64/-/sharp-win32-x64-0.35.3.tgz", - "integrity": "sha512-D4y1vNeZrIIJCN+uHaWVtH86B+aCrdMYYjicy9pXHvbGZeGYLLSd3wdVuC37FxVXlU1ARsk84eKWfWMXGYEqvA==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/@img/sharp-win32-x64/-/sharp-win32-x64-0.35.4.tgz", + "integrity": "sha512-XtmnYhBcrORsJ4XJngyzr/EWP0hRZLAZRFaApdKuviyqF78+ylxh2y06ZmtULAMOnObJ3ucpN0AcwSWnMowTRg==", "cpu": [ "x64" ], @@ -1955,16 +1990,6 @@ "url": "https://opencollective.com/libvips" } }, - "node_modules/@istanbuljs/schema": { - "version": "0.1.6", - "resolved": "https://registry.npmjs.org/@istanbuljs/schema/-/schema-0.1.6.tgz", - "integrity": "sha512-+Sg6GCR/wy1oSmQDFq4LQDAhm3ETKnorxN+y5nbLULOR3P0c14f2Wurzj3/xqPXtasLFfHd5iRFQ7AJt4KH2cw==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=8" - } - }, "node_modules/@jridgewell/gen-mapping": { "version": "0.3.13", "resolved": "https://registry.npmjs.org/@jridgewell/gen-mapping/-/gen-mapping-0.3.13.tgz", @@ -2035,25 +2060,26 @@ } }, "node_modules/@next/env": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/env/-/env-16.2.11.tgz", - "integrity": "sha512-0do5A3BJ2gxWr0ZCMcD6BhW+e595jyxdTl3rXTS6lOtD8ektMiW6CO+EPwt1Eca1DBnm90r/7GdiKWBKxH++DA==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/env/-/env-16.3.3.tgz", + "integrity": "sha512-U2eYQRwXj+dsqxV79zFqExDdatnNY/ZWc2nsJU1p/OgT7fd3dXwlF6OjYaFQCfMoeTA19PWq+wVmYgimVA+V+g==", "license": "MIT" }, "node_modules/@next/eslint-plugin-next": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/eslint-plugin-next/-/eslint-plugin-next-16.2.11.tgz", - "integrity": "sha512-vMEf/aXOpzFFdtIvFYOnIDPKb0xBbrXONsz83CcKdRrekfxNdL8PNkq5qHqAHSXVlIifnX68LOMaxr3z5PkeLQ==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/eslint-plugin-next/-/eslint-plugin-next-16.3.3.tgz", + "integrity": "sha512-pbEh30vvjKpDoTAmo1v3q2uM4JUi8QaEBpbmjWvGfoec2jLghy/WNtvzAT0bk+Ik9oz6etjt4YjXEk4BQnicCw==", "dev": true, "license": "MIT", "dependencies": { + "@eslint-community/eslint-utils": "4.9.1", "fast-glob": "3.3.1" } }, "node_modules/@next/swc-darwin-arm64": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.2.11.tgz", - "integrity": "sha512-wryL4pjKmDwGv2ox6+GZDFxvmtSRLqApBR8kL1j4+vhB7Z5vJC/zAnXpiR9Xkfzl0AS8WLMnsuGV/UKI67/rrw==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.3.3.tgz", + "integrity": "sha512-8Hiv32QJPwdV6KYJ8meR9SBA061tQqnIKTJDocvOXlEQqib0xMFpzArosuffFUUc0sslbh7QQ8a3Yey1QV8EIw==", "cpu": [ "arm64" ], @@ -2067,9 +2093,9 @@ } }, "node_modules/@next/swc-darwin-x64": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.2.11.tgz", - "integrity": "sha512-aZl2j4f/fLyjQvOhv0Oe9UaMAQHolYpKhctsoYzplSumKJKPUmgjcf6545aBtysLTcu994TREd0+pSgNE4ohmg==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.3.3.tgz", + "integrity": "sha512-A1lgKgwVchRYmSe467zdwhxT9040dd8lH+o65sL5Jet8fjB4kegw/rDyPIpYVRb6jAqwXFOJpjIXJLxQKLiE3A==", "cpu": [ "x64" ], @@ -2083,12 +2109,15 @@ } }, "node_modules/@next/swc-linux-arm64-gnu": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.2.11.tgz", - "integrity": "sha512-5jEriyEnH/LWFy27L2ZG0XaLlyEJIjhsImEsiS9P563PKEVp2BVups/xfOucIrsvVntp11oNcZwjHvaDPYVB5g==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.3.3.tgz", + "integrity": "sha512-bf0FIssMFueU2dm7vQEWWxk0c8UjKTdW0yzuh0sQsD8pf1+KCLDdaqhYZNMYGmXwEOiHAUzgBKudovIlcvvBjg==", "cpu": [ "arm64" ], + "libc": [ + "glibc" + ], "license": "MIT", "optional": true, "os": [ @@ -2099,12 +2128,15 @@ } }, "node_modules/@next/swc-linux-arm64-musl": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.2.11.tgz", - "integrity": "sha512-eIjcpx2fnnFSSkZDbTxy74KnokUXDjfoLClpWelfgHLf621aTqswhwXQ7GkD5K5rplrS6LZ/Bj+mVuvzluBOEg==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.3.3.tgz", + "integrity": "sha512-W7viwCk9JY/cAkdz/A273rd5bb3RgT/IHwR7Upv90tunjBWNtAAhGhoecHh+teRNRSinuAFmE+l7fwZ4YKkrXg==", "cpu": [ "arm64" ], + "libc": [ + "musl" + ], "license": "MIT", "optional": true, "os": [ @@ -2115,12 +2147,15 @@ } }, "node_modules/@next/swc-linux-x64-gnu": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.2.11.tgz", - "integrity": "sha512-8WgzpaWMs46qJT9kiV47cje86L0x/Mu9t8/Gwj+pnbgW3rETVfCnaScPjlYUwNScpOozdcIMHWmAvuZJUonR2w==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.3.3.tgz", + "integrity": "sha512-0W46zw1N3ODpI6n0GeivHvvob1pooozgZVqy65k0mh4/7vr+FbY9+WpHzNVXjHipJf/A3FDheBG19H1s5A25rA==", "cpu": [ "x64" ], + "libc": [ + "glibc" + ], "license": "MIT", "optional": true, "os": [ @@ -2131,12 +2166,15 @@ } }, "node_modules/@next/swc-linux-x64-musl": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.2.11.tgz", - "integrity": "sha512-I3UgPds7G4ZYnTb/H+5GBGuUT2DhAk6j0mL6A4s63RjFs74wB2hOWP0vaxsK+3NJraExt3eYEPQ/UtT0x/64Nw==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.3.3.tgz", + "integrity": "sha512-H4mBso8ZTMBPtdT0PN0pBx2ayTvQuTuvS6qT13d77yVFJXAPCxkyIhLTmdMaGTJs0krQYI/qpzdHijCeihXhbg==", "cpu": [ "x64" ], + "libc": [ + "musl" + ], "license": "MIT", "optional": true, "os": [ @@ -2147,9 +2185,9 @@ } }, "node_modules/@next/swc-win32-arm64-msvc": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.2.11.tgz", - "integrity": "sha512-n89CjtcThnjrwgJMAiI5xbqwLY51zvwC9tSlArmVndAJLYVl9T9UAdlkXTmZvE++idoXe8KdglQlhNRdUp1c6g==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.3.3.tgz", + "integrity": "sha512-cTMUJpcEGmeywofCUfhR+rSsoE33+rVPnPEYNTNdLNlsOeEg/vktOsKUSTb28vUGqD2jkm4Zaskcwn7OCI6FQg==", "cpu": [ "arm64" ], @@ -2163,9 +2201,9 @@ } }, "node_modules/@next/swc-win32-x64-msvc": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.2.11.tgz", - "integrity": "sha512-md8CLNggS1Dx9pUgApzps5uAf+N8GN9xywzmNx9vHAWo94HtBwCCqkSnhIrdfQe83Dhz8Lfo/20Nb1Zxal092w==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.3.3.tgz", + "integrity": "sha512-2VR4cTBzHXaBjnGsuH6GyJjENzQOmHeAh11uY1iUhjm3j5dEUrVJuUj+VL78jaGi/Dik8xS76zEj18BsFhlVZQ==", "cpu": [ "x64" ], @@ -2965,9 +3003,9 @@ "license": "MIT" }, "node_modules/@swc/helpers": { - "version": "0.5.15", - "resolved": "https://registry.npmjs.org/@swc/helpers/-/helpers-0.5.15.tgz", - "integrity": "sha512-JQ5TuMi45Owi4/BIMAJBoSQoOJu12oOk/gADqlcUL9JEdHB8vyjUSsxqeNXnmXHjYKMi2WcYtezGEEhqUI/E2g==", + "version": "0.5.23", + "resolved": "https://registry.npmjs.org/@swc/helpers/-/helpers-0.5.23.tgz", + "integrity": "sha512-5lSsMOTXURePglDfvuAQUqkGek9Hg2kksOYay2m0+XR++b2NWYL/4sWyuvVBIs8oKnJaxkdi9whaL/sqN13afw==", "license": "Apache-2.0", "dependencies": { "tslib": "^2.8.0" @@ -4325,32 +4363,29 @@ ] }, "node_modules/@vitest/coverage-v8": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-3.2.6.tgz", - "integrity": "sha512-LsAdmUapA0qSN306d8+zOyawM0hFm2m2Hg9IwVNIKBm+qJV8cijiq2c+gxKZcB1HCfIWAy+0qEZDCUQA58A1cw==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-4.1.11.tgz", + "integrity": "sha512-8MVGEFnJIcdGjcbfKmeq8z0pZHH0JlVtoVZH9Q/qwUp6wyFnEJUBMrw9DCaj+ra3vShGmhavjalMIhPNxZAUcw==", "dev": true, "license": "MIT", "dependencies": { - "@ampproject/remapping": "^2.3.0", "@bcoe/v8-coverage": "^1.0.2", - "ast-v8-to-istanbul": "^0.3.3", - "debug": "^4.4.1", + "@vitest/utils": "4.1.11", + "ast-v8-to-istanbul": "^1.0.0", "istanbul-lib-coverage": "^3.2.2", "istanbul-lib-report": "^3.0.1", - "istanbul-lib-source-maps": "^5.0.6", - "istanbul-reports": "^3.1.7", - "magic-string": "^0.30.17", - "magicast": "^0.3.5", - "std-env": "^3.9.0", - "test-exclude": "^7.0.1", - "tinyrainbow": "^2.0.0" + "istanbul-reports": "^3.2.0", + "magicast": "^0.5.2", + "obug": "^2.1.1", + "std-env": "^4.0.0-rc.1", + "tinyrainbow": "^3.1.0" }, "funding": { "url": "https://opencollective.com/vitest" }, "peerDependencies": { - "@vitest/browser": "3.2.6", - "vitest": "3.2.6" + "@vitest/browser": "4.1.11", + "vitest": "4.1.11" }, "peerDependenciesMeta": { "@vitest/browser": { @@ -4359,39 +4394,40 @@ } }, "node_modules/@vitest/expect": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-3.2.6.tgz", - "integrity": "sha512-1+7q9BtaKzEmO+fmNT3kYvoNn5Y71XWAx2Q5HRim4tTVRQVRv4uJFAQ5FbK0OPUeNP/WmVCpxYxoJdvuHVjzBQ==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-4.1.11.tgz", + "integrity": "sha512-VX2x5vNJXET47KAFzwERI+KRMtTTCSWTfSMKsW7JsUsXV4psq++e3DvZpuTDOpHcxytiDs6p2nhVb2tVDiiUYw==", "dev": true, "license": "MIT", "dependencies": { + "@standard-schema/spec": "^1.1.0", "@types/chai": "^5.2.2", - "@vitest/spy": "3.2.6", - "@vitest/utils": "3.2.6", - "chai": "^5.2.0", - "tinyrainbow": "^2.0.0" + "@vitest/spy": "4.1.11", + "@vitest/utils": "4.1.11", + "chai": "^6.2.2", + "tinyrainbow": "^3.1.0" }, "funding": { "url": "https://opencollective.com/vitest" } }, "node_modules/@vitest/mocker": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-3.2.6.tgz", - "integrity": "sha512-EZOrpDbkKotFAP7wPAQV1UIyoGOk4oX7ynWhBhLB7v+meMHbQhU16oPpIYGTTe4oFlhpryGpgpcZP/sin3hYuw==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-4.1.11.tgz", + "integrity": "sha512-2XJVD55d1o5AZous5CCGKS74g/riOj9odEt2bQpCVZeblHyHdnMeFl4jl0XjU21stf4mbjUkew2eXQZt65g5CQ==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/spy": "3.2.6", + "@vitest/spy": "4.1.11", "estree-walker": "^3.0.3", - "magic-string": "^0.30.17" + "magic-string": "^0.30.21" }, "funding": { "url": "https://opencollective.com/vitest" }, "peerDependencies": { "msw": "^2.4.9", - "vite": "^5.0.0 || ^6.0.0 || ^7.0.0-0" + "vite": "^6.0.0 || ^7.0.0 || ^8.0.0" }, "peerDependenciesMeta": { "msw": { @@ -4403,42 +4439,42 @@ } }, "node_modules/@vitest/pretty-format": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-3.2.6.tgz", - "integrity": "sha512-lb7XXXzmm2h2ASzFnRvQpDo6onT1NmMJA3tkGTWiBFtRJ9lxGY3d3mm/Apt36gej2bkkOVLL/yTOtufDaFa/jA==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-4.1.11.tgz", + "integrity": "sha512-yiZzPbGTS9Sr/JpFl8zHrcIkAofNbFV6k21vIgQN/cY/oxZeXhJv5sc/MBJ5jFKWmWs+oJHw0UXLZjmf931+Vw==", "dev": true, "license": "MIT", "dependencies": { - "tinyrainbow": "^2.0.0" + "tinyrainbow": "^3.1.0" }, "funding": { "url": "https://opencollective.com/vitest" } }, "node_modules/@vitest/runner": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-3.2.6.tgz", - "integrity": "sha512-HYcoSj1w5tcgUnzoF0HcyaAQjpA1gj9ftUJ7iSJSuipc02jW9gKkigwZbjFldAfYHA1fa8UZVRftdMY5msWM9Q==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-4.1.11.tgz", + "integrity": "sha512-LztvUgdwMNJMIkj3hQnnxiC2Xy1zNxq928W/xhjCLaNCzqTZOudjwbQf6v9IntZGPw132i2Lq2rgTRZHD3JHNw==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/utils": "3.2.6", - "pathe": "^2.0.3", - "strip-literal": "^3.0.0" + "@vitest/utils": "4.1.11", + "pathe": "^2.0.3" }, "funding": { "url": "https://opencollective.com/vitest" } }, "node_modules/@vitest/snapshot": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-3.2.6.tgz", - "integrity": "sha512-H+ZjNTWGpObenh0YnlBctAPnJSI20P81PL8BPzWpx54YXLLTm8hEsWawtcYLMrwvpK48hGxLLbCS+1KRXhsKhw==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-4.1.11.tgz", + "integrity": "sha512-pN7ikn1ON7h8ee4gIAp4AzyK+zBtJPzVbqOgu5LCEh4VaJVbPQcgYQYJIMGQPXVeJJq1fnfazis7a5pFNPahog==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "3.2.6", - "magic-string": "^0.30.17", + "@vitest/pretty-format": "4.1.11", + "@vitest/utils": "4.1.11", + "magic-string": "^0.30.21", "pathe": "^2.0.3" }, "funding": { @@ -4446,50 +4482,47 @@ } }, "node_modules/@vitest/spy": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-3.2.6.tgz", - "integrity": "sha512-oq6BbH68WzcWmwtBrU9nqLeaXTR4XwJF7FSLkKEZo4i6eoXcrxjcwSuTvWBIRUTC6VC72nXYunzqgZA+IKdtxg==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-4.1.11.tgz", + "integrity": "sha512-apNa/prQy2qCeywhnixOHPRCgGNhvg7T4Dapfl1GahLp/R+uhBm5cPyFoNVyqsNd2h1nJxL6BqqdIjiABL60YA==", "dev": true, "license": "MIT", - "dependencies": { - "tinyspy": "^4.0.3" - }, "funding": { "url": "https://opencollective.com/vitest" } }, "node_modules/@vitest/ui": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-3.2.6.tgz", - "integrity": "sha512-mATfG3zVdhobE9U1rIpvtYD3DGuSSxqZ3Aj/8ityGqKXy8YDJ9BoAjZmAz6dZ1IZ1xI5V+MerkCczvVa+3QK9Q==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-4.1.11.tgz", + "integrity": "sha512-r/rwyKoev21mWdRGSEkZOqkQ2BYy68mwjihg9M90nNRbf4NGrgzZ4cj6JNCEwlOGJkbKeMgsjlykvwKUbRr7gw==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/utils": "3.2.6", + "@vitest/utils": "4.1.11", "fflate": "^0.8.2", - "flatted": "^3.3.3", + "flatted": "^3.4.2", "pathe": "^2.0.3", - "sirv": "^3.0.1", - "tinyglobby": "^0.2.14", - "tinyrainbow": "^2.0.0" + "sirv": "^3.0.2", + "tinyglobby": "^0.2.15", + "tinyrainbow": "^3.1.0" }, "funding": { "url": "https://opencollective.com/vitest" }, "peerDependencies": { - "vitest": "3.2.6" + "vitest": "4.1.11" } }, "node_modules/@vitest/utils": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-3.2.6.tgz", - "integrity": "sha512-lI23nIs4bnT3T8NIoh+vFaz5s2/DdP0Jgt2jxwgWljvwn82cLJtyi/If+fjFyoLMGIOz0U/fKvWE0d4jsNQEfg==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-4.1.11.tgz", + "integrity": "sha512-zTCVGpyFsGWBhllOyKlTw/vnr6D9qxsfSDyfbyZmTyjHw5N/VuvzHpHoQjm2ZJzn4RJgx5w4r7V0er69CmLgPQ==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "3.2.6", - "loupe": "^3.1.4", - "tinyrainbow": "^2.0.0" + "@vitest/pretty-format": "4.1.11", + "convert-source-map": "^2.0.0", + "tinyrainbow": "^3.1.0" }, "funding": { "url": "https://opencollective.com/vitest" @@ -4800,9 +4833,9 @@ "license": "MIT" }, "node_modules/ast-v8-to-istanbul": { - "version": "0.3.12", - "resolved": "https://registry.npmjs.org/ast-v8-to-istanbul/-/ast-v8-to-istanbul-0.3.12.tgz", - "integrity": "sha512-BRRC8VRZY2R4Z4lFIL35MwNXmwVqBityvOIwETtsCSwvjl0IdgFsy9NhdaA6j74nUdtJJlIypeRhpDam19Wq3g==", + "version": "1.0.5", + "resolved": "https://registry.npmjs.org/ast-v8-to-istanbul/-/ast-v8-to-istanbul-1.0.5.tgz", + "integrity": "sha512-UPAgKJFSEGMWSDr3LX4tqnAb4f7KGT8O40Tyx8wbYmmZ/yn58lNCm8h3svs3eXgiGd5AXxz8NDOvXWvicq+rJA==", "dev": true, "license": "MIT", "dependencies": { @@ -4972,16 +5005,6 @@ "node": "^6 || ^7 || ^8 || ^9 || ^10 || ^11 || ^12 || >=13.7" } }, - "node_modules/cac": { - "version": "6.7.14", - "resolved": "https://registry.npmjs.org/cac/-/cac-6.7.14.tgz", - "integrity": "sha512-b6Ilus+c3RrdDk+JhLKUAQfzzgLEPy6wcXqS7f/xe1EETvsDP6GORG7SFuOs6cID5YkqchW/LXZbX5bc8j7ZcQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=8" - } - }, "node_modules/call-bind": { "version": "1.0.9", "resolved": "https://registry.npmjs.org/call-bind/-/call-bind-1.0.9.tgz", @@ -5072,18 +5095,11 @@ } }, "node_modules/chai": { - "version": "5.3.3", - "resolved": "https://registry.npmjs.org/chai/-/chai-5.3.3.tgz", - "integrity": "sha512-4zNhdJD/iOjSH0A05ea+Ke6MU5mmpQcbQsSOkgdaUMJ9zTlDTD/GYlwohmIE2u0gaxHYiVHEn1Fw9mZ/ktJWgw==", + "version": "6.2.2", + "resolved": "https://registry.npmjs.org/chai/-/chai-6.2.2.tgz", + "integrity": "sha512-NUPRluOfOiTKBKvWPtSD4PhFvWCqOi0BGStNWs57X9js7XGTprSmFoz5F0tWhR4WPjNeR9jXqdC7/UpSJTnlRg==", "dev": true, "license": "MIT", - "dependencies": { - "assertion-error": "^2.0.1", - "check-error": "^2.1.1", - "deep-eql": "^5.0.1", - "loupe": "^3.1.0", - "pathval": "^2.0.0" - }, "engines": { "node": ">=18" } @@ -5152,16 +5168,6 @@ "url": "https://github.com/sponsors/wooorm" } }, - "node_modules/check-error": { - "version": "2.1.3", - "resolved": "https://registry.npmjs.org/check-error/-/check-error-2.1.3.tgz", - "integrity": "sha512-PAJdDJusoxnwm1VwW07VWwUN1sl7smmC3OKggvndJFadxxDRyFJBX/ggnu/KE4kQAB7a3Dp8f/YXC1FlUprWmA==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 16" - } - }, "node_modules/class-variance-authority": { "version": "0.7.1", "resolved": "https://registry.npmjs.org/class-variance-authority/-/class-variance-authority-0.7.1.tgz", @@ -5577,16 +5583,6 @@ "url": "https://github.com/sponsors/wooorm" } }, - "node_modules/deep-eql": { - "version": "5.0.2", - "resolved": "https://registry.npmjs.org/deep-eql/-/deep-eql-5.0.2.tgz", - "integrity": "sha512-h5k/5U50IJJFpzfL6nO9jaaumfjO/f2NjK/oYB2Djzm4p9L+3T9qWpZqZ2hAbLPuuYq9wrU08WQyBTL5GbPk5Q==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=6" - } - }, "node_modules/deep-is": { "version": "0.1.4", "resolved": "https://registry.npmjs.org/deep-is/-/deep-is-0.1.4.tgz", @@ -5862,9 +5858,9 @@ } }, "node_modules/es-module-lexer": { - "version": "1.7.0", - "resolved": "https://registry.npmjs.org/es-module-lexer/-/es-module-lexer-1.7.0.tgz", - "integrity": "sha512-jEQoCwk8hyb2AZziIOLhDqpm5+2ww5uIE6lkO/6jcOCusfk6LhMHpXXfBLXTZ7Ydyt0j4VoUQv6uGNYbdW+kBA==", + "version": "2.3.2", + "resolved": "https://registry.npmjs.org/es-module-lexer/-/es-module-lexer-2.3.2.tgz", + "integrity": "sha512-poHGpORABojJJucnV9KbOavETW8lBVnphkW77ER5/BQ5Fz7oXSoCNek7IH3vR5nRjdsEz926ibFYX8KtLQmdyw==", "dev": true, "license": "MIT" }, @@ -6062,13 +6058,13 @@ } }, "node_modules/eslint-config-next": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/eslint-config-next/-/eslint-config-next-16.2.11.tgz", - "integrity": "sha512-FIpbK/dUyxUExchDB7eBg3k+VU8R2iR/Cx9/kqTBUTFv2bOIR9aRrpno4rvAQ9VhiPQAyFKNA2NlZwouGWtclA==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/eslint-config-next/-/eslint-config-next-16.3.3.tgz", + "integrity": "sha512-teqtsR26tnlfXFHfVLTM/4tzEzU8DMu6GS1sddZzhfGzgd2f2ofbgDUcsk6cssSCzX6Tk6fmWifJcdANSdPJrw==", "dev": true, "license": "MIT", "dependencies": { - "@next/eslint-plugin-next": "16.2.11", + "@next/eslint-plugin-next": "16.3.3", "eslint-import-resolver-node": "^0.3.6", "eslint-import-resolver-typescript": "^3.5.2", "eslint-plugin-import": "^2.32.0", @@ -6546,9 +6542,9 @@ "license": "MIT" }, "node_modules/expect-type": { - "version": "1.3.0", - "resolved": "https://registry.npmjs.org/expect-type/-/expect-type-1.3.0.tgz", - "integrity": "sha512-knvyeauYhqjOYvQ66MznSMs83wmHrCycNEN6Ao+2AeYEfxUIkuiVxdEa1qlGEPK+We3n0THiDciYSsCcgW/DoA==", + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/expect-type/-/expect-type-1.4.0.tgz", + "integrity": "sha512-KfYbmpRm0VbLjEvVa9yGwCi9GI34xvi7A/HXYWQO65CSD2u3MczUJSuwXKFIxlGsgBQizV9q5J9NHj4VG0n+pA==", "dev": true, "license": "Apache-2.0", "engines": { @@ -6952,24 +6948,6 @@ "url": "https://github.com/privatenumber/get-tsconfig?sponsor=1" } }, - "node_modules/glob": { - "version": "13.0.0", - "resolved": "https://registry.npmjs.org/glob/-/glob-13.0.0.tgz", - "integrity": "sha512-tvZgpqk6fz4BaNZ66ZsRaZnbHvP/jG3uKJvAZOwEVUL4RTA5nJeeLYfyN9/VA8NX/V3IBG+hkeuGpKjvELkVhA==", - "dev": true, - "license": "BlueOak-1.0.0", - "dependencies": { - "minimatch": "^10.1.1", - "minipass": "^7.1.2", - "path-scurry": "^2.0.0" - }, - "engines": { - "node": "20 || >=22" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/glob-parent": { "version": "6.0.2", "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-6.0.2.tgz", @@ -7941,21 +7919,6 @@ "node": ">=10" } }, - "node_modules/istanbul-lib-source-maps": { - "version": "5.0.6", - "resolved": "https://registry.npmjs.org/istanbul-lib-source-maps/-/istanbul-lib-source-maps-5.0.6.tgz", - "integrity": "sha512-yg2d+Em4KizZC5niWhQaIomgf5WlL4vOOjZ5xGCmF8SnPE/mDWWXgvRExdcpCgh9lLRRa1/fSYp2ymmbJ1pI+A==", - "dev": true, - "license": "BSD-3-Clause", - "dependencies": { - "@jridgewell/trace-mapping": "^0.3.23", - "debug": "^4.1.1", - "istanbul-lib-coverage": "^3.0.0" - }, - "engines": { - "node": ">=10" - } - }, "node_modules/istanbul-reports": { "version": "3.2.0", "resolved": "https://registry.npmjs.org/istanbul-reports/-/istanbul-reports-3.2.0.tgz", @@ -8015,9 +7978,9 @@ "license": "MIT" }, "node_modules/js-yaml": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.1.tgz", - "integrity": "sha512-CY6crGq313MX8GkwvB7tzgp99vjQxY1++5y10/BKN/GUfHqWaOGQMNZkBvqSzsZKWk/ijwHlWzzkLulsGHhjWQ==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.2.tgz", + "integrity": "sha512-SFNOvSJ+Dgf/9An904Yx+CgSlIPCkIpao4qo51lpee25TIRejdH3rhR4EZMGoNx3/TP3O+wzWuiTFl4sqbltzA==", "dev": true, "funding": [ { @@ -8607,13 +8570,6 @@ "loose-envify": "cli.js" } }, - "node_modules/loupe": { - "version": "3.2.1", - "resolved": "https://registry.npmjs.org/loupe/-/loupe-3.2.1.tgz", - "integrity": "sha512-CdzqowRJCeLU72bHvWqwRBBlLcMEtIvGrlvef74kMnV2AolS9Y8xUv1I0U/MNAWMhBlKIoyuEgoJ0t/bbwHbLQ==", - "dev": true, - "license": "MIT" - }, "node_modules/lowlight": { "version": "1.20.0", "resolved": "https://registry.npmjs.org/lowlight/-/lowlight-1.20.0.tgz", @@ -8668,15 +8624,15 @@ } }, "node_modules/magicast": { - "version": "0.3.5", - "resolved": "https://registry.npmjs.org/magicast/-/magicast-0.3.5.tgz", - "integrity": "sha512-L0WhttDl+2BOsybvEOLK7fW3UA0OQ0IQ2d6Zl2x/a6vVRs3bAY0ECOSHHeL5jD+SbOpOCUEi0y1DgHEn9Qn1AQ==", + "version": "0.5.4", + "resolved": "https://registry.npmjs.org/magicast/-/magicast-0.5.4.tgz", + "integrity": "sha512-llBEhWm1SacoRwgHUoQJYtwp4PBLF4faQi5TCpIGyGs9n4y5+juI0tDgyKIfpqxckRHaHzouUEph3THklWh03w==", "dev": true, "license": "MIT", "dependencies": { - "@babel/parser": "^7.25.4", - "@babel/types": "^7.25.4", - "source-map-js": "^1.2.0" + "@babel/parser": "^7.29.7", + "@babel/types": "^7.29.7", + "source-map-js": "^1.2.1" } }, "node_modules/make-dir": { @@ -9670,16 +9626,6 @@ "url": "https://github.com/sponsors/ljharb" } }, - "node_modules/minipass": { - "version": "7.1.3", - "resolved": "https://registry.npmjs.org/minipass/-/minipass-7.1.3.tgz", - "integrity": "sha512-tEBHqDnIoM/1rXME1zgka9g6Q2lcoCkxHLuc7ODJ5BxbP5d4c2Z5cGgtXAku59200Cx7diuHTOYfSBD8n6mm8A==", - "dev": true, - "license": "BlueOak-1.0.0", - "engines": { - "node": ">=16 || 14 >=14.17" - } - }, "node_modules/moment": { "version": "2.30.1", "resolved": "https://registry.npmjs.org/moment/-/moment-2.30.1.tgz", @@ -9747,16 +9693,16 @@ "license": "MIT" }, "node_modules/next": { - "version": "16.2.11", - "resolved": "https://registry.npmjs.org/next/-/next-16.2.11.tgz", - "integrity": "sha512-B339zaqbyK8cmxhoAvLrcwoabwCP1wz21zSzfqxqXAemTu2BXnH7tQnfcglKv1vnMUIDBc+Hth7XODQriTZiRQ==", + "version": "16.3.3", + "resolved": "https://registry.npmjs.org/next/-/next-16.3.3.tgz", + "integrity": "sha512-tuRTx1nQ/yVw83cwJBo9F+njGUgMn3UHQycreWHB8XsStvvAh1AthbI8/4IpKnFaF58F+iSiHejYOlMQ/eq83g==", "license": "MIT", "dependencies": { - "@next/env": "16.2.11", - "@swc/helpers": "0.5.15", + "@next/env": "16.3.3", + "@swc/helpers": "0.5.23", "baseline-browser-mapping": "^2.9.19", "caniuse-lite": "^1.0.30001579", - "postcss": "8.4.31", + "postcss": "8.5.23", "styled-jsx": "5.1.6" }, "bin": { @@ -9766,15 +9712,15 @@ "node": ">=20.9.0" }, "optionalDependencies": { - "@next/swc-darwin-arm64": "16.2.11", - "@next/swc-darwin-x64": "16.2.11", - "@next/swc-linux-arm64-gnu": "16.2.11", - "@next/swc-linux-arm64-musl": "16.2.11", - "@next/swc-linux-x64-gnu": "16.2.11", - "@next/swc-linux-x64-musl": "16.2.11", - "@next/swc-win32-arm64-msvc": "16.2.11", - "@next/swc-win32-x64-msvc": "16.2.11", - "sharp": "^0.34.5" + "@next/swc-darwin-arm64": "16.3.3", + "@next/swc-darwin-x64": "16.3.3", + "@next/swc-linux-arm64-gnu": "16.3.3", + "@next/swc-linux-arm64-musl": "16.3.3", + "@next/swc-linux-x64-gnu": "16.3.3", + "@next/swc-linux-x64-musl": "16.3.3", + "@next/swc-win32-arm64-msvc": "16.3.3", + "@next/swc-win32-x64-msvc": "16.3.3", + "sharp": "^0.35.3" }, "peerDependencies": { "@opentelemetry/api": "^1.1.0", @@ -10069,6 +10015,20 @@ "url": "https://github.com/sponsors/ljharb" } }, + "node_modules/obug": { + "version": "2.1.4", + "resolved": "https://registry.npmjs.org/obug/-/obug-2.1.4.tgz", + "integrity": "sha512-4a+OsYv9UktOJKE+l1A4OufDgdRF9PifWj+tJnHURo/P+WOxpG4GzUFL9qCalmWauao6ogiG+QvnCovwPoyAWA==", + "dev": true, + "funding": [ + "https://github.com/sponsors/sxzz", + "https://opencollective.com/debug" + ], + "license": "MIT", + "engines": { + "node": ">=12.20.0" + } + }, "node_modules/openai": { "version": "4.104.0", "resolved": "https://registry.npmjs.org/openai/-/openai-4.104.0.tgz", @@ -10378,23 +10338,6 @@ "dev": true, "license": "MIT" }, - "node_modules/path-scurry": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/path-scurry/-/path-scurry-2.0.2.tgz", - "integrity": "sha512-3O/iVVsJAPsOnpwWIeD+d6z/7PmqApyQePUtCndjatj/9I5LylHvt5qluFaBT3I5h3r1ejfR056c+FCv+NnNXg==", - "dev": true, - "license": "BlueOak-1.0.0", - "dependencies": { - "lru-cache": "^11.0.0", - "minipass": "^7.1.2" - }, - "engines": { - "node": "18 || 20 || >=22" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/pathe": { "version": "2.0.3", "resolved": "https://registry.npmjs.org/pathe/-/pathe-2.0.3.tgz", @@ -10402,16 +10345,6 @@ "dev": true, "license": "MIT" }, - "node_modules/pathval": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/pathval/-/pathval-2.0.1.tgz", - "integrity": "sha512-//nshmD55c46FuFw26xV/xFAaB5HF9Xdap7HJBBnrKdAd6/GxDBaNA1870O79+9ueg61cZLSVc+OaFlfmObYVQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 14.16" - } - }, "node_modules/picocolors": { "version": "1.1.1", "resolved": "https://registry.npmjs.org/picocolors/-/picocolors-1.1.1.tgz", @@ -11298,9 +11231,9 @@ } }, "node_modules/sharp": { - "version": "0.35.3", - "resolved": "https://registry.npmjs.org/sharp/-/sharp-0.35.3.tgz", - "integrity": "sha512-ej0zVHuZGHCiABXcNxeYhpRnPNPAcvbG8RMdBAhDAxLKkCRVSpK3Iyu7qbqw3JMzoj0REeM6f3tJLtVwl0023Q==", + "version": "0.35.4", + "resolved": "https://registry.npmjs.org/sharp/-/sharp-0.35.4.tgz", + "integrity": "sha512-n++8XWcj+jCOr2IOl7h8LbKnGBDY4aPbmprMONBNFdn0ImXqpGVv5zliDs0V9HbmbCQLpbuo2ej9rAoOQTvMDA==", "license": "Apache-2.0", "optional": true, "dependencies": { @@ -11315,31 +11248,31 @@ "url": "https://opencollective.com/libvips" }, "optionalDependencies": { - "@img/sharp-darwin-arm64": "0.35.3", - "@img/sharp-darwin-x64": "0.35.3", - "@img/sharp-freebsd-wasm32": "0.35.3", - "@img/sharp-libvips-darwin-arm64": "1.3.2", - "@img/sharp-libvips-darwin-x64": "1.3.2", - "@img/sharp-libvips-linux-arm": "1.3.2", - "@img/sharp-libvips-linux-arm64": "1.3.2", - "@img/sharp-libvips-linux-ppc64": "1.3.2", - "@img/sharp-libvips-linux-riscv64": "1.3.2", - "@img/sharp-libvips-linux-s390x": "1.3.2", - "@img/sharp-libvips-linux-x64": "1.3.2", - "@img/sharp-libvips-linuxmusl-arm64": "1.3.2", - "@img/sharp-libvips-linuxmusl-x64": "1.3.2", - "@img/sharp-linux-arm": "0.35.3", - "@img/sharp-linux-arm64": "0.35.3", - "@img/sharp-linux-ppc64": "0.35.3", - "@img/sharp-linux-riscv64": "0.35.3", - "@img/sharp-linux-s390x": "0.35.3", - "@img/sharp-linux-x64": "0.35.3", - "@img/sharp-linuxmusl-arm64": "0.35.3", - "@img/sharp-linuxmusl-x64": "0.35.3", - "@img/sharp-webcontainers-wasm32": "0.35.3", - "@img/sharp-win32-arm64": "0.35.3", - "@img/sharp-win32-ia32": "0.35.3", - "@img/sharp-win32-x64": "0.35.3" + "@img/sharp-darwin-arm64": "0.35.4", + "@img/sharp-darwin-x64": "0.35.4", + "@img/sharp-freebsd-wasm32": "0.35.4", + "@img/sharp-libvips-darwin-arm64": "1.3.3", + "@img/sharp-libvips-darwin-x64": "1.3.3", + "@img/sharp-libvips-linux-arm": "1.3.3", + "@img/sharp-libvips-linux-arm64": "1.3.3", + "@img/sharp-libvips-linux-ppc64": "1.3.3", + "@img/sharp-libvips-linux-riscv64": "1.3.3", + "@img/sharp-libvips-linux-s390x": "1.3.3", + "@img/sharp-libvips-linux-x64": "1.3.3", + "@img/sharp-libvips-linuxmusl-arm64": "1.3.3", + "@img/sharp-libvips-linuxmusl-x64": "1.3.3", + "@img/sharp-linux-arm": "0.35.4", + "@img/sharp-linux-arm64": "0.35.4", + "@img/sharp-linux-ppc64": "0.35.4", + "@img/sharp-linux-riscv64": "0.35.4", + "@img/sharp-linux-s390x": "0.35.4", + "@img/sharp-linux-x64": "0.35.4", + "@img/sharp-linuxmusl-arm64": "0.35.4", + "@img/sharp-linuxmusl-x64": "0.35.4", + "@img/sharp-webcontainers-wasm32": "0.35.4", + "@img/sharp-win32-arm64": "0.35.4", + "@img/sharp-win32-ia32": "0.35.4", + "@img/sharp-win32-x64": "0.35.4" }, "peerDependenciesMeta": { "@types/node": { @@ -11469,9 +11402,9 @@ } }, "node_modules/smol-toml": { - "version": "1.6.1", - "resolved": "https://registry.npmjs.org/smol-toml/-/smol-toml-1.6.1.tgz", - "integrity": "sha512-dWUG8F5sIIARXih1DTaQAX4SsiTXhInKf1buxdY9DIg4ZYPZK5nGM1VRIYmEbDbsHt7USo99xSLFu5Q1IqTmsg==", + "version": "1.8.0", + "resolved": "https://registry.npmjs.org/smol-toml/-/smol-toml-1.8.0.tgz", + "integrity": "sha512-kCZr2V3ch9i00x8zXRhjUNVcjG9ijES5dDudkXvUVCT5QlJNQWElSJdZqyPemffHoLNUYwOcou0Fy+ojN0uHSQ==", "dev": true, "license": "BSD-3-Clause", "engines": { @@ -11531,9 +11464,9 @@ "license": "MIT" }, "node_modules/std-env": { - "version": "3.10.0", - "resolved": "https://registry.npmjs.org/std-env/-/std-env-3.10.0.tgz", - "integrity": "sha512-5GS12FdOZNliM5mAOxFRg7Ir0pWz8MdpYm6AY6VPkGpbA7ZzmbzNcBJQ0GPvvyWgcY7QAhCgf9Uy89I03faLkg==", + "version": "4.2.0", + "resolved": "https://registry.npmjs.org/std-env/-/std-env-4.2.0.tgz", + "integrity": "sha512-oCUKSupKTHX53EyjDtuZQ64pjLJ6yYCtpmEw0goYxtjG9KpbRe8KAsl2tBUGU9DyMcJ0RwJ8GqJAFzMXcXW1Rw==", "dev": true, "license": "MIT" }, @@ -11714,26 +11647,6 @@ "url": "https://github.com/sponsors/sindresorhus" } }, - "node_modules/strip-literal": { - "version": "3.1.0", - "resolved": "https://registry.npmjs.org/strip-literal/-/strip-literal-3.1.0.tgz", - "integrity": "sha512-8r3mkIM/2+PpjHoOtiAW8Rg3jJLHaV7xPwG+YRGrv6FP0wwk/toTpATxWYOW0BKdWwl82VT2tFYi5DlROa0Mxg==", - "dev": true, - "license": "MIT", - "dependencies": { - "js-tokens": "^9.0.1" - }, - "funding": { - "url": "https://github.com/sponsors/antfu" - } - }, - "node_modules/strip-literal/node_modules/js-tokens": { - "version": "9.0.1", - "resolved": "https://registry.npmjs.org/js-tokens/-/js-tokens-9.0.1.tgz", - "integrity": "sha512-mxa9E9ITFOt0ban3j6L5MpjwegGz6lBQmM1IJkWeBZGcMxto50+eWdjC/52xDbS2vy0k7vIMK0Fe2wfL9OQSpQ==", - "dev": true, - "license": "MIT" - }, "node_modules/style-to-js": { "version": "1.1.21", "resolved": "https://registry.npmjs.org/style-to-js/-/style-to-js-1.1.21.tgz", @@ -11838,21 +11751,6 @@ "url": "https://opencollective.com/webpack" } }, - "node_modules/test-exclude": { - "version": "7.0.2", - "resolved": "https://registry.npmjs.org/test-exclude/-/test-exclude-7.0.2.tgz", - "integrity": "sha512-u9E6A+ZDYdp7a4WnarkXPZOx8Ilz46+kby6p1yZ8zsGTz9gYa6FIS7lj2oezzNKmtdyyJNNmmXDppga5GB7kSw==", - "dev": true, - "license": "ISC", - "dependencies": { - "@istanbuljs/schema": "^0.1.2", - "glob": "^10.4.1", - "minimatch": "^10.2.2" - }, - "engines": { - "node": ">=18" - } - }, "node_modules/tiny-invariant": { "version": "1.3.3", "resolved": "https://registry.npmjs.org/tiny-invariant/-/tiny-invariant-1.3.3.tgz", @@ -11867,11 +11765,14 @@ "license": "MIT" }, "node_modules/tinyexec": { - "version": "0.3.2", - "resolved": "https://registry.npmjs.org/tinyexec/-/tinyexec-0.3.2.tgz", - "integrity": "sha512-KQQR9yN7R5+OSwaK0XQoj22pwHoTlgYqmUscPYoknOoWCWfj/5/ABTMRi69FrKU5ffPVh5QcFikpWJI/P1ocHA==", + "version": "1.3.1", + "resolved": "https://registry.npmjs.org/tinyexec/-/tinyexec-1.3.1.tgz", + "integrity": "sha512-GCvB3aoys96IuDFBMcTB46JOR6mdMtAToqwiW8JlWhsoh1mhHi/xn9ss/Dg7N555GiJyEt2qzoG/NHCwM6h1EA==", "dev": true, - "license": "MIT" + "license": "MIT", + "engines": { + "node": ">=18" + } }, "node_modules/tinyglobby": { "version": "0.2.16", @@ -11890,30 +11791,10 @@ "url": "https://github.com/sponsors/SuperchupuDev" } }, - "node_modules/tinypool": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/tinypool/-/tinypool-1.1.1.tgz", - "integrity": "sha512-Zba82s87IFq9A9XmjiX5uZA/ARWDrB03OHlq+Vw1fSdt0I+4/Kutwy8BP4Y/y/aORMo61FQ0vIb5j44vSo5Pkg==", - "dev": true, - "license": "MIT", - "engines": { - "node": "^18.0.0 || >=20.0.0" - } - }, "node_modules/tinyrainbow": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/tinyrainbow/-/tinyrainbow-2.0.0.tgz", - "integrity": "sha512-op4nsTR47R6p0vMUUoYl/a+ljLFVtlfaXkLQmqfLR1qHma1h/ysYk4hEXZ880bf2CYgTskvTa/e196Vd5dDQXw==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=14.0.0" - } - }, - "node_modules/tinyspy": { - "version": "4.0.4", - "resolved": "https://registry.npmjs.org/tinyspy/-/tinyspy-4.0.4.tgz", - "integrity": "sha512-azl+t0z7pw/z958Gy9svOTuzqIk6xq+NSheJzn5MMWtWTFywIacg2wUlzKFGtt3cthx0r2SxMK0yzJOR0IES7Q==", + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/tinyrainbow/-/tinyrainbow-3.1.1.tgz", + "integrity": "sha512-yau8yJdTt989Mm0Bd/236QnzEiPf2xLLTqUZRUJOo/3CB078LSwzei343DgtJVmfJKJE3TMINY1u42SQsP6mXw==", "dev": true, "license": "MIT", "engines": { @@ -12547,29 +12428,6 @@ } } }, - "node_modules/vite-node": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/vite-node/-/vite-node-3.2.4.tgz", - "integrity": "sha512-EbKSKh+bh1E1IFxeO0pg1n4dvoOTt0UDiXMd/qn++r98+jPO1xtJilvXldeuQ8giIB5IkpjCgMleHMNEsGH6pg==", - "dev": true, - "license": "MIT", - "dependencies": { - "cac": "^6.7.14", - "debug": "^4.4.1", - "es-module-lexer": "^1.7.0", - "pathe": "^2.0.3", - "vite": "^5.0.0 || ^6.0.0 || ^7.0.0-0" - }, - "bin": { - "vite-node": "vite-node.mjs" - }, - "engines": { - "node": "^18.0.0 || ^20.0.0 || >=22.0.0" - }, - "funding": { - "url": "https://opencollective.com/vitest" - } - }, "node_modules/vite/node_modules/fsevents": { "version": "2.3.3", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz", @@ -12586,65 +12444,79 @@ } }, "node_modules/vitest": { - "version": "3.2.6", - "resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.6.tgz", - "integrity": "sha512-xejya+bT/j/+R/AGa1XOfRxLmNUlLtlwjRsFUILF+xHfzElmGcmFydy2gqqIrd62ptIEfwVMofd19uNWD9L7Nw==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/vitest/-/vitest-4.1.11.tgz", + "integrity": "sha512-fhACrNXUidIbGSBr5FlbuBkO7VWC1ZyLl0DO4CU2DrQoAPxX84Ysxs+HeGQpii5lZWV1Q4gBZTTu49mF+A6Edw==", "dev": true, "license": "MIT", "dependencies": { - "@types/chai": "^5.2.2", - "@vitest/expect": "3.2.6", - "@vitest/mocker": "3.2.6", - "@vitest/pretty-format": "^3.2.6", - "@vitest/runner": "3.2.6", - "@vitest/snapshot": "3.2.6", - "@vitest/spy": "3.2.6", - "@vitest/utils": "3.2.6", - "chai": "^5.2.0", - "debug": "^4.4.1", - "expect-type": "^1.2.1", - "magic-string": "^0.30.17", + "@vitest/expect": "4.1.11", + "@vitest/mocker": "4.1.11", + "@vitest/pretty-format": "4.1.11", + "@vitest/runner": "4.1.11", + "@vitest/snapshot": "4.1.11", + "@vitest/spy": "4.1.11", + "@vitest/utils": "4.1.11", + "es-module-lexer": "^2.0.0", + "expect-type": "^1.3.0", + "magic-string": "^0.30.21", + "obug": "^2.1.1", "pathe": "^2.0.3", - "picomatch": "^4.0.2", - "std-env": "^3.9.0", + "picomatch": "^4.0.3", + "std-env": "^4.0.0-rc.1", "tinybench": "^2.9.0", - "tinyexec": "^0.3.2", - "tinyglobby": "^0.2.14", - "tinypool": "^1.1.1", - "tinyrainbow": "^2.0.0", - "vite": "^5.0.0 || ^6.0.0 || ^7.0.0-0", - "vite-node": "3.2.4", + "tinyexec": "^1.0.2", + "tinyglobby": "^0.2.15", + "tinyrainbow": "^3.1.0", + "vite": "^6.0.0 || ^7.0.0 || ^8.0.0", "why-is-node-running": "^2.3.0" }, "bin": { "vitest": "vitest.mjs" }, "engines": { - "node": "^18.0.0 || ^20.0.0 || >=22.0.0" + "node": "^20.0.0 || ^22.0.0 || >=24.0.0" }, "funding": { "url": "https://opencollective.com/vitest" }, "peerDependencies": { "@edge-runtime/vm": "*", - "@types/debug": "^4.1.12", - "@types/node": "^18.0.0 || ^20.0.0 || >=22.0.0", - "@vitest/browser": "3.2.6", - "@vitest/ui": "3.2.6", + "@opentelemetry/api": "^1.9.0", + "@types/node": "^20.0.0 || ^22.0.0 || >=24.0.0", + "@vitest/browser-playwright": "4.1.11", + "@vitest/browser-preview": "4.1.11", + "@vitest/browser-webdriverio": "4.1.11", + "@vitest/coverage-istanbul": "4.1.11", + "@vitest/coverage-v8": "4.1.11", + "@vitest/ui": "4.1.11", "happy-dom": "*", - "jsdom": "*" + "jsdom": "*", + "vite": "^6.0.0 || ^7.0.0 || ^8.0.0" }, "peerDependenciesMeta": { "@edge-runtime/vm": { "optional": true }, - "@types/debug": { + "@opentelemetry/api": { "optional": true }, "@types/node": { "optional": true }, - "@vitest/browser": { + "@vitest/browser-playwright": { + "optional": true + }, + "@vitest/browser-preview": { + "optional": true + }, + "@vitest/browser-webdriverio": { + "optional": true + }, + "@vitest/coverage-istanbul": { + "optional": true + }, + "@vitest/coverage-v8": { "optional": true }, "@vitest/ui": { @@ -12655,6 +12527,9 @@ }, "jsdom": { "optional": true + }, + "vite": { + "optional": false } } }, diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index ededdfb4606..9786d5c1d6e 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -40,7 +40,7 @@ "jwt-decode": "4.0.0", "lucide-react": "0.513.0", "moment": "2.30.1", - "next": "16.2.11", + "next": "16.3.3", "next-themes": "^0.4.6", "nuqs": "^2.9.4", "openai": "4.104.0", @@ -74,10 +74,10 @@ "@types/react-copy-to-clipboard": "5.0.7", "@types/react-dom": "19.2.4", "@types/react-syntax-highlighter": "15.5.13", - "@vitest/coverage-v8": "3.2.6", - "@vitest/ui": "3.2.6", + "@vitest/coverage-v8": "4.1.11", + "@vitest/ui": "4.1.11", "eslint": "9.39.2", - "eslint-config-next": "16.2.11", + "eslint-config-next": "16.3.3", "eslint-config-prettier": "10.1.8", "eslint-plugin-jest-dom": "5.10.1", "eslint-plugin-testing-library": "7.16.2", @@ -91,11 +91,12 @@ "tw-animate-css": "1.4.0", "typescript": "5.9.3", "typescript-eslint": "8.60.1", - "vitest": "3.2.6" + "vite": "7.3.5", + "vitest": "4.1.11" }, "overrides": { "prismjs": "1.30.0", - "js-yaml": "4.3.1", + "js-yaml": "4.3.2", "brace-expansion": "5.0.9", "glob": "13.0.0", "minimatch": "10.2.4", @@ -104,7 +105,7 @@ "axios": "1.13.6", "postcss": "8.5.23", "esbuild": "0.28.1", - "sharp": "^0.35.0" + "sharp": "^0.35.4" }, "engines": { "node": ">=24.14.1", diff --git a/ui/litellm-dashboard/public/assets/logos/pointfive.png b/ui/litellm-dashboard/public/assets/logos/pointfive.png new file mode 100644 index 00000000000..4b7a6b8939e Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/pointfive.png differ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx index 1eeebe4ebba..118a8655d5a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx @@ -54,7 +54,7 @@ function ResourceBadge({ fallback, }: { resource: AccessGroupResource; - href: string; + href?: string; fallback: (id: string) => string; }) { const badge = ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.test.tsx index 7a006efce1e..1ea7286c686 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.test.tsx @@ -2,7 +2,7 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React from "react"; -import { describe, expect, it, vi } from "vitest"; +import { describe, expect, it, vi, type Mock } from "vitest"; vi.mock("@/components/ModelSelect/ModelSelect", () => ({ ModelSelect: ({ onChange }: { onChange: (values: string[]) => void }) => ( @@ -32,7 +32,7 @@ const Harness = ({ createAccessGroup }: { createAccessGroup: (body: unknown) => ); }; -const renderDialog = (overrides?: { createAccessGroup?: ReturnType }) => { +const renderDialog = (overrides?: { createAccessGroup?: Mock }) => { const createAccessGroup = overrides?.createAccessGroup ?? vi.fn().mockResolvedValue({}); const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); render( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts index 442dcd48f66..e0a5271366f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts @@ -313,6 +313,26 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { return agentData; }; +export const parseMcpPermissionsForForm = (agent: any) => ({ + allowed_mcp_servers_and_groups: { + servers: agent.object_permission?.mcp_servers ?? [], + accessGroups: agent.object_permission?.mcp_access_groups ?? [], + toolsets: agent.object_permission?.mcp_toolsets ?? [], + }, + mcp_tool_permissions: agent.object_permission?.mcp_tool_permissions ?? {}, +}); + +/** + * Always includes every MCP key (empty when cleared) so removals persist; + * the proxy merges object_permission per key, leaving non-MCP grants untouched. + */ +export const buildMcpObjectPermission = (values: any) => ({ + mcp_servers: values.allowed_mcp_servers_and_groups?.servers ?? [], + mcp_access_groups: values.allowed_mcp_servers_and_groups?.accessGroups ?? [], + mcp_toolsets: values.allowed_mcp_servers_and_groups?.toolsets ?? [], + mcp_tool_permissions: values.mcp_tool_permissions ?? {}, +}); + /** * Parse agent data for form fields */ @@ -356,5 +376,6 @@ export const parseAgentForForm = (agent: any) => { : [], // extra_headers: already an array of strings extra_headers: agent.extra_headers ?? [], + ...parseMcpPermissionsForForm(agent), }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx index de1eb153f6f..357f924cbe7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx @@ -1,4 +1,5 @@ import React from "react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; @@ -10,6 +11,12 @@ vi.mock("@/components/networking", () => ({ getAgentInfo: vi.fn(), patchAgentCall: vi.fn(), getAgentCreateMetadata: vi.fn(), + getProxyBaseUrl: vi.fn(() => ""), + getUiConfig: vi.fn(async () => ({})), + fetchMCPServers: vi.fn(async () => []), + fetchMCPAccessGroups: vi.fn(async () => []), + fetchMCPToolsets: vi.fn(async () => []), + listMCPTools: vi.fn(async () => ({ tools: [] })), })); vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ @@ -111,7 +118,14 @@ const bedrockAgentcoreInfo: AgentCreateInfo = { const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); -const renderView = () => render(); +const renderView = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render( + + + , + ); +}; const openEditor = async (user: ReturnType) => { await user.click(await screen.findByRole("tab", { name: "Settings" })); @@ -161,6 +175,7 @@ describe("AgentInfoView update payload", () => { rpm_limit: 222, session_tpm_limit: 333, session_rpm_limit: 444, + object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, }); }); @@ -201,6 +216,7 @@ describe("AgentInfoView update payload", () => { rpm_limit: 222, session_tpm_limit: 333, session_rpm_limit: 444, + object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, }); }); @@ -278,9 +294,30 @@ describe("AgentInfoView update payload", () => { api_base: "https://other.example.com", model: "langgraph/asst_1", }, + object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, }); }); + it("keeps the agent's existing MCP grants in the update payload", async () => { + const existingMcpGrants = { + mcp_servers: ["srv-1"], + mcp_access_groups: ["grp-a"], + mcp_toolsets: ["toolset-1"], + mcp_tool_permissions: { "srv-1": ["tool_x"] }, + }; + vi.mocked(networking.getAgentInfo).mockResolvedValue({ + ...A2A_AGENT, + object_permission: existingMcpGrants, + } as never); + const user = setup(); + renderView(); + await openEditor(user); + + await save(user); + + expect(patchedPayload().object_permission).toEqual(existingMcpGrants); + }); + it("preserves the full AgentCore runtime ARN (including the resource id after runtime/) across an unedited save", async () => { vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([bedrockAgentcoreInfo]); vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx index 0936b8e13db..29d97a20afe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx @@ -24,6 +24,18 @@ vi.mock("./agent_form_fields", () => ({ unmountedA2AFieldNames: () => [], })); +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: () => ({ data: [{ server_id: "srv-1", server_name: "github" }] }), +})); + +vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ + default: () =>
, +})); + +vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({ + default: () =>
, +})); + const agent = { agent_id: "agent-1", agent_name: "support-agent", @@ -62,5 +74,18 @@ describe("AgentInfoView settings", () => { expect(token).toBe("sk-test"); expect(agentId).toBe("agent-1"); expect(payload.tpm_limit).toBe(42); + const clearedMcpGrants = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }; + expect(payload.object_permission).toEqual(clearedMcpGrants); + }); + + it("shows MCP grants with server names on the overview tab", async () => { + vi.mocked(networking.getAgentInfo).mockResolvedValue({ + ...agent, + object_permission: { mcp_servers: ["srv-1"] }, + } as unknown as Agent); + + render(); + + expect(await screen.findByText("github (srv-1)")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx index eddeeec674b..6cb99e9692f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx @@ -15,16 +15,27 @@ import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } import { Agent } from "@/components/agents/types"; import { KeyResponse } from "@/components/key_team_helpers/key_list"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import KeyInfoView from "@/components/templates/key_info_view"; +import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; +import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; import AgentVirtualKeys from "./agent_virtual_keys"; import AgentFormFields, { unmountedA2AFieldNames } from "./agent_form_fields"; import DynamicAgentFormFields, { buildDynamicAgentData, unmountedDynamicFieldNames } from "./dynamic_agent_form_fields"; -import { AGENT_FORM_CONFIG, buildAgentDataFromForm, parseAgentForForm } from "./agent_config"; +import { + AGENT_FORM_CONFIG, + buildAgentDataFromForm, + buildMcpObjectPermission, + parseAgentForForm, + parseMcpPermissionsForForm, +} from "./agent_config"; import { AgentFormField, AgentFormValues, AgentNumberInput, AgentRequestPayload, + McpServerSelection, + labelWithHint, omitFieldValues, useCollapsiblePanels, } from "./AgentFormKit"; @@ -111,7 +122,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT } else { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset(parseDynamicAgentForForm(data, typeInfo)); + form.reset({ ...parseDynamicAgentForForm(data, typeInfo), ...parseMcpPermissionsForForm(data) }); } else { form.reset(parseAgentForForm(data)); } @@ -131,7 +142,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT if (agentType !== "a2a") { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset(parseDynamicAgentForForm(agent, typeInfo)); + form.reset({ ...parseDynamicAgentForForm(agent, typeInfo), ...parseMcpPermissionsForForm(agent) }); } } } @@ -139,6 +150,14 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT const selectedAgentTypeInfo = agentTypeMetadata.find((t) => t.agent_type === detectedAgentType); const watchedFormValues = useWatch({ control: form.control }); + const mcpSelection = useWatch({ control: form.control, name: "allowed_mcp_servers_and_groups" }); + const mcpToolPermissions = useWatch({ control: form.control, name: "mcp_tool_permissions" }); + const { data: mcpServers = [] } = useMCPServers(); + + const mcpServerLabel = (serverId: string) => { + const server = mcpServers.find((s) => s.server_id === serverId); + return server?.server_name ? `${server.server_name} (${serverId})` : serverId; + }; const discoveryRequest = useMemo( () => buildDiscoveryRequest(detectedAgentType, watchedFormValues || {}, selectedAgentTypeInfo), @@ -199,7 +218,10 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT ? overlayDiscoveredCardParams(built, appliedDiscoveredSelection.selected_card) : built; - await patchAgentCall(accessToken, agentId, updateData); + await patchAgentCall(accessToken, agentId, { + ...updateData, + object_permission: buildMcpObjectPermission(values), + }); toast.success("Agent updated successfully"); setIsEditing(false); fetchAgentInfo(); @@ -337,13 +359,20 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {agent.object_permission && (agent.object_permission.mcp_servers?.length || agent.object_permission.mcp_access_groups?.length || + agent.object_permission.mcp_toolsets?.length || (agent.object_permission.mcp_tool_permissions && Object.keys(agent.object_permission.mcp_tool_permissions).length > 0)) && (

MCP Tool Permissions

{agent.object_permission.mcp_servers && agent.object_permission.mcp_servers.length > 0 && ( - {agent.object_permission.mcp_servers.join(", ")} + +
+ {agent.object_permission.mcp_servers.map((serverId) => ( +
{mcpServerLabel(serverId)}
+ ))} +
+ )} {agent.object_permission.mcp_access_groups && agent.object_permission.mcp_access_groups.length > 0 && ( @@ -351,13 +380,16 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {agent.object_permission.mcp_access_groups.join(", ")} )} + {agent.object_permission.mcp_toolsets && agent.object_permission.mcp_toolsets.length > 0 && ( + {agent.object_permission.mcp_toolsets.join(", ")} + )} {agent.object_permission.mcp_tool_permissions && Object.keys(agent.object_permission.mcp_tool_permissions).length > 0 && (
{Object.entries(agent.object_permission.mcp_tool_permissions).map(([serverId, tools]) => (
- {serverId}:{" "} + {mcpServerLabel(serverId)}:{" "} {Array.isArray(tools) ? tools.join(", ") : String(tools)}
))} @@ -457,6 +489,41 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {rateLimitField("session_rpm_limit", "Session RPM Limit")}
+ +

MCP Servers

+ + + {({ value, onChange }) => ( + + )} + + +
+ ) => + form.setValue("mcp_tool_permissions", toolPerms) + } + /> +
+